From 46dada92fc9f8d373a4d8c4d8a2021f6c06b38c4 Mon Sep 17 00:00:00 2001
From: "K. Lange" <klange@toaruos.org>
Date: Thu, 18 Mar 2021 21:25:19 +0900
Subject: [PATCH] Support multiple context managers in a 'with' statement

---
 src/compiler.c                   | 17 ++++++++++----
 src/vm.c                         | 10 ++++++--
 test/testMultipleWith.krk        | 39 ++++++++++++++++++++++++++++++++
 test/testMultipleWith.krk.expect | 27 ++++++++++++++++++++++
 4 files changed, 86 insertions(+), 7 deletions(-)
 create mode 100644 test/testMultipleWith.krk
 create mode 100644 test/testMultipleWith.krk.expect

diff --git a/src/compiler.c b/src/compiler.c
index e8ba71a..c503493 100644
--- a/src/compiler.c
+++ b/src/compiler.c
@@ -1435,6 +1435,7 @@ static void withStatement() {
 
 	/* We only need this for block() */
 	size_t blockWidth = (parser.previous.type == TOKEN_INDENTATION) ? parser.previous.length : 0;
+	KrkToken myPrevious = parser.previous;
 
 	/* Collect the with token that started this statement */
 	advance();
@@ -1453,19 +1454,25 @@ static void withStatement() {
 		markInitialized();
 	}
 
-	consume(TOKEN_COLON, "Expected ':' after with statement");
-
 	/* Storage for return / exception */
 	addLocal(syntheticToken(""));
+	markInitialized();
 
 	/* Handler object */
 	addLocal(syntheticToken(""));
 	int withJump = emitJump(OP_PUSH_WITH);
 	markInitialized();
 
-	beginScope();
-	block(blockWidth,"with");
-	endScope();
+	if (check(TOKEN_COMMA)) {
+		parser.previous = myPrevious;
+		withStatement(); /* Keep nesting */
+	} else {
+		consume(TOKEN_COLON, "Expected ',' or ':' after with statement");
+
+		beginScope();
+		block(blockWidth,"with");
+		endScope();
+	}
 
 	patchJump(withJump);
 	emitByte(OP_CLEANUP_WITH);
diff --git a/src/vm.c b/src/vm.c
index bc647ab..dc115e6 100644
--- a/src/vm.c
+++ b/src/vm.c
@@ -300,7 +300,10 @@ inline void krk_push(KrkValue value) {
 inline KrkValue krk_pop() {
 	krk_currentThread.stackTop--;
 	if (unlikely(krk_currentThread.stackTop < krk_currentThread.stack)) {
-		fprintf(stderr, "Fatal error: stack underflow detected in VM.\n");
+		fprintf(stderr, "Fatal error: stack underflow detected in VM (krk_currentThread.stackTop = %p, krk_currentThread.stack = %p)\n",
+			(void*)krk_currentThread.stackTop,
+			(void*)krk_currentThread.stack);
+		abort();
 		return NONE_VAL();
 	}
 	return *krk_currentThread.stackTop;
@@ -738,8 +741,10 @@ _finishKwarg:
 		/* We can't have had any kwargs. */
 		if ((size_t)argCount > potentialPositionalArgs && closure->function->collectsArguments) {
 			krk_push(NONE_VAL()); krk_push(NONE_VAL()); krk_pop(); krk_pop();
-			startOfPositionals[offsetOfExtraArgs] = krk_list_of(argCount - potentialPositionalArgs,
+			KrkValue tmp = krk_list_of(argCount - potentialPositionalArgs,
 				&startOfPositionals[potentialPositionalArgs], 0);
+			startOfPositionals = &krk_currentThread.stackTop[-argCount];
+			startOfPositionals[offsetOfExtraArgs] = tmp;
 			argCount = closure->function->requiredArgs + 1;
 			argCountX = argCount - 1;
 			while (krk_currentThread.stackTop > startOfPositionals + argCount) krk_pop();
@@ -2004,6 +2009,7 @@ _resumeHook: (void)0;
 					krk_push(NONE_VAL());
 					krk_push(NONE_VAL());
 					krk_callSimple(OBJECT_VAL(type->_exit), 4, 0);
+					if (krk_currentThread.flags & KRK_THREAD_HAS_EXCEPTION) goto _finishException;
 				}
 				if (AS_HANDLER(handler).type != OP_RETURN) break;
 				krk_pop(); /* handler */
diff --git a/test/testMultipleWith.krk b/test/testMultipleWith.krk
new file mode 100644
index 0000000..32627de
--- /dev/null
+++ b/test/testMultipleWith.krk
@@ -0,0 +1,39 @@
+class Context:
+    def __init__(self,name):
+        self.name = name
+    def __enter__(self):
+        print("Entering",self.name)
+    def __exit__(self, *args):
+        print("Exiting",self.name,"with",[type(x) for x in args])
+
+def nested():
+    print("Before")
+    with Context('a') as a:
+        with Context('b') as b:
+            with Context('c') as c:
+                print("Hello, world.")
+    print("After")
+
+nested()
+
+def simple():
+    print("Before")
+    with Context('a') as a, Context('b') as b, Context('c') as c:
+        print("Hello, world.")
+    print("After")
+
+simple()
+
+
+def exceptional():
+    print("Before")
+    with Context('a') as a, Context('b') as b, Context('c') as c:
+        print("Prior")
+        raise ValueError("oh no")
+        print("Shouldn't happen")
+    print("After")
+
+try:
+    exceptional()
+except Exception as e:
+    print("Raised",repr(e))
diff --git a/test/testMultipleWith.krk.expect b/test/testMultipleWith.krk.expect
new file mode 100644
index 0000000..33690ce
--- /dev/null
+++ b/test/testMultipleWith.krk.expect
@@ -0,0 +1,27 @@
+Before
+Entering a
+Entering b
+Entering c
+Hello, world.
+Exiting c with [<class 'NoneType'>, <class 'NoneType'>, <class 'NoneType'>]
+Exiting b with [<class 'NoneType'>, <class 'NoneType'>, <class 'NoneType'>]
+Exiting a with [<class 'NoneType'>, <class 'NoneType'>, <class 'NoneType'>]
+After
+Before
+Entering a
+Entering b
+Entering c
+Hello, world.
+Exiting c with [<class 'NoneType'>, <class 'NoneType'>, <class 'NoneType'>]
+Exiting b with [<class 'NoneType'>, <class 'NoneType'>, <class 'NoneType'>]
+Exiting a with [<class 'NoneType'>, <class 'NoneType'>, <class 'NoneType'>]
+After
+Before
+Entering a
+Entering b
+Entering c
+Prior
+Exiting c with [<class 'type'>, <class 'ValueError'>, <class 'list'>]
+Exiting b with [<class 'type'>, <class 'ValueError'>, <class 'list'>]
+Exiting a with [<class 'type'>, <class 'ValueError'>, <class 'list'>]
+Raised ValueError('oh no')