Variable loop - Fixes #1566 (#1577)

* Variable loop wip

* Variable loop wip

* Variable loop wip

* Variable loop wip

* Fixed variable operator to work like jq

pkg/yqlib/lib.go

@@ -53,7 +53,7 @@ var subtractAssignOpType = &operationType{Type: "SUBTRACT_ASSIGN", NumArgs: 2, P
 
 var assignAttributesOpType = &operationType{Type: "ASSIGN_ATTRIBUTES", NumArgs: 2, Precedence: 40, Handler: assignAttributesOperator}
 var assignStyleOpType = &operationType{Type: "ASSIGN_STYLE", NumArgs: 2, Precedence: 40, Handler: assignStyleOperator}
-var assignVariableOpType = &operationType{Type: "ASSIGN_VARIABLE", NumArgs: 2, Precedence: 40, Handler: assignVariableOperator}
+var assignVariableOpType = &operationType{Type: "ASSIGN_VARIABLE", NumArgs: 2, Precedence: 40, Handler: useWithPipe}
 var assignTagOpType = &operationType{Type: "ASSIGN_TAG", NumArgs: 2, Precedence: 40, Handler: assignTagOperator}
 var assignCommentOpType = &operationType{Type: "ASSIGN_COMMENT", NumArgs: 2, Precedence: 40, Handler: assignCommentsOperator}
 var assignAnchorOpType = &operationType{Type: "ASSIGN_ANCHOR", NumArgs: 2, Precedence: 40, Handler: assignAnchorOperator}

pkg/yqlib/operator_add_test.go

@@ -34,7 +34,7 @@ var addOperatorScenarios = []expressionScenario{
 	{
 		skipDoc:    true,
 		document:   `{}`,
-		expression: "(.a + .b) as $x",
+		expression: "(.a + .b) as $x | .",
 		expected: []string{
 			"D0, P[], (doc)::{}\n",
 		},

pkg/yqlib/operator_alternative_test.go

@@ -16,7 +16,7 @@ var alternativeOperatorScenarios = []expressionScenario{
 	},
 	{
 		skipDoc:    true,
-		expression: `(.b // "hello") as $x`,
+		expression: `(.b // "hello") as $x | .`,
 		document:   `a: bridge`,
 		expected: []string{
 			"D0, P[], (doc)::a: bridge\n",

pkg/yqlib/operator_booleans_test.go

@@ -102,17 +102,25 @@ var booleanOperatorScenarios = []expressionScenario{
 	{
 		skipDoc:    true,
 		document:   `[{pet: cat}]`,
-		expression: `any_c(.name == "harry") as $c`,
+		expression: `any_c(.name == "harry") as $c | .`,
 		expected: []string{
 			"D0, P[], (doc)::[{pet: cat}]\n",
 		},
 	},
 	{
 		skipDoc:    true,
 		document:   `[{pet: cat}]`,
-		expression: `all_c(.name == "harry") as $c`,
+		expression: `any_c(.name == "harry") as $c | $c`,
 		expected: []string{
-			"D0, P[], (doc)::[{pet: cat}]\n",
+			"D0, P[], (!!bool)::false\n",
+		},
+	},
+	{
+		skipDoc:    true,
+		document:   `[{pet: cat}]`,
+		expression: `all_c(.name == "harry") as $c | $c`,
+		expected: []string{
+			"D0, P[], (!!bool)::false\n",
 		},
 	},
 	{
@@ -185,15 +193,15 @@ var booleanOperatorScenarios = []expressionScenario{
 	{
 		skipDoc:    true,
 		document:   `{}`,
-		expression: `(.a.b or .c) as $x`,
+		expression: `(.a.b or .c) as $x | .`,
 		expected: []string{
 			"D0, P[], (doc)::{}\n",
 		},
 	},
 	{
 		skipDoc:    true,
 		document:   `{}`,
-		expression: `(.a.b and .c) as $x`,
+		expression: `(.a.b and .c) as $x | .`,
 		expected: []string{
 			"D0, P[], (doc)::{}\n",
 		},

pkg/yqlib/operator_equals_test.go

@@ -47,7 +47,7 @@ var equalsOperatorScenarios = []expressionScenario{
 	{
 		skipDoc:    true,
 		document:   "{}",
-		expression: "(.a == .b) as $x",
+		expression: "(.a == .b) as $x | .",
 		expected: []string{
 			"D0, P[], (doc)::{}\n",
 		},
@@ -63,7 +63,7 @@ var equalsOperatorScenarios = []expressionScenario{
 	{
 		skipDoc:    true,
 		document:   "{}",
-		expression: "(.a != .b) as $x",
+		expression: "(.a != .b) as $x | .",
 		expected: []string{
 			"D0, P[], (doc)::{}\n",
 		},

pkg/yqlib/operator_has_test.go

@@ -16,7 +16,7 @@ var hasOperatorScenarios = []expressionScenario{
 	{
 		skipDoc:    true,
 		document:   `a: hello`,
-		expression: `has(.b) as $c`,
+		expression: `has(.b) as $c | .`,
 		expected: []string{
 			"D0, P[], (doc)::a: hello\n",
 		},

pkg/yqlib/operator_pipe.go

@@ -2,9 +2,9 @@ package yqlib
 
 func pipeOperator(d *dataTreeNavigator, context Context, expressionNode *ExpressionNode) (Context, error) {
 
-	//lhs may update the variable context, we should pass that into the RHS
-	// BUT we still return the original context back (see jq)
-	// https://stedolan.github.io/jq/manual/#Variable/SymbolicBindingOperator:...as$identifier|...
+	if expressionNode.LHS.Operation.OperationType == assignVariableOpType {
+		return variableLoop(d, context, expressionNode)
+	}
 
 	lhs, err := d.GetMatchingNodes(context, expressionNode.LHS)
 	if err != nil {

pkg/yqlib/operator_subtract_test.go

@@ -8,7 +8,7 @@ var subtractOperatorScenarios = []expressionScenario{
 	{
 		skipDoc:    true,
 		document:   `{}`,
-		expression: "(.a - .b) as $x",
+		expression: "(.a - .b) as $x | .",
 		expected: []string{
 			"D0, P[], (doc)::{}\n",
 		},

pkg/yqlib/operator_traverse_path_test.go

@@ -151,7 +151,7 @@ var traversePathOperatorScenarios = []expressionScenario{
 	{
 		skipDoc:    true,
 		document:   `c: dog`,
-		expression: `.[.a.b] as $x`,
+		expression: `.[.a.b] as $x | .`,
 		expected: []string{
 			"D0, P[], (doc)::c: dog\n",
 		},

pkg/yqlib/operator_union_test.go

@@ -8,7 +8,7 @@ var unionOperatorScenarios = []expressionScenario{
 	{
 		skipDoc:    true,
 		document:   "{}",
-		expression: `(.a, .b.c) as $x`,
+		expression: `(.a, .b.c) as $x | .`,
 		expected: []string{
 			"D0, P[], (doc)::{}\n",
 		},

pkg/yqlib/operator_variables.go

@@ -19,24 +19,81 @@ type assignVarPreferences struct {
 	IsReference bool
 }
 
-func assignVariableOperator(d *dataTreeNavigator, context Context, expressionNode *ExpressionNode) (Context, error) {
-	lhs, err := d.GetMatchingNodes(context.ReadOnlyClone(), expressionNode.LHS)
+func useWithPipe(d *dataTreeNavigator, context Context, originalExp *ExpressionNode) (Context, error) {
+	return Context{}, fmt.Errorf("must use variable with a pipe, e.g. `exp as $x | ...`")
+}
+
+// variables are like loops in jq
+// https://stedolan.github.io/jq/manual/#Variable
+func variableLoop(d *dataTreeNavigator, context Context, originalExp *ExpressionNode) (Context, error) {
+	log.Debug("variable loop!")
+	results := list.New()
+	var evaluateAllTogether = true
+	for matchEl := context.MatchingNodes.Front(); matchEl != nil; matchEl = matchEl.Next() {
+		evaluateAllTogether = evaluateAllTogether && matchEl.Value.(*CandidateNode).EvaluateTogether
+		if !evaluateAllTogether {
+			break
+		}
+	}
+	if evaluateAllTogether {
+		return variableLoopSingleChild(d, context, originalExp)
+	}
+
+	for el := context.MatchingNodes.Front(); el != nil; el = el.Next() {
+		result, err := variableLoopSingleChild(d, context.SingleChildContext(el.Value.(*CandidateNode)), originalExp)
+		if err != nil {
+			return Context{}, err
+		}
+		results.PushBackList(result.MatchingNodes)
+	}
+	return context.ChildContext(results), nil
+}
+
+func variableLoopSingleChild(d *dataTreeNavigator, context Context, originalExp *ExpressionNode) (Context, error) {
+
+	variableExp := originalExp.LHS
+	lhs, err := d.GetMatchingNodes(context.ReadOnlyClone(), variableExp.LHS)
 	if err != nil {
 		return Context{}, err
 	}
-	if expressionNode.RHS.Operation.OperationType.Type != "GET_VARIABLE" {
+	if variableExp.RHS.Operation.OperationType.Type != "GET_VARIABLE" {
 		return Context{}, fmt.Errorf("RHS of 'as' operator must be a variable name e.g. $foo")
 	}
-	variableName := expressionNode.RHS.Operation.StringValue
+	variableName := variableExp.RHS.Operation.StringValue
+
+	prefs := variableExp.Operation.Preferences.(assignVarPreferences)
 
-	prefs := expressionNode.Operation.Preferences.(assignVarPreferences)
+	results := list.New()
 
-	var variableValue *list.List
-	if prefs.IsReference {
-		variableValue = lhs.MatchingNodes
-	} else {
-		variableValue = lhs.DeepClone().MatchingNodes
+	// now we loop over lhs, set variable to each result and calculate originalExp.Rhs
+	for el := lhs.MatchingNodes.Front(); el != nil; el = el.Next() {
+		log.Debug("PROCESSING VARIABLE: ", NodeToString(el.Value.(*CandidateNode)))
+		var variableValue = list.New()
+		if prefs.IsReference {
+			variableValue.PushBack(el.Value)
+		} else {
+			candidateCopy, err := el.Value.(*CandidateNode).Copy()
+			if err != nil {
+				return Context{}, err
+			}
+			variableValue.PushBack(candidateCopy)
+		}
+		newContext := context.ChildContext(context.MatchingNodes)
+		newContext.SetVariable(variableName, variableValue)
+
+		rhs, err := d.GetMatchingNodes(newContext, originalExp.RHS)
+		log.Debug("PROCESSING VARIABLE DONE, got back: ", rhs.MatchingNodes.Len())
+		if err != nil {
+			return Context{}, err
+		}
+		results.PushBackList(rhs.MatchingNodes)
 	}
-	context.SetVariable(variableName, variableValue)
-	return context, nil
+
+	// if there is no LHS - then I guess we just calculate originalExp.Rhs
+	if lhs.MatchingNodes.Len() == 0 {
+		return d.GetMatchingNodes(context, originalExp.RHS)
+	}
+
+	return context.ChildContext(results), nil
+
 }

pkg/yqlib/operator_variables_test.go

@@ -8,15 +8,15 @@ var variableOperatorScenarios = []expressionScenario{
 	{
 		skipDoc:    true,
 		document:   `{}`,
-		expression: `.a.b as $foo`,
+		expression: `.a.b as $foo | .`,
 		expected: []string{
 			"D0, P[], (doc)::{}\n",
 		},
 	},
 	{
 		document:      "a: [cat]",
 		skipDoc:       true,
-		expression:    "(.[] | {.name: .}) as $item",
+		expression:    "(.[] | {.name: .}) as $item | .",
 		expectedError: `cannot index array with 'name' (strconv.ParseInt: parsing "name": invalid syntax)`,
 	},
 	{
@@ -36,6 +36,22 @@ var variableOperatorScenarios = []expressionScenario{
 			"D0, P[1], (!!str)::dog\n",
 		},
 	},
+	{
+		skipDoc:    true,
+		document:   `[1, 2]`,
+		expression: `.[] | . as $f | select($f == 2)`,
+		expected: []string{
+			"D0, P[1], (!!int)::2\n",
+		},
+	},
+	{
+		skipDoc:    true,
+		document:   `[1, 2]`,
+		expression: `[.[] | . as $f | $f + 1]`,
+		expected: []string{
+			"D0, P[], (!!seq)::- 2\n- 3\n",
+		},
+	},
 	{
 		description:    "Using variables as a lookup",
 		subdescription: "Example taken from [jq](https://stedolan.github.io/jq/manual/#Variable/SymbolicBindingOperator:...as$identifier|...)",