Fix object-literal generator method typing; extend generator .next() test coverage

tslang/lib/TypeScript/MLIRGenFunctions.cpp

@@ -528,11 +528,9 @@ namespace mlirgen
         return std::get<0>(res);
     }
 
-    std::tuple<mlir::LogicalResult, mlir_ts::FuncOp, std::string, bool> MLIRGenImpl::mlirGenFunctionGenerator(
-        FunctionLikeDeclarationBase functionLikeDeclarationBaseAST, const GenContext &genContext)
+    FunctionLikeDeclarationBase MLIRGenImpl::buildGeneratorWrapperDeclaration(
+        FunctionLikeDeclarationBase functionLikeDeclarationBaseAST, mlir::Location location)
     {
-        auto location = loc(functionLikeDeclarationBaseAST);
-
         auto fixThisReference = functionLikeDeclarationBaseAST == SyntaxKind::MethodDeclaration;
         if (functionLikeDeclarationBaseAST->parameters.size() > 0)
         {
@@ -653,15 +651,14 @@ namespace mlirgen
 
             // copy location info, to fix issue with names of anonymous functions
             methodOp->pos = functionLikeDeclarationBaseAST->pos;
-            methodOp->_end = functionLikeDeclarationBaseAST->_end;        
+            methodOp->_end = functionLikeDeclarationBaseAST->_end;
 
             // to ensure correct full name
             methodOp->parent = functionLikeDeclarationBaseAST->parent;
 
             LLVM_DEBUG(printDebug(methodOp););
 
-            auto genMethodOp = mlirGenFunctionLikeDeclaration(methodOp, genContext);
-            return genMethodOp;            
+            return methodOp;
         }
         else
         {
@@ -672,15 +669,22 @@ namespace mlirgen
 
             // copy location info, to fix issue with names of anonymous functions
             funcOp->pos = functionLikeDeclarationBaseAST->pos;
-            funcOp->_end = functionLikeDeclarationBaseAST->_end;        
+            funcOp->_end = functionLikeDeclarationBaseAST->_end;
 
             LLVM_DEBUG(printDebug(funcOp););
 
-            auto genFuncOp = mlirGenFunctionLikeDeclaration(funcOp, genContext);
-            return genFuncOp;
+            return funcOp;
         }
     }
 
