From cce3af001ac591c7e378dd8bbd7a1fb7b3b08646 Mon Sep 17 00:00:00 2001
From: Mike Farah <mikefarah@gmail.com>
Date: Wed, 7 Jul 2021 20:00:46 +1000
Subject: [PATCH] Extract out compound assign logic, use it add and subtract
 ops

---
 pkg/yqlib/operator_add.go      | 22 +---------------------
 pkg/yqlib/operator_subtract.go |  6 +-----
 pkg/yqlib/operators.go         | 26 ++++++++++++++++++++++++++
 3 files changed, 28 insertions(+), 26 deletions(-)

diff --git a/pkg/yqlib/operator_add.go b/pkg/yqlib/operator_add.go
index 4bd5045c..b785bc9c 100644
--- a/pkg/yqlib/operator_add.go
+++ b/pkg/yqlib/operator_add.go
@@ -15,27 +15,7 @@ func createAddOp(lhs *ExpressionNode, rhs *ExpressionNode) *ExpressionNode {
 }
 
 func addAssignOperator(d *dataTreeNavigator, context Context, expressionNode *ExpressionNode) (Context, error) {
-	lhs, err := d.GetMatchingNodes(context, expressionNode.Lhs)
-	if err != nil {
-		return Context{}, err
-	}
-
-	assignmentOp := &Operation{OperationType: assignOpType}
-	valueOp := &Operation{OperationType: valueOpType}
-
-	for el := lhs.MatchingNodes.Front(); el != nil; el = el.Next() {
-		candidate := el.Value.(*CandidateNode)
-		valueOp.CandidateNode = candidate
-		valueExpression := &ExpressionNode{Operation: valueOp}
-
-		assignmentOpNode := &ExpressionNode{Operation: assignmentOp, Lhs: valueExpression, Rhs: createAddOp(valueExpression, expressionNode.Rhs)}
-
-		_, err = d.GetMatchingNodes(context, assignmentOpNode)
-		if err != nil {
-			return Context{}, err
-		}
-	}
-	return context, nil
+	return compoundAssignFunction(d, context, expressionNode, createAddOp)
 }
 
 func toNodes(candidate *CandidateNode) []*yaml.Node {
diff --git a/pkg/yqlib/operator_subtract.go b/pkg/yqlib/operator_subtract.go
index af152771..9c5bcc0f 100644
--- a/pkg/yqlib/operator_subtract.go
+++ b/pkg/yqlib/operator_subtract.go
@@ -15,11 +15,7 @@ func createSubtractOp(lhs *ExpressionNode, rhs *ExpressionNode) *ExpressionNode
 }
 
 func subtractAssignOperator(d *dataTreeNavigator, context Context, expressionNode *ExpressionNode) (Context, error) {
-	assignmentOp := &Operation{OperationType: assignOpType}
-	assignmentOp.UpdateAssign = true
-	selfExpression := &ExpressionNode{Operation: &Operation{OperationType: selfReferenceOpType}}
-	assignmentOpNode := &ExpressionNode{Operation: assignmentOp, Lhs: expressionNode.Lhs, Rhs: createSubtractOp(selfExpression, expressionNode.Rhs)}
-	return d.GetMatchingNodes(context, assignmentOpNode)
+	return compoundAssignFunction(d, context, expressionNode, createSubtractOp)
 }
 
 func subtractOperator(d *dataTreeNavigator, context Context, expressionNode *ExpressionNode) (Context, error) {
diff --git a/pkg/yqlib/operators.go b/pkg/yqlib/operators.go
index 5b138a2b..70d9d8d5 100644
--- a/pkg/yqlib/operators.go
+++ b/pkg/yqlib/operators.go
@@ -10,6 +10,32 @@ import (
 
 type operatorHandler func(d *dataTreeNavigator, context Context, expressionNode *ExpressionNode) (Context, error)
 
+type compoundCalculation func(lhs *ExpressionNode, rhs *ExpressionNode) *ExpressionNode
+
+func compoundAssignFunction(d *dataTreeNavigator, context Context, expressionNode *ExpressionNode, calculation compoundCalculation) (Context, error) {
+	lhs, err := d.GetMatchingNodes(context, expressionNode.Lhs)
+	if err != nil {
+		return Context{}, err
+	}
+
+	assignmentOp := &Operation{OperationType: assignOpType}
+	valueOp := &Operation{OperationType: valueOpType}
+
+	for el := lhs.MatchingNodes.Front(); el != nil; el = el.Next() {
+		candidate := el.Value.(*CandidateNode)
+		valueOp.CandidateNode = candidate
+		valueExpression := &ExpressionNode{Operation: valueOp}
+
+		assignmentOpNode := &ExpressionNode{Operation: assignmentOp, Lhs: valueExpression, Rhs: calculation(valueExpression, expressionNode.Rhs)}
+
+		_, err = d.GetMatchingNodes(context, assignmentOpNode)
+		if err != nil {
+			return Context{}, err
+		}
+	}
+	return context, nil
+}
+
 func unwrapDoc(node *yaml.Node) *yaml.Node {
 	if node.Kind == yaml.DocumentNode {
 		return node.Content[0]