Fix generators losing state across manual .next() calls

tslang/include/TypeScript/MLIRLogic/MLIRCodeLogic.h

@@ -1150,8 +1150,18 @@ class MLIRPropertyAccessCodeLogic
     mlir::Attribute fieldId;
     mlir::Value argument;
     CompileOptions& compileOptions;
+    llvm::ScopedHashTable<mlir::Value, mlir::Value> *boundRefMaterializedCache = nullptr;
 
   public:
+    // optional: lets a bound-method property access (e.g. `g.next` on a storage-less
+    // `const` binding) reuse the ref it materialized on a previous access instead of
+    // minting a fresh one seeded from the pristine, never-mutated value each time. See
+    // TupleNoError() and MLIRGenImpl::boundRefMaterializedCache.
+    void setBoundRefMaterializedCache(llvm::ScopedHashTable<mlir::Value, mlir::Value> *cache)
+    {
+        boundRefMaterializedCache = cache;
+    }
+
     MLIRPropertyAccessCodeLogic(CompileOptions& compileOptions, mlir::OpBuilder &builder, mlir::Location location, mlir::Value expression,
                                 StringRef name)
         : builder(builder), location(location), expression(expression), name(name), compileOptions(compileOptions)
@@ -1239,11 +1249,31 @@ class MLIRPropertyAccessCodeLogic
         auto elementType = mth.isBoundReference(elementTypeForRef, isBoundRef);
 
         auto refValue = getExprLoadRefValue(location);
+        if (isBoundRef && !refValue && boundRefMaterializedCache)
+        {
+            // only reuse a ref materialized in the SAME block as this access: a ref
+            // minted inside a nested block (e.g. a `{ }` scope) does not dominate uses
+            // outside that block, and checking real dominance would need MLIR's
+            // DominanceInfo, which isn't otherwise used in this codebase. Same-block is
+            // a conservative, cheap approximation -- it misses some reuse opportunities
+            // across block boundaries but never returns a ref that fails to dominate.
+            auto cached = boundRefMaterializedCache->lookup(expression);
+            if (cached && cached.getParentBlock() == builder.getInsertionBlock())
+            {
+                refValue = cached;
+            }
+        }
+
         if (isBoundRef && !refValue)
         {
             // allocate in stack
             refValue =
                 builder.create<mlir_ts::VariableOp>(location, mlir_ts::RefType::get(expression.getType()), expression);
+
+            if (boundRefMaterializedCache)
+            {
+                boundRefMaterializedCache->insert(expression, refValue);
+            }
         }
 
         if (refValue)
@@ -1260,7 +1290,7 @@ class MLIRPropertyAccessCodeLogic
             location, elementTypeForRef, expression, MLIRHelper::getStructIndex(builder, fieldIndex));
     }
 
-    template <typename T> ValueOrLogicalResult TupleGetSetAccessor(T tupleType, mlir::Attribute fieldId) 
+    template <typename T> ValueOrLogicalResult TupleGetSetAccessor(T tupleType, mlir::Attribute fieldId)
     {
         MLIRCodeLogic mcl(builder, compileOptions);
 
@@ -1380,11 +1410,31 @@ class MLIRPropertyAccessCodeLogic
         auto elementType = mth.isBoundReference(elementTypeForRef, isBoundRef);
 
         auto refValue = getExprLoadRefValue(location);
+        if (isBoundRef && !refValue && boundRefMaterializedCache)
+        {
+            // only reuse a ref materialized in the SAME block as this access: a ref
+            // minted inside a nested block (e.g. a `{ }` scope) does not dominate uses
+            // outside that block, and checking real dominance would need MLIR's
+            // DominanceInfo, which isn't otherwise used in this codebase. Same-block is
+            // a conservative, cheap approximation -- it misses some reuse opportunities
+            // across block boundaries but never returns a ref that fails to dominate.
+            auto cached = boundRefMaterializedCache->lookup(expression);
+            if (cached && cached.getParentBlock() == builder.getInsertionBlock())
+            {
+                refValue = cached;
+            }
+        }
+
         if (isBoundRef && !refValue)
         {
             // allocate in stack
             refValue =
                 builder.create<mlir_ts::VariableOp>(location, mlir_ts::RefType::get(expression.getType()), expression);
+
+            if (boundRefMaterializedCache)
+            {
+                boundRefMaterializedCache->insert(expression, refValue);
+            }
         }
 
         if (refValue)

tslang/include/TypeScript/MLIRLogic/MLIRDefines.h

@@ -46,6 +46,9 @@
 using VariablePairT = std::pair<mlir::Value, ts::VariableDeclarationDOM::TypePtr>;
 using SymbolTableScopeT = llvm::ScopedHashTableScope<StringRef, VariablePairT>;
 
+// see MLIRGenImpl::boundRefMaterializedCache
+using BoundRefCacheScopeT = llvm::ScopedHashTableScope<mlir::Value, mlir::Value>;
+
 typedef std::pair<mlir::Type, StringRef> SafeTypeKeyType;
 using SafeTypesMapScopeT = llvm::ScopedHashTableScope<SafeTypeKeyType, mlir::Value>;
 

tslang/lib/TypeScript/MLIRGenAccessCall.cpp

@@ -32,6 +32,10 @@ namespace mlirgen
     {
         assert(objectValue);
         MLIRPropertyAccessCodeLogic cl(compileOptions, builder, location, objectValue, name);
+        if (!genContext.dummyRun && !genContext.allowPartialResolve)
+        {
+            cl.setBoundRefMaterializedCache(&boundRefMaterializedCache);
+        }
         return mlirGenPropertyAccessExpressionLogic(location, objectValue, false, cl, genContext);
     }
 
@@ -41,13 +45,21 @@ namespace mlirgen
     {
         assert(objectValue);
         MLIRPropertyAccessCodeLogic cl(compileOptions, builder, location, objectValue, name);
+        if (!genContext.dummyRun && !genContext.allowPartialResolve)
+        {
+            cl.setBoundRefMaterializedCache(&boundRefMaterializedCache);
+        }
         return mlirGenPropertyAccessExpressionLogic(location, objectValue, isConditional, cl, genContext);
     }
 
     ValueOrLogicalResult MLIRGenImpl::mlirGenPropertyAccessExpression(mlir::Location location, mlir::Value objectValue,
                                                          mlir::Attribute id, const GenContext &genContext)
     {
         MLIRPropertyAccessCodeLogic cl(compileOptions, builder, location, objectValue, id);
+        if (!genContext.dummyRun && !genContext.allowPartialResolve)
+        {
+            cl.setBoundRefMaterializedCache(&boundRefMaterializedCache);
+        }
         return mlirGenPropertyAccessExpressionLogic(location, objectValue, false, cl, genContext);
     }
 
@@ -56,6 +68,10 @@ namespace mlirgen
                                                          const GenContext &genContext)
     {
         MLIRPropertyAccessCodeLogic cl(compileOptions, builder, location, objectValue, id);
+        if (!genContext.dummyRun && !genContext.allowPartialResolve)
+        {
+            cl.setBoundRefMaterializedCache(&boundRefMaterializedCache);
+        }
         return mlirGenPropertyAccessExpressionLogic(location, objectValue, isConditional, cl, genContext);
     }
 
