Fix singleton unwrapping for $each, $spread, and $match (RecoLabs#3)

These three functions returned []any instead of *Sequence,
preventing CollapseSequence from unwrapping single-element
results. This caused parity divergence from jsonata-js:

- $each({"a":1}, fn) returned ["result"] instead of "result"
- $spread({"a":1}) returned [{"a":1}] instead of {"a":1}
- $match("hello", /^h/) returned [{...}] instead of {...}

Return *Sequence from all three so evalFunction properly
singleton-unwraps single-element results. Also introduces
CollapseAndKeep helper in value.go (with shallow struct copy
for alias safety) to centralize KeepArray/singleton logic,
and collapses intermediate results in function composition
to prevent raw *Sequence leaking between composed functions.

Signed-off-by: NirBarak-RecoLabs <nirb@recolabs.ai>

Author: NirBarak-RecoLabs

evaluator_test.go

@@ -495,6 +495,79 @@ func TestRegressionExpr(t *testing.T) {
 	}
 }
 
+// TestSingletonUnwrapEdgeCases covers scenarios the JSON suite cannot:
+// error expectations and chain operator (~>) paths.
+func TestSingletonUnwrapEdgeCases(t *testing.T) {
+	tests := []struct {
+		name      string
+		expr      string
+		data      string
+		want      any
+		wantError bool
+	}{
+		{
+			name: "join_nested_each_single_key_errors",
+			expr: `$join($each({"x": {"p": "hi", "q": "lo"}},` +
+				` function($obj, $name) { $each($obj,` +
+				` function($v, $k) { $name & "." & $k & "=" & $v }) })[], ', ')`,
+			data:      `{}`,
+			wantError: true,
+		},
+		{
+			name: "chain_each_single_key",
+			expr: `{"a": 1} ~> $each(function($v, $k) { $k })`,
+			data: `{}`,
+			want: "a",
+		},
+		{
+			name: "chain_spread_single_key",
+			expr: `{"a": 1} ~> $spread()`,
+			data: `{}`,
+			want: map[string]any{"a": float64(1)},
+		},
+		{
+			name: "chain_bare_spread_single_key",
+			expr: `{"a": 1} ~> $spread`,
+			data: `{}`,
+			want: map[string]any{"a": float64(1)},
+		},
+		{
+			name: "composition_spread_merge",
+			expr: `($spread ~> $merge)({"a": 1, "b": 2})`,
+			data: `{}`,
+			want: map[string]any{"a": float64(1), "b": float64(2)},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			var data any
+			if err := json.Unmarshal([]byte(tt.data), &data); err != nil {
+				t.Fatalf("unmarshal: %v", err)
+			}
+			e, err := gnata.Compile(tt.expr)
+			if err != nil {
+				t.Fatalf("compile %q: %v", tt.expr, err)
+			}
+			result, err := e.Eval(context.Background(), data)
+			if tt.wantError {
+				if err == nil {
+					t.Fatalf("expected error, got %v", result)
+				}
+				return
+			}
+			if err != nil {
+				t.Fatalf("eval %q: %v", tt.expr, err)
+			}
+			if tt.want != nil && !gnata.DeepEqual(result, tt.want) {
+				got, _ := json.Marshal(result)
+				want, _ := json.Marshal(tt.want)
+				t.Fatalf("got %s, want %s", got, want)
+			}
+		})
+	}
+}
+
 // TestRegressionBytes covers regressions requiring direct EvalBytes
 // access (e.g. large integers that exceed float64 precision).
 func TestRegressionBytes(t *testing.T) {

functions/object_funcs.go

@@ -113,7 +113,7 @@ func fnSpread(args []any, _ any) (any, error) {
 		return result
 	}
 	if evaluator.IsMap(args[0]) {
-		return spreadOne(args[0]), nil
+		return &evaluator.Sequence{Values: spreadOne(args[0])}, nil
 	}
 	if arr, ok := args[0].([]any); ok {
 		var result []any
@@ -241,18 +241,18 @@ func makeFnEach(evalFn EvalFn) evaluator.EnvAwareBuiltin {
 		}
 
 		keys := evaluator.MapKeys(objVal)
-		result := make([]any, 0, len(keys))
+		seq := evaluator.CreateSequence()
 		for _, ks := range keys {
 			val, _ := evaluator.MapGet(objVal, ks)
 			res, err := evalFn(fn, []any{val, ks}, focus, env)
 			if err != nil {
 				return nil, err
 			}
 			if res != nil {
-				result = append(result, res)
+				seq.Values = append(seq.Values, res)
 			}
 		}
-		return result, nil
+		return seq, nil
 	}
 }
 

functions/string_match_replace.go

@@ -73,11 +73,15 @@ func makeFnMatch(evalFn EvalFn) evaluator.EnvAwareBuiltin {
 				return nil, &evaluator.JSONataError{Code: "D3137", Message: fmt.Sprintf("regex error: %v", matchErr)}
 			}
 		}
