From 9ea28e77f115267ac0b238a6723c5c37bc02403f Mon Sep 17 00:00:00 2001
From: K Lange <klange@dakko.us>
Date: Sat, 13 Mar 2021 14:22:32 +0900
Subject: [PATCH] Add class and module member annotations

---
 src/compiler.c                      | 42 +++++++++++++++++++++++++----
 src/vm.c                            | 13 ++++++---
 test/testClassAnnotation.krk        | 10 +++++++
 test/testClassAnnotation.krk.expect |  2 ++
 4 files changed, 58 insertions(+), 9 deletions(-)
 create mode 100644 test/testClassAnnotation.krk
 create mode 100644 test/testClassAnnotation.krk.expect

diff --git a/src/compiler.c b/src/compiler.c
index 2de2019..18c5ef0 100644
--- a/src/compiler.c
+++ b/src/compiler.c
@@ -137,6 +137,7 @@ typedef struct Compiler {
 typedef struct ClassCompiler {
 	struct ClassCompiler * enclosing;
 	KrkToken name;
+	int hasAnnotations;
 } ClassCompiler;
 
 static Parser parser;
@@ -739,8 +740,23 @@ static void letDeclaration(void) {
 		if (current->scopeDepth > 0) {
 			/* Need locals space */
 			args[argCount++] = current->localCount - 1;
+			if (match(TOKEN_COLON)) {
+				error("Annotation on scoped variable declaration is meaningless.");
+				goto _letDone;
+			}
 		} else {
 			args[argCount++] = ind;
+			if (check(TOKEN_COLON)) {
+				KrkToken name = parser.previous;
+				match(TOKEN_COLON);
+				/* Get __annotations__ from globals */
+				KrkToken annotations = syntheticToken("__annotations__");
+				size_t ind = identifierConstant(&annotations);
+				EMIT_CONSTANT_OP(OP_GET_GLOBAL, ind);
+				emitConstant(OBJECT_VAL(krk_copyString(name.start, name.length)));
+				parsePrecedence(PREC_TERNARY);
+				emitBytes(OP_INVOKE_SETTER, OP_POP);
+			}
 		}
 	} while (match(TOKEN_COMMA));
 
@@ -1094,6 +1110,24 @@ static void method(size_t blockWidth) {
 		decorator(0, TYPE_METHOD);
 	} else if (match(TOKEN_IDENTIFIER)) {
 		size_t ind = identifierConstant(&parser.previous);
+		if (check(TOKEN_COLON)) {
+			KrkToken name = parser.previous;
+			match(TOKEN_COLON);
+			/* Get __annotations__ from class */
+			emitBytes(OP_DUP, 0);
+			KrkToken annotations = syntheticToken("__annotations__");
+			size_t ind = identifierConstant(&annotations);
+			if (!currentClass->hasAnnotations) {
+				EMIT_CONSTANT_OP(OP_MAKE_DICT, 0);
+				EMIT_CONSTANT_OP(OP_SET_PROPERTY, ind);
+				currentClass->hasAnnotations = 1;
+			} else {
+				EMIT_CONSTANT_OP(OP_GET_PROPERTY, ind);
+			}
+			emitConstant(OBJECT_VAL(krk_copyString(name.start, name.length)));
+			parsePrecedence(PREC_TERNARY);
+			emitBytes(OP_INVOKE_SETTER, OP_POP);
+		}
 		consume(TOKEN_EQUAL, "Class field must have value.");
 		expression();
 		rememberClassProperty(ind);
@@ -1156,6 +1190,7 @@ static KrkToken classDeclaration() {
 	classCompiler.enclosing = currentClass;
 	currentClass = &classCompiler;
 	int hasSuperclass = 0;
+	classCompiler.hasAnnotations = 0;
 
 	if (match(TOKEN_LEFT_PAREN)) {
 		startEatingWhitespace();
@@ -1177,13 +1212,10 @@ static KrkToken classDeclaration() {
 	addLocal(syntheticToken("super"));
 	defineVariable(0);
 
-	if (hasSuperclass) {
-		namedVariable(className, 0);
-		emitByte(OP_INHERIT);
-	}
-
 	namedVariable(className, 0);
 
+	emitByte(OP_INHERIT);
+
 	consume(TOKEN_COLON, "Expected colon after class");
 
 	/* Set Class.__module__ to the value of __name__, which is the string
diff --git a/src/vm.c b/src/vm.c
index 396d133..1db0151 100644
--- a/src/vm.c
+++ b/src/vm.c
@@ -2157,7 +2157,6 @@ _resumeHook: (void)0;
 				subclass->allocSize = AS_CLASS(superclass)->allocSize;
 				subclass->_ongcsweep = AS_CLASS(superclass)->_ongcsweep;
 				subclass->_ongcscan = AS_CLASS(superclass)->_ongcscan;
-				krk_pop();
 				break;
 			}
 			case OP_DOCSTRING: {
@@ -2204,12 +2203,17 @@ _resumeHook: (void)0;
 				return result;
 			}
 			case OP_ANNOTATE: {
-				if (!IS_CLOSURE(krk_peek(0))) {
+				if (IS_CLOSURE(krk_peek(0))) {
+					krk_swap(1);
+					AS_CLOSURE(krk_peek(1))->annotations = krk_pop();
+				} else if (IS_NONE(krk_peek(0))) {
+					krk_swap(1);
+					fprintf(stderr, "TODO: Global annotation.\n");
+					krk_pop();
+				} else {
 					krk_runtimeError(vm.exceptions->typeError, "Can not annotate '%s'.", krk_typeName(krk_peek(0)));
 					goto _finishException;
 				}
-				krk_swap(1);
-				AS_CLOSURE(krk_peek(1))->annotations = krk_pop();
 				break;
 			}
 
@@ -2626,6 +2630,7 @@ KrkInstance * krk_startModule(const char * name) {
 	krk_attachNamedObject(&vm.modules, name, (KrkObj*)module);
 	krk_attachNamedObject(&module->fields, "__builtins__", (KrkObj*)vm.builtins);
 	krk_attachNamedObject(&module->fields, "__name__", (KrkObj*)krk_copyString(name,strlen(name)));
+	krk_attachNamedValue(&module->fields, "__annotations__", krk_dict_of(0,NULL,0));
 	return module;
 }
 
diff --git a/test/testClassAnnotation.krk b/test/testClassAnnotation.krk
new file mode 100644
index 0000000..42500ca
--- /dev/null
+++ b/test/testClassAnnotation.krk
@@ -0,0 +1,10 @@
+class Foo(object):
+    a:int = 42
+    b:float = 96.8
+    s:str = "test"
+    def method(self, foo: type(a)) -> type(a):
+        return self.a + foo
+    d:bool = False
+
+print(Foo.__annotations__)
+print(Foo().method(42))
diff --git a/test/testClassAnnotation.krk.expect b/test/testClassAnnotation.krk.expect
new file mode 100644
index 0000000..9477a80
--- /dev/null
+++ b/test/testClassAnnotation.krk.expect
@@ -0,0 +1,2 @@
+{'a': <class 'int'>, 'b': <class 'float'>, 's': <class 'str'>, 'd': <class 'bool'>}
+84