yield keyword; generator functions

src/chunk.h

@@ -67,7 +67,8 @@ typedef enum {
 	OP_INVOKE_ITER,
 	OP_INVOKE_CONTAINS,
 	OP_BREAKPOINT, /* NEVER output this instruction in the compiler or bad things can happen */
-	/* current highest: 43 */
+	OP_YIELD,
+	/* current highest: 44 */
 
 	OP_CALL = 64,
 	OP_CLASS,

src/compiler.c

@@ -1588,6 +1588,20 @@ 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->function->isGenerator = 1;
+	if (check(TOKEN_EOL) || check(TOKEN_EOF)) {
+		emitByte(OP_NONE);
+	} else {
+		expression();
+	}
+	emitByte(OP_YIELD);
+}
+
 static void tryStatement() {
 	size_t blockWidth = (parser.previous.type == TOKEN_INDENTATION) ? parser.previous.length : 0;
 	advance();
@@ -1780,6 +1794,8 @@ _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)) {
@@ -2425,6 +2441,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_AT,            NULL,     NULL,   PREC_NONE),
 

src/obj_gen.c

@@ -0,0 +1,138 @@
+#include <string.h>
+#include "vm.h"
+#include "value.h"
+#include "memory.h"
+#include "util.h"
+
+static KrkClass * generator;
+struct generator {
+	KrkInstance inst;
+	KrkClosure * closure;
+	KrkValue * args;
+	size_t argCount;
+	uint8_t * ip;
+	int running;
+};
+
+#define AS_generator(o) ((struct generator *)AS_OBJECT(o))
+#define IS_generator(o) (krk_isInstanceOf(o, generator))
+
+#define CURRENT_CTYPE struct generator *
+#define CURRENT_NAME  self
+
+static void _generator_gcscan(KrkInstance * _self) {
+	struct generator * self = (struct generator*)_self;
+	krk_markObject((KrkObj*)self->closure);
+	for (size_t i = 0; i < self->argCount; ++i) {
+		krk_markValue(self->args[i]);
+	}
+}
+
+static void _generator_gcsweep(KrkInstance * self) {
+	free(((struct generator*)self)->args);
+}
+
+static void _set_generator_done(struct generator * self) {
+	self->ip = NULL;
+}
+
+KrkInstance * krk_buildGenerator(KrkClosure * closure, KrkValue * argsIn, size_t argCount) {
+	/* Copy the args */
+	KrkValue * args = malloc(sizeof(KrkValue) * (argCount));
+	memcpy(args, argsIn, sizeof(KrkValue) * argCount);
+
+	/* Create a generator object */
+	struct generator * self = (struct generator *)krk_newInstance(generator);
+	self->args = args;
+	self->argCount = argCount;
+	self->closure = closure;
+	self->ip = self->closure->function->chunk.code;
+	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 * tmp = malloc(estimatedLength);
+	size_t lenActual = snprintf(tmp, estimatedLength, "<generator object %s at %p>",
+		self->closure->function->name->chars,
+		(void*)self);
+
+	return OBJECT_VAL(krk_takeString(tmp,lenActual));
+})
+
+KRK_METHOD(generator,__iter__,{
+	METHOD_TAKES_NONE();
+	return OBJECT_VAL(self);
+})
+
+KRK_METHOD(generator,__call__,{
+	METHOD_TAKES_NONE();
+	if (!self->ip) return OBJECT_VAL(self);
+	/* Prepare frame */
+	KrkCallFrame * frame = &krk_currentThread.frames[krk_currentThread.frameCount++];
+	frame->closure = self->closure;
+	frame->ip      = self->ip;
+	frame->slots   = krk_currentThread.stackTop - krk_currentThread.stack;
+	frame->outSlots = frame->slots;
+	frame->globals = &self->closure->function->globalsContext->fields;
+
+	/* Stick our stack on their stack */
+	for (size_t i = 0; i < self->argCount; ++i) {
+		krk_push(self->args[i]);
+	}
+
+	/* Jump into the iterator */
+	self->running = 1;
+	size_t stackBefore = krk_currentThread.stackTop - krk_currentThread.stack;
+	KrkValue result = krk_runNext();
+	size_t stackAfter = krk_currentThread.stackTop - krk_currentThread.stack;
+	self->running = 0;
+
+	/* Was there an exception? */
+	if (krk_currentThread.flags & KRK_THREAD_HAS_EXCEPTION) {
+		_set_generator_done(self);
+		return NONE_VAL();
+	}
+
+	/* Determine the stack state */
+	if (stackAfter > stackBefore) {
+		size_t newArgs = stackAfter - stackBefore;
+		self->args = realloc(self->args, sizeof(KrkValue) * (self->argCount + newArgs));
+		self->argCount += newArgs;
+	} else if (stackAfter <= stackBefore) {
+		_set_generator_done(self);
+		return OBJECT_VAL(self);
+	}
+
+	/* Save stack entries */
+	memcpy(self->args, krk_currentThread.stackTop - self->argCount, sizeof(KrkValue) * self->argCount);
+	self->ip      = frame->ip;
+
+	krk_currentThread.stackTop = krk_currentThread.stack + frame->slots;
+
+	return result;
+})
+
+/*
+ * For compatibility with Python...
+ */
+KRK_METHOD(generator,gi_running,{
+	METHOD_TAKES_NONE();
+	return BOOLEAN_VAL(self->running);
+})
+
+_noexport
+void _createAndBind_generatorClass(void) {
+	generator = ADD_BASE_CLASS(vm.baseClasses->generatorClass, "generator", vm.baseClasses->objectClass);
+	generator->allocSize = sizeof(KrkList);
+	generator->_ongcscan = _generator_gcscan;
+	generator->_ongcsweep = _generator_gcsweep;
+	BIND_METHOD(generator,__iter__);
+	BIND_METHOD(generator,__call__);
+	BIND_METHOD(generator,__repr__);
+	BIND_PROP(generator,gi_running);
+	krk_defineNative(&generator->methods, "__str__", FUNC_NAME(generator,__repr__));
+	krk_finalizeClass(generator);
+}

