Support 'return' from within a 'with' block

README.md

@@ -1026,7 +1026,35 @@ print(f)
 # → <closed file 'README.md' ...>
 ```
 
-_**Note:** The implementation of `with` blocks is incomplete; exceptions raised from within a `with` that are not caught within the block will cause `__exit__` to not be called. The same is also true of `return` statements within a `with` block._
+If an early return is encountered inside of a `with` block, the `__exit__` method for the context manager will be called before the function returns.
+
+```py
+class ContextManager:
+    def __init__(title):
+        self.title = title
+    def __enter__():
+        print("Enter context manager", self.title)
+    def __exit__():
+        print("Exit context manager", self.title)
+def doesANestedThing():
+    with ContextManager('outer'):
+        with ContextManager('inner'):
+            with ContextManager('triple'):
+                return 42
+            print('Should not print')
+        print('Should not print')
+    print('Should not print')
+print(doesANestedThing())
+# → Enter context manager outer
+#   Enter context manager inner
+#   Enter context manager triple
+#   Exit context manager triple
+#   Exit context manager inner
+#   Exit context manager outer
+#   42
+```
+
+_**Note:** The implementation of `with` blocks is incomplete; exceptions raised from within a `with` that are not caught within the block will cause `__exit__` to not be called._
 
 ## About the REPL
 

chunk.h

@@ -72,6 +72,8 @@ typedef enum {
 	OP_FINALIZE,
 	OP_TUPLE,
 	OP_UNPACK_TUPLE,
+	OP_PUSH_WITH,
+	OP_CLEANUP_WITH,
 
 	OP_CONSTANT_LONG = 128,
 	OP_DEFINE_GLOBAL_LONG,

compiler.c

@@ -1067,33 +1067,22 @@ static void withStatement() {
 		defineVariable(ind);
 	} else {
 		/* Otherwise we want an unnamed local; TODO: Wait, can't we do this for iterable counts? */
-		Local * local = &current->locals[current->localCount++];
-		local->depth = current->scopeDepth;
-		local->isCaptured = 0;
-		local->name.start = "";
-		local->name.length = 0;
+		addLocal(syntheticToken(""));
+		markInitialized();
 	}
 
 	consume(TOKEN_COLON, "Expected ':' after with statement");
 
-	/* Call enter */
-	emitBytes(OP_DUP, 0);
-	KrkToken _enter = syntheticToken("__enter__");
-	ssize_t ind = identifierConstant(&_enter);
-	EMIT_CONSTANT_OP(OP_GET_PROPERTY, ind);
-	emitBytes(OP_CALL, 0);
-	emitByte(OP_POP); /* We don't care about the result */
+	addLocal(syntheticToken(""));
+	int withJump = emitJump(OP_PUSH_WITH);
+	markInitialized();
 
 	beginScope();
 	block(blockWidth,"with");
 	endScope();
 
-	emitBytes(OP_DUP, 0);
-	KrkToken _exit = syntheticToken("__exit__");
-	ind = identifierConstant(&_exit);
-	EMIT_CONSTANT_OP(OP_GET_PROPERTY, ind);
-	emitBytes(OP_CALL, 0);
-	emitByte(OP_POP); /* We don't care about the result */
+	patchJump(withJump);
+	emitByte(OP_CLEANUP_WITH);
 
 	/* Scope exit pops context manager */
 	endScope();
@@ -1235,7 +1224,6 @@ static void forStatement() {
 	int exitJump;
 
 	if (match(TOKEN_IN)) {
-		defineVariable(loopInd);
 
 		/* ITERABLE.__iter__() */
 		beginScope();

debug.c

@@ -146,6 +146,8 @@ size_t krk_disassembleInstruction(FILE * f, KrkFunction * func, size_t offset) {
 		JUMP(OP_JUMP_IF_TRUE,+)
 		JUMP(OP_LOOP,-)
 		JUMP(OP_PUSH_TRY,+)
+		JUMP(OP_PUSH_WITH,+)
+		SIMPLE(OP_CLEANUP_WITH)
 		default:
 			fprintf(f, "Unknown opcode: %02x", opcode);
 	}

object.h

@@ -111,6 +111,8 @@ typedef struct KrkClass {
 	KrkObj * _init;
 	KrkObj * _eq;
 	KrkObj * _len;
+	KrkObj * _enter;
+	KrkObj * _exit;
 } KrkClass;
 
 typedef struct KrkInstance {

test/testContextManagerEarlyReturn.krk

@@ -0,0 +1,35 @@
+from fileio import open
+
+class ContextManager:
+    def __init__(title):
+        self.title = title
+    def __enter__():
+        print("Enter context manager", self.title)
+    def __exit__():
+        print("Exit context manager", self.title)
+
+
+def doesThing():
+    print("Creating CM")
+    let cm = ContextManager('simple')
+
+    with cm:
+        print("In context manager")
+        return 42
+
+    print("Not exiting correctly?")
+    return "oh no"
+
+print(doesThing())
+
+
+def doesANestedThing():
+    with ContextManager('outer'):
+        with ContextManager('inner'):
+            with ContextManager('triple'):
+                return 42
+            print('Should not print')
+        print('Should not print')
+    print('Should not print')
+
+print(doesANestedThing())

test/testContextManagerEarlyReturn.krk.expect

@@ -0,0 +1,12 @@
+Creating CM
+Enter context manager simple
+In context manager
+Exit context manager simple
+42
+Enter context manager outer
+Enter context manager inner
+Enter context manager triple
+Exit context manager triple
+Exit context manager inner
+Exit context manager outer
+42

value.c

@@ -34,7 +34,7 @@ void krk_printValue(FILE * f, KrkValue printable) {
 			case VAL_BOOLEAN:  fprintf(f, "%s", AS_BOOLEAN(printable) ? "True" : "False"); break;
 			case VAL_FLOATING: fprintf(f, "%g", AS_FLOATING(printable)); break;
 			case VAL_NONE:     fprintf(f, "None"); break;
-			case VAL_HANDLER:  fprintf(f, "{try->%ld}", AS_HANDLER(printable)); break;
+			case VAL_HANDLER:  fprintf(f, "{%s->%d}", AS_HANDLER(printable).type == OP_PUSH_TRY ? "try" : "with", (int)AS_HANDLER(printable).target); break;
 			case VAL_KWARGS: {
 				if (AS_INTEGER(printable) == LONG_MAX) {
 					fprintf(f, "{unpack single}");

value.h

@@ -17,13 +17,18 @@ typedef enum {
 	/* More here later */
 } KrkValueType;
 
+typedef struct {
+	unsigned short type;
+	unsigned short target;
+} KrkJumpTarget;
+
 typedef struct {
 	KrkValueType type;
 	union {
 		char boolean;
 		long integer;
 		double  floating;
-		long handler;
+		KrkJumpTarget handler;
 		KrkObj *   object;
 	} as;
 } KrkValue;
@@ -32,7 +37,7 @@ typedef struct {
 #define NONE_VAL(value)     ((KrkValue){VAL_NONE,    {.integer = 0}})
 #define INTEGER_VAL(value)  ((KrkValue){VAL_INTEGER, {.integer = value}})
 #define FLOATING_VAL(value) ((KrkValue){VAL_FLOATING,{.floating = value}})
-#define HANDLER_VAL(value)  ((KrkValue){VAL_HANDLER, {.handler = value}})
+#define HANDLER_VAL(ty,ta)  ((KrkValue){VAL_HANDLER, {.handler = (KrkJumpTarget){.type = ty, .target = ta}}})
 #define OBJECT_VAL(value)   ((KrkValue){VAL_OBJECT,  {.object = (KrkObj*)value}})
 #define KWARGS_VAL(value)   ((KrkValue){VAL_KWARGS,  {.integer = value}})
 
@@ -50,6 +55,9 @@ typedef struct {
 #define IS_OBJECT(value)    ((value).type == VAL_OBJECT)
 #define IS_KWARGS(value)    ((value).type == VAL_KWARGS)
 
+#define IS_TRY_HANDLER(value)  (IS_HANDLER(value) && AS_HANDLER(value).type == OP_PUSH_TRY)
+#define IS_WITH_HANDLER(value) (IS_HANDLER(value) && AS_HANDLER(value).type == OP_PUSH_WITH)
+
 typedef struct {
 	size_t capacity;
 	size_t count;

vm.c

@@ -298,6 +298,8 @@ void krk_finalizeClass(KrkClass * _class) {
 		{&_class->_init, METHOD_INIT},
 		{&_class->_eq, METHOD_EQ},
 		{&_class->_len, METHOD_LEN},
+		{&_class->_enter, METHOD_ENTER},
+		{&_class->_exit, METHOD_EXIT},
 		{NULL, 0},
 	};
 
@@ -2637,6 +2639,8 @@ void krk_initVM(int flags) {
 	vm.specialMethodNames[METHOD_DICT_INT] = OBJECT_VAL(S("__dict"));
 	vm.specialMethodNames[METHOD_INREPR] = OBJECT_VAL(S("__inrepr"));
 	vm.specialMethodNames[METHOD_EQ] = OBJECT_VAL(S("__eq__"));
+	vm.specialMethodNames[METHOD_ENTER] = OBJECT_VAL(S("__enter__"));
+	vm.specialMethodNames[METHOD_EXIT] = OBJECT_VAL(S("__exit__"));
 
 	/* Create built-in class `object` */
 	vm.objectClass = krk_newClass(S("object"));
@@ -3011,7 +3015,7 @@ static void addObjects() {
 static int handleException() {
 	int stackOffset, frameOffset;
 	int exitSlot = (vm.exitOnFrame >= 0) ? vm.frames[vm.exitOnFrame].outSlots : 0;
-	for (stackOffset = (int)(vm.stackTop - vm.stack - 1); stackOffset >= exitSlot && !IS_HANDLER(vm.stack[stackOffset]); stackOffset--);
+	for (stackOffset = (int)(vm.stackTop - vm.stack - 1); stackOffset >= exitSlot && !IS_TRY_HANDLER(vm.stack[stackOffset]); stackOffset--);
 	if (stackOffset < exitSlot) {
 		if (exitSlot == 0) {
 			/*
@@ -3288,9 +3292,33 @@ static KrkValue run() {
 		int operandWidth = (opcode & (1 << 7)) ? 3 : 1;
 
 		switch (opcode) {
+			case OP_CLEANUP_WITH: {
+				/* Top of stack is a HANDLER that should have had something loaded into it if it was still valid */
+				KrkValue handler = krk_peek(0);
+				KrkValue contextManager = krk_peek(1);
+				KrkClass * type = AS_CLASS(krk_typeOf(1, &contextManager));
+				krk_push(contextManager);
+				krk_callSimple(OBJECT_VAL(type->_exit), 1, 0);
+				/* Top of stack is now either someone else's problem or a return value */
+				if (AS_HANDLER(handler).type != OP_RETURN) break;
+				krk_pop(); /* handler */
+				krk_pop(); /* context manager */
+			} /* fallthrough */
 			case OP_RETURN: {
 				KrkValue result = krk_pop();
 				closeUpvalues(frame->slots);
+				/* See if this frame had a thing */
+				int stackOffset;
+				for (stackOffset = (int)(vm.stackTop - vm.stack - 1); stackOffset >= (int)frame->slots && !IS_WITH_HANDLER(vm.stack[stackOffset]); stackOffset--);
+				if (stackOffset >= (int)frame->slots) {
+					vm.stackTop = &vm.stack[stackOffset + 1];
+					krk_push(result);
+					krk_swap(2);
+					krk_swap(1);
+					frame->ip = frame->closure->function->chunk.code + AS_HANDLER(krk_peek(0)).target;
+					AS_HANDLER(vm.stackTop[-1]).type = OP_RETURN;
+					break;
+				}
 				vm.frameCount--;
 				if (vm.frameCount == 0) {
 					krk_pop();
@@ -3423,7 +3451,7 @@ static KrkValue run() {
 			}
 			case OP_PUSH_TRY: {
 				uint16_t tryTarget = readBytes(frame, 2) + (frame->ip - frame->closure->function->chunk.code);
-				KrkValue handler = HANDLER_VAL(tryTarget);
+				KrkValue handler = HANDLER_VAL(OP_PUSH_TRY, tryTarget);
 				krk_push(handler);
 				break;
 			}
@@ -3647,12 +3675,27 @@ static KrkValue run() {
 				vm.stackTop[-count] = AS_TUPLE(tuple)->values.values[0];
 				break;
 			}
+			case OP_PUSH_WITH: {
+				uint16_t cleanupTarget = readBytes(frame, 2) + (frame->ip - frame->closure->function->chunk.code);
+				KrkValue contextManager = krk_peek(0);
+				KrkClass * type = AS_CLASS(krk_typeOf(1, &contextManager));
+				if (!type->_enter || !type->_exit) {
+					krk_runtimeError(vm.exceptions.attributeError, "Can not use '%s' as context manager.", krk_typeName(contextManager));
+					goto _finishException;
+				}
+				krk_push(contextManager);
+				krk_callSimple(OBJECT_VAL(type->_enter), 1, 0);
+				/* Ignore result; don't need to pop */
+				KrkValue handler = HANDLER_VAL(OP_PUSH_WITH, cleanupTarget);
+				krk_push(handler);
+				break;
+			}
 		}
 		if (!(vm.flags & KRK_HAS_EXCEPTION)) continue;
 _finishException:
 		if (!handleException()) {
 			frame = &vm.frames[vm.frameCount - 1];
-			frame->ip = frame->closure->function->chunk.code + AS_HANDLER(krk_peek(0));
+			frame->ip = frame->closure->function->chunk.code + AS_HANDLER(krk_peek(0)).target;
 			/* Replace the exception handler with the exception */
 			krk_pop();
 			krk_push(vm.currentException);

vm.h

@@ -39,6 +39,8 @@ typedef enum {
 	METHOD_ORD,
 	METHOD_CALL,
 	METHOD_EQ,
+	METHOD_ENTER,
+	METHOD_EXIT,
 
 	METHOD__MAX,
 } KrkSpecialMethods;