-mm=rc: a parameter the body assigns owns what it is assigned

tslang/lib/TypeScript/MLIRGenFunctions.cpp

@@ -1418,6 +1418,122 @@ namespace mlirgen
         return mlir::success();
     }    
 
+    // The names a function body assigns: `x = ...` and every compound assignment, the targets of a destructuring
+    // assignment, and the variable of `for (x of ...)` / `for (x in ...)`. Nested functions are walked as well, and
+    // a destructuring target counts every name in it: naming one too many costs a retain and a release, missing
+    // one leaves a parameter pointing at freed memory.
+    static void collectAssignedNames(ts::Node body, llvm::StringSet<> &names)
+    {
+        auto addTargets = [&](ts::Node target) {
+            while (target == SyntaxKind::ParenthesizedExpression)
+            {
+                target = target.as<ts::ParenthesizedExpression>()->expression;
+            }
+
+            if (target == SyntaxKind::Identifier)
+            {
+                names.insert(MLIRHelper::getName(target.as<ts::Identifier>()));
+            }
+            else if (target == SyntaxKind::ArrayLiteralExpression || target == SyntaxKind::ObjectLiteralExpression)
+            {
+                ts::VisitorAST targetNames([&](ts::Node node) {
+                    if (node == SyntaxKind::Identifier)
+                    {
+                        names.insert(MLIRHelper::getName(node.as<ts::Identifier>()));
+                    }
+                });
+                targetNames.visit(target);
+            }
+        };
+
+        ts::VisitorAST visitor([&](ts::Node node) {
+            if (node == SyntaxKind::BinaryExpression)
+            {
+                auto binaryExpression = node.as<ts::BinaryExpression>();
+                if (ts::isAssignmentOperator(binaryExpression->operatorToken))
+                {
+                    addTargets(binaryExpression->left);
+                }
+            }
+            else if (node == SyntaxKind::ForOfStatement)
+            {
+                auto initializer = node.as<ts::ForOfStatement>()->initializer;
+                if (initializer != SyntaxKind::VariableDeclarationList)
+                {
+                    addTargets(initializer);
+                }
+            }
+            else if (node == SyntaxKind::ForInStatement)
+            {
+                auto initializer = node.as<ts::ForInStatement>()->initializer;
+                if (initializer != SyntaxKind::VariableDeclarationList)
+                {
+                    addTargets(initializer);
+                }
+            }
+        });
+        visitor.visit(body);
+    }
+
+    // -mm=rc: a parameter is borrowed - the caller owns its value (takeOwnershipOfLocal), so its slot neither
+    // retains nor releases. One the body assigns holds values the frame made, though: `x = new B(1)` stored the
+    // instance with no reference of its own, the block it was made in released it, and `x` pointed at freed
+    // memory - in a loop, the condition read it on the next iteration (#512). A parameter the body assigns is
+    // copied into a local of the same name, which owns what it holds as any local does: a reference taken on the
+    // caller's value here, handed over by each assignment, given back at the function's exit.
+    //
+    // -mm=own refuses the assignment itself (OwnershipInferencePass), and a collector needs none of this.
+    mlir::LogicalResult MLIRGenImpl::mlirGenFunctionOwnAssignedParams(mlir::Location location,
+                                                                      FunctionLikeDeclarationBase functionLikeDeclarationBaseAST,
+                                                                      FunctionPrototypeDOM::TypePtr funcProto,
+                                                                      const GenContext &genContext)
+    {
+        if (!compileOptions.isRefCounted())
+        {
+            return mlir::success();
+        }
+
+        llvm::StringSet<> assignedNames;
+        auto scanned = false;
+        for (const auto &param : funcProto->getParams())
+        {
+            if (param->getBindingPattern() || !mth.ownsHeapMemory(location, param->getType()))
+            {
+                continue;
+            }
+
+            if (!scanned)
+            {
+                collectAssignedNames(functionLikeDeclarationBaseAST->body, assignedNames);
+                scanned = true;
+            }
+
+            if (!assignedNames.contains(param->getName()))
+            {
+                continue;
+            }
+
+            auto paramValue = resolveIdentifier(location, param->getName(), genContext);
+            if (!paramValue)
+            {
+                return mlir::failure();
+            }
+
+            auto paramType = param->getType();
+            if (!registerVariable(
+                    location, param->getName(), false, VariableType::Let,
+                    [&](mlir::Location, const GenContext &) -> TypeValueInitType {
+                        return {paramType, paramValue, TypeProvided::Yes};
+                    },
+                    genContext))
+            {
+                return mlir::failure();
+            }
+        }
+
+        return mlir::success();
+    }
+
     mlir::LogicalResult MLIRGenImpl::mlirGenFunctionParamsBindings(int firstIndex, FunctionPrototypeDOM::TypePtr funcProto,
                                                       mlir::Block::BlockArgListType arguments,
                                                       const GenContext &genContext)
