From 19e5beba055cb6a80182cf4e2a1e946399a370c0 Mon Sep 17 00:00:00 2001
From: "K. Lange" <klange@toaruos.org>
Date: Fri, 5 Feb 2021 19:54:46 +0900
Subject: [PATCH] Implement class decorators and move class definitions into
 pseudo-functions

---
 src/compiler.c                     | 31 ++++++++++++++++++++++++++----
 src/vm.c                           |  1 -
 test/testClassDecorator.krk        | 18 +++++++++++++++++
 test/testClassDecorator.krk.expect |  2 ++
 4 files changed, 47 insertions(+), 5 deletions(-)
 create mode 100644 test/testClassDecorator.krk
 create mode 100644 test/testClassDecorator.krk.expect

diff --git a/src/compiler.c b/src/compiler.c
index 8c264c9..f0601d2 100644
--- a/src/compiler.c
+++ b/src/compiler.c
@@ -188,7 +188,7 @@ static void declaration();
 static void or_(int canAssign);
 static void ternary(int canAssign);
 static void and_(int canAssign);
-static void classDeclaration();
+static KrkToken classDeclaration();
 static void declareVariable();
 static void namedVariable(KrkToken name, int canAssign);
 static void addLocal(KrkToken name);
@@ -761,7 +761,11 @@ static void declaration() {
 	} else if (match(TOKEN_LET)) {
 		letDeclaration();
 	} else if (check(TOKEN_CLASS)) {
-		classDeclaration();
+		KrkToken className = classDeclaration();
+		size_t classConst = identifierConstant(&className);
+		parser.previous = className;
+		declareVariable();
+		defineVariable(classConst);
 	} else if (check(TOKEN_AT)) {
 		decorator(0, TYPE_FUNCTION);
 	} else if (match(TOKEN_EOL) || match(TOKEN_EOF)) {
@@ -1018,11 +1022,17 @@ static void method(size_t blockWidth) {
 	}
 }
 
-static void classDeclaration() {
+static KrkToken classDeclaration() {
 	size_t blockWidth = (parser.previous.type == TOKEN_INDENTATION) ? parser.previous.length : 0;
 	advance(); /* Collect the `class` */
 
 	consume(TOKEN_IDENTIFIER, "Expected class name.");
+	Compiler subcompiler;
+	initCompiler(&subcompiler, TYPE_LAMBDA);
+	subcompiler.function->chunk.filename = subcompiler.enclosing->function->chunk.filename;
+
+	beginScope();
+
 	KrkToken className = parser.previous;
 	size_t constInd = identifierConstant(&parser.previous);
 	declareVariable();
@@ -1094,8 +1104,15 @@ static void classDeclaration() {
 	} /* else empty class (and at end of file?) we'll allow it for now... */
 _pop_class:
 	emitByte(OP_FINALIZE);
-	endScope();
 	currentClass = currentClass->enclosing;
+	KrkFunction * makeclass = endCompiler();
+	size_t indFunc = krk_addConstant(currentChunk(), OBJECT_VAL(makeclass));
+	EMIT_CONSTANT_OP(OP_CLOSURE, indFunc);
+	doUpvalues(&subcompiler, makeclass);
+	freeCompiler(&subcompiler);
+	emitBytes(OP_CALL, 0);
+
+	return className;
 }
 
 static void markInitialized() {
@@ -1187,6 +1204,12 @@ static KrkToken decorator(size_t level, FunctionType type) {
 		function(type, blockWidth);
 	} else if (check(TOKEN_AT)) {
 		funcName = decorator(level+1, type);
+	} else if (check(TOKEN_CLASS)) {
+		if (type != TYPE_FUNCTION) {
+			error("Invalid decorator applied to class");
+			return funcName;
+		}
+		funcName = classDeclaration();
 	} else {
 		error("Expected a function declaration or another decorator.");
 		return funcName;
diff --git a/src/vm.c b/src/vm.c
index 7996477..cc6662d 100644
--- a/src/vm.c
+++ b/src/vm.c
@@ -2066,7 +2066,6 @@ static KrkValue run() {
 				KrkClass * _class = AS_CLASS(krk_peek(0));
 				/* Store special methods for quick access */
 				krk_finalizeClass(_class);
-				krk_pop(); /* Pop the class as we're done attaching methods */
 				break;
 			}
 			case OP_INHERIT: {
diff --git a/test/testClassDecorator.krk b/test/testClassDecorator.krk
new file mode 100644
index 0000000..5bf1f83
--- /dev/null
+++ b/test/testClassDecorator.krk
@@ -0,0 +1,18 @@
+def classDecorator(cls):
+    class NewClass(cls):
+        def extraMethod(self):
+            print("decorator-added method",self.__class__.__name__)
+    return NewClass
+
+
+@classDecorator
+class FooBar(object):
+    def __init__(self):
+        self.bar = "baz"
+    def myMethod(self):
+        print(self.bar)
+
+
+let l = FooBar()
+l.extraMethod()
+l.myMethod()
diff --git a/test/testClassDecorator.krk.expect b/test/testClassDecorator.krk.expect
new file mode 100644
index 0000000..90b72a1
--- /dev/null
+++ b/test/testClassDecorator.krk.expect
@@ -0,0 +1,2 @@
+decorator-added method NewClass
+baz