diff --git a/tslang/lib/TypeScript/OwnedReturnConsumptionPass.cpp b/tslang/lib/TypeScript/OwnedReturnConsumptionPass.cpp index 9b78786e9..243e4f020 100644 --- a/tslang/lib/TypeScript/OwnedReturnConsumptionPass.cpp +++ b/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 generators; module.walk([&](mlir_ts::FuncOp funcOp) { auto isGenerator = false; @@ -856,7 +858,7 @@ class OwnedReturnConsumptionPass if (auto funcOp = op->getParentOfType()) { - 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(*rangeEnd)) + { + return false; + } + for (auto *user : op->getResult(0).getUsers()) { if (user->hasTrait() || !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(*it)) + auto resumes = it->walk([&](mlir::Operation *nested) { + if (mlir::isa(nested)) + { + return mlir::WalkResult::skip(); + } + + return mlir::isa(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(*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(user)) + { + if (storeOp.getValue() == value) + { + return true; + } + + continue; + } + + if (mlir::isa(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(*it)) + if (mlir::isa(*it)) { return &*it; } diff --git a/tslang/test/tester/CMakeLists.txt b/tslang/test/tester/CMakeLists.txt index 6960525fe..aeccd35ee 100644 --- a/tslang/test/tester/CMakeLists.txt +++ b/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=$ + "-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 (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). diff --git a/tslang/test/tester/generator-temporaries.cmake b/tslang/test/tester/generator-temporaries.cmake new file mode 100644 index 000000000..35f45f00f --- /dev/null +++ b/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() diff --git a/tslang/test/tester/tests/00generator_temporaries.ts b/tslang/test/tester/tests/00generator_temporaries.ts new file mode 100644 index 000000000..c7efa8293 --- /dev/null +++ b/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."); +}