Non local returns.

"return" statements inside a function which is inside a method will
return from the surrounding method.

The implementation is surprisingly clean. It turns out using an upvalue
for "this", which we already do, give us all the state we need to find
the enclosing method on the stack.

src/vm/wren_compiler.c

@@ -470,11 +470,11 @@ static void initCompiler(Compiler* compiler, Parser* parser, Compiler* parent,
   }
   else
   {
-    // Declare a fake local variable for the receiver so that it's slot in the
+    // Declare a fake local variable for the receiver so that its slot in the
     // stack is taken. For methods, we call this "this", so that we can resolve
     // references to that like a normal variable. For functions, they have no
-    // explicit "this". So we pick a bogus name. That way references to "this"
-    // inside a function will try to walk up the parent chain to find a method
+    // explicit "this", so we pick a bogus name. That way, references to "this"
+    // inside a function will walk up the parent chain to find a method
     // enclosing the function whose "this" we can close over.
     compiler->numLocals = 1;
     if (isFunction)
@@ -501,6 +501,15 @@ static void initCompiler(Compiler* compiler, Parser* parser, Compiler* parent,
   wrenIntBufferInit(&compiler->debugSourceLines);
 }
 
+// Returns `true` if [compiler] is compiling a function, and not a method or
+// top level code.
+static bool isFunction(Compiler* compiler)
+{
+  // Functions store a null name in their first local slot since they don't
+  // have their own "this".
+  return compiler->numLocals > 0 && compiler->locals[0].name == NULL;
+}
+
 // Lexing ----------------------------------------------------------------------
 
 typedef struct
@@ -2736,6 +2745,7 @@ static int getNumArguments(const uint8_t* bytecode, const Value* constants,
     case CODE_STORE_FIELD_THIS:
     case CODE_LOAD_FIELD:
     case CODE_STORE_FIELD:
+    case CODE_NONLOCAL_RETURN:
     case CODE_CLASS:
       return 1;
 
@@ -2973,6 +2983,7 @@ static void whileStatement(Compiler* compiler)
 // curly blocks. Unlike expressions, these do not leave a value on the stack.
 void statement(Compiler* compiler)
 {
+  // TODO: Allow "break <value>" inside functions for an explicit local return?
   if (match(compiler, TOKEN_BREAK))
   {
     if (compiler->loop == NULL)
@@ -3046,7 +3057,21 @@ void statement(Compiler* compiler)
       expression(compiler);
     }
 
-    emitOp(compiler, CODE_RETURN);
+    // If we are in a function inside a method, then "return" performs a
+    // non-local return.
+    if (getEnclosingClass(compiler) != NULL && isFunction(compiler))
+    {
+      Code loadInstruction;
+      int index = resolveNonmodule(compiler, "this", 4, &loadInstruction);
+      ASSERT(loadInstruction == CODE_LOAD_UPVALUE,
+             "Should have upvalue for 'this'.");
+
+      emitByteArg(compiler, CODE_NONLOCAL_RETURN, index);
+    }
+    else
+    {
+      emitOp(compiler, CODE_RETURN);
+    }
     return;
   }
 

src/vm/wren_opcodes.h

@@ -140,10 +140,13 @@ OPCODE(OR, -1)
 // Close the upvalue for the local on the top of the stack, then pop it.
 OPCODE(CLOSE_UPVALUE, -1)
 
-// Exit from the current function and return the value on the top of the
-// stack.
+// Return from the current method and return the value on top of the stack.
 OPCODE(RETURN, 0)
 
+// Perform a non-local return from the current function up to the enclosing
+// method and return the value on top of the stack.
+OPCODE(NONLOCAL_RETURN, 0)
+
 // Creates a closure for the function stored at [arg] in the constant table.
 //
 // Following the function argument is a number of arguments, two for each

src/vm/wren_value.c

@@ -900,7 +900,7 @@ ObjUpvalue* wrenNewUpvalue(WrenVM* vm, Value* value)
   initObj(vm, &upvalue->obj, OBJ_UPVALUE, NULL);
 
   upvalue->value = value;
-  upvalue->closed = NULL_VAL;
+  upvalue->closed = UNDEFINED_VAL;
   upvalue->next = NULL;
   return upvalue;
 }

src/vm/wren_vm.c

@@ -700,6 +700,59 @@ void wrenFinalizeForeign(WrenVM* vm, ObjForeign* foreign)
   finalizer(foreign->data);
 }
 
