3.7 Evaluation (return statements)

evaluator/evaluator.go

@@ -33,17 +33,22 @@ func Eval(node ast.Node) object.Object {
 	// Statements
 	case *ast.Program:
 		// Traverse the tree and evaluate every statement of the *ast.Program.
-		return evalStatements(node.Statements)
+		return evalProgram(node)
 
 	case *ast.BlockStatement:
-		return evalStatements(node.Statements)
+		return evalBlockStatement(node)
 
 	case *ast.ExpressionStatement:
 		// If the statement is an *ast.ExpressionStatement we evaluate its
 		// expression. An expression statement (not a return statement and not
 		// a let statement).
 		return Eval(node.Expression)
 
+	case *ast.ReturnStatement:
+		// Evaluate the expression associated with the return statement.
+		val := Eval(node.ReturnValue)
+		return &object.ReturnValue{Value: val}
+
 	// Expressions
 	case *ast.IntegerLiteral:
 		return &object.Integer{Value: node.Value}
@@ -69,11 +74,42 @@ func Eval(node ast.Node) object.Object {
 	return nil
 }
 
-func evalStatements(stmts []ast.Statement) object.Object {
+func evalProgram(program *ast.Program) object.Object {
+	// evalProgram was renamed from evalStatements and make less generic because
+	// we can’t reuse evalStatements function for evaluating block statements.
+	// We are using evalBlockStatement for evaluating block statements.
+
 	var result object.Object
 
-	for _, statement := range stmts {
+	for _, statement := range program.Statements {
 		result = Eval(statement)
+
+		// Check if the last evaluation result is such an object.ReturnValue and
+		// if so, we stop the evaluation and return the unwrapped value.
+		if returnValue, ok := result.(*object.ReturnValue); ok {
+			return returnValue.Value
+		}
+	}
+
+	return result
+}
+
+func evalBlockStatement(block *ast.BlockStatement) object.Object {
+	// Evaluate an *ast.BlockStatement.
+
+	var result object.Object
+
+	for _, statement := range block.Statements {
+		result = Eval(statement)
+
+		// Here we explicitly don't unwrap the return value and only check the
+		// Type() of each evaluation result. If it's object.RETURN_VALUE_OBJ we
+		// simply return the *object.ReturnValue, without unwrapping its .Value,
+		// so it stops execution in a possible outer block statement and bubbles
+		// up to evalProgram, where it finally get's unwrapped.
+		if result != nil && result.Type() == object.RETURN_VALUE_OBJ {
+			return result
+		}
 	}
 
 	return result

evaluator/evaluator_test.go

@@ -132,6 +132,36 @@ func TestIfElseExpressions(t *testing.T) {
 	}
 }
 
+func TestReturnStatements(t *testing.T) {
+	tests := []struct {
+		input    string
+		expected int64
+	}{
+		{"return 10;", 10},
+		{"return 10; 9;", 10},
+		{"return 2 * 5; 9;", 10},
+		{"9; return 2 * 5; 9;", 10},
+		{"if (10 > 1) { return 10; }", 10},
+		{
+			`
+if (10 > 1) {
+  if (10 > 1) {
+    return 10;
+  }
+
+  return 1;
+}
+`,
+			10,
+		},
+	}
+
+	for _, tt := range tests {
+		evaluated := testEval(tt.input)
+		testIntegerObject(t, evaluated, tt.expected)
+	}
+}
+
 func testEval(input string) object.Object {
 	l := lexer.New(input)
 	p := parser.New(l)

object/object.go

@@ -15,6 +15,9 @@ const (
 
 	// NULL_OBJ is the Null object type.
 	NULL_OBJ = "NULL"
+
+	// RETURN_VALUE_OBJ is the Return value object type.
+	RETURN_VALUE_OBJ = "RETURN_VALUE"
 )
 
 // ObjectType represents the type of an object.
@@ -66,3 +69,17 @@ func (n *Null) Type() ObjectType { return NULL_OBJ }
 
 // Inspect returns a stringified version of the object for debugging.
 func (n *Null) Inspect() string { return "null" }
+
+// ReturnValue is the return value type and used to hold the value of another
+// object. This is used for `return` statements and this object is tracked
+// through the evaluator and when encountered stops evaluation of the program,
+// or body of a function.
+type ReturnValue struct {
+	Value Object
+}
+
+// Type returns the type of the object.
+func (rv *ReturnValue) Type() ObjectType { return RETURN_VALUE_OBJ }
+
+// Inspect returns a stringified version of the object for debugging.
+func (rv *ReturnValue) Inspect() string { return rv.Value.Inspect() }