'yield' as an expression with a value; 'yield from'

src/builtins.c

@@ -846,6 +846,13 @@ KRK_FUNC(hash,{
 	return INTEGER_VAL(krk_hashValue(argv[0]));
 })
 
+KRK_FUNC(next,{
+	FUNCTION_TAKES_EXACTLY(1);
+	krk_push(argv[0]);
+	krk_push(krk_callSimple(argv[0], 0, 0));
+	return krk_pop();
+})
+
 _noexport
 void _createAndBind_builtins(void) {
 	vm.baseClasses->objectClass = krk_newClass(S("object"), NULL);
@@ -1026,5 +1033,6 @@ void _createAndBind_builtins(void) {
 	BUILTIN_FUNCTION("enumerate", FUNC_NAME(krk,enumerate), "Return an iterator that produces a tuple with a count the iterated values of the passed iteratable.");
 	BUILTIN_FUNCTION("bin", FUNC_NAME(krk,bin), "Convert an integer value to a binary string.");
 	BUILTIN_FUNCTION("zip", FUNC_NAME(krk,zip), "Returns an iterator that produces tuples of the nth element of each passed iterable.");
+	BUILTIN_FUNCTION("next", FUNC_NAME(krk,next), "Compatibility function. Calls an iterable.");
 }
 

src/chunk.h

@@ -109,6 +109,7 @@ typedef enum {
 	OP_LOOP,
 	OP_PUSH_TRY,
 	OP_PUSH_WITH,
+	OP_YIELD_FROM,
 
 	OP_CALL_LONG = 192,
 	OP_CLASS_LONG,

src/compiler.c

@@ -1631,9 +1631,7 @@ static void forStatement() {
 		expression();
 		endScope();
 
-		KrkToken _it = syntheticToken("");
-		size_t indLoopIter = current->localCount;
-		addLocal(_it);
+		size_t indLoopIter = addLocal(syntheticToken(""));
 		defineVariable(indLoopIter);
 
 		emitByte(OP_INVOKE_ITER);
@@ -1721,20 +1719,6 @@ static void returnStatement() {
 	}
 }
 
-static void yieldStatement() {
-	if (current->type == TYPE_MODULE || current->type == TYPE_INIT || current->type == TYPE_CLASS) {
-		error("'yield' outside function");
-		return;
-	}
-	current->codeobject->isGenerator = 1;
-	if (check(TOKEN_EOL) || check(TOKEN_EOF)) {
-		emitByte(OP_NONE);
-	} else {
-		parsePrecedence(PREC_ASSIGNMENT);
-	}
-	emitBytes(OP_YIELD, OP_POP);
-}
-
 static void tryStatement() {
 	size_t blockWidth = (parser.previous.type == TOKEN_INDENTATION) ? parser.previous.length : 0;
 	advance();
@@ -1994,8 +1978,6 @@ _anotherSimpleStatement:
 			raiseStatement();
 		} else if (match(TOKEN_RETURN)) {
 			returnStatement();
-		} else if (match(TOKEN_YIELD)) {
-			yieldStatement();
 		} else if (match(TOKEN_IMPORT)) {
 			importStatement();
 		} else if (match(TOKEN_FROM)) {
@@ -2020,6 +2002,32 @@ _anotherSimpleStatement:
 	}
 }
 
+static void yield(int canAssign) {
+	if (current->type == TYPE_MODULE ||
+		current->type == TYPE_INIT ||
+		current->type == TYPE_CLASS) {
+		error("'yield' outside function");
+		return;
+	}
+	current->codeobject->isGenerator = 1;
+	if (match(TOKEN_FROM)) {
+		parsePrecedence(PREC_ASSIGNMENT);
+		emitByte(OP_INVOKE_ITER);
+		emitByte(OP_NONE);
+		size_t loopContinue = currentChunk()->count;
+		size_t exitJump = emitJump(OP_YIELD_FROM);
+		emitByte(OP_YIELD);
+		emitLoop(loopContinue);
+		patchJump(exitJump);
+	} else if (check(TOKEN_EOL) || check(TOKEN_EOF) || check(TOKEN_RIGHT_PAREN) || check(TOKEN_RIGHT_BRACE)) {
+		emitByte(OP_NONE);
+		emitByte(OP_YIELD);
+	} else {
+		parsePrecedence(PREC_ASSIGNMENT);
+		emitByte(OP_YIELD);
+	}
+}
+
 static void unot_(int canAssign) {
 	parsePrecedence(PREC_NOT);
 	emitByte(OP_NOT);
@@ -2711,7 +2719,7 @@ ParseRule krk_parseRules[] = {
 	RULE(TOKEN_CONTINUE,      NULL,     NULL,   PREC_NONE),
 	RULE(TOKEN_IMPORT,        NULL,     NULL,   PREC_NONE),
 	RULE(TOKEN_RAISE,         NULL,     NULL,   PREC_NONE),
-	RULE(TOKEN_YIELD,         NULL,     NULL,   PREC_NONE),
+	RULE(TOKEN_YIELD,         yield,    NULL,   PREC_NONE),
 
 	RULE(TOKEN_AT,            NULL,     NULL,   PREC_NONE),
 

src/obj_gen.c

@@ -11,6 +11,7 @@
 #include "value.h"
 #include "memory.h"
 #include "util.h"
+#include "debug.h"
 
 static KrkClass * generator;
 /**
@@ -24,6 +25,8 @@ struct generator {
 	size_t argCount;
 	uint8_t * ip;
 	int running;
+	int started;
+	KrkValue result;
 };
 
 #define AS_generator(o) ((struct generator *)AS_OBJECT(o))
@@ -38,6 +41,7 @@ static void _generator_gcscan(KrkInstance * _self) {
 	for (size_t i = 0; i < self->argCount; ++i) {
 		krk_markValue(self->args[i]);
 	}
+	krk_markValue(self->result);
 }
 
 static void _generator_gcsweep(KrkInstance * self) {
@@ -70,6 +74,7 @@ KrkInstance * krk_buildGenerator(KrkClosure * closure, KrkValue * argsIn, size_t
 	self->argCount = argCount;
 	self->closure = closure;
 	self->ip = self->closure->function->chunk.code;
+	self->result = NONE_VAL();
 	return (KrkInstance *)self;
 }
 
@@ -91,7 +96,7 @@ KRK_METHOD(generator,__iter__,{
 })
 
 KRK_METHOD(generator,__call__,{
-	METHOD_TAKES_NONE();
+	METHOD_TAKES_AT_MOST(1);
 	if (!self->ip) return OBJECT_VAL(self);
 	/* Prepare frame */
 	KrkCallFrame * frame = &krk_currentThread.frames[krk_currentThread.frameCount++];
@@ -106,6 +111,15 @@ KRK_METHOD(generator,__call__,{
 		krk_push(self->args[i]);
 	}
 
+	if (self->started) {
+		krk_pop();
+		if (argc > 1) {
+			krk_push(argv[1]);
+		} else {
+			krk_push(NONE_VAL());
+		}
+	}
+
 	/* Jump into the iterator */
 	self->running = 1;
 	size_t stackBefore = krk_currentThread.stackTop - krk_currentThread.stack;
@@ -113,7 +127,10 @@ KRK_METHOD(generator,__call__,{
 	size_t stackAfter = krk_currentThread.stackTop - krk_currentThread.stack;
 	self->running = 0;
 
+	self->started = 1;
+
 	if (IS_KWARGS(result) && AS_INTEGER(result) == 0) {
+		self->result = krk_pop();
 		_set_generator_done(self);
 		return OBJECT_VAL(self);
 	}
@@ -144,6 +161,19 @@ KRK_METHOD(generator,__call__,{
 	return result;
 })
 
+KRK_METHOD(generator,send,{
+	METHOD_TAKES_EXACTLY(1);
+	if (!self->started && !IS_NONE(argv[1])) {
+		return krk_runtimeError(vm.exceptions->typeError, "Can not send non-None value to just-started generator");
+	}
+	return FUNC_NAME(generator,__call__)(argc,argv,0);
+})
+
+KRK_METHOD(generator,__finish__,{
+	METHOD_TAKES_NONE();
+	return self->result;
+})
+
 /*
  * For compatibility with Python...
  */
@@ -161,6 +191,8 @@ void _createAndBind_generatorClass(void) {
 	BIND_METHOD(generator,__iter__);
 	BIND_METHOD(generator,__call__);
 	BIND_METHOD(generator,__repr__);
+	BIND_METHOD(generator,__finish__);
+	BIND_METHOD(generator,send);
 	BIND_PROP(generator,gi_running);
 	krk_defineNative(&generator->methods, "__str__", FUNC_NAME(generator,__repr__));
 	krk_finalizeClass(generator);

src/opcodes.h

@@ -80,3 +80,4 @@ JUMP(OP_JUMP_IF_TRUE,+)
 JUMP(OP_LOOP,-)
 JUMP(OP_PUSH_TRY,+)
 JUMP(OP_PUSH_WITH,+)
+JUMP(OP_YIELD_FROM,+)

src/vm.c

@@ -2041,6 +2041,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->isGenerator) {
+						krk_push(result);
 						return KWARGS_VAL(0);
 					}
 					return result;
@@ -2341,6 +2342,34 @@ _finishReturn: (void)0;
 				krk_push(handler);
 				break;
 			}
+			case OP_YIELD_FROM: {
+				uint8_t * exitIp = frame->ip + OPERAND;
+				/* Stack has [iterator] [sent value] */
+				/* Is this a generator or something with a 'send' method? */
+				KrkValue method = krk_valueGetAttribute_default(krk_peek(1), "send", NONE_VAL());
+				if (!IS_NONE(method)) {
+					krk_push(method);
+					krk_swap(1);
+					krk_push(krk_callSimple(krk_peek(1),1,0));
+				} else {
+					krk_pop();
+					krk_push(krk_peek(0));
+					krk_push(krk_callSimple(krk_peek(0),0,0));
+				}
+				if (!krk_valuesSame(krk_peek(0), krk_peek(1))) break;
+
+				/* Does it have a final value? */
+				method = krk_valueGetAttribute_default(krk_peek(0), "__finish__", NONE_VAL());
+				if (!IS_NONE(method)) {
+					krk_push(method);
+					krk_push(krk_callSimple(krk_peek(0),0,0));
+				} else {
+					krk_pop();
+					krk_push(NONE_VAL());
+				}
+				frame->ip = exitIp;
+				break;
+			}
 
 			case OP_CONSTANT_LONG:
 			case OP_CONSTANT: {

test/testYieldFrom.krk

@@ -0,0 +1,61 @@
+def gen(x=0):
+    x = yield from [1,2,3]
+    print("checkpoint 1", x)
+    x = yield 4
+    print("checkpoint 2", x)
+    yield from [5,6,7]
+
+def main():
+    for i in gen():
+        print(i)
+
+main()
+
+def foo(x=-1):
+    print("Entering")
+    x = yield 42
+    print("Got",x)
+    x = yield 1024
+    print("got",x)
+    return 1234
+
+def main(f=None):
+    f = foo()
+    print('yield 1 got',f.send(None))
+    print('yield 2 got',f.send('hello'))
+    try:
+        f.send('bye') # End of iterator
+    except:
+        pass # On Python, this throws stop-iteration with 1234
+
+main()
+
+def accumulate(tally=0,next=None):
+    print("Entering accumulate")
+    while True:
+        print("Yielding None")
+        next = yield
+        print("Got",next)
+        if next is None:
+            print("Returning tally",tally)
+            return tally
+        tally += next
+
+def gather_tallies(tallies,tally=None):
+    while True:
+        tally = yield from accumulate()
+        print("Appending",tally)
+        tallies.append(tally)
+
+def main(tallies=None,acc=None):
+    tallies = []
+    acc = gather_tallies(tallies)
+    print(next(acc))
+    for i in range(4):
+        print(acc.send(i))
+    acc.send(None)
+    for i in range(5):
+        acc.send(i)
+    acc.send(None)
+    print(tallies)
+main()

test/testYieldFrom.krk.expect

@@ -0,0 +1,50 @@
+1
+2
+3
+checkpoint 1 None
+4
+checkpoint 2 None
+5
+6
+7
+Entering
+yield 1 got 42
+Got hello
+yield 2 got 1024
+got bye
+Entering accumulate
+Yielding None
+None
+Got 0
+Yielding None
+None
+Got 1
+Yielding None
+None
+Got 2
+Yielding None
+None
+Got 3
+Yielding None
+None
+Got None
+Returning tally 6
+Appending 6
+Entering accumulate
+Yielding None
+Got 0
+Yielding None
+Got 1
+Yielding None
+Got 2
+Yielding None
+Got 3
+Yielding None
+Got 4
+Yielding None
+Got None
+Returning tally 10
+Appending 10
+Entering accumulate
+Yielding None
+[6, 10]