+// Executes a non-local return from the current call frame which is assumed to
+// be for a function.
+//
+// [thisUpvalue] is the upvalue that has closed over the "this" of the
+// enclosing method.
+//
+// If the enclosing method is still on the stack, this returns to it. Otherwise,
+// it sets the fiber to a runtime error.
+static void nonlocalReturn(WrenVM* vm, ObjUpvalue* thisUpvalue,
+                           Value returnValue)
+{
+  // Very conveniently, upvalues already track exactly what we need to know to
+  // see if the method's call frame is still on the stack. When compiling a
+  // non-local return, we load "this" and store the index of its upvalue as an
+  // argument to the instruction.
+  //
+  // This ensures we have an upvalue pointing to the enclosing method's
+  // receiver. This upvalue starts off open and closes when the method's frame
+  // returns. So, if the upvalue is still open (its closed value is still
+  // undefined), we know the method hasn't returned yet.
+  if (!IS_UNDEFINED(thisUpvalue->closed))
+  {
+    vm->fiber->error = wrenStringFormat(vm,
+        "Non-local return from method that already returned.");
+    return;
+  }
+
+  // Figure out how many call frames we'll be unwinding. We know "this" is
+  // stored in the first stack slot of the method's frame, so find that frame.
+  //
+  // Since we handled the method not being on the stack above, we can be
+  // certain it will be found.
+  ObjFiber* fiber = vm->fiber;
+  while (fiber->frames[fiber->numFrames - 1].stackStart != thisUpvalue->value) {
+    fiber->numFrames--;
+  }
+  
+  // We are on the method frame now, so make it return too.
+  fiber->numFrames--;
+  
+  // Close any upvalues still in scope.
+  CallFrame* frame = &fiber->frames[fiber->numFrames];
+  closeUpvalues(fiber, frame->stackStart);
+
+  // Store the result of the function in the first slot, which is where the
+  // caller expects it.
+  frame->stackStart[0] = returnValue;
+  
+  // Discard the stack slots for the discarded frames (leaving one slot for the
+  // result).
+  fiber->stackTop = frame->stackStart + 1;
+}
+
 // The main bytecode interpreter loop. This is where the magic happens. It is
 // also, as you can imagine, highly performance critical. Returns `true` if the
 // fiber completed without error.
@@ -1124,7 +1177,7 @@ static WrenInterpretResult runInterpreter(WrenVM* vm, register ObjFiber* fiber)
       }
       else
       {
-        // Store the result of the block in the first slot, which is where the
+        // Store the result of the method in the first slot, which is where the
         // caller expects it.
         stackStart[0] = result;
 
@@ -1136,6 +1189,19 @@ static WrenInterpretResult runInterpreter(WrenVM* vm, register ObjFiber* fiber)
       LOAD_FRAME();
       DISPATCH();
     }
+    
+    CASE_CODE(NONLOCAL_RETURN):
+    {
+      Value result = POP();
+      ObjUpvalue** upvalues = ((ObjClosure*)frame->fn)->upvalues;
+      ObjUpvalue* thisUpvalue = upvalues[READ_BYTE()];
+
+      nonlocalReturn(vm, thisUpvalue, result);
+      if (!IS_NULL(fiber->error)) RUNTIME_ERROR();
+
+      LOAD_FRAME();
+      DISPATCH();
+    }
 
     CASE_CODE(CONSTRUCT):
       ASSERT(IS_CLASS(stackStart[0]), "'this' should be a class.");

test/language/return/nonlocal.wren

@@ -0,0 +1,11 @@
+class Foo {
+  static test() {
+    Fn.new {
+      return "ok"
+    }.call()
+
+    System.print("bad")
+  }
+}
+
+System.print(Foo.test()) // expect: ok

test/language/return/nonlocal_after_method_returns.wren

@@ -0,0 +1,14 @@
+class Foo {
+  static setup() {
+    __field = Fn.new {
+      return // expect runtime error: Non-local return from method that already returned.
+    }
+  }
+
+  static test() {
+    __field.call()
+  }
+}
+
+Foo.setup()
+System.print(Foo.test())

test/language/return/nonlocal_nested_functions.wren

@@ -0,0 +1,23 @@
+class Foo {
+  static test() {
+    System.print("1")
+    Fn.new {
+      System.print("2")
+      Fn.new {
+        System.print("3")
+        Fn.new {
+          return "ok"
+        }.call()
+        System.print("bad")
+      }.call()
+      System.print("bad")
+    }.call()
+    System.print("bad")
+  }
+}
+
+System.print(Foo.test())
+// expect: 1
+// expect: 2
+// expect: 3
+// expect: ok

test/language/return/nonlocal_return_past_methods.wren

@@ -0,0 +1,32 @@
+class Foo {
+  static test() {
+    a {
+      return "ok"
+    }
+    System.print("bad")
+  }
+
+  static a(fn) {
+    System.print("a")
+    b(fn)
+    System.print("bad")
+  }
+
+  static b(fn) {
+    System.print("b")
+    c(fn)
+    System.print("bad")
+  }
+
+  static c(fn) {
+    System.print("c")
+    fn.call()
+    System.print("bad")
+  }
+}
+
+System.print(Foo.test())
+// expect: a
+// expect: b
+// expect: c
+// expect: ok