@@ -65,8 +81,12 @@ namespace mlirgen
                                                          const GenContext &genContext)
     {
         MLIRPropertyAccessCodeLogic cl(compileOptions, builder, location, objectValue, id, argument);
+        if (!genContext.dummyRun && !genContext.allowPartialResolve)
+        {
+            cl.setBoundRefMaterializedCache(&boundRefMaterializedCache);
+        }
         return mlirGenPropertyAccessExpressionLogic(location, objectValue, isConditional, cl, genContext);
-    }    
+    }
 
     ValueOrLogicalResult MLIRGenImpl::mlirGenPropertyAccessExpressionLogic(mlir::Location location, mlir::Value objectValue,
                                                               bool isConditional, MLIRPropertyAccessCodeLogic &cl,

tslang/lib/TypeScript/MLIRGenFunctions.cpp

@@ -1172,6 +1172,7 @@ namespace mlirgen
         }
 
         SymbolTableScopeT varScope(symbolTable);
+        BoundRefCacheScopeT boundRefCacheScope(boundRefMaterializedCache);
 
         auto location = loc(functionLikeDeclarationBaseAST);
 
@@ -1282,6 +1283,7 @@ namespace mlirgen
         LLVM_DEBUG(llvm::dbgs() << "\n!! >>>> SYNTH. FUNCTION: '" << fullFuncName << "' ~~~ " << (genContext.dummyRun ? "dummy run" : "") <<  (genContext.allowPartialResolve ? " allowed partial resolve" : "") << "\n";);
 
         SymbolTableScopeT varScope(symbolTable);
+        BoundRefCacheScopeT boundRefCacheScope(boundRefMaterializedCache);
 
         SmallVector<mlir::NamedAttribute> attrs;
         processFunctionAttributes(attrs, genContext);

tslang/lib/TypeScript/MLIRGenImpl.h

@@ -10744,6 +10744,20 @@ class MLIRGenImpl
 
     llvm::ScopedHashTable<StringRef, VariablePairT> symbolTable;
 
+    // Caches the stack-allocated ref materialized for a storage-less value (e.g. a
+    // `const` binding with no backing storage) the first time a bound-method property
+    // access needs an address for it (see MLIRPropertyAccessCodeLogic::TupleNoError).
+    // Without this, each access re-materializes a fresh copy seeded from the pristine,
+    // never-mutated SSA value, so repeated calls like `g.next()` on a `const`-bound
+    // generator never observe state changes made by earlier calls. Keyed by mlir::Value
+    // identity, which is stable and unique within a function. Scoped (not just cleared)
+    // at each mlirGenFunctionBody entry via BoundRefCacheScopeT, mirroring symbolTable's
+    // own scoping -- codegen for a nested closure recurses into mlirGenFunctionBody
+    // while the enclosing function's generation is still on the call stack, so a plain
+    // clear-on-entry would permanently drop the outer function's cache entries instead
+    // of restoring them when the nested closure's generation finishes.
+    llvm::ScopedHashTable<mlir::Value, mlir::Value> boundRefMaterializedCache;
+
     NamespaceInfo::TypePtr rootNamespace;
 
     NamespaceInfo::TypePtr currentNamespace;

tslang/test/tester/CMakeLists.txt

@@ -123,6 +123,7 @@ add_test(NAME test-compile-00-numbers COMMAND test-runner "${PROJECT_SOURCE_DIR}
 add_test(NAME test-compile-00-equals COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00equals.ts")
 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-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")
@@ -467,6 +468,7 @@ add_test(NAME test-jit-00-numbers COMMAND test-runner -jit "${PROJECT_SOURCE_DIR
 add_test(NAME test-jit-00-equals COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00equals.ts")
 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-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")

docs/bugs/00generator_manual_next.ts → tslang/test/tester/tests/00generator_manual_next.ts

@@ -8,9 +8,6 @@
 // original -- on every single call site. So every manual `.next()` call restarted the
 // generator instead of resuming it. `for...of` happened to work because its lowering
 // materializes the generator object into one persistent local up front and reuses it.
-//
-// fix: a const whose value is a tuple with a bound-method field (e.g. a generator or
-// closure object) now gets real stack storage, matching what for...of already relied on.
 
 function* gen() {
     for (let i = 0; i < 5; i++) {