rc: a const chosen by a conditional expression is released at the end of its block

tslang/lib/TypeScript/OwnedReturnConsumptionPass.cpp

@@ -1610,7 +1610,8 @@ class OwnedReturnConsumptionPass
     //    path this cannot see - the body of a statement written in the block is not outside it
     //    (runsWithinBlock);
     //  - a user that is a terminator, since a value handed to a successor as a block argument is
-    //    still live after the point this would release it;
+    //    still live after the point this would release it - save a conditional expression's
+    //    branch, whose `ts.If` result is judged in the value's place (allUsesReleasableInOwnBlock);
     //  - a user at or after a `ts.StateLabel` that follows the definition, which is a generator's
     //    resume point: the state machine re-enters the block THERE, so the end of the block is
     //    reachable on a path that never ran the op that produced the value. Nothing about that is
@@ -1923,6 +1924,13 @@ class OwnedReturnConsumptionPass
                     continue;
                 }
 
+                // `n = b ? h : h2` stores what the branch hands on
+                if (auto handedOn = handedOnByConditional(user))
+                {
+                    values.push_back(handedOn);
+                    continue;
+                }
+
                 auto storeOp = mlir::dyn_cast<mlir_ts::StoreOp>(user);
                 auto varOp = storeOp && storeOp.getValue() == value
                                  ? storeOp.getReference().getDefiningOp<mlir_ts::VariableOp>()
@@ -1946,32 +1954,63 @@ class OwnedReturnConsumptionPass
     // `const a = new C(); if (x) { let b = a; }` - is exactly that, and leaked before this. Any
     // other region is: a nested function or an `async.execute` body runs later, from somewhere
     // else.
+    //
+    // A conditional expression's branch is the one terminator that is followed: `b ? h : h2`
+    // yields `h` through `ts.Result`, and from there the value is the `ts.If`'s result, so it is
+    // that result's users that decide (handedOnByConditional). A receiver of a result nobody owns
+    // retains, and a branch that makes it owned retained `h` for it (consumeConditionalResults),
+    // so `h`'s own reference is still given back at the end of its block either way.
     static bool allUsesReleasableInOwnBlock(mlir::Operation *op)
     {
         auto *block = op->getBlock();
         auto rangeEnd = releaseRangeEnd(op);
 
-        for (auto *user : op->getResult(0).getUsers())
+        llvm::SmallVector<mlir::Value> values{op->getResult(0)};
+        while (!values.empty())
         {
-            if (user->hasTrait<mlir::OpTrait::IsTerminator>() || !runsWithinBlock(user, block))
+            auto value = values.pop_back_val();
+            for (auto *user : value.getUsers())
             {
-                return false;
-            }
+                auto handedOn = handedOnByConditional(user);
+                if ((user->hasTrait<mlir::OpTrait::IsTerminator>() && !handedOn) || !runsWithinBlock(user, block))
+                {
+                    return false;
+                }
 
-            // A use at or after the end of the range runs on a path the release has already left:
-            // after the yield, or inside the statement that holds one (the resume point there
-            // re-enters past the definition). Such a value leaks. MLIRGen keeps a generator's
-            // locals in its state object, so a temporary used across a yield does not occur.
-            auto *inBlock = block->findAncestorOpInBlock(*user);
-            if (rangeEnd != block->end() && !inBlock->isBeforeInBlock(&*rangeEnd))
-            {
-                return false;
+                // A use at or after the end of the range runs on a path the release has already
+                // left: after the yield, or inside the statement that holds one (the resume point
+                // there re-enters past the definition). Such a value leaks. MLIRGen keeps a
+                // generator's locals in its state object, so a temporary used across a yield does
+                // not occur.
+                auto *inBlock = block->findAncestorOpInBlock(*user);
+                if (rangeEnd != block->end() && !inBlock->isBeforeInBlock(&*rangeEnd))
+                {
+                    return false;
+                }
+
+                if (handedOn)
+                {
+                    values.push_back(handedOn);
+                }
             }
         }
 
         return true;
     }
 
+    // The value a branch of a conditional expression hands on: the result of the `ts.If` whose
+    // branch this `ts.Result` ends, when that `ts.If` has the one result. None for any other user.
+    static mlir::Value handedOnByConditional(mlir::Operation *user)
+    {
+        if (!mlir::isa<mlir_ts::ResultOp>(user) || user->getNumOperands() != 1)
+        {
+            return mlir::Value();
+        }
+
+        auto ifOp = mlir::dyn_cast_or_null<mlir_ts::IfOp>(user->getParentOp());
+        return ifOp && ifOp.getNumResults() == 1 ? ifOp.getResult(0) : mlir::Value();
+    }
+
     // Where the stretch of `op`'s block that a release of it may cover ends: at the first op after
     // it that is a yield or a resume point or has one inside 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
