Refactor loop aggregation logic; optimize destination register handling, adjust collector allocation, and improve resource management in nested aggregation scenarios. Update integration tests with consistent formatting for nested FOR cases.

2025-08-15 20:02:56 +02:00 · 2025-06-25 16:23:39 -04:00
parent dd817e7115
commit e60a02ec12
2 changed files with 431 additions and 245 deletions
--- a/pkg/compiler/internal/loop_collect_agg.go
+++ b/pkg/compiler/internal/loop_collect_agg.go
@@ -49,14 +49,12 @@ func (c *LoopCollectCompiler) compileGroupedAggregation(ctx fql.ICollectAggregat
 func (c *LoopCollectCompiler) compileGlobalAggregation(ctx fql.ICollectAggregatorContext) {
 	parentLoop := c.ctx.Loops.Current()
 	// we create a custom collector for aggregators
-	c.ctx.Emitter.PatchSwapAx(parentLoop.Pos, vm.OpDataSetCollector, parentLoop.Dst, int(core.CollectorTypeKeyGroup))
-
+	dst := parentLoop.PatchDestinationAx(c.ctx.Registers, c.ctx.Emitter, vm.OpDataSetCollector, int(core.CollectorTypeKeyGroup))
 	// Nested scope for aggregators
 	c.ctx.Symbols.EnterScope()
 	// Now we add value selectors to the collector
 	selectors := ctx.AllCollectAggregateSelector()
-	argsPkg := c.compileAggregationFuncArgs(selectors, parentLoop.Dst)
-
+	argsPkg := c.compileAggregationFuncArgs(selectors, dst)
 	parentLoop.EmitFinalization(c.ctx.Emitter)
 	c.ctx.Loops.Pop()
 	c.ctx.Symbols.ExitScope()
@@ -66,14 +64,18 @@ func (c *LoopCollectCompiler) compileGlobalAggregation(ctx fql.ICollectAggregato
 	c.ctx.Emitter.EmitA(vm.OpLoadZero, zero)
 	// We move the aggregator to a temporary register to access it later from the new loop
 	aggregator := c.ctx.Registers.Allocate(core.Temp)
-	c.ctx.Emitter.EmitAB(vm.OpMove, aggregator, parentLoop.Dst)
+	c.ctx.Emitter.EmitAB(vm.OpMove, aggregator, dst)
+
+	if parentLoop.Dst != dst && !parentLoop.Allocate {
+		c.ctx.Registers.Free(dst)
+	}

 	// CreateFor new loop with 1 iteration only
 	c.ctx.Symbols.EnterScope()
 	c.ctx.Emitter.EmitABC(vm.OpRange, parentLoop.Src, zero, zero)
 	loop := c.ctx.Loops.CreateFor(core.TemporalLoop, parentLoop.Src, parentLoop.Distinct)
 	loop.Dst = parentLoop.Dst
-	loop.Allocate = true
+	loop.Allocate = parentLoop.Allocate
 	c.ctx.Loops.Push(loop)
 	loop.EmitInitialization(c.ctx.Registers, c.ctx.Emitter)

--- a/test/integration/vm/vm_for_nested_test.go
+++ b/test/integration/vm/vm_for_nested_test.go
@@ -8,223 +8,223 @@ import (

 func TestForNested(t *testing.T) {
 	RunUseCases(t, []UseCase{
-		CaseArray(
-			`FOR prop IN ["a"]
-							FOR val IN [1, 2, 3]
-								RETURN {[prop]: val}`,
-			[]any{map[string]any{"a": 1}, map[string]any{"a": 2}, map[string]any{"a": 3}},
-		),
-		CaseArray(
-			`FOR val IN 1..3
-							FOR prop IN ["a"]
-								RETURN {[prop]: val}`,
-			[]any{map[string]any{"a": 1}, map[string]any{"a": 2}, map[string]any{"a": 3}},
-		),
-		CaseArray(
-			`FOR prop IN ["a"]
-							FOR val IN 1..3
-								RETURN {[prop]: val}`,
-			[]any{map[string]any{"a": 1}, map[string]any{"a": 2}, map[string]any{"a": 3}},
-		),
-		CaseArray(
-			`FOR prop IN ["a"]
-							FOR val IN [1, 2, 3]
-								FOR val2 IN [1, 2, 3]
-									RETURN { [prop]: [val, val2] }`,
-			[]any{map[string]any{"a": []int{1, 1}}, map[string]any{"a": []int{1, 2}}, map[string]any{"a": []int{1, 3}}, map[string]any{"a": []int{2, 1}}, map[string]any{"a": []int{2, 2}}, map[string]any{"a": []int{2, 3}}, map[string]any{"a": []int{3, 1}}, map[string]any{"a": []int{3, 2}}, map[string]any{"a": []int{3, 3}}},
-		),
-		CaseArray(
-			`FOR val IN [1, 2, 3]
-							RETURN (
-								FOR prop IN ["a", "b", "c"]
-									RETURN { [prop]: val }
-							)`,
-			[]any{[]any{map[string]any{"a": 1}, map[string]any{"b": 1}, map[string]any{"c": 1}}, []any{map[string]any{"a": 2}, map[string]any{"b": 2}, map[string]any{"c": 2}}, []any{map[string]any{"a": 3}, map[string]any{"b": 3}, map[string]any{"c": 3}}},
-		),
-		CaseArray(
-			`FOR val IN [1, 2, 3]
-							LET sub = (
-								FOR prop IN ["a", "b", "c"]
-									RETURN { [prop]: val }
-							)
-		
-							RETURN sub`,
-			[]any{[]any{map[string]any{"a": 1}, map[string]any{"b": 1}, map[string]any{"c": 1}}, []any{map[string]any{"a": 2}, map[string]any{"b": 2}, map[string]any{"c": 2}}, []any{map[string]any{"a": 3}, map[string]any{"b": 3}, map[string]any{"c": 3}}},
+		CaseArray(`
+			FOR prop IN ["a"]
+				FOR val IN [1, 2, 3]
+					RETURN {[prop]: val}
+`, []any{map[string]any{"a": 1}, map[string]any{"a": 2}, map[string]any{"a": 3}},
 		),
 		CaseArray(`
-LET strs = ["foo", "bar", "qaz", "abc"]
+			FOR val IN 1..3
+				FOR prop IN ["a"]
+					RETURN {[prop]: val}
+`, []any{map[string]any{"a": 1}, map[string]any{"a": 2}, map[string]any{"a": 3}},
+		),
+		CaseArray(`
+			FOR prop IN ["a"]
+				FOR val IN 1..3
+					RETURN {[prop]: val}
+`, []any{map[string]any{"a": 1}, map[string]any{"a": 2}, map[string]any{"a": 3}},
+		),
+		CaseArray(`
+			FOR prop IN ["a"]
+				FOR val IN [1, 2, 3]
+					FOR val2 IN [1, 2, 3]
+						RETURN { [prop]: [val, val2] }
+`, []any{map[string]any{"a": []int{1, 1}}, map[string]any{"a": []int{1, 2}}, map[string]any{"a": []int{1, 3}}, map[string]any{"a": []int{2, 1}}, map[string]any{"a": []int{2, 2}}, map[string]any{"a": []int{2, 3}}, map[string]any{"a": []int{3, 1}}, map[string]any{"a": []int{3, 2}}, map[string]any{"a": []int{3, 3}}},
+		),
+		CaseArray(`
+			FOR val IN [1, 2, 3]
+				RETURN (
+					FOR prop IN ["a", "b", "c"]
+						RETURN { [prop]: val }
+				)
+`, []any{[]any{map[string]any{"a": 1}, map[string]any{"b": 1}, map[string]any{"c": 1}}, []any{map[string]any{"a": 2}, map[string]any{"b": 2}, map[string]any{"c": 2}}, []any{map[string]any{"a": 3}, map[string]any{"b": 3}, map[string]any{"c": 3}}},
+		),
+		CaseArray(`
+			FOR val IN [1, 2, 3]
+			LET sub = (
+				FOR prop IN ["a", "b", "c"]
+					RETURN { [prop]: val }
+			)

-FOR s IN strs
-	SORT s
-	FOR n IN 0..1
-		RETURN CONCAT(s, n)
+			RETURN sub
+`, []any{[]any{map[string]any{"a": 1}, map[string]any{"b": 1}, map[string]any{"c": 1}}, []any{map[string]any{"a": 2}, map[string]any{"b": 2}, map[string]any{"c": 2}}, []any{map[string]any{"a": 3}, map[string]any{"b": 3}, map[string]any{"c": 3}}},
+		),
+		CaseArray(`
+			LET strs = ["foo", "bar", "qaz", "abc"]
+			
+			FOR s IN strs
+				SORT s
+				FOR n IN 0..1
+					RETURN CONCAT(s, n)
 `, []any{"abc0", "abc1", "bar0", "bar1", "foo0", "foo1", "qaz0", "qaz1"}),
 		CaseArray(`
-LET strs = ["foo", "bar", "qaz", "abc"]
+			LET strs = ["foo", "bar", "qaz", "abc"]
 			
-FOR n IN 0..1
-	FOR s IN strs
-		SORT s
-		RETURN CONCAT(s, n)
+			FOR n IN 0..1
+				FOR s IN strs
+					SORT s
+					RETURN CONCAT(s, n)
 `, []any{"abc0", "bar0", "foo0", "qaz0", "abc1", "bar1", "foo1", "qaz1"}),
 		CaseArray(`
-LET users = [
-	{
-		name: "Ron",
-		age: 31,
-		gender: "m"
-	},
-	{
-		name: "Angela",
-		age: 29,
-		gender: "f"
-	},
-	{
-		name: "Bob",
-		age: 36,
-		gender: "m"
-	}
-]
+			LET users = [
+				{
+					name: "Ron",
+					age: 31,
+					gender: "m"
+				},
+				{
+					name: "Angela",
+					age: 29,
+					gender: "f"
+				},
+				{
+					name: "Bob",
+					age: 36,
+					gender: "m"
+				}
+			]
 			
-FOR n IN 0..1
-	FOR u IN users
-		SORT u.gender, u.age
-		RETURN CONCAT(u.name, n)
+			FOR n IN 0..1
+				FOR u IN users
+					SORT u.gender, u.age
+					RETURN CONCAT(u.name, n)
 `, []any{"Angela0", "Ron0", "Bob0", "Angela1", "Ron1", "Bob1"}),
 		CaseArray(`
-LET strs = ["foo", "bar", "qaz", "abc"]
+			LET strs = ["foo", "bar", "qaz", "abc"]
 			
-FOR n IN 0..1
-	FOR m IN 0..1
-		FOR s IN strs
-			SORT s
-			RETURN CONCAT(s, n, m)
+			FOR n IN 0..1
+				FOR m IN 0..1
+					FOR s IN strs
+						SORT s
+						RETURN CONCAT(s, n, m)
 `, []any{"abc00", "bar00", "foo00", "qaz00", "abc01", "bar01", "foo01", "qaz01", "abc10", "bar10", "foo10", "qaz10", "abc11", "bar11", "foo11", "qaz11"}),
 		CaseArray(`
-LET strs = ["foo", "bar", "qaz", "abc"]
+			LET strs = ["foo", "bar", "qaz", "abc"]
 			
-FOR n IN 0..1
-	FOR s IN strs
-		SORT s
-		FOR m IN 0..1
-			RETURN CONCAT(s, n, m)
+			FOR n IN 0..1
+				FOR s IN strs
+					SORT s
+					FOR m IN 0..1
+						RETURN CONCAT(s, n, m)
 `, []any{"abc00", "abc01", "bar00", "bar01", "foo00", "foo01", "qaz00", "qaz01", "abc10", "abc11", "bar10", "bar11", "foo10", "foo11", "qaz10", "qaz11"}),
 		CaseArray(`
-LET users = [
-	{
-		active: true,
-		married: true,
-		age: 31,
-		gender: "m"
-	},
-	{
-		active: true,
-		married: false,
-		age: 25,
-		gender: "f"
-	},
-	{
-		active: true,
-		married: false,
-		age: 36,
-		gender: "m"
-	},
-	{
-		active: false,
-		married: true,
-		age: 69,
-		gender: "m"
-	},
-	{
-		active: true,
-		married: true,
-		age: 45,
-		gender: "f"
-	}
-]
-FOR n IN 0..1
-	FOR i IN users
-		COLLECT gender = i.gender
-		RETURN CONCAT(gender, n)
+			LET users = [
+				{
+					active: true,
+					married: true,
+					age: 31,
+					gender: "m"
+				},
+				{
+					active: true,
+					married: false,
+					age: 25,
+					gender: "f"
+				},
+				{
+					active: true,
+					married: false,
+					age: 36,
+					gender: "m"
+				},
+				{
+					active: false,
+					married: true,
+					age: 69,
+					gender: "m"
+				},
+				{
+					active: true,
+					married: true,
+					age: 45,
+					gender: "f"
+				}
+			]
+			FOR n IN 0..1
+				FOR i IN users
+					COLLECT gender = i.gender
+					RETURN CONCAT(gender, n)
 `, []any{"f0", "m0", "f1", "m1"}),
 		CaseArray(`
-LET users = [
-	{
-		active: true,
-		married: true,
-		age: 31,
-		gender: "m"
-	},
-	{
-		active: true,
-		married: false,
-		age: 25,
-		gender: "f"
-	},
-	{
-		active: true,
-		married: false,
-		age: 36,
-		gender: "m"
-	},
-	{
-		active: false,
-		married: true,
-		age: 69,
-		gender: "m"
-	},
-	{
-		active: true,
-		married: true,
-		age: 45,
-		gender: "f"
-	}
-]
-FOR i IN users
-	COLLECT gender = i.gender
-	FOR n IN 0..1
-		RETURN CONCAT(gender, n)
+			LET users = [
+				{
+					active: true,
+					married: true,
+					age: 31,
+					gender: "m"
+				},
+				{
+					active: true,
+					married: false,
+					age: 25,
+					gender: "f"
+				},
+				{
+					active: true,
+					married: false,
+					age: 36,
+					gender: "m"
+				},
+				{
+					active: false,
+					married: true,
+					age: 69,
+					gender: "m"
+				},
+				{
+					active: true,
+					married: true,
+					age: 45,
+					gender: "f"
+				}
+			]
+			FOR i IN users
+				COLLECT gender = i.gender
+				FOR n IN 0..1
+					RETURN CONCAT(gender, n)
 `, []any{"f0", "f1", "m0", "m1"}),
 		CaseArray(`
-LET users = [
-	{
-		active: true,
-		married: true,
-		age: 31,
-		gender: "m"
-	},
-	{
-		active: true,
-		married: false,
-		age: 25,
-		gender: "f"
-	},
-	{
-		active: true,
-		married: false,
-		age: 36,
-		gender: "m"
-	},
-	{
-		active: false,
-		married: true,
-		age: 69,
-		gender: "m"
-	},
-	{
-		active: true,
-		married: true,
-		age: 45,
-		gender: "f"
-	}
-]
-FOR n IN 0..1
-	FOR i IN users
-		COLLECT gender = i.gender INTO genders
-		RETURN {
-			gender: CONCAT(gender, n),
-			values: genders
-		}
+			LET users = [
+				{
+					active: true,
+					married: true,
+					age: 31,
+					gender: "m"
+				},
+				{
+					active: true,
+					married: false,
+					age: 25,
+					gender: "f"
+				},
+				{
+					active: true,
+					married: false,
+					age: 36,
+					gender: "m"
+				},
+				{
+					active: false,
+					married: true,
+					age: 69,
+					gender: "m"
+				},
+				{
+					active: true,
+					married: true,
+					age: 45,
+					gender: "f"
+				}
+			]
+			FOR n IN 0..1
+				FOR i IN users
+					COLLECT gender = i.gender INTO genders
+					RETURN {
+						gender: CONCAT(gender, n),
+						values: genders
+					}
 `, []any{map[string]any{
 			"gender": "f0",
 			"values": []map[string]any{
@@ -327,45 +327,45 @@ FOR n IN 0..1
 			},
 		}),
 		CaseArray(`
-LET users = [
-	{
-		active: true,
-		married: true,
-		age: 31,
-		gender: "m"
-	},
-	{
-		active: true,
-		married: false,
-		age: 25,
-		gender: "f"
-	},
-	{
-		active: true,
-		married: false,
-		age: 36,
-		gender: "m"
-	},
-	{
-		active: false,
-		married: true,
-		age: 69,
-		gender: "m"
-	},
-	{
-		active: true,
-		married: true,
-		age: 45,
-		gender: "f"
-	}
-]
-FOR i IN users
-	COLLECT gender = i.gender INTO genders
-	FOR n IN 0..1
-		RETURN {
-			gender: CONCAT(gender, n),
-			values: genders
-		}
+			LET users = [
+				{
+					active: true,
+					married: true,
+					age: 31,
+					gender: "m"
+				},
+				{
+					active: true,
+					married: false,
+					age: 25,
+					gender: "f"
+				},
+				{
+					active: true,
+					married: false,
+					age: 36,
+					gender: "m"
+				},
+				{
+					active: false,
+					married: true,
+					age: 69,
+					gender: "m"
+				},
+				{
+					active: true,
+					married: true,
+					age: 45,
+					gender: "f"
+				}
+			]
+			FOR i IN users
+				COLLECT gender = i.gender INTO genders
+				FOR n IN 0..1
+					RETURN {
+						gender: CONCAT(gender, n),
+						values: genders
+					}
 `, []any{map[string]any{
 			"gender": "f0",
 			"values": []map[string]any{
@@ -467,5 +467,189 @@ FOR i IN users
 				},
 			},
 		}),
+		CaseArray(`
+			LET users = [
+				{
+					active: true,
+					married: true,
+					age: 31,
+					gender: "m"
+				},
+				{
+					active: true,
+					married: false,
+					age: 25,
+					gender: "f"
+				},
+				{
+					active: true,
+					married: false,
+					age: 36,
+					gender: "m"
+				},
+				{
+					active: false,
+					married: true,
+					age: 69,
+					gender: "m"
+				},
+				{
+					active: true,
+					married: true,
+					age: 45,
+					gender: "f"
+				}
+			]
+			FOR n IN 0..1
+				FOR u IN users
+				  COLLECT gender = u.gender
+					AGGREGATE minAge = MIN(u.age), maxAge = MAX(u.age)
+				
+				  RETURN {
+					gender: CONCAT(gender, n),
+					minAge,
+					maxAge
+				  }
+`, []any{
+			map[string]any{
+				"gender": "f0",
+				"maxAge": 45,
+				"minAge": 25,
+			},
+			map[string]any{
+				"gender": "m0",
+				"maxAge": 69,
+				"minAge": 31,
+			},
+			map[string]any{
+				"gender": "f1",
+				"maxAge": 45,
+				"minAge": 25,
+			},
+			map[string]any{
+				"gender": "m1",
+				"maxAge": 69,
+				"minAge": 31,
+			},
+		}),
+		CaseArray(`
+			LET users = [
+				{
+					active: true,
+					married: true,
+					age: 31,
+					gender: "m"
+				},
+				{
+					active: true,
+					married: false,
+					age: 25,
+					gender: "f"
+				},
+				{
+					active: true,
+					married: false,
+					age: 36,
+					gender: "m"
+				},
+				{
+					active: false,
+					married: true,
+					age: 69,
+					gender: "m"
+				},
+				{
+					active: true,
+					married: true,
+					age: 45,
+					gender: "f"
+				}
+			]
+			
+			FOR u IN users
+			  COLLECT gender = u.gender
+				AGGREGATE minAge = MIN(u.age), maxAge = MAX(u.age)
+				FOR n IN 0..1
+				  RETURN {
+					gender: CONCAT(gender, n),
+					minAge,
+					maxAge
+				  }
+`, []any{
+			map[string]any{
+				"gender": "f0",
+				"maxAge": 45,
+				"minAge": 25,
+			},
+			map[string]any{
+				"gender": "f1",
+				"maxAge": 45,
+				"minAge": 25,
+			},
+			map[string]any{
+				"gender": "m0",
+				"maxAge": 69,
+				"minAge": 31,
+			},
+			map[string]any{
+				"gender": "m1",
+				"maxAge": 69,
+				"minAge": 31,
+			},
+		}),
+		CaseArray(`
+			LET users = [
+				{
+					active: true,
+					married: true,
+					age: 31,
+					gender: "m"
+				},
+				{
+					active: true,
+					married: false,
+					age: 25,
+					gender: "f"
+				},
+				{
+					active: true,
+					married: false,
+					age: 36,
+					gender: "m"
+				},
+				{
+					active: false,
+					married: true,
+					age: 69,
+					gender: "m"
+				},
+				{
+					active: true,
+					married: true,
+					age: 45,
+					gender: "f"
+				}
+			]
+			FOR n IN 0..1
+				FOR u IN users
+  				COLLECT AGGREGATE minAge = MIN(u.age), maxAge = MAX(u.age)
+				
+				  RETURN {
+					iteration: n,
+					minAge,
+					maxAge
+				  }
+`, []any{
+			map[string]any{
+				"iteration": 0,
+				"maxAge":    69,
+				"minAge":    25,
+			},
+			map[string]any{
+				"iteration": 1,
+				"maxAge":    69,
+				"minAge":    25,
+			},
+		}),
 	})
 }