A generator body's temporaries are released before the yield

tslang/lib/TypeScript/OwnedReturnConsumptionPass.cpp

@@ -819,19 +819,21 @@ class OwnedReturnConsumptionPass
     //    never ran the op that produced the value. Nothing about that is visible while the
     //    generator is still one block - it only becomes a use before definition once the state
     //    machine is expanded, and it surfaces as a dominance failure in the affine lowering
-    //    rather than as anything this pass could notice.
+    //    rather than as anything this pass could notice. A `yield` in front of the resume point
+    //    is different: the release goes before it, inside the stretch the definition dominates.
     //
     // That bias is the same one the rest of this arc takes: an unreleased reference is invisible,
     // a released one that was still owned is a use-after-free.
     void releaseDiscardedTemporaries(MLIRTypeHelper &mth, mlir::ModuleOp module)
     {
-        // Whole generators are left alone, the same way functionReturnsOwned leaves them alone
-        // and for the same reason: what a generator's body looks like here is not what runs.
-        // The state machine cuts its blocks at every resume point afterwards, so a position that
-        // is plainly after a definition now may not be on the path that reaches it later - and
-        // the failure is a dominance error in the affine lowering, not something visible here.
-        // Their temporaries leak, which is the side of the line this arc keeps everything
-        // uncertain on.
+        // What a generator's body looks like here is not what runs: the state machine cuts its
+        // blocks at every resume point afterwards, so a position that is plainly after a
+        // definition now may not be on the path that reaches it later - and the failure is a
+        // dominance error in the affine lowering, not something visible here. A temporary in a
+        // generator is released only inside the stretch of its block that no resume point cuts
+        // (releaseRangeEnd), and only when it does not go out with what the generator yields
+        // (escapesThroughStore). Every other one leaks, which is the side of the line this arc
+        // keeps everything uncertain on.
         llvm::DenseSet<mlir::Operation *> generators;
         module.walk([&](mlir_ts::FuncOp funcOp) {
             auto isGenerator = false;
@@ -856,7 +858,7 @@ class OwnedReturnConsumptionPass
 
             if (auto funcOp = op->getParentOfType<mlir_ts::FuncOp>())
             {
-                if (generators.contains(funcOp.getOperation()))
+                if (generators.contains(funcOp.getOperation()) && escapesThroughStore(op->getResult(0)))
                 {
                     return;
                 }
@@ -874,6 +876,7 @@ class OwnedReturnConsumptionPass
             }
 
             auto *block = op->getBlock();
+            auto rangeEnd = releaseRangeEnd(op);
 
             // The end of the block is not the only way out of it. A `return` written inside an
             // `if` - which is where returns are usually written - leaves from a nested region and
@@ -900,7 +903,7 @@ class OwnedReturnConsumptionPass
             //
             // jumpLeavesBlock answers it by walking from the jump up to this block and asking
             // whether anything on the way catches it.
-            for (auto it = std::next(mlir::Block::iterator(op)); it != block->end(); ++it)
+            for (auto it = std::next(mlir::Block::iterator(op)); it != rangeEnd; ++it)
             {
                 if (it->getNumRegions() == 0)
                 {
@@ -1026,18 +1029,46 @@ class OwnedReturnConsumptionPass
     static bool allUsesReleasableInOwnBlock(mlir::Operation *op)
     {
         auto *block = op->getBlock();
+        auto rangeEnd = releaseRangeEnd(op);
+
+        // a resume point between the definition and the end of the block with no yield in front
+        // of it to release before - see above
+        if (rangeEnd != block->end() && mlir::isa<mlir_ts::StateLabelOp>(*rangeEnd))
+        {
+            return false;
+        }
+
         for (auto *user : op->getResult(0).getUsers())
         {
             if (user->hasTrait<mlir::OpTrait::IsTerminator>() || !runsWithinBlock(user, block))
             {
                 return false;
             }
+
+            // a use after the yield runs on the resume path, which the release has already left
+            auto *inBlock = block->findAncestorOpInBlock(*user);
+            if (rangeEnd != block->end() && !inBlock->isBeforeInBlock(&*rangeEnd))
+            {
+                return false;
+            }
         }
 
-        // a resume point between the definition and the end of the block - see above
-        for (auto it = std::next(mlir::Block::iterator(op)); it != block->end(); ++it)
+        // A resume point inside a statement written before the yield - `if (c) { yield 1; }` -
+        // re-enters the block in the middle of that statement and falls through to the release
+        // without having run the definition.
+        for (auto it = std::next(mlir::Block::iterator(op)); it != rangeEnd; ++it)
         {
-            if (mlir::isa<mlir_ts::StateLabelOp>(*it))
+            auto resumes = it->walk([&](mlir::Operation *nested) {
+                if (mlir::isa<mlir_ts::FuncOp>(nested))
+                {
+                    return mlir::WalkResult::skip();
+                }
+
+                return mlir::isa<mlir_ts::StateLabelOp>(nested) ? mlir::WalkResult::interrupt()
+                                                                : mlir::WalkResult::advance();
+            });
+
+            if (resumes.wasInterrupted())
             {
                 return false;
             }
@@ -1046,6 +1077,56 @@ class OwnedReturnConsumptionPass
         return true;
     }
 
+    // Where the stretch of `op`'s block that a release of it may cover ends: at the first yield
+    // or resume point after it, or at the end of the block. The yield ends the call (the affine
+    // lowering turns it into a store and a jump to the exit block), and the next call re-enters
+    // at the resume point after it, past the definition - so a release has to come before the
+    // yield, and nothing after it is on the path that ran the definition. A plain function has
+    // neither and gets the whole block.
+    static mlir::Block::iterator releaseRangeEnd(mlir::Operation *op)
+    {
+        auto *block = op->getBlock();
+        for (auto it = std::next(mlir::Block::iterator(op)); it != block->end(); ++it)
+        {
+            if (mlir::isa<mlir_ts::YieldReturnValOp, mlir_ts::StateLabelOp>(*it))
+            {
+                return it;
+            }
+        }
+
+        return block->end();
+    }
+
+    // Does this generator temporary, or a view of it, get stored somewhere? The `{ value, done }`
+    // record a `yield` hands back is filled by a store that takes no count of its own - the
+    // caller retains on receipt - so `yield "a" + i` keeps the concatenation alive only through
+    // the +1 it was born with, and a release before the yield would free what the caller reads.
+    // Any store is treated that way: it leaks, as before, rather than telling the record apart
+    // from a place that does count.
+    static bool escapesThroughStore(mlir::Value value)
+    {
+        for (auto *user : value.getUsers())
+        {
+            if (auto storeOp = mlir::dyn_cast<mlir_ts::StoreOp>(user))
+            {
+                if (storeOp.getValue() == value)
+                {
+                    return true;
+                }
+
+                continue;
+            }
+
+            if (mlir::isa<mlir_ts::CastOp, mlir_ts::OptionalValueOp, mlir_ts::CreateUnionInstanceOp>(user) &&
+                user->getNumResults() == 1 && escapesThroughStore(user->getResult(0)))
+            {
+                return true;
+            }
+        }
+
+        return false;
+    }
+
     // The first op after `op` that ends this block's execution, if there is one.
     //
     // "End of the block" is not the end of what runs. `ts.ReturnVal` is not a terminator - the
@@ -1061,12 +1142,13 @@ class OwnedReturnConsumptionPass
     //
     // Releasing before the return is right for a discarded temporary by definition: what is
     // being returned was consumed by the return's own retain (§9.24) and so is not in this set.
+    // A generator's `yield` ends the call the same way (releaseRangeEnd).
     static mlir::Operation *firstExitingOpAfter(mlir::Operation *op)
     {
         auto *block = op->getBlock();
         for (auto it = std::next(mlir::Block::iterator(op)); it != block->end(); ++it)
         {
-            if (mlir::isa<mlir_ts::ReturnValOp, mlir_ts::ReturnOp, mlir_ts::ExitOp>(*it))
+            if (mlir::isa<mlir_ts::ReturnValOp, mlir_ts::ReturnOp, mlir_ts::ExitOp, mlir_ts::YieldReturnValOp>(*it))
             {
                 return &*it;
             }

tslang/test/tester/CMakeLists.txt

@@ -389,6 +389,7 @@ tslang_add_test(NAME test-compile-00-string-concat-bounds COMMAND test-runner "$
 tslang_add_test(NAME test-compile-00-parse-number-range COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00parse_number_range.ts")
 tslang_add_test(NAME test-compile-00-narrowed-assign-other-member COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00narrowed_assign_other_member.ts")
 tslang_add_test(NAME test-compile-00-const-record-owned-fields COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00const_record_owned_fields.ts")
+tslang_add_test(NAME test-compile-00-generator-temporaries COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00generator_temporaries.ts")
 tslang_add_test(NAME test-compile-00-array-length-valid COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00array_length_valid.ts")
 tslang_add_test(NAME test-compile-00-string-view-copy COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00string_view_copy.ts")
 tslang_add_test(NAME test-compile-00-c-string-copy COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00c_string_copy.ts")
@@ -950,6 +951,7 @@ tslang_add_test(NAME test-jit-00-string-concat-bounds COMMAND test-runner -jit "
 tslang_add_test(NAME test-jit-00-parse-number-range COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00parse_number_range.ts")
 tslang_add_test(NAME test-jit-00-narrowed-assign-other-member COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00narrowed_assign_other_member.ts")
 tslang_add_test(NAME test-jit-00-const-record-owned-fields COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00const_record_owned_fields.ts")
+tslang_add_test(NAME test-jit-00-generator-temporaries COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00generator_temporaries.ts")
 tslang_add_test(NAME test-jit-00-array-length-valid COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00array_length_valid.ts")
 tslang_add_test(NAME test-jit-00-string-view-copy COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00string_view_copy.ts")
 tslang_add_test(NAME test-jit-00-c-string-copy COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00c_string_copy.ts")
@@ -1871,6 +1873,7 @@ set(TSLANG_CORPUS
     00parse_number_range.ts
     00narrowed_assign_other_member.ts
     00const_record_owned_fields.ts
+    00generator_temporaries.ts
     00array_length_valid.ts
     00string_view_copy.ts
     00c_string_copy.ts
@@ -3457,6 +3460,15 @@ add_test(NAME test-rc-const-record-release
                  "-DFILE=${PROJECT_SOURCE_DIR}/test/tester/tests/00const_record_owned_fields.ts"
                  -P "${CMAKE_CURRENT_SOURCE_DIR}/const-record-release.cmake")
 
+# A generator body's temporaries are released before the yield that ends the call, where no resume
+# point comes between, and the file runs under own as well.
+add_test(NAME test-rc-generator-temporaries
+         COMMAND ${CMAKE_COMMAND} -DTSLANG=$<TARGET_FILE:tslang>
+                 "-DFILE=${PROJECT_SOURCE_DIR}/test/tester/tests/00generator_temporaries.ts"
+                 -P "${CMAKE_CURRENT_SOURCE_DIR}/generator-temporaries.cmake")
+tslang_add_test(NAME test-jit-own-00-generator-temporaries COMMAND test-runner -jit -mm=own "${PROJECT_SOURCE_DIR}/test/tester/tests/00generator_temporaries.ts")
+tslang_add_test(NAME test-compile-own-00-generator-temporaries COMMAND test-runner -mm=own "${PROJECT_SOURCE_DIR}/test/tester/tests/00generator_temporaries.ts")
+
 # Shared<T> (spec 23): counted handles. These are not in the own_test list, whose
 # test-own-no-counting check forbids __tslang_inc_ref: a handle is counted under own by design.
 # own_shared_user_named: a program's own class named Shared wins over the built-in (spec 23.1).

tslang/test/tester/generator-temporaries.cmake

@@ -0,0 +1,74 @@
+# A temporary inside a generator body is released before the yield that ends the call, and only
+# where no resume point comes between (00generator_temporaries.ts). No run shows which ones leak,
+# so the IR is read, after OwnedReturnConsumptionPass (--emit=mlir), under rc:
+#  - lengths:     `yield ("a" + i).length` releases the number's string and the concatenation;
+#  - strings:     `yield "s" + i` releases the number's string but not the concatenation, which
+#                 the yielded record holds;
+#  - branchYield: a temporary in front of an `if` with a yield inside is not released;
+#  - twoSegments: each of the two yields has its own two releases in front of it.
+# Under own the file must also pass the ownership verifier.
+cmake_minimum_required(VERSION 3.17.3)
+execute_process(COMMAND "${TSLANG}" --emit=mlir --no-default-lib -mm=rc "${FILE}"
+                OUTPUT_VARIABLE out ERROR_VARIABLE err RESULT_VARIABLE result)
+# the IR is printed to stderr
+string(APPEND out "${err}")
+if(NOT result EQUAL 0)
+    message(FATAL_ERROR "tslang failed (${result}):\n${out}")
+endif()
+
+# the body of a generator's `next`, from its `ts.Func` to the next one
+function(next_body name outVar)
+    string(FIND "${out}" "ts.Func @\".f_${name}.." begin)
+    if(begin EQUAL -1)
+        message(FATAL_ERROR "no generator `${name}` in:\n${out}")
+    endif()
+    string(SUBSTRING "${out}" ${begin} -1 body)
+    string(SUBSTRING "${body}" 1 -1 rest)
+    string(FIND "${rest}" "ts.Func @" end)
+    if(NOT end EQUAL -1)
+        math(EXPR end "${end} + 1")
+        string(SUBSTRING "${body}" 0 ${end} body)
+    endif()
+    set(${outVar} "${body}" PARENT_SCOPE)
+endfunction()
+
+function(count_releases body outVar)
+    string(REGEX MATCHALL "\"ts[.]Release\"[(]%[0-9]+[)] : [(]!ts[.]string[)]" releases "${body}")
+    list(LENGTH releases count)
+    set(${outVar} ${count} PARENT_SCOPE)
+endfunction()
+
+# (CMake's regex has no {n})
+set(release "\"ts[.]Release\"[(]%[0-9]+[)] : [(]!ts[.]string[)] -> [(][)][\r\n]+ *")
+set(releasesBeforeYield "${release}${release}\"ts[.]YieldReturnVal\"")
+
+next_body(lengths body)
+if(NOT body MATCHES "${releasesBeforeYield}")
+    message(FATAL_ERROR "`lengths` does not release its two temporaries before the yield:\n${body}")
+endif()
+
+next_body(strings body)
+count_releases("${body}" count)
+string(REGEX MATCH "%([0-9]+) = \"ts[.]ArithmeticBinary\"[^\r\n]*-> !ts[.]string" concat "${body}")
+if(NOT count EQUAL 1 OR NOT concat OR body MATCHES "\"ts[.]Release\"[(]%${CMAKE_MATCH_1}[)]")
+    message(FATAL_ERROR "`strings` must release the number's string and keep the yielded one:\n${body}")
+endif()
+
+next_body(branchYield body)
+count_releases("${body}" count)
+if(NOT count EQUAL 0)
+    message(FATAL_ERROR "`branchYield` releases a temporary a resume point cuts off:\n${body}")
+endif()
+
+next_body(twoSegments body)
+string(REGEX MATCHALL "${releasesBeforeYield}" segments "${body}")
+list(LENGTH segments count)
+if(NOT count EQUAL 2)
+    message(FATAL_ERROR "`twoSegments` does not release two temporaries before each yield:\n${body}")
+endif()
+
+execute_process(COMMAND "${TSLANG}" --emit=mlir-affine --verify-ownership --no-default-lib -mm=own "${FILE}"
+                OUTPUT_QUIET ERROR_VARIABLE err RESULT_VARIABLE result)
+if(NOT result EQUAL 0)
+    message(FATAL_ERROR "own rejects it (${result}):\n${err}")
+endif()

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

@@ -0,0 +1,74 @@
+// A temporary made inside a generator body is given back before the yield that ends the call.
+//
+// OwnedReturnConsumptionPass used to skip whole generators: the state machine cuts their blocks
+// at every resume point later, so a release at the end of a block could stop being dominated by
+// the definition. Under rc and own `yield ("a" + i).length` leaked the number's string and the
+// concatenation on every call: 500000 runs of `lengths` below grew by about 165 MB.
+//
+// Now a temporary is released when no resume point lies between it and the release, which goes
+// in front of the first yield after it. Three cases stay leaking, as before:
+//  - a value that goes out with the yield (`yield "a" + i`): the `{ value, done }` record holds
+//    it without a count of its own, and the caller reads it after the call returns;
+//  - a value made in front of a statement that has a yield inside it (`branchYield`): resuming
+//    there reaches the end of the block without having made it;
+//  - a value used after a yield.
+//
+// test-rc-generator-temporaries reads the IR for which is which.
+
+function* lengths() {
+    let i = 0;
+    while (true) {
+        i++;
+        if (i > 3) break;
+        yield ("a" + i).length;
+    }
+}
+
+function* strings() {
+    let i = 0;
+    while (true) {
+        i++;
+        if (i > 3) break;
+        yield "s" + i;
+    }
+}
+
+function* branchYield() {
+    let i = 0;
+    while (i < 3) {
+        i++;
+        const n = ("b" + i).length;
+        if (n > 0) {
+            yield n;
+        }
+    }
+}
+
+// a temporary between two yields belongs to the stretch after the first one's resume point
+function* twoSegments() {
+    yield ("x" + 1).length;
+    yield ("yy" + 22).length;
+}
+
+function main() {
+    let total = 0;
+    for (let n = 0; n < 1000; n++) {
+        for (const v of lengths()) total += v;
+    }
+
+    assert(total == 6000, "lengths");
+
+    let joined = "";
+    for (const s of strings()) joined += s;
+    assert(joined == "s1s2s3", "strings");
+
+    let branchTotal = 0;
+    for (const v of branchYield()) branchTotal += v;
+    assert(branchTotal == 6, "branchYield");
+
+    let segments = 0;
+    for (const v of twoSegments()) segments += v;
+    assert(segments == 6, "twoSegments");
+
+    print("done.");
+}