diff --git a/tslang/lib/TypeScript/MLIRGenFunctions.cpp b/tslang/lib/TypeScript/MLIRGenFunctions.cpp index 7ba92ea8b..5c5199bb0 100644 --- a/tslang/lib/TypeScript/MLIRGenFunctions.cpp +++ b/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()->expression; + } + + if (target == SyntaxKind::Identifier) + { + names.insert(MLIRHelper::getName(target.as())); + } + else if (target == SyntaxKind::ArrayLiteralExpression || target == SyntaxKind::ObjectLiteralExpression) + { + ts::VisitorAST targetNames([&](ts::Node node) { + if (node == SyntaxKind::Identifier) + { + names.insert(MLIRHelper::getName(node.as())); + } + }); + targetNames.visit(target); + } + }; + + ts::VisitorAST visitor([&](ts::Node node) { + if (node == SyntaxKind::BinaryExpression) + { + auto binaryExpression = node.as(); + if (ts::isAssignmentOperator(binaryExpression->operatorToken)) + { + addTargets(binaryExpression->left); + } + } + else if (node == SyntaxKind::ForOfStatement) + { + auto initializer = node.as()->initializer; + if (initializer != SyntaxKind::VariableDeclarationList) + { + addTargets(initializer); + } + } + else if (node == SyntaxKind::ForInStatement) + { + auto initializer = node.as()->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 ¶m : 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))) diff --git a/tslang/lib/TypeScript/MLIRGenImpl.h b/tslang/lib/TypeScript/MLIRGenImpl.h index 8e6b9693f..868806cff 100644 --- a/tslang/lib/TypeScript/MLIRGenImpl.h +++ b/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, diff --git a/tslang/test/tester/CMakeLists.txt b/tslang/test/tester/CMakeLists.txt index d0f0a113b..45df54e6e 100644 --- a/tslang/test/tester/CMakeLists.txt +++ b/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 diff --git a/tslang/test/tester/tests/00for_condition_narrowing.ts b/tslang/test/tester/tests/00for_condition_narrowing.ts index 03c0b0d7b..50f594e94 100644 --- a/tslang/test/tester/tests/00for_condition_narrowing.ts +++ b/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; diff --git a/tslang/test/tester/tests/00param_assigned_owned.ts b/tslang/test/tester/tests/00param_assigned_owned.ts new file mode 100644 index 000000000..ab1829837 --- /dev/null +++ b/tslang/test/tester/tests/00param_assigned_owned.ts @@ -0,0 +1,236 @@ +// -mm=rc: a parameter the body assigns owns what it is assigned (#512). The slot of a parameter is borrowed +// from the caller and took no reference, so `x = new B(1)` stored an instance the block that made it then +// released: read after that block - the next loop condition, a statement after an `if` - it was freed memory. +// `while (x instanceof A) { x = new B(1); }` crashed; the rest read garbage once the heap was reused. +class A { + constructor(public a: number) {} +} + +class B { + constructor(public b: number) {} +} + +// reuses freed blocks, so a read of one sees another value rather than the old one +function churn() { + let fill: B[] = []; + for (let i = 0; i < 200; i++) { + fill.push(new B(1000 + i)); + } + + return fill.length; +} + +function loopInstanceOf(x: A | B) { + let n = 0; + while (x instanceof A) { + x = new B(1); + n++; + } + + return n; +} + +function forInstanceOf(x: A | B) { + let n = 0; + for (; x instanceof A; x = new B(1)) { + n++; + } + + return n; +} + +function inIf(x: B, c: boolean) { + if (c) { + x = new B(7); + } + + churn(); + return x.b; +} + +function inLoop(x: B) { + for (let i = 0; i < 3; i++) { + x = new B(i); + } + + churn(); + return x.b; +} + +function captured(x: B) { + const g = () => x.b; + { + x = new B(9); + } + + churn(); + return g(); +} + +function inClosure(x: B) { + const set = () => { + x = new B(8); + }; + set(); + churn(); + return x.b; +} + +function optional(x?: B) { + if (!x) { + x = new B(7); + } + + churn(); + return x.b; +} + +function withDefault(x: B = new B(1)) { + { + x = new B(6); + } + + churn(); + return x.b; +} + +function returned(x: B) { + { + x = new B(5); + } + + return x; +} + +function destructured(x: B) { + { + [x] = [new B(4)]; + } + + churn(); + return x.b; +} + +function forOf(x: B) { + const arr = [new B(1), new B(3)]; + for (x of arr) { + } + + churn(); + return x.b; +} + +class M { + m(x: B) { + for (let i = 0; i < 2; i++) { + x = new B(10 + i); + } + + churn(); + return x.b; + } +} + +const arrow = (x: B) => { + { + x = new B(12); + } + + churn(); + return x.b; +}; + +function nullish(x: B | undefined) { + { + x ??= new B(13); + } + + churn(); + return x.b; +} + +function strings(s: string) { + for (let i = 0; i < 3; i++) { + s = s + i; + } + + churn(); + return s; +} + +function* generator(x: B) { + { + x = new B(1); + } + + churn(); + yield x.b; + { + x = new B(2); + } + + churn(); + yield x.b; +} + +async function asynchronous(x: B) { + { + x = new B(3); + } + + await 0; + churn(); + return x.b; +} + +function rest(...xs: B[]) { + { + xs = [new B(4)]; + } + + churn(); + return xs[0].b; +} + +function unassigned(x: B) { + churn(); + return x.b; +} + +async function main() { + assert(loopInstanceOf(new A(5)) == 1, "while on instanceof"); + assert(loopInstanceOf(new B(5)) == 0, "while on instanceof, not entered"); + assert(forInstanceOf(new A(5)) == 1, "for on instanceof"); + assert(inIf(new B(0), true) == 7, "assigned in an if"); + assert(inIf(new B(3), false) == 3, "not assigned"); + assert(inLoop(new B(0)) == 2, "assigned in a loop"); + assert(captured(new B(0)) == 9, "captured, assigned in a block"); + assert(inClosure(new B(0)) == 8, "assigned in a closure"); + assert(optional() == 7, "optional, assigned"); + assert(optional(new B(2)) == 2, "optional, given"); + assert(withDefault() == 6, "default, assigned"); + assert(withDefault(new B(0)) == 6, "default given, assigned"); + + const r = returned(new B(0)); + churn(); + assert(r.b == 5, "returned after assignment"); + + assert(destructured(new B(0)) == 4, "destructuring assignment"); + assert(forOf(new B(0)) == 3, "for...of into the parameter"); + assert(new M().m(new B(0)) == 11, "method"); + assert(arrow(new B(0)) == 12, "arrow function"); + assert(nullish(undefined) == 13, "??="); + assert(strings("s") == "s012", "string"); + assert(unassigned(new B(14)) == 14, "unassigned"); + + let yielded = 0; + for (const v of generator(new B(0))) { + yielded = yielded * 10 + v; + } + + assert(yielded == 12, "generator"); + assert((await asynchronous(new B(0))) == 3, "async function"); + assert(rest(new B(0)) == 4, "rest parameter"); + + print("done."); +}