@@ -2041,6 +2080,12 @@ class OwnedReturnConsumptionPass
             {
                 return true;
             }
+
+            // `yield b ? h : h2` stores what the branch hands on
+            if (auto handedOn = handedOnByConditional(user); handedOn && escapesThroughStore(handedOn))
+            {
+                return true;
+            }
         }
 
         return false;

tslang/test/tester/CMakeLists.txt

@@ -395,6 +395,7 @@ tslang_add_test(NAME test-compile-00-const-view-of-field COMMAND test-runner "${
 tslang_add_test(NAME test-compile-00-const-view-of-element COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00const_view_of_element.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-conditional-named-const-release COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00conditional_named_const_release.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-const-interface-of-local COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00const_interface_of_local.ts")
 tslang_add_test(NAME test-compile-00-string-view-copy COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00string_view_copy.ts")
@@ -993,6 +994,7 @@ tslang_add_test(NAME test-jit-00-const-view-of-field COMMAND test-runner -jit "$
 tslang_add_test(NAME test-jit-00-const-view-of-element COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00const_view_of_element.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-conditional-named-const-release COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00conditional_named_const_release.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-const-interface-of-local COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00const_interface_of_local.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")
@@ -1981,6 +1983,7 @@ set(TSLANG_CORPUS
     00const_view_of_element.ts
     00const_record_owned_fields.ts
     00generator_temporaries.ts
+    00conditional_named_const_release.ts
     00array_length_valid.ts
     00const_interface_of_local.ts
     00string_view_copy.ts
@@ -3823,6 +3826,13 @@ 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 `const` chosen by a conditional expression - `b ? h : h2` - is released at the end of its block
+# under rc; no run can see the leak, so the IR is read.
+add_test(NAME test-rc-conditional-named-const-release
+         COMMAND ${CMAKE_COMMAND} -DTSLANG=$<TARGET_FILE:tslang>
+                 "-DFILE=${PROJECT_SOURCE_DIR}/test/tester/tests/00conditional_named_const_release.ts"
+                 -P "${CMAKE_CURRENT_SOURCE_DIR}/conditional-named-const-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

tslang/test/tester/conditional-named-const-release.cmake

@@ -0,0 +1,76 @@
+# A folded `const` made by `new` and chosen by a conditional expression is released at the end of
+# its block (00conditional_named_const_release.ts). The branch yields it through `ts.Result`, a
+# terminator, and that kept it from any release under rc (and own, which reads rc's IR). No run
+# shows the leak, so the IR is read, after OwnedReturnConsumptionPass (--emit=mlir). Each `new`
+# bound to a `const` in a function is released exactly once - except, in `assignedOutside`, the
+# one chosen into a local that owns nothing and outlives the `if` - and in `chosenOrNew` the
+# `ts.If`, which owns the side it made, is released once as well.
+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 function, from its `ts.Func` to the next one
+function(func_body function outVar)
+    string(FIND "${out}" "ts.Func @${function} " begin)
+    if(begin EQUAL -1)
+        message(FATAL_ERROR "no function `${function}` 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(releases value body outVar)
+    string(REGEX MATCHALL "\"ts[.]Release\"[(]%${value}[)]" matches "${body}")
+    list(LENGTH matches n)
+    set(${outVar} ${n} PARENT_SCOPE)
+endfunction()
+
+# function name, then how many times each of its named `new`s is released, in order
+set(expected
+    "constChosen:1:1"
+    "letChosen:1:1"
+    "chosenOrNew:1"
+    "nestedChosen:1:1"
+    "chosenInLoop:1:1"
+    "returnChosen:1:1"
+    "assignedOutside:1:0")
+foreach(entry IN LISTS expected)
+    string(REPLACE ":" ";" entry "${entry}")
+    list(POP_FRONT entry function)
+    func_body("${function}" body)
+    string(REGEX MATCHALL "%[0-9]+ = \"ts[.]New\"[^\r\n]*__owned_result_named" news "${body}")
+    list(LENGTH news count)
+    list(LENGTH entry wanted)
+    if(NOT count EQUAL wanted)
+        message(FATAL_ERROR "`${function}` has ${count} named `new`s, not ${wanted}:\n${body}")
+    endif()
+    foreach(new IN LISTS news)
+        list(POP_FRONT entry times)
+        string(REGEX MATCH "^%([0-9]+)" ignored "${new}")
+        releases("${CMAKE_MATCH_1}" "${body}" n)
+        if(NOT n EQUAL times)
+            message(FATAL_ERROR "`${function}` releases %${CMAKE_MATCH_1} ${n} times, not ${times}:\n${body}")
+        endif()
+    endforeach()
+endforeach()
+
+func_body("chosenOrNew" body)
+string(REGEX MATCH "%([0-9]+) = \"ts[.]If\"[^\r\n]*[\r\n]+ *\"ts[.]Retain\"" ignored "${body}")
+if(NOT CMAKE_MATCH_1)
+    message(FATAL_ERROR "`chosenOrNew` does not retain the borrowed side of its `ts.If`:\n${body}")
+endif()
+releases("${CMAKE_MATCH_1}" "${body}" n)
+if(NOT n EQUAL 1)
+    message(FATAL_ERROR "`chosenOrNew` releases its `ts.If` %${CMAKE_MATCH_1} ${n} times, not once:\n${body}")
+endif()

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

@@ -0,0 +1,100 @@
+// A `const` made by `new` and then chosen by a conditional expression - `b ? h : h2` - is given
+// back at the end of its block. The branch hands it on through `ts.Result`, a terminator, and that
+// kept it from ever being released: under rc and own neither side was freed. What reads the chosen
+// value retains it, so the release stays balanced; conditional-named-const-release.cmake checks the
+// releases under rc. Here each value is read after allocations of the same size have reused
+// anything freed too early.
+class C {
+    v: number[] = [];
+    constructor(public dist: number) {}
+}
+
+function churn() {
+    const junk: C[] = [];
+    for (let i = 0; i < 64; i++) {
+        const c = new C(999);
+        c.v.push(99);
+        junk.push(c);
+    }
+}
+
+function constChosen(b: boolean) {
+    const h = new C(3);
+    const h2 = new C(30);
+    const each: C = b ? h : h2;
+    churn();
+    return each.dist + h.dist + h2.dist;
+}
+
+function letChosen(b: boolean) {
+    const h = new C(3);
+    const h2 = new C(30);
+    let each = b ? h : h2;
+    churn();
+    return each.dist + h.dist + h2.dist;
+}
+
+function chosenOrNew(b: boolean) {
+    const h = new C(3);
+    const each: C = b ? h : new C(30);
+    churn();
+    return each.dist + h.dist;
+}
+
+function nestedChosen(a: boolean, b: boolean) {
+    const h = new C(3);
+    const h2 = new C(30);
+    const each: C = a ? h : b ? h2 : h;
+    churn();
+    return each.dist + h.dist + h2.dist;
+}
+
+function chosenInLoop() {
+    const h = new C(3);
+    const h2 = new C(30);
+    let sum = 0;
+    for (let i = 0; i < 10; i++) {
+        const each: C = i % 2 == 0 ? h : h2;
+        sum += each.dist;
+    }
+
+    churn();
+    return sum + h.dist + h2.dist;
+}
+
+function returnChosen(b: boolean) {
+    const h = new C(3);
+    const h2 = new C(30);
+    return b ? h : h2;
+}
+
+// the local declared without a value owns nothing, so `h` is what keeps the value it is given
+// alive after the `if`: it is not released at the end of the `if`
+function assignedOutside(b: boolean) {
+    const h2 = new C(30);
+    let n: C;
+    if (b) {
+        const h = new C(3);
+        n = b ? h : h2;
+    } else {
+        n = h2;
+    }
+
+    churn();
+    return n.dist;
+}
+
+function main() {
+    assert(constChosen(true) == 3 + 33 && constChosen(false) == 30 + 33, "a const chosen from two");
+    assert(letChosen(true) == 3 + 33 && letChosen(false) == 30 + 33, "a let chosen from two");
+    assert(chosenOrNew(true) == 3 + 3 && chosenOrNew(false) == 30 + 3, "one side made by the choice");
+    assert(nestedChosen(false, true) == 30 + 33 && nestedChosen(false, false) == 3 + 33, "a choice inside a choice");
+    assert(chosenInLoop() == 5 * 3 + 5 * 30 + 33, "a choice in a loop");
+
+    const r = returnChosen(false);
+    churn();
+    assert(r.dist == 30, "the chosen one returned");
+
+    assert(assignedOutside(true) == 3 && assignedOutside(false) == 30, "a choice given to a local outside");
+    print("done.");
+}