+    std::tuple<mlir::LogicalResult, mlir_ts::FuncOp, std::string, bool> MLIRGenImpl::mlirGenFunctionGenerator(
+        FunctionLikeDeclarationBase functionLikeDeclarationBaseAST, const GenContext &genContext)
+    {
+        auto location = loc(functionLikeDeclarationBaseAST);
+        auto wrapperDecl = buildGeneratorWrapperDeclaration(functionLikeDeclarationBaseAST, location);
+        return mlirGenFunctionLikeDeclaration(wrapperDecl, genContext);
+    }
+
     bool MLIRGenImpl::registerFunctionOp(FunctionPrototypeDOM::TypePtr funcProto, mlir_ts::FuncOp funcOp)
     {
         auto name = funcProto->getNameWithoutNamespace();

tslang/lib/TypeScript/MLIRGenImpl.h

@@ -1806,6 +1806,13 @@ class MLIRGenImpl
     std::tuple<mlir::LogicalResult, mlir_ts::FuncOp, std::string, bool> mlirGenFunctionGenerator(
         FunctionLikeDeclarationBase functionLikeDeclarationBaseAST, const GenContext &genContext);
 
+    // Builds the synthetic non-generator declaration (method/function whose body just returns the
+    // generator wrapper object literal with its `next` method) that mlirGenFunctionGenerator generates
+    // from. Factored out so callers that only need the correctly-typed prototype (e.g. object-literal
+    // method prototype registration) can reuse it without running full body codegen twice.
+    FunctionLikeDeclarationBase buildGeneratorWrapperDeclaration(
+        FunctionLikeDeclarationBase functionLikeDeclarationBaseAST, mlir::Location location);
+
     std::pair<mlir::LogicalResult, std::string> registerGenericFunctionLike(
         FunctionLikeDeclarationBase functionLikeDeclarationBaseAST, bool ignoreFunctionArgsDetection,
         const GenContext &genContext)
@@ -7458,7 +7465,16 @@ class MLIRGenImpl
 
         funcLikeDecl->parent = oli.objectLiteral;
 
-        auto [funcOp, funcProto, result, isGeneric] = mlirGenFunctionPrototype(funcLikeDecl, funcGenContext);
+        // generator methods/properties must resolve to the generator wrapper type (the object with
+        // `.next()`, not the bare yielded-value tuple), same as mlirGenFunctionLikeDeclaration does for
+        // top-level function* and class generator methods; otherwise the field type registered here
+        // (used for the object literal's own type) is wrong and later access like `obj.gen().next()`
+        // fails to resolve.
+        auto protoDecl = funcLikeDecl->asteriskToken
+            ? buildGeneratorWrapperDeclaration(funcLikeDecl, loc(funcLikeDecl))
+            : funcLikeDecl;
+
+        auto [funcOp, funcProto, result, isGeneric] = mlirGenFunctionPrototype(protoDecl, funcGenContext);
         if (mlir::failed(result) || !funcOp)
         {
             return mlir::failure();

tslang/test/tester/CMakeLists.txt

@@ -124,6 +124,7 @@ add_test(NAME test-compile-00-equals COMMAND test-runner "${PROJECT_SOURCE_DIR}/
 add_test(NAME test-compile-00-bool-arith-ops COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00bool_arith_ops.ts")
 add_test(NAME test-compile-00-mixed-type-ops COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00mixed_type_ops.ts")
 add_test(NAME test-compile-00-generator-manual-next COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00generator_manual_next.ts")
+add_test(NAME test-compile-00-generator-manual-next-2 COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00generator_manual_next2.ts")
 add_test(NAME test-compile-00-funcs COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00funcs.ts")
 add_test(NAME test-compile-00-funcs-capture COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00funcs_capture.ts")
 add_test(NAME test-compile-00-funcs-vararg COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00funcs_vararg.ts")
@@ -294,6 +295,7 @@ add_test(NAME test-compile-00-generator-3 COMMAND test-runner "${PROJECT_SOURCE_
 add_test(NAME test-compile-00-generator-4 COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00generator4.ts")
 add_test(NAME test-compile-00-generator-5 COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00generator5.ts")
 add_test(NAME test-compile-00-generator-6 COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00generator6.ts")
+add_test(NAME test-compile-00-generator-7 COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00generator7.ts")
 add_test(NAME test-compile-00-safe-cast COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00safe_cast.ts")
 add_test(NAME test-compile-00-safe-cast-2 COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00safe_cast2.ts")
 add_test(NAME test-compile-00-safe-cast-typeof COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00safe_cast_typeof.ts")
@@ -469,6 +471,7 @@ add_test(NAME test-jit-00-equals COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}
 add_test(NAME test-jit-00-bool-arith-ops COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00bool_arith_ops.ts")
 add_test(NAME test-jit-00-mixed-type-ops COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00mixed_type_ops.ts")
 add_test(NAME test-jit-00-generator-manual-next COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00generator_manual_next.ts")
+add_test(NAME test-jit-00-generator-manual-next-2 COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00generator_manual_next2.ts")
 add_test(NAME test-jit-00-funcs COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00funcs.ts")
 add_test(NAME test-jit-00-funcs-capture COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00funcs_capture.ts")
 add_test(NAME test-jit-00-funcs-vararg COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00funcs_vararg.ts")
@@ -638,6 +641,7 @@ add_test(NAME test-jit-00-generator-3 COMMAND test-runner -jit "${PROJECT_SOURCE
 add_test(NAME test-jit-00-generator-4 COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00generator4.ts")
 add_test(NAME test-jit-00-generator-5 COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00generator5.ts")
 add_test(NAME test-jit-00-generator-6 COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00generator6.ts")
+add_test(NAME test-jit-00-generator-7 COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00generator7.ts")
 add_test(NAME test-jit-00-safe-cast COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00safe_cast.ts")
 add_test(NAME test-jit-00-safe-cast-2 COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00safe_cast2.ts")
 add_test(NAME test-jit-00-safe-cast-typeof COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00safe_cast_typeof.ts")

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

@@ -0,0 +1,93 @@
+// regression coverage: generator *methods* (as opposed to top-level function* declarations)
+// on object literals and classes. These share the same asteriskToken dispatch in
+// mlirGenFunctionLikeDeclaration but were never exercised by the existing 00generator*.ts
+// suite, which only uses top-level `function*`.
+
+const objGen = {
+    *gen() {
+        yield 1;
+        yield 2;
+        yield 3;
+    },
+};
+
+function main1() {
+    let count = 0;
+    let t = 1;
+    for (const o of objGen.gen()) {
+        assert(t++ == o);
+        count++;
+    }
+
+    assert(count == 3);
+}
+
+class Counter {
+    constructor(private limit: number) {}
+
+    *gen() {
+        for (let i = 0; i < this.limit; i++) {
+            yield i;
+        }
+    }
+}
+
+function main2() {
+    const c = new Counter(4);
+
+    let count = 0;
+    let t = 0;
+    for (const o of c.gen()) {
+        assert(t++ == o);
+        count++;
+    }
+
+    assert(count == 4);
+}
+
+// manual .next() driving on a class generator method, mirroring the top-level
+// function* manual-next regression (00generator_manual_next.ts).
+function main3() {
+    const c = new Counter(3);
+    const it = c.gen();
+
+    let r = it.next();
+    assert(!r.done);
+    assert(r.value == 0);
+
+    r = it.next();
+    assert(!r.done);
+    assert(r.value == 1);
+
+    r = it.next();
+    assert(!r.done);
+    assert(r.value == 2);
+
+    r = it.next();
+    assert(r.done);
+}
+
+// two independently-constructed instances must not share generator state.
+function main4() {
+    const a = new Counter(2);
+    const b = new Counter(2);
+
+    const ia = a.gen();
+    const ib = b.gen();
+
+    assert(ia.next().value == 0);
+    assert(ib.next().value == 0);
+    assert(ia.next().value == 1);
+    assert(ib.next().value == 1);
+    assert(ia.next().done);
+    assert(ib.next().done);
+}
+
+function main() {
+    main1();
+    main2();
+    main3();
+    main4();
+
+    print("done.");
+}

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

@@ -0,0 +1,122 @@
+// further coverage around manual .next() driving, extending 00generator_manual_next.ts:
+// interleaving two independent generator instances by hand, manually driving a generator
+// from inside a plain function (not for...of), and a generator body that closes over and
+// mutates outer state across suspensions.
+
+function* gen(start: number, count: number) {
+    for (let i = 0; i < count; i++) {
+        yield start + i;
+    }
+}
+
+// two independent instances, manually interleaved -- each must keep its own state.
+function main1() {
+    const a = gen(0, 3);
+    const b = gen(100, 3);
+
+    let ra = a.next();
+    let rb = b.next();
+    assert(ra.value == 0);
+    assert(rb.value == 100);
+
+    rb = b.next();
+    ra = a.next();
+    assert(rb.value == 101);
+    assert(ra.value == 1);
+
+    ra = a.next();
+    rb = b.next();
+    assert(ra.value == 2);
+    assert(rb.value == 102);
+
+    ra = a.next();
+    rb = b.next();
+    assert(ra.done);
+    assert(rb.done);
+}
+
+// manually drain a generator entirely from inside a plain (non-generator) helper
+// function that receives it as a parameter.
+//
+// NOTE: driving the SAME iterator further from the caller after it has been passed
+// into and mutated by a helper function is a known, separate bug (not covered here):
+// generator objects have value semantics and are copied across a function-parameter
+// boundary, so .next() calls made inside the callee do not advance the caller's
+// binding. That is unlike a same-function const local (see main1/00generator_manual_next.ts),
+// which works via an alloca-caching mechanism scoped to a single function body. Fixing
+// this would require pass-by-reference semantics for generator-typed parameters at the
+// ABI level in mlirGenFunctionParams -- out of scope for this regression file.
+function drainTwo(it: ReturnType<typeof gen>) {
+    const first = it.next();
+    const second = it.next();
+    return [first.value, second.value];
+}
+
+function main2() {
+    const it = gen(5, 4);
+
+    const [v0, v1] = drainTwo(it);
+    assert(v0 == 5);
+    assert(v1 == 6);
+}
+
+// generator closing over outer mutable state; manual .next() calls interleaved with
+// mutation of that outer state between resumptions.
+let outer = 0;
+
+function* counterFromOuter() {
+    while (outer < 5) {
+        yield outer;
+        outer++;
+    }
+}
+
+function main3() {
+    outer = 0;
+    const it = counterFromOuter();
+
+    let r = it.next();
+    assert(r.value == 0);
+
+    outer = 3; // mutate captured state between manual resumptions;
+               // resuming re-enters the loop body right after the yield, so
+               // outer++ (-> 4) runs before the while condition is rechecked
+
+    r = it.next();
+    assert(r.value == 4);
+
+    r = it.next();
+    assert(r.done);
+}
+
+// calling .next() again after the generator has already completed keeps returning done.
+function* short() {
+    yield 42;
+}
+
+function main4() {
+    const it = short();
+
+    let r = it.next();
+    assert(!r.done);
+    assert(r.value == 42);
+
+    r = it.next();
+    assert(r.done);
+
+    // extra calls past completion must stay done, not restart or crash
+    r = it.next();
+    assert(r.done);
+
+    r = it.next();
+    assert(r.done);
+}
+
+function main() {
+    main1();
+    main2();
+    main3();
+    main4();
+
+    print("done.");
+}