Refactor FOR-WHILE loop compilation to ensure compatibility with FOR-IN loops, update value register resolution, and add integration tests for sorting scenarios.

pkg/compiler/internal/loop_sort.go

@@ -91,15 +91,19 @@ func (c *LoopSortCompiler) compileSingleSortKey(clause fql.ISortClauseExpression
 // If the loop already has a value name, reuse it; otherwise, allocate a new register
 // and load the value from the iterator.
 func (c *LoopSortCompiler) resolveValueRegister(loop *core.Loop) vm.Operand {
-	// If value is already used in the loop body, reuse the existing register
+	if loop.Kind == core.ForInLoop {
-	if loop.ValueName != "" {
+		// If value is already used in the loop body, reuse the existing register
-		return loop.Value
+		if loop.ValueName != "" {
+			return loop.Value
+		}
+
+		// Otherwise, allocate a new register and load the value from iterator
+		kvValReg := c.ctx.Registers.Allocate(core.Temp)
+		loop.EmitValue(kvValReg, c.ctx.Emitter)
+		return kvValReg
 	}
 
-	// Otherwise, allocate a new register and load the value from iterator
+	return loop.Key
-	kvValReg := c.ctx.Registers.Allocate(core.Temp)
-	loop.EmitValue(kvValReg, c.ctx.Emitter)
-	return kvValReg
 }
 
 // compileSorter configures a sorter for a loop based on provided sort clauses and directions.
@@ -140,6 +144,15 @@ func (c *LoopSortCompiler) finalizeSorting(loop *core.Loop, kv *core.KV, sorter
 		c.ctx.Registers.Free(sorter)
 	}
 
+	if loop.Kind != core.ForInLoop {
+		// We switched from a ForWhileLoop to a ForInLoop because the underlying data is Iterable now.
+		loop.Kind = core.ForInLoop
+		loop.ValueName = loop.KeyName
+		loop.Value = loop.Key
+		loop.Key = vm.NoopOperand
+		loop.KeyName = ""
+	}
+
 	// Reinitialize the loop to iterate over sorted data
 	loop.EmitInitialization(c.ctx.Registers, c.ctx.Emitter)
 }

test/integration/vm/vm_for_while_sort_test.go

@@ -0,0 +1,239 @@
+package vm_test
+
+import (
+	"context"
+	"testing"
+
+	"github.com/MontFerret/ferret/pkg/runtime"
+	"github.com/MontFerret/ferret/pkg/vm"
+)
+
+func TestForWhileSort(t *testing.T) {
+	RunUseCases(t, []UseCase{
+		CaseArray(`
+			FOR i WHILE UNTIL(5)
+				SORT i DESC
+				RETURN i
+`, []any{4, 3, 2, 1, 0}),
+		CaseArray(`
+			LET strs = ["foo", "bar", "qaz", "abc"]
+
+			FOR i WHILE UNTIL(4)
+				SORT strs[i]
+				RETURN i
+`, []any{3, 1, 0, 2}),
+		CaseArray(`
+			LET strs = ["foo", "bar", "qaz", "abc"]
+
+			FOR i WHILE UNTIL(4)
+				SORT i DESC
+				RETURN strs[i]
+`, []any{"abc", "qaz", "bar", "foo"}),
+		CaseArray(`
+			LET users = [
+				{
+					name: "Ron",
+					age: 31,
+					gender: "m"
+				},
+				{
+					name: "Angela",
+					age: 29,
+					gender: "f"
+				},
+				{
+					name: "Bob",
+					age: 36,
+					gender: "m"
+				}
+			]
+			FOR i WHILE UNTIL(3)
+				LET u = users[i]
+				SORT u.name
+				RETURN users[i]
+`, []any{
+			map[string]any{"name": "Angela", "age": 29, "gender": "f"},
+			map[string]any{"name": "Bob", "age": 36, "gender": "m"},
+			map[string]any{"name": "Ron", "age": 31, "gender": "m"},
+		}),
+		CaseArray(`
+			LET users = [
+				{
+					active: true,
+					age: 31,
+					gender: "m"
+				},
+				{
+					active: true,
+					age: 29,
+					gender: "f"
+				},
+				{
+					active: true,
+					age: 36,
+					gender: "m"
+				}
+			]
+			FOR i WHILE UNTIL(3)
+				LET u = users[i]
+				SORT u.age DESC
+				RETURN users[i]
+		`, []any{
+			map[string]any{"active": true, "age": 36, "gender": "m"},
+			map[string]any{"active": true, "age": 31, "gender": "m"},
+			map[string]any{"active": true, "age": 29, "gender": "f"},
+		}),
+		CaseArray(`			
+			LET users = [
+				{
+					active: true,
+					age: 31,
+					gender: "m"
+				},
+				{
+					active: true,
+					age: 29,
+					gender: "f"
+				},
+				{
+					active: true,
+					age: 31,
+					gender: "f"
+				},
+				{
+					active: true,
+					age: 36,
+					gender: "m"
+				}
+			]
+			FOR i WHILE UNTIL(4)
+				LET u = users[i]
+				SORT u.age, u.gender
+				RETURN users[i]`,
+			[]any{
+				map[string]any{"active": true, "age": 29, "gender": "f"},
+				map[string]any{"active": true, "age": 31, "gender": "f"},
+				map[string]any{"active": true, "age": 31, "gender": "m"},
+				map[string]any{"active": true, "age": 36, "gender": "m"},
+			}),
+		CaseArray(`
+			LET users = [
+				{
+					active: true,
+					age: 31,
+					gender: "m"
+				},
+				{
+					active: true,
+					age: 29,
+					gender: "f"
+				},
+				{
+					active: true,
+					age: 31,
+					gender: "f"
+				},
+				{
+					active: true,
+					age: 36,
+					gender: "m"
+				}
+			]
+			FOR i WHILE UNTIL(4)
+				LET u = users[i]
+				LET x = "foo"
+				TEST(x)
+				SORT u.age, u.gender
+				RETURN users[i]
+		`, []any{
+			map[string]any{"active": true, "age": 29, "gender": "f"},
+			map[string]any{"active": true, "age": 31, "gender": "f"},
+			map[string]any{"active": true, "age": 31, "gender": "m"},
+			map[string]any{"active": true, "age": 36, "gender": "m"},
+		}),
+		CaseArray(`
+			LET users = [
+				{
+					active: true,
+					age: 31,
+					gender: "m"
+				},
+				{
+					active: true,
+					age: 29,
+					gender: "f"
+				},
+				{
+					active: true,
+					age: 36,
+					gender: "m"
+				}
+			]
+			FOR i WHILE UNTIL(3)
+				LET u = users[i]
+				FILTER u.gender == "m"
+				SORT u.age
+				RETURN users[i]
+		`, []any{
+			map[string]any{"active": true, "age": 31, "gender": "m"},
+			map[string]any{"active": true, "age": 36, "gender": "m"},
+		}),
+		CaseArray(`
+			LET users = [
+				{
+					active: true,
+					age: 31,
+					gender: "m"
+				},
+				{
+					active: true,
+					age: 29,
+					gender: "f"
+				},
+				{
+					active: true,
+					age: 36,
+					gender: "m"
+				}
+			]
+			FOR i WHILE UNTIL(3)
+				SORT users[i].age
+				FILTER users[i].gender == "m"
+				RETURN users[i]
+		`, []any{
+			map[string]any{"active": true, "age": 31, "gender": "m"},
+			map[string]any{"active": true, "age": 36, "gender": "m"},
+		}),
+		CaseObject(`
+			LET users = [
+				{
+					active: true,
+					age: 31,
+					gender: "m"
+				},
+				{
+					active: true,
+					age: 29,
+					gender: "f"
+				},
+				{
+					active: true,
+					age: 36,
+					gender: "m"
+				}
+			]
+			LET sorted = (FOR i WHILE UNTIL(3)
+				SORT users[i].age
+				FILTER users[i].gender == "m"
+				RETURN users[i])
+			
+			RETURN sorted[0]
+		`, map[string]any{"active": true, "age": 31, "gender": "m"}),
+	}, vm.WithFunctionSetter(func(fns runtime.Functions) {
+		fns.F().Set("TEST", func(ctx context.Context, args ...runtime.Value) (runtime.Value, error) {
+			return runtime.None, nil
+		})
+
+		fns.SetAll(ForWhileHelpers())
+	}))
+}