return in a generator finishes it

tslang/lib/TypeScript/MLIRGenFunctions.cpp

@@ -670,12 +670,32 @@ namespace mlirgen
         // add main switcher
         auto stepAccess = nf.createPropertyAccessExpression(nf.createToken(SyntaxKind::ThisKeyword), stepIdent);
 
+        // a finished generator - `return` sets step to -1 - gives {value: undefined, done: true} from then on
+        auto finishedCheck = nf.createIfStatement(
+            nf.createBinaryExpression(stepAccess, nf.createToken(SyntaxKind::LessThanToken), nf.createNumericLiteral(S("0"), TokenFlags::None)),
+            nf.createReturnStatement(getYieldReturnObject(nf, location, nf.createIdentifier(S(UNDEFINED_NAME)), true)),
+            undefined);
+        nextStatements.push_back(finishedCheck);
+
         // call stateswitch
         auto callStat = nf.createExpressionStatement(
             nf.createCallExpression(nf.createIdentifier(S(GENERATOR_SWITCHSTATE)), undefined, {stepAccess}));
 
         nextStatements.push_back(callStat);
 
+        // the statements of next() that use the generator object's own `this`, not the body's
+        auto ownStatementsCount = nextStatements.size();
+
+        // a `return` of the body finishes the generator; one of a nested function is that function's own
+        FilterVisitorSkipFuncsAST<ReturnStatement> returnsVisitor(SyntaxKind::ReturnStatement, [&](auto returnStatement) {
+            returnStatement->internalFlags |= InternalFlags::GeneratorReturn;
+        });
+
+        if (functionLikeDeclarationBaseAST->body)
+        {
+            returnsVisitor.visit(functionLikeDeclarationBaseAST->body);
+        }
+
         // add function body to statements to first step
         if (functionLikeDeclarationBaseAST->body == SyntaxKind::Block)
         {
@@ -715,7 +735,7 @@ namespace mlirgen
                 thisNode->internalFlags |= InternalFlags::ThisArgAlias;
             });
 
-            for (auto it = begin(nextStatements) + 1; it != end(nextStatements); ++it)
+            for (auto it = begin(nextStatements) + ownStatementsCount; it != end(nextStatements); ++it)
             {
                 visitor.visit(*it);
             }

tslang/lib/TypeScript/MLIRGenStatements.cpp

