From 70268ba2c7c23bc5418a1f9513bea9b4fcc99f3b Mon Sep 17 00:00:00 2001
From: Tim Voronov <ziflex@gmail.com>
Date: Mon, 1 Apr 2024 12:44:22 -0400
Subject: [PATCH] Optimized runtime DataSet usage

---
 pkg/compiler/compiler_func_test.go | 32 ++++++++++--------------------
 pkg/compiler/visitor.go            | 15 +++++++++++++-
 pkg/runtime/dataset.go             | 24 ++++++++++++++++++----
 pkg/runtime/opcodes.go             |  1 +
 pkg/runtime/vm.go                  |  8 +++++---
 5 files changed, 50 insertions(+), 30 deletions(-)

diff --git a/pkg/compiler/compiler_func_test.go b/pkg/compiler/compiler_func_test.go
index 316b6cff..ba862153 100644
--- a/pkg/compiler/compiler_func_test.go
+++ b/pkg/compiler/compiler_func_test.go
@@ -42,27 +42,15 @@ func TestFunctionCall(t *testing.T) {
 			[]int{2, 4, 6, 8},
 			ShouldEqualJSON,
 		},
-		//{
-		//	`RETURN FIRST((FOR i IN 1..10 RETURN i * 2))`,
-		//	2,
-		//	nil,
-		//},
+		{
+			`RETURN FIRST((FOR i IN 1..10 RETURN i * 2))`,
+			2,
+			nil,
+		},
+		{
+			`RETURN UNION((FOR i IN 0..5 RETURN i), (FOR i IN 6..10 RETURN i))`,
+			[]int{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10},
+			ShouldEqualJSON,
+		},
 	})
-
-	//
-	//Convey("Should be able to use FOR as arguments", t, func() {
-	//	c := compiler.New()
-	//
-	//	p, err := c.Compile(`
-	//		RETURN UNION((FOR i IN 0..5 RETURN i), (FOR i IN 6..10 RETURN i))
-	//	`)
-	//
-	//	So(err, ShouldBeNil)
-	//
-	//	out, err := p.Run(context.Background())
-	//
-	//	So(err, ShouldBeNil)
-	//
-	//	So(string(out), ShouldEqual, `[0,1,2,3,4,5,6,7,8,9,10]`)
-	//})
 }
diff --git a/pkg/compiler/visitor.go b/pkg/compiler/visitor.go
index 305030d3..17b08e6d 100644
--- a/pkg/compiler/visitor.go
+++ b/pkg/compiler/visitor.go
@@ -231,7 +231,6 @@ func (v *visitor) VisitForExpression(ctx *fql.ForExpressionContext) interface{}
 	v.emitLoop(loopJump)
 	v.patchJump(exitJump)
 	v.endScope()
-	v.endLoop()
 	// pop the boolean value from the stack
 	v.emitPop()
 
@@ -243,6 +242,8 @@ func (v *visitor) VisitForExpression(ctx *fql.ForExpressionContext) interface{}
 		v.emitPop()
 	}
 
+	v.endLoop()
+
 	return nil
 }
 
@@ -850,6 +851,18 @@ func (v *visitor) resolveLoopResult() int {
 
 func (v *visitor) endLoop() {
 	v.loops = v.loops[:len(v.loops)-1]
+
+	var unwrap bool
+
+	if len(v.loops) == 0 {
+		unwrap = true
+	} else if !v.loops[len(v.loops)-1].passThrough {
+		unwrap = true
+	}
+
+	if unwrap {
+		v.emit(runtime.OpLoopUnwrapOutput)
+	}
 }
 
 func (v *visitor) resolveLocalVariable(name string) int {
diff --git a/pkg/runtime/dataset.go b/pkg/runtime/dataset.go
index f1031c92..5d4476b5 100644
--- a/pkg/runtime/dataset.go
+++ b/pkg/runtime/dataset.go
@@ -23,6 +23,26 @@ func NewDataSet(distinct bool) *DataSet {
 	}
 }
 
+func (ds *DataSet) String() string {
+	return ds.values.String()
+}
+
+func (ds *DataSet) Unwrap() interface{} {
+	return ds.values
+}
+
+func (ds *DataSet) Hash() uint64 {
+	return ds.values.Hash()
+}
+
+func (ds *DataSet) Copy() core.Value {
+	return ds.values.Copy()
+}
+
+func (ds *DataSet) MarshalJSON() ([]byte, error) {
+	return ds.values.MarshalJSON()
+}
+
 func (ds *DataSet) Push(item core.Value) {
 	if ds.hashmap != nil {
 		hash := item.Hash()
@@ -39,10 +59,6 @@ func (ds *DataSet) Push(item core.Value) {
 	ds.values.Push(item)
 }
 
-func (ds *DataSet) MarshalJSON() ([]byte, error) {
-	return ds.values.MarshalJSON()
-}
-
 func (ds *DataSet) ToArray() *values.Array {
 	return ds.values
 }
diff --git a/pkg/runtime/opcodes.go b/pkg/runtime/opcodes.go
index e95883c7..7c6f2a5c 100644
--- a/pkg/runtime/opcodes.go
+++ b/pkg/runtime/opcodes.go
@@ -61,6 +61,7 @@ const (
 	OpJump
 	OpJumpBackward
 	OpLoopInitOutput
+	OpLoopUnwrapOutput
 	OpForLoopInitInput
 	OpForLoopHasNext
 	OpForLoopNext
diff --git a/pkg/runtime/vm.go b/pkg/runtime/vm.go
index 9004c520..5081e6b4 100644
--- a/pkg/runtime/vm.go
+++ b/pkg/runtime/vm.go
@@ -382,9 +382,11 @@ loop:
 				return nil, err
 			}
 		case OpLoopInitOutput:
-			output := NewDataSet(arg == 1)
+			stack.Push(NewDataSet(arg == 1))
 
-			stack.Push(values.NewBoxedValue(output))
+		case OpLoopUnwrapOutput:
+			ds := stack.Pop().(*DataSet)
+			stack.Push(ds.ToArray())
 
 		case OpForLoopInitInput:
 			// start a new iteration
@@ -466,7 +468,7 @@ loop:
 			// pop the return value from the stack
 			res := stack.Pop()
 			pos := stack.Len() - arg
-			ds := stack.Get(pos).(*values.Boxed).Unwrap().(*DataSet)
+			ds := stack.Get(pos).(*DataSet)
 			ds.Push(res)
 
 		case OpReturn: