Refactored stack management

pkg/compiler/visitor.go

@@ -882,7 +882,7 @@ func (v *visitor) beginLoopScope(passThrough, distinct bool) {
 		}
 
 		resultPos = v.operandsStackTracker
-		v.emit(runtime.OpLoopInitOutput, arg)
+		v.emit(runtime.OpLoopInit, arg)
 	} else {
 		resultPos = prevResult
 	}
@@ -917,7 +917,7 @@ func (v *visitor) endLoopScope() {
 	}
 
 	if unwrap {
-		v.emit(runtime.OpLoopUnwrapOutput)
+		v.emit(runtime.OpLoopFin)
 	}
 }
 
@@ -1148,10 +1148,10 @@ func (v *visitor) updateStackTracker(op runtime.Opcode, arg int) {
 	case runtime.OpRange:
 		v.operandsStackTracker--
 
-	case runtime.OpLoopInitOutput:
+	case runtime.OpLoopInit:
 		v.operandsStackTracker++
 
-	case runtime.OpLoopUnwrapOutput, runtime.OpForLoopInitInput:
+	case runtime.OpLoopFin, runtime.OpForLoopInitInput:
 		break
 
 	case runtime.OpForLoopHasNext:

pkg/runtime/frame.go

@@ -0,0 +1,16 @@
+package runtime
+
+type Frame struct {
+	Operands  *Stack
+	Variables *Stack
+	State     *Stack
+	ReturnPC  int
+}
+
+func NewFrame(size int) *Frame {
+	return &Frame{
+		Operands:  NewStack(size),
+		Variables: NewStack(size),
+		State:     NewStack(size),
+	}
+}

pkg/runtime/opcodes.go

@@ -3,7 +3,14 @@ package runtime
 type Opcode byte
 
 const (
-	OpNone Opcode = iota
+	OpPush Opcode = iota
+	OpPop
+	OpPopClose
+	OpJumpIfFalse
+	OpJumpIfTrue
+	OpJump
+	OpJumpBackward
+	OpNone
 	OpCastBool
 	OpTrue
 	OpFalse
@@ -52,15 +59,8 @@ const (
 	OpCall4Safe
 	OpCallN
 	OpCallNSafe
-	OpPush
-	OpPop
-	OpPopClose
-	OpJumpIfFalse
-	OpJumpIfTrue
-	OpJump
-	OpJumpBackward
-	OpLoopInitOutput
-	OpLoopUnwrapOutput
+	OpLoopInit
+	OpLoopFin
 	OpForLoopInitInput
 	OpForLoopHasNext
 	OpForLoopNext

pkg/runtime/stack.go

@@ -5,61 +5,41 @@ import (
 )
 
 type Stack struct {
-	operands  []core.Value
-	variables []core.Value
+	values []core.Value
 }
 
-func NewStack(operands, variables int) *Stack {
+func NewStack(size int) *Stack {
 	return &Stack{
-		make([]core.Value, 0, operands),
-		make([]core.Value, 0, variables),
+		values: make([]core.Value, 0, size),
 	}
 }
 
 func (s *Stack) Len() int {
-	return len(s.operands)
+	return len(s.values)
 }
 
 func (s *Stack) Peek() core.Value {
-	return s.operands[len(s.operands)-1]
+	return s.values[len(s.values)-1]
 }
 
 func (s *Stack) Push(value core.Value) {
-	s.operands = append(s.operands, value)
+	s.values = append(s.values, value)
 }
 
 func (s *Stack) Pop() core.Value {
-	value := s.operands[len(s.operands)-1]
-	s.operands = s.operands[:len(s.operands)-1]
+	value := s.values[len(s.values)-1]
+	s.values = s.values[:len(s.values)-1]
 	return value
 }
 
 func (s *Stack) Get(index int) core.Value {
-	return s.operands[index]
+	return s.values[index]
 }
 
 func (s *Stack) Set(index int, value core.Value) {
-	s.operands[index] = value
-}
-
-func (s *Stack) GetVariable(index int) core.Value {
-	return s.variables[index]
-}
-
-func (s *Stack) SetVariable(index int, value core.Value) {
-	// TODO: Calculate in advance the number of variables
-	if index >= len(s.variables) {
-		s.variables = append(s.variables, value)
-		return
+	if index < len(s.values) {
+		s.values[index] = value
+	} else {
+		s.values = append(s.values, value)
 	}
-
-	s.variables[index] = value
-}
-
-func (s *Stack) PopVariable() {
-	if len(s.variables) == 0 {
-		return
-	}
-
-	s.variables = s.variables[:len(s.variables)-1]
 }

pkg/runtime/vm.go

@@ -11,7 +11,7 @@ import (
 
 type VM struct {
 	pc      int
-	stack   *Stack
+	frames  []*Frame
 	globals map[string]core.Value
 	env     *Environment
 }
@@ -35,13 +35,16 @@ func (vm *VM) Run(ctx context.Context, program *Program) ([]byte, error) {
 	}
 
 	// TODO: Add code analysis to calculate the number of operands and variables
-	stack := NewStack(len(program.Bytecode), 8)
-	vm.stack = stack
+	frame := NewFrame(32)
+	vm.frames = []*Frame{frame}
 	vm.globals = make(map[string]core.Value)
 	vm.pc = 0
 
 loop:
 	for vm.pc < len(program.Bytecode) {
+		stack := frame.Operands
+		variables := frame.Variables
+		state := frame.State
 		op := program.Bytecode[vm.pc]
 		arg := program.Arguments[vm.pc]
 		vm.pc++
@@ -67,13 +70,16 @@ loop:
 			stack.Push(vm.globals[program.Constants[arg].String()])
 
 		case OpStoreLocal:
-			stack.SetVariable(arg, stack.Pop())
+			variables.Set(arg, stack.Pop())
 
 		case OpPopLocal:
-			stack.PopVariable()
+			// TODO: Figure out how to remove the check. Added to handle variable cleanup after an empty iteration
+			if variables.Len() > 0 {
+				variables.Pop()
+			}
 
 		case OpLoadLocal:
-			stack.Push(stack.GetVariable(arg))
+			stack.Push(variables.Get(arg))
 
 		case OpNone:
 			stack.Push(values.None)
@@ -92,7 +98,7 @@ loop:
 			arr := values.NewSizedArray(size)
 
 			// iterate from the end to the beginning
-			// because stack is LIFO
+			// because frame is LIFO
 			for i := size - 1; i >= 0; i-- {
 				arr.MustSet(values.Int(i), stack.Pop())
 			}
@@ -358,17 +364,17 @@ loop:
 				return nil, err
 			}
 		case OpCallN, OpCallNSafe:
-			// pop arguments from the stack
+			// pop arguments from the frame
 			// and push them to the arguments array
-			// in reverse order because stack is LIFO and arguments array is FIFO
+			// in reverse order because frame is LIFO and arguments array is FIFO
 			argCount := arg
 			args := make([]core.Value, argCount)
 
 			for i := argCount - 1; i >= 0; i-- {
 				args[i] = stack.Pop()
 			}
 
-			// pop the function name from the stack
+			// pop the function name from the frame
 			fnName := stack.Pop().String()
 
 			// call the function
@@ -392,11 +398,11 @@ loop:
 			} else {
 				return nil, err
 			}
-		case OpLoopInitOutput:
-			stack.Push(NewDataSet(arg == 1))
+		case OpLoopInit:
+			state.Push(NewDataSet(arg == 1))
 
-		case OpLoopUnwrapOutput:
-			ds := stack.Pop().(*DataSet)
+		case OpLoopFin:
+			ds := state.Pop().(*DataSet)
 			stack.Push(ds.ToArray())
 
 		case OpForLoopInitInput:
@@ -476,15 +482,15 @@ loop:
 			}
 
 		case OpLoopReturn:
-			// pop the return value from the stack
+			// pop the return value from the frame
 			res := stack.Pop()
-			ds := stack.Get(arg).(*DataSet)
+			ds := state.Peek().(*DataSet)
 			ds.Push(res)
 
 		case OpReturn:
 			break loop
 		}
 	}
 
-	return stack.Pop().MarshalJSON()
+	return frame.Operands.Pop().MarshalJSON()
 }