src/object.h

@@ -125,6 +125,7 @@ typedef struct {
 	unsigned char collectsArguments:1;
 	unsigned char collectsKeywords:1;
 	unsigned char isClassMethod:1;
+	unsigned char isGenerator:1;
 	struct KrkInstance * globalsContext;
 } KrkFunction;
 

src/opcodes.h

@@ -41,6 +41,7 @@ SIMPLE(OP_CLEANUP_WITH)
 SIMPLE(OP_FILTER_EXCEPT)
 SIMPLE(OP_CREATE_CLASSMETHOD)
 SIMPLE(OP_BREAKPOINT)
+SIMPLE(OP_YIELD)
 CONSTANT(OP_DEFINE_GLOBAL,(void)0)
 CONSTANT(OP_CONSTANT,(void)0)
 CONSTANT(OP_GET_GLOBAL,(void)0)

src/rline.c

@@ -544,7 +544,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",
+	"pass","assert","yield",
 	NULL
 };
 

src/scanner.c

@@ -257,6 +257,7 @@ static KrkTokenType identifierType() {
 			case 'h': return checkKeyword(2, "ile", TOKEN_WHILE);
 			case 'i': return checkKeyword(2, "th", TOKEN_WITH);
 		} break;
+		case 'y': return checkKeyword(1, "ield", TOKEN_YIELD);
 	}
 	return TOKEN_IDENTIFIER;
 }

src/scanner.h

@@ -91,6 +91,7 @@ typedef enum {
 	TOKEN_FROM,
 	TOKEN_LAMBDA,
 	TOKEN_ASSERT,
+	TOKEN_YIELD,
 	TOKEN_WITH,
 
 	TOKEN_PREFIX_B,

src/vm.c

@@ -1,3 +1,4 @@
+#include <assert.h>
 #include <limits.h>
 #include <stdarg.h>
 #include <string.h>
@@ -744,6 +745,12 @@ _finishKwarg:
 		krk_push(KWARGS_VAL(0));
 		argCount++;
 	}
