Demo changes

vm/frame.go

@@ -21,3 +21,7 @@ func NewFrame(cl *object.Closure, basePointer int) *Frame {
 func (f *Frame) Instructions() code.Instructions {
 	return f.cl.Fn.Instructions
 }
+
+func (f *Frame) NextOp() code.Opcode {
+	return code.Opcode(f.Instructions()[f.ip+1])
+}

vm/vm.go

@@ -119,7 +119,11 @@ func (vm *VM) Run() error {
 			}
 
 		case code.OpPop:
-			vm.pop()
+			// This makes things like this work:
+			// >> let x = 1; if (x == 1) { x = 2 }
+			if vm.sp > 0 {
+				vm.pop()
+			}
 
 		case code.OpTrue:
 			err := vm.push(True)
@@ -563,6 +567,19 @@ func (vm *VM) callClosure(cl *object.Closure, numArgs int) error {
 		return fmt.Errorf("wrong number of arguments: want=%d, got=%d", cl.Fn.NumParameters, numArgs)
 	}
 
+	// Optimize tail calls and avoid a new frame
+	if cl.Fn == vm.currentFrame().cl.Fn {
+		nextOP := vm.currentFrame().NextOp()
+		if nextOP == code.OpReturn {
+			for p := 0; p < numArgs; p++ {
+				vm.stack[vm.currentFrame().basePointer+p] = vm.stack[vm.sp-numArgs+p]
+			}
+			vm.sp -= numArgs + 1
+			vm.currentFrame().ip = -1 // reset IP to the beginning of the frame
+			return nil
+		}
+	}
+
 	frame := NewFrame(cl, vm.sp-numArgs)
 	vm.pushFrame(frame)
 	vm.sp = frame.basePointer + cl.Fn.NumLocals

vm/vm_test.go

@@ -833,6 +833,40 @@ func TestAssignmentStatements(t *testing.T) {
 	runVmTests(t, tests)
 }
 
+func TestTailCalls(t *testing.T) {
+	tests := []vmTestCase{
+		{
+			input: `
+			let fact = fn(n, a) {
+			  if (n == 0) {
+				return a
+			  }
+			  return fact(n - 1, a * n)
+			}
+
+			fact(5, 1)
+        	`,
+			expected: 120,
+		},
+
+		// without tail recursion optimization this will cause a stack overflow
+		{
+			input: `
+			let iter = fn(n, max) {
+				if (n == max) {
+					return n
+				}
+				return iter(n + 1, max)
+			}
+			iter(0, 9999)
+			`,
+			expected: 9999,
+		},
+	}
+
+	runVmTests(t, tests)
+}
+
 func TestIntegration(t *testing.T) {
 	matches, err := filepath.Glob("../testdata/*.monkey")
 	if err != nil {
@@ -874,42 +908,79 @@ func TestIntegration(t *testing.T) {
 	}
 }
 
-func TestExamples(t *testing.T) {
-	matches, err := filepath.Glob("../examples/*.monkey")
-	if err != nil {
-		t.Error(err)
+func BenchmarkFibonacci(b *testing.B) {
+	tests := map[string]string{
+		"iterative": `
+		let fib = fn(n) {
+		   if (n < 3) {
+			 return 1
+		   }
+		   let a = 1
+		   let b = 1
+		   let c = 0
+		   let i = 0
+		   while (i < n - 2) {
+			 c = a + b
+			 b = a
+			 a = c
+			 i = i + 1
+		   }
+		   return a
+		}
+
+		fib(35)
+		`,
+		"recursive": `
+		let fib = fn(x) {
+		  if (x == 0) {
+			return 0
+		  }
+		  if (x == 1) {
+			return 1
+		  }
+		  return fib(x-1) + fib(x-2)
+		}
+
+		fib(35)
+		`,
+		"tail-recursive": `
+		let fib = fn(n, a, b) {
+		  if (n == 0) {
+			return a
+		  }
+		  if (n == 1) {
+			return b
+		  }
+		  return fib(n - 1, b, a + b)
+		}
+
+		fib(35, 0, 1)
+		`,
 	}
 
-	for _, match := range matches {
-		basename := path.Base(match)
-		name := strings.TrimSuffix(basename, filepath.Ext(basename))
+	for name, input := range tests {
+		b.Run(name, func(b *testing.B) {
+			for i := 0; i < b.N; i++ {
+				program := parse(input)
 
-		t.Run(name, func(t *testing.T) {
-			b, err := os.ReadFile(match)
-			if err != nil {
-				t.Error(err)
-			}
+				c := compiler.New()
+				err := c.Compile(program)
+				if err != nil {
+					b.Log(input)
+					b.Fatalf("compiler error: %s", err)
+				}
 
-			input := string(b)
-			program := parse(input)
+				vm := New(c.Bytecode())
 
-			c := compiler.New()
-			err = c.Compile(program)
-			if err != nil {
-				t.Log(input)
-				t.Fatalf("compiler error: %s", err)
-			}
-
-			vm := New(c.Bytecode())
-
-			err = vm.Run()
-			if err != nil {
-				t.Log(input)
-				t.Fatalf("vm error: %s", err)
-			}
-			if vm.sp != 0 {
-				t.Log(input)
-				t.Fatal("vm stack pointer non-zero")
+				err = vm.Run()
+				if err != nil {
+					b.Log(input)
+					b.Fatalf("vm error: %s", err)
+				}
+				if vm.sp != 0 {
+					b.Log(input)
+					b.Fatal("vm stack pointer non-zero")
+				}
 			}
 		})
 	}