-		if len(result) == 0 {
-			return nil, nil
-		}
-		return result, nil
+		return matchResultSeq(result), nil
+	}
+}
+
+func matchResultSeq(result []any) any {
+	if len(result) == 0 {
+		return nil
 	}
+	return &evaluator.Sequence{Values: result}
 }
 
 func matchWithCustomMatcher(s string, matcherFn any, limit int, evalFn EvalFn, env *evaluator.Environment) (any, error) {
@@ -119,10 +123,7 @@ func matchWithCustomMatcher(s string, matcherFn any, limit int, evalFn EvalFn, e
 		}
 	}
 
-	if len(result) == 0 {
-		return nil, nil
-	}
-	return result, nil
+	return matchResultSeq(result), nil
 }
 
 // ── $replace ──────────────────────────────────────────────────────────────────

internal/evaluator/eval_chain.go

@@ -38,20 +38,7 @@ func evalChain(right *parser.Node, piped, input any, env *Environment) (any, err
 		if err != nil {
 			return nil, err
 		}
-		if right.KeepArray {
-			if seq, ok := result.(*Sequence); ok {
-				result = CollapseSequence(seq)
-			}
-			switch result.(type) {
-			case []any:
-				return result, nil
-			case nil:
-				return nil, nil
-			default:
-				return []any{result}, nil
-			}
-		}
-		return result, nil
+		return CollapseAndKeep(result, right.KeepArray), nil
 	}
 	// Right is a function reference or other expression.
 	fn, err := Eval(right, input, env)
@@ -87,10 +74,19 @@ func evalChain(right *parser.Node, piped, input any, env *Environment) (any, err
 			if err != nil {
 				return nil, err
 			}
-			return callFunction(fn, []any{intermediate}, focus, env)
+			intermediate = CollapseAndKeep(intermediate, false)
+			res, err := callFunction(fn, []any{intermediate}, focus, env)
+			if err != nil {
+				return nil, err
+			}
+			return CollapseAndKeep(res, false), nil
 		}), nil
 	}
-	return callFunction(fn, []any{piped}, input, env)
+	result, err := callFunction(fn, []any{piped}, input, env)
+	if err != nil {
+		return nil, err
+	}
+	return CollapseAndKeep(result, false), nil
 }
 
 func evalBlock(node *parser.Node, input any, env *Environment) (any, error) {

internal/evaluator/eval_function.go

@@ -76,23 +76,7 @@ func evalFunction(node *parser.Node, input any, env *Environment) (any, error) {
 	if err != nil {
 		return nil, err
 	}
-	// Normalize Sequence results: jsonata-js collapses Sequence values
-	// after every function call. Single-element → scalar, multi → []any.
-	// Functions returning *Sequence ($keys, $distinct) rely on this.
-	if seq, ok := result.(*Sequence); ok {
-		result = CollapseSequence(seq)
-	}
-	if node.KeepArray {
-		switch result.(type) {
-		case []any:
-			return result, nil
-		case nil:
-			return nil, nil
-		default:
-			return []any{result}, nil
-		}
-	}
-	return result, nil
+	return CollapseAndKeep(result, node.KeepArray), nil
 }
 
 func evalLambda(node *parser.Node, input any, env *Environment) (any, error) {

internal/evaluator/value.go

@@ -65,6 +65,32 @@ func CollapseSequence(s *Sequence) any {
 	}
 }
 
+// CollapseAndKeep normalizes a function call result for callers that need
+// KeepArray (the [] suffix) support. Builtins returning *Sequence rely on
+// CollapseSequence; when keepArray is true, singletons are preserved as
+// one-element arrays instead of being unwrapped.
+func CollapseAndKeep(result any, keepArray bool) any {
+	if seq, ok := result.(*Sequence); ok {
+		if keepArray {
+			s := *seq
+			s.KeepSingleton = true
+			seq = &s
+		}
+		result = CollapseSequence(seq)
+	}
+	if keepArray {
+		switch result.(type) {
+		case []any:
+			return result
+		case nil:
+			return nil
+		default:
+			return []any{result}
+		}
+	}
+	return result
+}
+
 // IsArray returns true for []any values (not *Sequence).
 func IsArray(v any) bool {
 	_, ok := v.([]any)

testdata/groups/function-each/case003.json

@@ -0,0 +1,6 @@
+{
+    "expr": "$each({'a': 1}, function($v, $k) { $k })",
+    "data": {},
+    "bindings": {},
+    "result": "a"
+}

testdata/groups/function-each/case004.json

@@ -0,0 +1,6 @@
+{
+    "expr": "$each({'a': 1}, function($v, $k) { [1, 2, 3] })",
+    "data": {},
+    "bindings": {},
+    "result": [1, 2, 3]
+}

testdata/groups/function-each/case005.json

@@ -0,0 +1,6 @@
+{
+    "expr": "$each({'a': 1}, function($v, $k) { $k })[]",
+    "data": {},
+    "bindings": {},
+    "result": ["a"]
+}

testdata/groups/function-each/case006.json

@@ -0,0 +1,7 @@
+{
+    "expr": "$each({'x': {'p': 1, 'q': 2}}, function($v, $k) { $each($v, function($val, $key) { $k & '.' & $key }) })",
+    "data": {},
+    "bindings": {},
+    "result": ["x.p", "x.q"],
+    "unordered": true
+}