@@ -425,6 +425,32 @@ namespace mlirgen
     mlir::LogicalResult MLIRGenImpl::mlirGen(ReturnStatement returnStatementAST, const GenContext &genContext)
     {
         auto location = loc(returnStatementAST);
+
+        // `return e` in a generator: {value: e, done: true}, and every next() after it is done (step -1)
+        if (!!(returnStatementAST->internalFlags & InternalFlags::GeneratorReturn))
+        {
+            NodeFactory nf(NodeFactoryFlags::None);
+
+            // as yield sets its resume point: `this.step` of the generator object, or `step` of a captured one
+            auto stepTarget = evaluateProperty(nf.createToken(SyntaxKind::ThisKeyword), GENERATOR_STEP, genContext)
+                ? static_cast<Expression>(nf.createPropertyAccessExpression(nf.createToken(SyntaxKind::ThisKeyword), nf.createIdentifier(S(GENERATOR_STEP))))
+                : static_cast<Expression>(nf.createIdentifier(S(GENERATOR_STEP)));
+            auto finishExpr = nf.createBinaryExpression(
+                stepTarget, nf.createToken(SyntaxKind::EqualsToken),
+                nf.createPrefixUnaryExpression(SyntaxKind::MinusToken, nf.createNumericLiteral(S("1"), TokenFlags::None)));
+            auto finishResult = mlirGen(finishExpr, genContext);
+            EXIT_IF_FAILED(finishResult)
+
+            Expression value = returnStatementAST->expression;
+            if (!value)
+            {
+                value = nf.createIdentifier(S(UNDEFINED_NAME));
+            }
+            auto finishedReturn = nf.createReturnStatement(getYieldReturnObject(nf, location, value, true));
+            finishedReturn->pos = returnStatementAST->pos;
+            finishedReturn->_end = returnStatementAST->_end;
+            return mlirGen(finishedReturn, genContext);
+        }
         if (auto expression = returnStatementAST->expression)
         {
             GenContext receiverTypeGenContext(genContext);

tslang/test/tester/CMakeLists.txt

@@ -394,6 +394,7 @@ tslang_add_test(NAME test-compile-00-strict-equals-mismatched-types COMMAND test
 tslang_add_test(NAME test-compile-00-length-with-negative-merge COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00length_with_negative_merge.ts")
 tslang_add_test(NAME test-compile-00-length-signed-arithmetic COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00length_signed_arithmetic.ts")
 tslang_add_test(NAME test-compile-00-optional-fields COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00optional_fields.ts")
+tslang_add_test(NAME test-compile-00-generator-return COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00generator_return.ts")
 tslang_add_test(NAME test-compile-00-intersection-type COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00intersection_type.ts")
 tslang_add_test(NAME test-compile-00-intersection-type-generic COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00intersection_type_generic.ts")
 tslang_add_test(NAME test-compile-00-length COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00length.ts")
@@ -871,6 +872,7 @@ tslang_add_test(NAME test-jit-00-strict-equals-mismatched-types COMMAND test-run
 tslang_add_test(NAME test-jit-00-length-with-negative-merge COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00length_with_negative_merge.ts")
 tslang_add_test(NAME test-jit-00-length-signed-arithmetic COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00length_signed_arithmetic.ts")
 tslang_add_test(NAME test-jit-00-optional-fields COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00optional_fields.ts")
+tslang_add_test(NAME test-jit-00-generator-return COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00generator_return.ts")
 tslang_add_test(NAME test-jit-00-union-global-init COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00union_global_init.ts")
 tslang_add_test(NAME test-jit-00-intersection-type COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00intersection_type.ts")
 tslang_add_test(NAME test-jit-00-intersection-type-generic COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00intersection_type_generic.ts")

tslang/test/tester/tests/00generator_return.ts

@@ -0,0 +1,77 @@
+// `return` in a generator finishes it: next() gives {value: <returned>, done: true}, every later next() is
+// done, and for..of stops - a `return` of a function nested in the body is that function's own
+function* withValue() {
+    yield "a";
+    return "end";
+    yield "never";
+}
+
+function* early(stop: boolean) {
+    yield 1;
+    if (stop) return;
+    yield 2;
+}
+
+function* inLoop(limit: number) {
+    for (let i = 0; ; i++) {
+        if (i >= limit) return;
+        yield i;
+    }
+}
+
+function* nestedFunctionReturn() {
+    const twice = (x: number) => { return x * 2; };
+    yield twice(1);
+    yield twice(2);
+}
+
+class Counter {
+    count = 2;
+    *values() {
+        yield this.count;
+        if (this.count > 1) return;
+        yield 0;
+    }
+}
+
+function main() {
+    const it = withValue();
+    let r = it.next();
+    assert(r.value == "a" && !r.done, "first yield");
+    r = it.next();
+    assert(r.value == "end" && r.done, "return value, done");
+    r = it.next();
+    assert(r.done, "done after return");
+    r = it.next();
+    assert(r.done, "still done");
+
+    let seen = 0;
+    for (const v of withValue()) {
+        seen++;
+        assert(v == "a", "for..of skips the returned value");
+    }
+    assert(seen == 1, "for..of stops at return");
+
+    const e = early(true);
+    assert(!e.next().done, "early: yield");
+    assert(e.next().done, "early: return");
+    assert(e.next().done, "early: done after return");
+
+    let sum = 0;
+    for (const v of early(false)) sum += v;
+    assert(sum == 3, "no return taken");
+
+    let count = 0;
+    for (const v of inLoop(3)) count++;
+    assert(count == 3, "return from a loop");
+
+    let nested = 0;
+    for (const v of nestedFunctionReturn()) nested += v;
+    assert(nested == 6, "nested function return");
+
+    let fromMethod = 0;
+    for (const v of new Counter().values()) fromMethod++;
+    assert(fromMethod == 1, "method generator return");
+
+    print("done.");
+}

tslang/ts-new-parser/enums.h

@@ -560,7 +560,9 @@ enum class InternalFlags : number
     // object literal must be heap-boxed into a reference-typed ObjectType value
     // instead of the default value-typed tuple (used for generator wrappers,
     // whose mutable `step` state must be shared across aliases)
-    BoxAsObject = 1 << 11
+    BoxAsObject = 1 << 11,
+    // a `return` of a generator's own body: it finishes the generator (mlirGen(ReturnStatement))
+    GeneratorReturn = 1 << 12
 };
 
 ENUM_OPS(InternalFlags)