@@ -1612,6 +1728,12 @@ namespace mlirgen
         auto discoverParamsOnly = funcGenContext.allowPartialResolve && funcGenContext.discoverParamsOnly;
         if (!discoverParamsOnly)
         {
+            // after the list is wired in: the copy is the function's to give back at exit
+            if (failed(mlirGenFunctionOwnAssignedParams(location, functionLikeDeclarationBaseAST, funcProto, funcGenContext)))
+            {
+                return mlir::failure();
+            }
+
             // we need it to skip lexical block
             functionLikeDeclarationBaseAST->body->parent = functionLikeDeclarationBaseAST->body;
             if (failed(mlirGenBody(functionLikeDeclarationBaseAST->body, funcGenContext)))

tslang/lib/TypeScript/MLIRGenImpl.h

@@ -3214,6 +3214,10 @@ class MLIRGenImpl
 
     mlir::LogicalResult mlirGenFunctionParams(mlir::Location location, int firstIndex, mlir::Block::BlockArgListType arguments, const GenContext &genContext);
 
+    mlir::LogicalResult mlirGenFunctionOwnAssignedParams(mlir::Location location,
+                                                         FunctionLikeDeclarationBase functionLikeDeclarationBaseAST,
+                                                         FunctionPrototypeDOM::TypePtr funcProto, const GenContext &genContext);
+
     mlir::LogicalResult mlirGenFunctionCaptures(mlir::Location location, FunctionPrototypeDOM::TypePtr funcProto, const GenContext &genContext);
 
     mlir::LogicalResult mlirGenFunctionBody(FunctionLikeDeclarationBase functionLikeDeclarationBaseAST,

tslang/test/tester/CMakeLists.txt

@@ -357,6 +357,7 @@ tslang_add_test(NAME test-compile-00-for-await-sequential COMMAND test-runner "$
 tslang_add_test(NAME test-compile-00-for-await-control COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00for_await_control.ts")
 tslang_add_test(NAME test-compile-00-for-optional-class-condition COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00for_optional_class_condition.ts")
 tslang_add_test(NAME test-compile-00-for-condition-narrowing COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00for_condition_narrowing.ts")
+tslang_add_test(NAME test-compile-00-param-assigned-owned COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00param_assigned_owned.ts")
 tslang_add_test(NAME test-compile-00-string-empty-falsy COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00string_empty_falsy.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")
@@ -873,6 +874,7 @@ tslang_add_test(NAME test-jit-00-for-await-sequential COMMAND test-runner -jit "
 tslang_add_test(NAME test-jit-00-for-await-control COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00for_await_control.ts")
 tslang_add_test(NAME test-jit-00-for-optional-class-condition COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00for_optional_class_condition.ts")
 tslang_add_test(NAME test-jit-00-for-condition-narrowing COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00for_condition_narrowing.ts")
+tslang_add_test(NAME test-jit-00-param-assigned-owned COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00param_assigned_owned.ts")
 tslang_add_test(NAME test-jit-00-string-empty-falsy COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00string_empty_falsy.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")
@@ -1718,6 +1720,7 @@ set(TSLANG_CORPUS
     00for_await_control.ts
     00for_optional_class_condition.ts
     00for_condition_narrowing.ts
+    00param_assigned_owned.ts
     00string_empty_falsy.ts
     00narrowed_assign_other_member.ts
     00const_record_owned_fields.ts

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

@@ -53,9 +53,8 @@ function typeGuard(v: string | number) {
 
 function instanceOf(x: A | B) {
     let n = 0;
-    // a local: reassigning the parameter itself in this loop crashes under -mm=rc, narrowed or not (#512)
-    for (let y = x; y instanceof A; y = new B(1)) {
-        n += y.onlyA();
+    for (; x instanceof A; x = new B(1)) {
+        n += x.onlyA();
     }
 
     return n;