+	if (closure->function->isGenerator) {
+		KrkInstance * gen = krk_buildGenerator(closure, krk_currentThread.stackTop - argCount, argCount);
+		krk_currentThread.stackTop = krk_currentThread.stackTop - argCount - extra;
+		krk_push(OBJECT_VAL(gen));
+		return 1;
+	}
 	if (krk_currentThread.frameCount == KRK_CALL_FRAMES_MAX) {
 		krk_runtimeError(vm.exceptions->baseException, "Too many call frames.");
 		return 0;
@@ -1178,6 +1185,7 @@ void krk_initVM(int flags) {
 	_createAndBind_timeMod();
 	_createAndBind_osMod();
 	_createAndBind_fileioMod();
+	_createAndBind_generatorClass();
 #ifdef ENABLE_THREADING
 	_createAndBind_threadsMod();
 #endif
@@ -2084,6 +2092,13 @@ _resumeHook: (void)0;
 				if (krk_currentThread.flags & KRK_THREAD_HAS_EXCEPTION) goto _finishException;
 				goto _resumeHook;
 			}
+			case OP_YIELD: {
+				KrkValue result = krk_peek(0);
+				krk_currentThread.frameCount--;
+				assert(krk_currentThread.frameCount == (size_t)krk_currentThread.exitOnFrame);
+				/* Do NOT restore the stack */
+				return result;
+			}
 
 			/*
 			 * Two-byte operands

src/vm.h

@@ -157,6 +157,7 @@ struct BaseClasses {
 	KrkClass * bytesiteratorClass;   /**< Iterator over the integer byte values of a bytes object. */
 	KrkClass * propertyClass;        /**< Magic object that calls a function when accessed from an instance through the dot operator. */
 	KrkClass * codeobjectClass;      /**< Static compiled bytecode container. */
+	KrkClass * generatorClass;       /**< Generator object. */
 };
 
 /**
@@ -773,4 +774,5 @@ extern void _createAndBind_threadsMod(void);
 extern KrkValue krk_operator_lt(KrkValue,KrkValue);
 extern KrkValue krk_operator_gt(KrkValue,KrkValue);
 
-
+extern void _createAndBind_generatorClass(void);
+extern KrkInstance * krk_buildGenerator(KrkClosure *, KrkValue *, size_t);

test/testGeneratorFunctions.krk

@@ -0,0 +1,45 @@
+def infinite_sequence(num=0):
+    print("Called")
+    yield
+    while num < 20:
+        print("Yielding")
+        yield num
+        if num == 15:
+            raise ValueError()
+        num += 1
+    print("All out of yields?")
+
+print('function infinite_sequence' in str(infinite_sequence))
+print('generator object infinite_sequence' in str(infinite_sequence()))
+
+def foo(o=None):
+    try:
+        o = infinite_sequence()
+        for i in o:
+            print(i)
+    except:
+        pass
+    print('===')
+    print(o.gi_running)
+    print('---')
+    # It's empty now.
+    for i in o:
+        print(i)
+
+foo()
+
+# We should also be able to make generator methods
+
+class Foo():
+    def __init__(self):
+        self.bar = "baz"
+    def generatorMethod(self):
+        for c in self.bar:
+            print('yielding',c)
+            yield c
+
+let f = Foo()
+for entry in f.generatorMethod():
+    print('got',entry)
+for entry in f.generatorMethod():
+    print('got',entry)

test/testGeneratorFunctions.krk.expect

@@ -0,0 +1,51 @@
+True
+True
+Called
+None
+Yielding
+0
+Yielding
+1
+Yielding
+2
+Yielding
+3
+Yielding
+4
+Yielding
+5
+Yielding
+6
+Yielding
+7
+Yielding
+8
+Yielding
+9
+Yielding
+10
+Yielding
+11
+Yielding
+12
+Yielding
+13
+Yielding
+14
+Yielding
+15
+===
+False
+---
+yielding b
+got b
+yielding a
+got a
+yielding z
+got z
+yielding b
+got b
+yielding a
+got a
+yielding z
+got z