From 03167dcb83ec4e7899052856062d1cff81455dfd Mon Sep 17 00:00:00 2001
From: Tim Voronov <ziflex@gmail.com>
Date: Fri, 1 Nov 2024 17:40:52 -0400
Subject: [PATCH] Fixed LOOP body statements

---
 pkg/compiler/compiler_exec_test.go  | 23 +++++++++++------------
 pkg/compiler/compiler_setup_test.go | 18 +++++++++++++-----
 pkg/compiler/visitor.go             | 16 ++++++++--------
 pkg/runtime/vm.go                   |  2 +-
 4 files changed, 33 insertions(+), 26 deletions(-)

diff --git a/pkg/compiler/compiler_exec_test.go b/pkg/compiler/compiler_exec_test.go
index 51b20757..d546c5a8 100644
--- a/pkg/compiler/compiler_exec_test.go
+++ b/pkg/compiler/compiler_exec_test.go
@@ -325,7 +325,7 @@ func TestUnaryOperators(t *testing.T) {
 			RETURN { enabled: !val }
 		`)
 
-		v1, err := runtime.NewVM(p1).Run(context.Background())
+		v1, err := runtime.NewVM(p1).Run(context.Background(), nil)
 
 		So(err, ShouldBeNil)
 
@@ -339,7 +339,7 @@ func TestUnaryOperators(t *testing.T) {
 			RETURN { enabled: !!val }
 		`)
 
-		v2, err := runtime.NewVM(p2).Run(context.Background())
+		v2, err := runtime.NewVM(p2).Run(context.Background(), nil)
 
 		So(err, ShouldBeNil)
 
@@ -1121,20 +1121,19 @@ func TestFor(t *testing.T) {
 	//	ShouldEqualJSON,
 	//},
 	RunUseCases(t, []UseCase{
-		{
-			"FOR i IN 1..5 RETURN i",
-			[]any{1, 2, 3, 4, 5},
-			ShouldEqualJSON,
-		},
 		//{
-		//	`FOR i IN 1..5
-		//		                           LET x = i
-		//		                           PRINT(x)
-		//									RETURN i
-		//		`,
+		//	"FOR i IN 1..5 RETURN i",
 		//	[]any{1, 2, 3, 4, 5},
 		//	ShouldEqualJSON,
 		//},
+		{
+			`FOR i IN 1..5
+				                           LET x = i * 2
+											RETURN x
+				`,
+			[]any{2, 4, 6, 8, 10},
+			ShouldEqualJSON,
+		},
 		//{
 		//	`FOR val, counter IN 1..5
 		//                    LET x = val
diff --git a/pkg/compiler/compiler_setup_test.go b/pkg/compiler/compiler_setup_test.go
index 01530f57..40e5461a 100644
--- a/pkg/compiler/compiler_setup_test.go
+++ b/pkg/compiler/compiler_setup_test.go
@@ -41,7 +41,7 @@ func Compile(expression string) (*runtime.Program, error) {
 func Run(p *runtime.Program, opts ...runtime.EnvironmentOption) ([]byte, error) {
 	vm := runtime.NewVM(p)
 
-	out, err := vm.Run(context.Background(), opts...)
+	out, err := vm.Run(context.Background(), opts)
 
 	if err != nil {
 		return nil, err
@@ -126,7 +126,12 @@ func RunAsmUseCases(t *testing.T, useCases []ByteCodeUseCase) {
 
 func RunUseCasesWith(t *testing.T, c *compiler.Compiler, useCases []UseCase, opts ...runtime.EnvironmentOption) {
 	for _, useCase := range useCases {
-		t.Run(useCase.Expression, func(t *testing.T) {
+		name := strings.TrimSpace(useCase.Expression)
+		name = strings.Replace(name, "\n", " ", -1)
+		name = strings.Replace(name, "\t", " ", -1)
+		// Replace multiple spaces with a single space
+		name = strings.Join(strings.Fields(name), " ")
+		t.Run(name, func(t *testing.T) {
 			Convey(useCase.Expression, t, func() {
 				// catch panic
 				//defer func() {
@@ -188,10 +193,13 @@ func RunBenchmarkWith(b *testing.B, c *compiler.Compiler, expression string, opt
 	}
 	options = append(options, opts...)
 
-	for n := 0; n < b.N; n++ {
-		vm := runtime.NewVM(prog)
+	ctx := context.Background()
+	vm := runtime.NewVM(prog)
+
+	b.ResetTimer()
 
-		_, err := vm.Run(context.Background(), opts...)
+	for n := 0; n < b.N; n++ {
+		_, err := vm.Run(ctx, opts)
 
 		if err != nil {
 			panic(err)
diff --git a/pkg/compiler/visitor.go b/pkg/compiler/visitor.go
index f4631112..635925f5 100644
--- a/pkg/compiler/visitor.go
+++ b/pkg/compiler/visitor.go
@@ -311,15 +311,15 @@ func (v *visitor) VisitFilterClause(ctx *fql.FilterClauseContext) interface{} {
 }
 
 func (v *visitor) VisitForExpressionStatement(ctx *fql.ForExpressionStatementContext) interface{} {
-	//if c := ctx.VariableDeclaration(); c != nil {
-	//	c.Accept(v)
-	//} else if c := ctx.FunctionCallExpression(); c != nil {
-	//	c.Accept(v)
-	//	// remove un-used return value
-	//	v.emitPop()
-	//}
+	if c := ctx.VariableDeclaration(); c != nil {
+		return c.Accept(v)
+	}
 
-	return nil
+	if c := ctx.FunctionCallExpression(); c != nil {
+		return c.Accept(v)
+	}
+
+	panic(core.Error(ErrUnexpectedToken, ctx.GetText()))
 }
 
 func (v *visitor) VisitFunctionCallExpression(ctx *fql.FunctionCallExpressionContext) interface{} {
diff --git a/pkg/runtime/vm.go b/pkg/runtime/vm.go
index 882b3695..1fb72286 100644
--- a/pkg/runtime/vm.go
+++ b/pkg/runtime/vm.go
@@ -25,7 +25,7 @@ func NewVM(program *Program) *VM {
 	return vm
 }
 
-func (vm *VM) Run(ctx context.Context, opts ...EnvironmentOption) (core.Value, error) {
+func (vm *VM) Run(ctx context.Context, opts []EnvironmentOption) (core.Value, error) {
 	tryCatch := func(pos int) bool {
 		for _, pair := range vm.program.CatchTable {
 			if pos >= pair[0] && pos <= pair[1] {