Improved loop scoping

benchmark/benchmark.go

@@ -13,8 +13,9 @@ import (
 )
 
 var engine = flag.String("engine", "vm", "use 'vm' or 'eval'")
+var input = flag.String("input", "recursive", "recursive or iterative benchmark")
 
-var input = `
+var recursive = `
 let fibonacci = fn(x) {
 	if (x == 0) {
 		0
@@ -29,13 +30,31 @@ let fibonacci = fn(x) {
 fibonacci(35);
 `
 
+var iterative = `
+let fibonacci = fn(x) {
+	let sequence = [0, 1];
+	for (let i = 2; i <= x; i += 1) {
+		sequence = push(sequence, sequence[i - 1] + sequence[i - 2]);
+	}
+	return sequence[x];
+};
+fibonacci(35);
+`
+
 func main() {
 	flag.Parse()
 
 	var duration time.Duration
 	var result object.Object
+	var benchmark string
+
+	if *input == "recursive" {
+		benchmark = recursive
+	} else {
+		benchmark = iterative
+	}
 
-	l := lexer.New(input)
+	l := lexer.New(benchmark)
 	p := parser.New(l)
 	program := p.ParseProgram()
 

code/code.go

@@ -113,6 +113,7 @@ const (
 	OpCall       // Call top n+1 elements of the stack as a function // n is the number of arguments // last element is the function
 
 	OpCurrentClosure // Push the current closure as a variable // recursion
+	OpSetFree        // Pop the top element of the stack and set it to a free scope variable
 	OpGetFree        // Push a variable from the free scope
 
 	OpReturnValue // Return from a function with a value
@@ -193,6 +194,7 @@ var definitions = map[Opcode]*Definition{
 	OpReturn:      {"OpReturn", []int{}},      // No operands, 1 byte in total
 
 	OpCurrentClosure: {"OpCurrentClosure", []int{}}, // No operands, 1 byte in total
+	OpSetFree:        {"OpSetFree", []int{1}},       // Single operand of 1 byte, 2 bytes in total
 	OpGetFree:        {"OpGetFree", []int{1}},       // Single operand of 1 byte, 2 bytes in total
 
 	// Loop Opcodes

compiler/compiler.go

@@ -126,9 +126,12 @@ func (c *Compiler) Compile(node ast.Node) error {
 			c.emit(code.OpMod)
 		}
 
-		if symbol.Scope == GlobalScope {
+		switch scope := symbol.Scope; scope {
+		case GlobalScope:
 			c.emit(code.OpSetGlobal, symbol.Index)
-		} else {
+		case FreeScope:
+			c.emit(code.OpSetFree, symbol.Index)
+		default:
 			c.emit(code.OpSetLocal, symbol.Index)
 		}
 
@@ -201,12 +204,19 @@ func (c *Compiler) Compile(node ast.Node) error {
 
 		c.emit(code.OpBreak)
 
+		freeSymbols := c.symbolTable.FreeSymbols
 		numLoc := c.symbolTable.numDefinitions
 		ins := c.leaveScope()
 
+		free := make([]int, len(freeSymbols))
+		for i, s := range freeSymbols {
+			free[i] = s.Index
+		}
+
 		compiled := &object.CompiledFor{
 			Instructions: ins,
 			NumLocals:    numLoc,
+			Free:         free,
 		}
 
 		idx := c.addConstant(compiled)

compiler/symbol_table.go

@@ -11,9 +11,10 @@ const (
 )
 
 type Symbol struct {
-	Name  string
-	Scope SymbolScope
-	Index int
+	Name       string
+	Scope      SymbolScope
+	Index      int
+	StackIndex int
 }
 
 type SymbolTable struct {

evaluator/evaluator.go

@@ -31,7 +31,7 @@ func Eval(node ast.Node, env *object.Environment) object.Object {
 	// Statements
 	case *ast.LetStatement:
 		_, ok := env.GetNoRecursion(node.Name.Value)
-		if ok {
+		if ok && !env.IsLoop() {
 			return newError("identifier already declared: " + node.Name.Value)
 		}
 
@@ -77,6 +77,7 @@ func Eval(node ast.Node, env *object.Environment) object.Object {
 	case *ast.ForLoopStatement:
 		// Create a new environment for the for loop
 		forEnv := object.NewEnclosedEnvironment(env)
+		forEnv.SetLoop(true)
 
 		// Evaluate the init expression
 		if node.Initializer != nil {

evaluator/evaluator_test.go

@@ -664,6 +664,22 @@ func TestEvalEmptyForLoop(t *testing.T) {
 	testIntegerObject(t, testEval(input), 5)
 }
 
+func TestEvalDoubleForLoop(t *testing.T) {
+	input := `
+	let x = 0;
+	for (let i = 0; i < 10; i += 1) {
+		let sum = 0;
+		for (let j = 0; j < 10; j += 1) {
+			sum = sum + j;
+		}
+		x = sum;
+	}
+	x;
+	`
+
+	testIntegerObject(t, testEval(input), 45)
+}
+
 func TestEvalCompoundAssignment(t *testing.T) {
 	input := `
 	let x = 10;
@@ -681,3 +697,37 @@ func TestEvalModulo(t *testing.T) {
 
 	testIntegerObject(t, testEval(input), 2)
 }
+
+func TestRecursiveFibonacci(t *testing.T) {
+	input := `
+	let fibonacci = fn(x) {
+		if (x == 0) {
+			return 0;
+		} else {
+			if (x == 1) {
+				return 1;
+			} else {
+				fibonacci(x - 1) + fibonacci(x - 2);
+			}
+		}
+	};
+	fibonacci(15);
+	`
+
+	testIntegerObject(t, testEval(input), 610)
+}
+
+func TestIterativeFibonacci(t *testing.T) {
+	input := `
+	let fibonacci = fn(x) {
+		let sequence = [0, 1];
+		for (let i = 2; i <= x; i += 1) {
+			sequence = push(sequence, sequence[i - 1] + sequence[i - 2]);
+		}
+		return sequence[x];
+	};
+	fibonacci(15);
+	`
+
+	testIntegerObject(t, testEval(input), 610)
+}

object/environment.go

@@ -3,6 +3,7 @@ package object
 type Environment struct {
 	store map[string]Object
 	outer *Environment
+	loop  bool
 }
 
 func NewEnvironment() *Environment {
@@ -35,7 +36,10 @@ func (e *Environment) Set(name string, val Object) Object {
 	return val
 }
 
-func SetInEnv(name string, val Object, env *Environment) Object {
-	env.store[name] = val
-	return val
+func (e *Environment) IsLoop() bool {
+	return e.loop
+}
+
+func (e *Environment) SetLoop(loop bool) {
+	e.loop = loop
 }

object/object.go

@@ -219,6 +219,7 @@ func (c *Closure) Inspect() string {
 type CompiledFor struct {
 	Instructions code.Instructions
 	NumLocals    int
+	Free         []int
 }
 
 func (cf *CompiledFor) Type() ObjectType { return COMPILED_FOR_OBJ }

vm/vm.go

@@ -259,16 +259,42 @@ func (vm *VM) Run() error {
 				return err
 			}
 
+		case code.OpSetFree:
+			freeIndex := code.ReadUint8(ins[ip+1:])
+			vm.currentFrame().ip++
+
+			currentLoop, ok := vm.currentFrame().obj.(*object.CompiledFor)
+			if !ok {
+				return fmt.Errorf("not a compiled for loop: %+v", vm.currentFrame().obj)
+			}
+
+			idx := currentLoop.Free[freeIndex]
+			vm.stack[idx] = vm.pop()
+
 		case code.OpGetFree:
 			freeIndex := code.ReadUint8(ins[ip+1:])
 			vm.currentFrame().ip++
 
-			currentClosure := vm.currentFrame().obj.(*object.Closure)
-			err := vm.push(currentClosure.Free[freeIndex])
-			if err != nil {
-				return err
+			switch current := vm.currentFrame().obj.(type) {
+			case *object.Closure:
+				err := vm.push(current.Free[freeIndex])
+				if err != nil {
+					return err
+				}
+			case *object.CompiledFor:
+				free := vm.stack[current.Free[freeIndex]]
+				err := vm.push(free)
+				if err != nil {
+					return err
+				}
 			}
 
+			// currentClosure := vm.currentFrame().obj.(*object.Closure)
+			// err := vm.push(currentClosure.Free[freeIndex])
+			// if err != nil {
+			// 	return err
+			// }
+
 		case code.OpCurrentClosure:
 			currentClosure := vm.currentFrame().obj.(*object.Closure)
 			err := vm.push(currentClosure)
@@ -285,6 +311,10 @@ func (vm *VM) Run() error {
 				return fmt.Errorf("not a compiled for loop: %+v", vm.constants[constIndex])
 			}
 
+			for i, freeIndex := range compiledFor.Free {
+				compiledFor.Free[i] = vm.currentFrame().basePointer + int(freeIndex)
+			}
+
 			vm.push(compiledFor)
 
 			loopFrame := NewFrame(compiledFor, vm.sp-1)

vm/vm_test.go

@@ -635,6 +635,26 @@ func TestRecursiveFibonacci(t *testing.T) {
 	runVmTests(t, tests)
 }
 
+func TestIterativeFibonacci(t *testing.T) {
+	tests := []vmTestCase{
+		{
+			input: `
+			let fibonacci = fn(x) {
+				let sequence = [0, 1];
+				for (let i = 2; i <= x; i += 1) {
+					sequence = push(sequence, sequence[i - 1] + sequence[i - 2]);
+				}
+				return sequence[x];
+			};
+			fibonacci(15);
+			`,
+			expected: 610,
+		},
+	}
+
+	runVmTests(t, tests)
+}
+
 func TestFloatingPointNumbers(t *testing.T) {
 	tests := []vmTestCase{
 		{"3.14", 3.14},
@@ -685,6 +705,20 @@ func TestForLoop(t *testing.T) {
 			`,
 			expected: 5,
 		},
+		{
+			input: `
+			let x = 0;
+			for (let i = 0; i < 10; i += 1) {
+				let sum = 0;
+				for (let j = 0; j < 10; j += 1) {
+					sum = sum + j;
+				}
+				x = sum;
+			}
+			x;
+			`,
+			expected: 45,
+		},
 	}
 
 	runVmTests(t, tests)

vm/vm_test_helpers.go

@@ -31,6 +31,7 @@ func runVmTests(t *testing.T, tests []vmTestCase) {
 		}
 
 		vm := New(comp.Bytecode())
+
 		err = vm.Run()
 		if err != nil {
 			t.Fatalf("vm error: %s", err)