diff --git a/tslang/lib/TypeScript/MLIRGenFunctions.cpp b/tslang/lib/TypeScript/MLIRGenFunctions.cpp index 224b60f4e..bf993f0ba 100644 --- a/tslang/lib/TypeScript/MLIRGenFunctions.cpp +++ b/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 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); } diff --git a/tslang/lib/TypeScript/MLIRGenStatements.cpp b/tslang/lib/TypeScript/MLIRGenStatements.cpp index 12c59d678..bae026750 100644 --- a/tslang/lib/TypeScript/MLIRGenStatements.cpp +++ b/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(nf.createPropertyAccessExpression(nf.createToken(SyntaxKind::ThisKeyword), nf.createIdentifier(S(GENERATOR_STEP)))) + : static_cast(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); diff --git a/tslang/test/tester/CMakeLists.txt b/tslang/test/tester/CMakeLists.txt index 0a625dbd1..87bf71c67 100644 --- a/tslang/test/tester/CMakeLists.txt +++ b/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") diff --git a/tslang/test/tester/tests/00generator_return.ts b/tslang/test/tester/tests/00generator_return.ts new file mode 100644 index 000000000..3078792f7 --- /dev/null +++ b/tslang/test/tester/tests/00generator_return.ts @@ -0,0 +1,77 @@ +// `return` in a generator finishes it: next() gives {value: , 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."); +} diff --git a/tslang/ts-new-parser/enums.h b/tslang/ts-new-parser/enums.h index 5f01aff0e..65fb6a741 100644 --- a/tslang/ts-new-parser/enums.h +++ b/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)