From 98ff66490e53c6deef10e1925b84c7c8f4ac8578 Mon Sep 17 00:00:00 2001
From: Tim Voronov <ziflex@gmail.com>
Date: Sun, 4 Feb 2024 14:02:49 -0500
Subject: [PATCH] for-loop body wip

---
 pkg/compiler/compiler_for_test.go | 30 ++++++++++++-------
 pkg/compiler/visitor.go           | 48 +++++++++++++++----------------
 pkg/runtime/values/range_iter.go  | 11 ++++---
 3 files changed, 51 insertions(+), 38 deletions(-)

diff --git a/pkg/compiler/compiler_for_test.go b/pkg/compiler/compiler_for_test.go
index 1a5a83c9..96a46914 100644
--- a/pkg/compiler/compiler_for_test.go
+++ b/pkg/compiler/compiler_for_test.go
@@ -10,18 +10,28 @@ import (
 
 func TestFor(t *testing.T) {
 	RunUseCases(t, compiler.New(), []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
+		//`,
+		//			[]any{1, 2, 3, 4, 5},
+		//			ShouldEqualJSON,
+		//		},
 		{
-			"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 val, counter IN 1..5
+                            LET x = val
+                            PRINT(counter)
+							LET y = counter
+							RETURN [x, y]
 `,
-			[]any{1, 2, 3, 4, 5},
+			[]any{[]any{1, 0}, []any{2, 1}, []any{3, 2}, []any{4, 3}, []any{5, 4}},
 			ShouldEqualJSON,
 		},
 	})
diff --git a/pkg/compiler/visitor.go b/pkg/compiler/visitor.go
index e471507b..8b48b6cc 100644
--- a/pkg/compiler/visitor.go
+++ b/pkg/compiler/visitor.go
@@ -199,30 +199,12 @@ func (v *visitor) VisitForExpression(ctx *fql.ForExpressionContext) interface{}
 		v.defineVariable(index)
 	}
 
-	//// filter
-	//if c := ctx.ForExpressionFilter(); c != nil {
-	//	c.Accept(v)
-	//}
-	//
-	//// sort
-	//if c := ctx.ForExpressionSortClause(); c != nil {
-	//	c.Accept(v)
-	//}
-	//
-	//// limit
-	//if c := ctx.ForExpressionLimitClause(); c != nil {
-	//	c.Accept(v)
-	//}
-	//
-	//// collect
-	//if c := ctx.ForExpressionCollectClause(); c != nil {
-	//	c.Accept(v)
-	//}
-	//
-	//// return
-	//if c := ctx.ForExpressionReturnClause(); c != nil {
-	//	c.Accept(v)
-	//}
+	// body
+	if body := ctx.AllForExpressionBody(); body != nil && len(body) > 0 {
+		for _, b := range body {
+			b.Accept(v)
+		}
+	}
 
 	// return
 	if c := ctx.ForExpressionReturn(); c != nil {
@@ -263,6 +245,24 @@ func (v *visitor) VisitForExpressionSource(ctx *fql.ForExpressionSourceContext)
 	return nil
 }
 
+func (v *visitor) VisitForExpressionBody(ctx *fql.ForExpressionBodyContext) interface{} {
+	if c := ctx.ForExpressionStatement(); c != nil {
+		c.Accept(v)
+	}
+
+	return nil
+}
+
+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)
+	}
+
+	return nil
+}
+
 func (v *visitor) VisitForExpressionReturn(ctx *fql.ForExpressionReturnContext) interface{} {
 	if c := ctx.ReturnExpression(); c != nil {
 		c.Accept(v)
diff --git a/pkg/runtime/values/range_iter.go b/pkg/runtime/values/range_iter.go
index 69fd9499..c7004e7b 100644
--- a/pkg/runtime/values/range_iter.go
+++ b/pkg/runtime/values/range_iter.go
@@ -6,9 +6,10 @@ import (
 )
 
 type RangeIterator struct {
-	values *Range
-	dir    int64
-	pos    int64
+	values  *Range
+	dir     int64
+	pos     int64
+	counter int64
 }
 
 func NewRangeIterator(values *Range) core.Iterator {
@@ -29,8 +30,10 @@ func (iterator *RangeIterator) HasNext(_ context.Context) (bool, error) {
 
 func (iterator *RangeIterator) Next(_ context.Context) (value core.Value, key core.Value, err error) {
 	val := NewInt64(iterator.pos)
+	counter := NewInt64(iterator.counter)
 
 	iterator.pos += iterator.dir
+	iterator.counter++
 
-	return val, val, nil
+	return val, counter, nil
 }