for-loop body wip

MontFerret · Feb 4, 2024 · 98ff664 · 98ff664
1 parent 2b3fd4f
commit 98ff664
Show file tree

Hide file tree

Showing 3 changed files with 51 additions and 38 deletions.
diff --git 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
@@ -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
@@ -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
 }