Fix up edge cases in exception handling

src/compiler.c

@@ -2254,12 +2254,17 @@ static void tryStatement(struct GlobalState * state) {
 
 	if (state->parser.hadError) return;
 
-#define EXIT_JUMP_MAX 32
-	int exitJumps = 1;
+#define EXIT_JUMP_MAX 64
+	int exitJumps = 2;
 	int exitJumpOffsets[EXIT_JUMP_MAX] = {0};
 
+	/* Jump possibly to `else` */
 	exitJumpOffsets[0] = emitJump(OP_JUMP);
+
+	/* Except entry point; ENTER_EXCEPT jumps to `finally` or continues to
+	 * first `except` expression test; may end up redundant if there is only an 'else'. */
 	patchJump(tryJump);
+	exitJumpOffsets[1] = emitJump(OP_ENTER_EXCEPT);
 
 	int firstJump = 0;
 	int nextJump = -1;
@@ -2275,16 +2280,14 @@ static void tryStatement(struct GlobalState * state) {
 		if (exitJumps && !firstJump && match(TOKEN_EXCEPT)) {
 			if (nextJump != -1) {
 				patchJump(nextJump);
-				emitByte(OP_POP);
 			}
 			/* Match filter expression (should be class or tuple) */
 			if (!check(TOKEN_COLON) && !check(TOKEN_AS)) {
 				expression(state);
 			} else {
 				emitByte(OP_NONE);
 			}
-			emitByte(OP_FILTER_EXCEPT);
-			nextJump = emitJump(OP_JUMP_IF_FALSE_OR_POP);
+			nextJump = emitJump(OP_FILTER_EXCEPT);
 
 			/* Match 'as' to rename exception */
 			if (match(TOKEN_AS)) {
@@ -2336,9 +2339,7 @@ static void tryStatement(struct GlobalState * state) {
 			emitByte(OP_BEGIN_FINALLY);
 			exitJumps = 0;
 			if (nextJump != -1) {
-				emitByte(OP_NONE);
 				patchJump(nextJump);
-				emitByte(OP_POP);
 			}
 			beginScope(state);
 			block(state,blockWidth,"finally");
@@ -2363,9 +2364,7 @@ static void tryStatement(struct GlobalState * state) {
 
 	if (nextJump >= 0) {
 		emitByte(OP_BEGIN_FINALLY);
-		emitByte(OP_NONE);
 		patchJump(nextJump);
-		emitByte(OP_POP);
 		emitByte(OP_END_FINALLY);
 	}
 

src/kuroko/value.h

@@ -219,7 +219,7 @@ static inline uintptr_t _krk_sanitize(uintptr_t input) {
 #define INTEGER_VAL(value)  ((KrkValue)(((uint64_t)(value) & KRK_VAL_MASK_LOW) | KRK_VAL_MASK_INTEGER))
 #define KWARGS_VAL(value)   ((KrkValue)((uint32_t)(value) | KRK_VAL_MASK_KWARGS))
 #define OBJECT_VAL(value)   ((KrkValue)((_krk_sanitize((uintptr_t)(value)) & KRK_VAL_MASK_LOW) | KRK_VAL_MASK_OBJECT))
-#define HANDLER_VAL(ty,ta)  ((KrkValue)((uint32_t)((((uint16_t)ty) << 16) | ((uint16_t)ta)) | KRK_VAL_MASK_HANDLER))
+#define HANDLER_VAL(ty,ta)  ((KrkValue)((uint64_t)((((uint64_t)ty) << 32) | ((uint32_t)ta)) | KRK_VAL_MASK_HANDLER))
 #define FLOATING_VAL(value) (((KrkValueDbl){.dbl = (value)}).val)
 
 #define KRK_VAL_TYPE(value) ((value) >> 48)
@@ -230,7 +230,7 @@ static inline uintptr_t _krk_sanitize(uintptr_t input) {
 #define AS_BOOLEAN(value) AS_INTEGER(value)
 
 #define AS_NOTIMPL(value)   ((krk_integer_type)((value) & KRK_VAL_MASK_LOW))
-#define AS_HANDLER(value)   ((uint32_t)((value) & KRK_VAL_MASK_LOW))
+#define AS_HANDLER(value)   ((uint64_t)((value) & KRK_VAL_MASK_LOW))
 #define AS_OBJECT(value)    ((KrkObj*)(uintptr_t)(((value) & KRK_VAL_MASK_LOW) | KRK_HEAP_TAG))
 #define AS_FLOATING(value)  (((KrkValueDbl){.val = (value)}).dbl)
 
@@ -249,8 +249,8 @@ static inline uintptr_t _krk_sanitize(uintptr_t input) {
 /* ... and as we said above, if any of the MASK_NAN bits are unset, it's a float. */
 #define IS_FLOATING(value)  (((value) & KRK_VAL_MASK_NAN) != KRK_VAL_MASK_NAN)
 
-#define AS_HANDLER_TYPE(value)    (AS_HANDLER(value) >> 16)
-#define AS_HANDLER_TARGET(value)  (AS_HANDLER(value) & 0xFFFF)
+#define AS_HANDLER_TYPE(value)    (AS_HANDLER(value) >> 32)
+#define AS_HANDLER_TARGET(value)  (AS_HANDLER(value) & 0xFFFFFFFF)
 #define IS_HANDLER_TYPE(value,type) (IS_HANDLER(value) && AS_HANDLER_TYPE(value) == type)
 
 #define KWARGS_SINGLE (INT32_MAX)

src/opcodes.h

@@ -49,7 +49,7 @@ SIMPLE(OP_RAISE_FROM)
 SIMPLE(OP_SHIFTLEFT)
 JUMP(OP_POP_JUMP_IF_FALSE,+)
 SIMPLE(OP_ANNOTATE)
-SIMPLE(OP_FILTER_EXCEPT)
+JUMP(OP_FILTER_EXCEPT,+)
 SIMPLE(OP_BITAND)
 SIMPLE(OP_NONE)
 SIMPLE(OP_POP)
@@ -120,3 +120,4 @@ SIMPLE(OP_SET_ADD_TOP)
 SIMPLE(OP_TUPLE_FROM_LIST)
 
 OPERAND(OP_UNPACK_EX,NOOP)
+JUMP(OP_ENTER_EXCEPT,+)

src/vm.c

@@ -1131,7 +1131,8 @@ static int handleException(void) {
 		stackOffset >= exitSlot &&
 		!IS_HANDLER_TYPE(krk_currentThread.stack[stackOffset], OP_PUSH_TRY) &&
 		!IS_HANDLER_TYPE(krk_currentThread.stack[stackOffset], OP_PUSH_WITH) &&
-		!IS_HANDLER_TYPE(krk_currentThread.stack[stackOffset], OP_FILTER_EXCEPT)
+		!IS_HANDLER_TYPE(krk_currentThread.stack[stackOffset], OP_FILTER_EXCEPT) &&
+		!IS_HANDLER_TYPE(krk_currentThread.stack[stackOffset], OP_RAISE)
 		; stackOffset--);
 	if (stackOffset < exitSlot) {
 		if (exitSlot == 0) {
@@ -2161,6 +2162,7 @@ _resumeHook: (void)0;
 				KrkClass * type = krk_getType(contextManager);
 				krk_push(contextManager);
 				if (AS_HANDLER_TYPE(handler) == OP_RAISE) {
+					krk_currentThread.stackTop[-2] = HANDLER_VAL(OP_CLEANUP_WITH,AS_HANDLER_TARGET(krk_peek(1)));
 					krk_push(OBJECT_VAL(krk_getType(exceptionObject)));
 					krk_push(exceptionObject);
 					KrkValue tracebackEntries = NONE_VAL();
@@ -2193,7 +2195,6 @@ _resumeHook: (void)0;
 			case OP_RETURN: {
 _finishReturn: (void)0;
 				KrkValue result = krk_pop();
-				closeUpvalues(frame->slots);
 				/* See if this frame had a thing */
 				int stackOffset;
 				for (stackOffset = (int)(krk_currentThread.stackTop - krk_currentThread.stack - 1);
@@ -2203,12 +2204,14 @@ _finishReturn: (void)0;
 					!IS_HANDLER_TYPE(krk_currentThread.stack[stackOffset],OP_FILTER_EXCEPT)
 					; stackOffset--);
 				if (stackOffset >= (int)frame->slots) {
+					closeUpvalues(stackOffset);
 					krk_currentThread.stackTop = &krk_currentThread.stack[stackOffset + 1];
 					frame->ip = frame->closure->function->chunk.code + AS_HANDLER_TARGET(krk_peek(0));
 					krk_currentThread.stackTop[-1] = HANDLER_VAL(OP_RETURN,AS_HANDLER_TARGET(krk_peek(0)));
 					krk_currentThread.stackTop[-2] = result;
 					break;
 				}
+				closeUpvalues(frame->slots);
 				FRAME_OUT(frame);
 				krk_currentThread.frameCount--;
 				if (krk_currentThread.frameCount == 0) {
@@ -2311,33 +2314,6 @@ _finishReturn: (void)0;
 			case OP_SWAP:
 				krk_swap(1);
 				break;
-			case OP_FILTER_EXCEPT: {
-				int isMatch = 0;
-				if (AS_HANDLER_TYPE(krk_peek(1)) == OP_RETURN) {
-					isMatch = 0;
-				} else if (AS_HANDLER_TYPE(krk_peek(1)) == OP_END_FINALLY) {
-					isMatch = 0;
-				} else if (AS_HANDLER_TYPE(krk_peek(1)) == OP_EXIT_LOOP) {
-					isMatch = 0;
-				} else if (IS_CLASS(krk_peek(0)) && krk_isInstanceOf(krk_peek(2), AS_CLASS(krk_peek(0)))) {
-					isMatch = 1;
-				} else if (IS_TUPLE(krk_peek(0))) {
-					for (size_t i = 0; i < AS_TUPLE(krk_peek(0))->values.count; ++i) {
-						if (IS_CLASS(AS_TUPLE(krk_peek(0))->values.values[i]) && krk_isInstanceOf(krk_peek(2), AS_CLASS(AS_TUPLE(krk_peek(0))->values.values[i]))) {
-							isMatch = 1;
-							break;
-						}
-					}
-				} else if (IS_NONE(krk_peek(0))) {
-					isMatch = !IS_NONE(krk_peek(2));
-				}
-				if (isMatch) {
-					krk_currentThread.stackTop[-2] = HANDLER_VAL(OP_FILTER_EXCEPT,AS_HANDLER_TARGET(krk_peek(1)));
-				}
-				krk_pop();
-				krk_push(BOOLEAN_VAL(isMatch));
-				break;
-			}
 			case OP_TRY_ELSE: {
 				if (IS_HANDLER(krk_peek(0))) {
 					krk_currentThread.stackTop[-1] = HANDLER_VAL(OP_FILTER_EXCEPT,AS_HANDLER_TARGET(krk_peek(0)));
@@ -2573,6 +2549,43 @@ _finishReturn: (void)0;
 				if (krk_pop() != KWARGS_VAL(0)) frame->ip += OPERAND;
 				break;
 			}
+			case OP_FILTER_EXCEPT: {
+				TWO_BYTE_OPERAND;
+				/* "Pop exception to match with and jump if not a match" */
+				int isMatch = 0;
+				if (IS_CLASS(krk_peek(0)) && krk_isInstanceOf(krk_peek(2), AS_CLASS(krk_peek(0)))) {
+					isMatch = 1;
+				} else if (IS_TUPLE(krk_peek(0))) {
+					for (size_t i = 0; i < AS_TUPLE(krk_peek(0))->values.count; ++i) {
+						if (IS_CLASS(AS_TUPLE(krk_peek(0))->values.values[i]) && krk_isInstanceOf(krk_peek(2), AS_CLASS(AS_TUPLE(krk_peek(0))->values.values[i]))) {
+							isMatch = 1;
+							break;
+						}
+					}
+				} else if (IS_NONE(krk_peek(0))) {
+					isMatch = !IS_NONE(krk_peek(2));
+				}
+				if (isMatch) {
+					/* If exception matched, set handler state. */
+					krk_currentThread.stackTop[-2] = HANDLER_VAL(OP_FILTER_EXCEPT,AS_HANDLER_TARGET(krk_peek(1)));
+				} else {
+					/* If exception did not match, jump to next 'except' or 'finally' */
+					frame->ip += OPERAND;
+				}
+				krk_pop();
+				break;
+			}
+			case OP_ENTER_EXCEPT: {
+				TWO_BYTE_OPERAND;
+				switch (AS_HANDLER_TYPE(krk_peek(0))) {
+					case OP_RETURN:
+					case OP_END_FINALLY:
+					case OP_EXIT_LOOP:
+						frame->ip += OPERAND;
+						break;
+				}
+				break;
+			}
 
 			case OP_CONSTANT_LONG:
 				THREE_BYTE_OPERAND;
@@ -2897,8 +2910,7 @@ _finishReturn: (void)0;
 				THREE_BYTE_OPERAND;
 			case OP_EXIT_LOOP: {
 				ONE_BYTE_OPERAND;
-_finishPopBlock:
-				closeUpvalues(frame->slots + OPERAND);
+_finishPopBlock: (void)0;
 				int stackOffset;
 				for (stackOffset = (int)(krk_currentThread.stackTop - krk_currentThread.stack - 1);
 					stackOffset >= (int)(frame->slots + OPERAND) && 
@@ -2909,11 +2921,14 @@ _finishReturn: (void)0;
 
 				/* Do the handler. */
 				if (stackOffset >= (int)(frame->slots + OPERAND)) {
+					closeUpvalues(stackOffset);
 					uint16_t popTarget = (frame->ip - frame->closure->function->chunk.code);
 					krk_currentThread.stackTop = &krk_currentThread.stack[stackOffset + 1];
 					frame->ip = frame->closure->function->chunk.code + AS_HANDLER_TARGET(krk_peek(0));
 					krk_currentThread.stackTop[-1] = HANDLER_VAL(OP_EXIT_LOOP, popTarget);
 					krk_currentThread.stackTop[-2] = INTEGER_VAL(OPERAND);
+				} else {
+					closeUpvalues(frame->slots + OPERAND);
 				}
 
 				/* Continue normally */
@@ -3145,10 +3160,14 @@ _finishReturn: (void)0;
 				frame = &krk_currentThread.frames[krk_currentThread.frameCount - 1];
 				frame->ip = frame->closure->function->chunk.code + AS_HANDLER_TARGET(krk_peek(0));
 				/* Stick the exception into the exception slot */
-				if (AS_HANDLER_TYPE(krk_currentThread.stackTop[-1])== OP_FILTER_EXCEPT) {
-					krk_currentThread.stackTop[-1] = HANDLER_VAL(OP_END_FINALLY,AS_HANDLER_TARGET(krk_peek(0)));
-				} else {
-					krk_currentThread.stackTop[-1] = HANDLER_VAL(OP_RAISE,AS_HANDLER_TARGET(krk_peek(0)));
+				switch (AS_HANDLER_TYPE(krk_currentThread.stackTop[-1])) {
+					case OP_RAISE:
+					case OP_FILTER_EXCEPT:
+						krk_currentThread.stackTop[-1] = HANDLER_VAL(OP_END_FINALLY,AS_HANDLER_TARGET(krk_peek(0)));
+						break;
+					default:
+						krk_currentThread.stackTop[-1] = HANDLER_VAL(OP_RAISE,AS_HANDLER_TARGET(krk_peek(0)));
+						break;
 				}
 				krk_currentThread.stackTop[-2] = krk_currentThread.currentException;
 				krk_currentThread.currentException = NONE_VAL();

test/testElseWithNoExcept.krk

@@ -0,0 +1,9 @@
+try:
+    try:
+        raise ValueError()
+    else:
+        print("oh no (else should not run)")
+    finally:
+        print('finally')
+except ValueError:
+    print("value error caught in outer")

test/testElseWithNoExcept.krk.expect

@@ -0,0 +1,2 @@
+finally
+value error caught in outer

test/testExceptionElseFinally.krk

@@ -0,0 +1,13 @@
+for i in range(3):
+    try:
+        print('hello world')
+        if i == 1:
+            break
+    except ValueError:
+        print('value error')
+    else:
+        print('no error')
+    finally:
+        print('and done')
+
+print('finished')

test/testExceptionElseFinally.krk.expect

@@ -0,0 +1,6 @@
+hello world
+no error
+and done
+hello world
+and done
+finished

test/testExceptionOnExceptExpression.krk

@@ -0,0 +1,14 @@
+try:
+    try:
+        print('hello world')
+        raise ValueError()
+    except 1/0:
+        print("oh no")
+    except ValueError:
+        print("value error")
+    finally:
+        print('this is the finally')
+except ZeroDivisionError:
+    print('caught the zero div error')
+
+print('done')

test/testExceptionOnExceptExpression.krk.expect

@@ -0,0 +1,4 @@
+hello world
+this is the finally
+caught the zero div error
+done

test/testUpvalueClosureOnEarlyReturnOrLoopExit.krk

@@ -0,0 +1,23 @@
+let closures = []
+for i in range(3):
+    let a
+    try:
+        a = 'bad'
+        closures.append(lambda: print(a))
+        continue
+    finally:
+        a = f'good {i}'
+
+def foo():
+    let a
+    try:
+        a = 'bad'
+        closures.append(lambda: print(a))
+        return
+    finally:
+        a = 'good'
+
+foo()
+
+for i in closures:
+    i()

test/testUpvalueClosureOnEarlyReturnOrLoopExit.krk.expect

@@ -0,0 +1,4 @@
+good 0
+good 1
+good 2
+good