Initial support for async/await

src/compiler.c

@@ -106,6 +106,8 @@ typedef enum {
 	TYPE_PROPERTY,
 	TYPE_CLASS,
 	TYPE_CLASSMETHOD,
+	TYPE_COROUTINE,
+	TYPE_COROUTINE_METHOD,
 } FunctionType;
 
 struct IndexWithNext {
@@ -161,7 +163,11 @@ static int inDel = 0;
 	else { emitBytes(opc ## _LONG, arg >> 16); emitBytes(arg >> 8, arg); } } while (0)
 
 static int isMethod(int type) {
-	return type == TYPE_METHOD || type == TYPE_INIT || type == TYPE_PROPERTY;
+	return type == TYPE_METHOD || type == TYPE_INIT || type == TYPE_PROPERTY || type == TYPE_COROUTINE_METHOD;
+}
+
+static int isCoroutine(int type) {
+	return type == TYPE_COROUTINE || type == TYPE_COROUTINE_METHOD;
 }
 
 static char * calculateQualName(void) {
@@ -248,6 +254,7 @@ static ssize_t identifierConstant(KrkToken * name);
 static ssize_t resolveLocal(Compiler * compiler, KrkToken * name);
 static ParseRule * getRule(KrkTokenType type);
 static void defDeclaration();
+static void asyncDeclaration(int);
 static void expression();
 static void statement();
 static void declaration();
@@ -904,6 +911,7 @@ static void synchronize() {
 			case TOKEN_IF:
 			case TOKEN_WHILE:
 			case TOKEN_RETURN:
+			case TOKEN_ASYNC:
 				return;
 			default: break;
 		}
@@ -925,6 +933,8 @@ static void declaration() {
 		defineVariable(classConst);
 	} else if (check(TOKEN_AT)) {
 		decorator(0, TYPE_FUNCTION);
+	} else if (check(TOKEN_ASYNC)) {
+		asyncDeclaration(1);
 	} else if (match(TOKEN_EOL) || match(TOKEN_EOF)) {
 		return;
 	} else if (check(TOKEN_INDENTATION)) {
@@ -1047,6 +1057,7 @@ static void function(FunctionType type, size_t blockWidth) {
 	beginScope();
 
 	if (isMethod(type)) current->codeobject->requiredArgs = 1;
+	if (isCoroutine(type)) current->codeobject->flags |= KRK_CODEOBJECT_FLAGS_IS_COROUTINE;
 
 	int hasCollectors = 0;
 	KrkToken self = syntheticToken("self");
@@ -1225,14 +1236,27 @@ static void method(size_t blockWidth) {
 		/* bah */
 		consume(TOKEN_EOL, "Expected linefeed after 'pass' in class body.");
 	} else {
+		FunctionType type = TYPE_METHOD;
+		if (match(TOKEN_ASYNC)) {
+			type = TYPE_COROUTINE_METHOD;
+		}
 		consume(TOKEN_DEF, "expected a definition, got nothing");
 		consume(TOKEN_IDENTIFIER, "expected method name");
 		size_t ind = identifierConstant(&parser.previous);
-		FunctionType type = TYPE_METHOD;
 
 		if (parser.previous.length == 8 && memcmp(parser.previous.start, "__init__", 8) == 0) {
+			if (type == TYPE_COROUTINE_METHOD) {
+				error("'%.*s' can not be a coroutine",
+					(int)parser.previous.length, parser.previous.start);
+				return;
+			}
 			type = TYPE_INIT;
 		} else if (parser.previous.length == 17 && memcmp(parser.previous.start, "__class_getitem__", 17) == 0) {
+			if (type == TYPE_COROUTINE_METHOD) {
+				error("'%.*s' can not be a coroutine",
+					(int)parser.previous.length, parser.previous.start);
+				return;
+			}
 			/* This magic method is implicitly always a class method,
 			 * so mark it as such so we don't do implicit self for it. */
 			type = TYPE_CLASSMETHOD;
@@ -1370,7 +1394,42 @@ static void defDeclaration() {
 	defineVariable(global);
 }
 
+static void asyncDeclaration(int declarationLevel) {
+	size_t blockWidth = (parser.previous.type == TOKEN_INDENTATION) ? parser.previous.length : 0;
+	advance(); /* 'async' */
+
+	if (match(TOKEN_DEF)) {
+		if (!declarationLevel) {
+			error("'async def' not valid here");
+			return;
+		}
+		ssize_t global = parseVariable("Expected coroutine name after 'async def'");
+		markInitialized();
+		function(TYPE_COROUTINE, blockWidth);
+		defineVariable(global);
+	} else if (match(TOKEN_FOR)) {
+		if (!isCoroutine(current->type)) {
+			error("'async for' outside of async function");
+			return;
+		}
+		error("'async for' unsupported");
+		return;
+	} else if (match(TOKEN_WITH)) {
+		if (!isCoroutine(current->type)) {
+			error("'async with' outside of async function");
+			return;
+		}
+		error("'async with' unsupported");
+		return;
+	} else {
+		errorAtCurrent("'%.*s' can not be prefixed with 'async'",
+			(int)parser.current.length, parser.current.start);
+		return;
+	}
+}
+
 static KrkToken decorator(size_t level, FunctionType type) {
+	int inType = type;
 	size_t blockWidth = (parser.previous.type == TOKEN_INDENTATION) ? parser.previous.length : 0;
 	advance(); /* Collect the `@` */
 
@@ -1379,8 +1438,10 @@ static KrkToken decorator(size_t level, FunctionType type) {
 	KrkToken at_staticmethod = syntheticToken("staticmethod");
 	KrkToken at_classmethod  = syntheticToken("classmethod");
 
-	if (identifiersEqual(&at_staticmethod, &parser.current)) type = TYPE_STATIC;
+	if (type == TYPE_METHOD) {
-	if (identifiersEqual(&at_classmethod, &parser.current)) type = TYPE_CLASSMETHOD;
+		if (identifiersEqual(&at_staticmethod, &parser.current)) type = TYPE_STATIC;
+		if (identifiersEqual(&at_classmethod, &parser.current)) type = TYPE_CLASSMETHOD;
+	}
 
 	expression();
 
@@ -1399,6 +1460,14 @@ static KrkToken decorator(size_t level, FunctionType type) {
 			type = TYPE_INIT;
 		}
 		function(type, blockWidth);
+	} else if (match(TOKEN_ASYNC)) {
+		if (!match(TOKEN_DEF)) {
+			errorAtCurrent("Expected 'def' after 'async' with decorator, not '%*.s'",
+				(int)parser.current.length, parser.current.start);
+		}
+		consume(TOKEN_IDENTIFIER, "Expected coroutine name.");
+		funcName = parser.previous;
+		function(type == TYPE_METHOD ? TYPE_COROUTINE_METHOD : TYPE_COROUTINE, blockWidth);
 	} else if (check(TOKEN_AT)) {
 		funcName = decorator(level+1, type);
 	} else if (check(TOKEN_CLASS)) {
@@ -1415,7 +1484,7 @@ static KrkToken decorator(size_t level, FunctionType type) {
 	emitBytes(OP_CALL, 1);
 
 	if (level == 0) {
-		if (type == TYPE_FUNCTION) {
+		if (inType == TYPE_FUNCTION) {
 			parser.previous = funcName;
 			declareVariable();
 			size_t ind = (current->scopeDepth > 0) ? 0 : identifierConstant(&funcName);
@@ -1977,6 +2046,8 @@ static void statement() {
 		whileStatement();
 	} else if (check(TOKEN_FOR)) {
 		forStatement();
+	} else if (check(TOKEN_ASYNC)) {
+		asyncDeclaration(0);
 	} else if (check(TOKEN_TRY)) {
 		tryStatement();
 	} else if (check(TOKEN_WITH)) {
@@ -2048,6 +2119,22 @@ static void yield(int canAssign) {
 	}
 }
 
+static void await(int canAssign) {
+	if (!isCoroutine(current->type)) {
+		error("'await' outside async function");
+		return;
+	}
+
+	parsePrecedence(PREC_ASSIGNMENT);
+	emitByte(OP_INVOKE_AWAIT);
+	emitByte(OP_NONE);
+	size_t loopContinue = currentChunk()->count;
+	size_t exitJump = emitJump(OP_YIELD_FROM);
+	emitByte(OP_YIELD);
+	emitLoop(loopContinue);
+	patchJump(exitJump);
+}
+
 static void unot_(int canAssign) {
 	parsePrecedence(PREC_NOT);
 	emitByte(OP_NOT);
@@ -2741,6 +2828,8 @@ ParseRule krk_parseRules[] = {
 	RULE(TOKEN_IMPORT,        NULL,     NULL,   PREC_NONE),
 	RULE(TOKEN_RAISE,         NULL,     NULL,   PREC_NONE),
 	RULE(TOKEN_YIELD,         yield,    NULL,   PREC_NONE),
+	RULE(TOKEN_AWAIT,         await,    NULL,   PREC_NONE),
+	RULE(TOKEN_ASYNC,         NULL,     NULL,   PREC_NONE),
 
 	RULE(TOKEN_AT,            NULL,     NULL,   PREC_NONE),
 

src/kuroko/chunk.h

@@ -70,7 +70,8 @@ typedef enum {
 	OP_END_FINALLY,
 	OP_GREATER_EQUAL,
 	OP_LESS_EQUAL,
-	/* current highest: 47 */
+	OP_INVOKE_AWAIT,
+	/* current highest: 48 */
 
 	OP_CALL = 64,
 	OP_CLASS,

src/kuroko/object.h

@@ -147,6 +147,7 @@ typedef struct {
 #define KRK_CODEOBJECT_FLAGS_COLLECTS_ARGS 0x0001
 #define KRK_CODEOBJECT_FLAGS_COLLECTS_KWS  0x0002
 #define KRK_CODEOBJECT_FLAGS_IS_GENERATOR  0x0004
+#define KRK_CODEOBJECT_FLAGS_IS_COROUTINE  0x0008
 
 /**
  * @brief Function object.
@@ -410,6 +411,11 @@ extern size_t krk_codepointToBytes(krk_integer_type value, unsigned char * out);
  */
 extern KrkInstance * krk_buildGenerator(KrkClosure * function, KrkValue * arguments, size_t argCount);
 
+/**
+ * @brief Calls __await__
+ */
+extern int krk_getAwaitable(void);
+
 /**
  * @brief Special value for type hint expressions.
  *

src/kuroko/scanner.h

@@ -94,6 +94,8 @@ typedef enum {
 	TOKEN_LAMBDA,
 	TOKEN_ASSERT,
 	TOKEN_YIELD,
+	TOKEN_ASYNC,
+	TOKEN_AWAIT,
 	TOKEN_WITH,
 
 	TOKEN_PREFIX_B,

src/obj_gen.c

@@ -27,6 +27,7 @@ struct generator {
 	int running;
 	int started;
 	KrkValue result;
+	int type;
 };
 
 #define AS_generator(o) ((struct generator *)AS_OBJECT(o))
@@ -75,15 +76,25 @@ KrkInstance * krk_buildGenerator(KrkClosure * closure, KrkValue * argsIn, size_t
 	self->closure = closure;
 	self->ip = self->closure->function->chunk.code;
 	self->result = NONE_VAL();
+	self->type = closure->function->flags & (KRK_CODEOBJECT_FLAGS_IS_GENERATOR | KRK_CODEOBJECT_FLAGS_IS_COROUTINE);
 	return (KrkInstance *)self;
 }
 
 KRK_METHOD(generator,__repr__,{
 	METHOD_TAKES_NONE();
 
-	size_t estimatedLength = sizeof("<generator object  at 0x1234567812345678>") + 1 + self->closure->function->name->length;
+	char * typeStr = "generator";
+	if (self->type == KRK_CODEOBJECT_FLAGS_IS_COROUTINE) {
+		/* Regular coroutine */
+		typeStr = "coroutine";
+	} else if (self->type == (KRK_CODEOBJECT_FLAGS_IS_COROUTINE | KRK_CODEOBJECT_FLAGS_IS_GENERATOR)) {
+		typeStr = "async_generator";
+	}
+
+	size_t estimatedLength = sizeof("< object at 0x1234567812345678>") + strlen(typeStr) + 1 + self->closure->function->name->length;
 	char * tmp = malloc(estimatedLength);
-	size_t lenActual = snprintf(tmp, estimatedLength, "<generator object %s at %p>",
+	size_t lenActual = snprintf(tmp, estimatedLength, "<%s object %s at %p>",
+		typeStr,
 		self->closure->function->name->chars,
 		(void*)self);
 
@@ -183,6 +194,32 @@ KRK_METHOD(generator,gi_running,{
 	return BOOLEAN_VAL(self->running);
 })
 
+int krk_getAwaitable(void) {
+	if (IS_generator(krk_peek(0)) && AS_generator(krk_peek(0))->type == KRK_CODEOBJECT_FLAGS_IS_COROUTINE) {
+		/* Good to go */
+		return 1;
+	}
+
+	/* Need to try for __await__ */
+	KrkValue method = krk_valueGetAttribute_default(krk_peek(0), "__await__", NONE_VAL());
+	if (!IS_NONE(method)) {
+		krk_push(method);
+		krk_swap(1);
+		krk_pop();
+		krk_push(krk_callSimple(krk_peek(0),0,0));
+		KrkClass * _type = krk_getType(krk_peek(0));
+		if (!_type || !_type->_iter) {
+			krk_runtimeError(vm.exceptions->attributeError, "__await__ returned non-iterator of type '%s'", krk_typeName(krk_peek(0)));
+			return 0;
+		}
+	} else {
+		krk_runtimeError(vm.exceptions->attributeError, "'%s' object is not awaitable", krk_typeName(krk_peek(0)));
+		return 0;
+	}
+
+	return 1;
+}
+
 _noexport
 void _createAndBind_generatorClass(void) {
 	generator = ADD_BASE_CLASS(vm.baseClasses->generatorClass, "generator", vm.baseClasses->objectClass);

src/opcodes.h

@@ -31,6 +31,7 @@ SIMPLE(OP_INVOKE_SETSLICE)
 SIMPLE(OP_INVOKE_DELSLICE)
 SIMPLE(OP_INVOKE_ITER)
 SIMPLE(OP_INVOKE_CONTAINS)
+SIMPLE(OP_INVOKE_AWAIT)
 SIMPLE(OP_SWAP)
 SIMPLE(OP_FINALIZE)
 SIMPLE(OP_IS)

src/scanner.c

@@ -197,8 +197,13 @@ static KrkTokenType identifierType() {
 	switch (*scanner.start) {
 		case 'a': if (MORE(1)) switch(scanner.start[1]) {
 			case 'n': return checkKeyword(2, "d", TOKEN_AND);
+			case 'w': return checkKeyword(2, "ait", TOKEN_AWAIT);
 			case 's': if (MORE(2)) {
-					return checkKeyword(2, "sert", TOKEN_ASSERT);
+					switch (scanner.start[2]) {
+						case 's': return checkKeyword(3, "ert", TOKEN_ASSERT);
+						case 'y': return checkKeyword(3, "nc", TOKEN_ASYNC);
+					}
+					break;
 				} else {
 					return checkKeyword(2, "", TOKEN_AS);
 				}

src/vendor/rline.c

@@ -539,7 +539,7 @@ char * syn_krk_keywords[] = {
 	"and","class","def","else","for","if","in","import","del",
 	"let","not","or","return","while","try","except","raise",
 	"continue","break","as","from","elif","lambda","with","is",
-	"pass","assert","yield","finally",
+	"pass","assert","yield","finally","async","await",
 	NULL
 };
 

src/vm.c

@@ -753,7 +753,7 @@ _finishKwarg:
 		krk_push(KWARGS_VAL(0));
 		argCount++;
 	}
-	if (unlikely(closure->function->flags & KRK_CODEOBJECT_FLAGS_IS_GENERATOR)) {
+	if (unlikely(closure->function->flags & (KRK_CODEOBJECT_FLAGS_IS_GENERATOR | KRK_CODEOBJECT_FLAGS_IS_COROUTINE))) {
 		KrkInstance * gen = krk_buildGenerator(closure, krk_currentThread.stackTop - argCount, argCount);
 		krk_currentThread.stackTop = krk_currentThread.stackTop - argCount - extra;
 		krk_push(OBJECT_VAL(gen));
@@ -2133,7 +2133,7 @@ _finishReturn: (void)0;
 				}
 				krk_currentThread.stackTop = &krk_currentThread.stack[frame->outSlots];
 				if (krk_currentThread.frameCount == (size_t)krk_currentThread.exitOnFrame) {
-					if (frame->closure->function->flags & KRK_CODEOBJECT_FLAGS_IS_GENERATOR) {
+					if (frame->closure->function->flags & (KRK_CODEOBJECT_FLAGS_IS_GENERATOR | KRK_CODEOBJECT_FLAGS_IS_COROUTINE)) {
 						krk_push(result);
 						return KWARGS_VAL(0);
 					}
@@ -2285,6 +2285,10 @@ _finishReturn: (void)0;
 				}
 				break;
 			}
+			case OP_INVOKE_AWAIT: {
+				if (!krk_getAwaitable()) goto _finishException;
+				break;
+			}
 			case OP_FINALIZE: {
 				KrkClass * _class = AS_CLASS(krk_peek(0));
 				/* Store special methods for quick access */

test/testAsyncAwait.krk

@@ -0,0 +1,66 @@
+
+if not hasattr(__builtins__,'StopIteration'):
+    class StopIteration(Exception):
+        pass
+    __builtins__.StopIteration = StopIteration
+
+class Awaiter:
+    def __await__(self):
+        print("    __await__ called")
+        yield "    Awaiter(): awaitable returns an iterator"
+
+def decorate(func):
+    print("Decorating", func.__qualname__)
+    return func
+
+class Baz:
+    def __init__(self, fromVal):
+        self.identifier = fromVal
+    async def asyncMethod(self):
+        print("   async method on", self.identifier)
+        return 'a value'
+    @decorate
+    async def decoratedAsyncMethod(self):
+        print("   decorated async method on", self.identifier)
+        return 3.141519
+
+async def foo(i):
+    print('    foo(): hi')
+    print('      Awaiting result 1:', await i())
+    print('      Awaiting result 2:', await Awaiter())
+    print('      Awaiting result 3:', await i())
+    print('      Awaiting result 4:', await Baz('<baz>').asyncMethod())
+    print('      Awaiting result 5:', await decoratedAsync())
+    print('      Awaiting result 6:', await Baz('<foo>').decoratedAsyncMethod())
+    print('    foo(): bye')
+    return "done"
+
+async def bar():
+    print("    bar(): hello, there, I'm an async function")
+    return 42
+
+@decorate
+async def decoratedAsync():
+    print("     I am a decorated async function")
+    return 'decorated result'
+
+def run(coro, scheduled=None, next=None, result=None):
+    # Okay, let's see.
+    scheduled = [coro]
+    print("Starting run loop.")
+    while scheduled:
+        print("  Popping from scheduled list.")
+        next = scheduled.pop(0) # Yes, that's slow, I know.
+        try:
+            print("  Calling",type(next))
+            result = next.send(None)
+            if result == next:
+                raise StopIteration(result.__finish__())
+            scheduled.append(next)
+            print("  Returned with",result)
+        except StopIteration as e:
+            # Stop iteration value should be return value from foo()
+            print('Exception:', type(e), e)
+    print('Done with run loop.')
+
+run(foo(bar))

test/testAsyncAwait.krk.expect

@@ -0,0 +1,24 @@
+Decorating Baz.decoratedAsyncMethod
+Decorating decoratedAsync
+Starting run loop.
+  Popping from scheduled list.
+  Calling <class 'generator'>
+    foo(): hi
+    bar(): hello, there, I'm an async function
+      Awaiting result 1: 42
+    __await__ called
+  Returned with     Awaiter(): awaitable returns an iterator
+  Popping from scheduled list.
+  Calling <class 'generator'>
+      Awaiting result 2: None
+    bar(): hello, there, I'm an async function
+      Awaiting result 3: 42
+   async method on <baz>
+      Awaiting result 4: a value
+     I am a decorated async function
+      Awaiting result 5: decorated result
+   decorated async method on <foo>
+      Awaiting result 6: 3.141519
+    foo(): bye
+Exception: <class '__main__.StopIteration'> done
+Done with run loop.