new C() of a local class allocates in place instead of calling C..new

tslang/include/TypeScript/Defines.h

@@ -25,10 +25,13 @@
 #define INTEGER_LITERAL_INFERRED_ATTR_NAME "__int_literal_inferred"
 
 // Marks an operation whose result already carries a reference the receiver is expected to take
-// over, rather than one it must retain for itself. `new C()` is the case that matters: it lowers
-// to a call of the generated `C..new`, and every function retains its result before returning
-// (§9.24), so the value arrives owned. A receiver that retained it again would be one owner
-// above the truth, which is exactly the leak §9.25 removes.
+// over, rather than one it must retain for itself. `new C()` is the case that matters. Of a class
+// from another module it lowers to a call of that class's generated `C..new`, and every function
+// retains its result before returning (§9.24), so the value arrives owned. Of a class defined in
+// the module it is the `ts.New` itself that carries this mark, and that is what makes the block
+// start with one reference under rc (NewOpLowering) - so a pass that put this mark on any other
+// `ts.New` would change what its block is born with. A receiver that retained the value again
+// would be one owner above the truth, which is exactly the leak §9.25 removes.
 //
 // Only set where the producer is known to retain - never inferred from a call being a call. A
 // runtime or builtin helper, or a function imported from a module built before that convention,

tslang/lib/TypeScript/LowerToLLVM.cpp

@@ -3380,6 +3380,19 @@ struct NewOpLowering : public TsLlvmPattern<mlir_ts::NewOp>
         else
         {
             value = ch.MemoryAlloc(storageType, MemoryAllocSet::Zero);
+
+            // `new C()` of a class defined here (MLIRGen's NewClassInstanceAsMethodCallOp): the
+            // block is handed over owned, as `C..new`'s retained return hands it, so under rc it
+            // starts with that one reference. A block is otherwise born with none, and whoever
+            // keeps it retains it. Under own the block starts as it did from `C..new`, whose
+            // retain ownership inference erased, and under gc or none there is no count.
+            if (tsLlvmContext->compileOptions.isRefCounted() && newOp->hasAttr(OWNED_RESULT_ATTR_NAME))
+            {
+                auto llvmIndexType = tch.convertType(th.getIndexType());
+                auto blockPtr = ch.getBlockPtrFromPayloadPtr(loc, value, llvmIndexType);
+                auto one = rewriter.create<LLVM::ConstantOp>(loc, llvmIndexType, rewriter.getIntegerAttr(llvmIndexType, 1));
+                rewriter.create<LLVM::StoreOp>(loc, one, blockPtr);
+            }
         }
 
         rewriter.replaceOp(newOp, ValueRange{value});

tslang/lib/TypeScript/MLIRGenAccessCall.cpp

@@ -2367,7 +2367,9 @@ namespace mlirgen
             resultType = getValueRefType(type);
         }
 
-        // if true, will call Class..new method, otheriwise ts::NewOp which we need to implement Class..new method 
+        // true: allocate as `new C()` does (NewClassInstanceAsMethodCallOp - in place for a class
+        // defined here, else a call of its `C..new`), then call the constructor; false: only the
+        // bare allocation, which is the body of `C..new` itself
         auto methodCallWay = !suppressConstructorCall;
 
         mlir::Value newOp;
@@ -2539,6 +2541,22 @@ namespace mlirgen
                                              const GenContext &genContext)
     {
 #ifdef USE_NEW_AS_METHOD
+        // A class defined in this module allocates in place, with the body `C..new` would run:
+        // the `ts.New` and the vtable store. The instance arrives owned, as from the call, whose
+        // return retained it (§9.24): the `ts.New` carries the mark, and the block it makes starts
+        // with that one reference (NewOpLowering) rather than gaining it through a `ts.Retain`
+        // the ownership passes would have to account for in the caller. `C..new` is still
+        // generated, for the vtable's `.new` slot and for other modules. A class from elsewhere -
+        // a declaration, an import, a dynamic import - goes through its `C..new`, which allocates
+        // in the module that knows its layout and vtable. A @cpp class keeps the call as well.
+        if (asMethodCall && !classInfo->isDeclaration && !classInfo->isImport && !classInfo->isDynamicImport &&
+            !classInfo->isCppVTable)
+        {
+            auto newOp = NewClassInstanceLogicAsOp(location, classInfo, false, genContext);
+            newOp.getDefiningOp()->setAttr(OWNED_RESULT_ATTR_NAME, builder.getUnitAttr());
+            return newOp;
+        }
+
         if (asMethodCall)
         {
             auto classRefVal = builder.create<mlir_ts::ClassRefOp>(

tslang/lib/TypeScript/MLIRGenVariables.cpp

@@ -200,7 +200,7 @@ namespace mlirgen
         varOp->setAttr(OWNED_LOCAL_ATTR_NAME, builder.getUnitAttr());
 
         // When the initializer already carries a reference for its receiver - `let x = new C()`,
-        // whose `C..new` retained the instance on the way out (§9.24) - taking another would put
+        // whose instance arrives owned (OWNED_RESULT_ATTR_NAME) - taking another would put
         // the slot one owner above the truth and nothing would ever be freed. Consume that
         // reference instead: no retain here, and the scope-exit release is what gives it back.
         // The pair stays balanced, so this cannot over-release; it is only the retain that moves.

tslang/lib/TypeScript/OwnedReturnConsumptionPass.cpp

@@ -169,8 +169,9 @@ class OwnedReturnConsumptionPass
                 return;
             }
 
-            // Already settled at the point the call was built - `new C()` is marked there,
-            // where the callee is known outright. Nothing left to remove.
+            // Already settled at the point the call was built - `new C()` of another module's
+            // class calls its `C..new`, and is marked there, where the callee is known outright.
+            // Nothing left to remove.
             if (callOp->hasAttr(OWNED_RESULT_ATTR_NAME))
             {
                 return;

tslang/lib/TypeScript/StackPromotionPass.cpp

@@ -1,5 +1,4 @@
 #include "mlir/Pass/Pass.h"
-#include "mlir/IR/IRMapping.h"
 #include "mlir/IR/Dominance.h"
 #include "mlir/Analysis/CFGLoopInfo.h"
 
@@ -33,9 +32,9 @@ constexpr llvm::StringLiteral NewMethodSuffix("..new");
 // Every other allocation stays on the heap, silently. Runs at the affine level, where control
 // flow is an ordinary CFG and every allocation and ownership op is still present.
 //
-// An allocation site is a `ts.New` (a boxed object literal, a class expression) or a call of a
-// class's `C..new` defined in this module, which is how `new C()` reaches the caller. A promoted
-// `C..new` call is replaced by its body: a `ts.New` on the stack plus the vtable store.
+// An allocation site is a `ts.New`: a boxed object literal, a class expression, or `new C()` of a
+// class defined in this module, which MLIRGen allocates in place rather than through `C..new`.
+// `new C()` of a class from elsewhere calls that module's `C..new`, an unknown callee: a heap site.
 class StackPromotionPass : public mlir::PassWrapper<StackPromotionPass, mlir::OperationPass<mlir::ModuleOp>>
 {
     CompileOptions compileOptions;
@@ -158,28 +157,7 @@ class StackPromotionPass : public mlir::PassWrapper<StackPromotionPass, mlir::Op
             return !op->hasAttr(STACK_ALLOC_ATTR_NAME);
         }
 
-        return static_cast<bool>(newMethodOf(op));
-    }
-
-    // The `C..new` a call allocates through, when it is one defined here with a body. A class
-    // imported as a declaration has none: an unknown callee, so a heap site.
-    mlir_ts::FuncOp newMethodOf(mlir::Operation *op)
-    {
-        auto callOp = dyn_cast<mlir_ts::SymbolCallInternalOp>(op);
-        if (!callOp || !callOp.getCallOperands().empty() || callOp->getNumResults() != 1 ||
-            !isa<mlir_ts::ClassType>(callOp->getResult(0).getType()) ||
-            !callOp.getCallee().ends_with(NewMethodSuffix))
-        {
-            return {};
-        }
-
-        auto funcOp = getOperation().lookupSymbol<mlir_ts::FuncOp>(callOp.getCalleeAttr());
-        if (!funcOp || funcOp.getBody().empty())
-        {
-            return {};
-        }
-
-        return funcOp;
+        return false;
     }
 
     // The storage the object's block holds, which is what goes on the stack.
@@ -422,11 +400,6 @@ class StackPromotionPass : public mlir::PassWrapper<StackPromotionPass, mlir::Op
             return "too-big";
         }
 
-        if (isa<mlir_ts::SymbolCallInternalOp>(site) && newMethodBody(newMethodOf(site)).empty())
-        {
-            return "unclassified";
-        }
-
         stack_escape::StackEscapeFacts::Walk walk;
         auto reason = facts->escapes(value, walk);
         if (!reason.empty())
@@ -534,56 +507,6 @@ class StackPromotionPass : public mlir::PassWrapper<StackPromotionPass, mlir::Op
         return *entry;
     }
 
-    // The ops of `C..new` that make the object: its `ts.New` and what follows it up to the point
-    // the object is handed to the return slot - the vtable store. Empty when the body is not of
-    // that shape, which keeps the site on the heap.
-    llvm::SmallVector<mlir::Operation *> newMethodBody(mlir_ts::FuncOp newMethod)
-    {
-        llvm::SmallVector<mlir::Operation *> ops;
-        mlir_ts::NewOp newOp;
-        for (auto &op : newMethod.getBody().front())
-        {
-            if (!newOp)
-            {
-                newOp = dyn_cast<mlir_ts::NewOp>(&op);
-                if (newOp)
-                {
-                    ops.push_back(&op);
-                }
-
-                continue;
-            }
-
-            if (isa<mlir_ts::RetainOp, mlir_ts::ReturnInternalOp>(&op))
-            {
-                return ops;
-            }
-
-            if (auto storeOp = dyn_cast<mlir_ts::StoreOp>(&op); storeOp && storeOp.getValue() == newOp.getResult())
-            {
-                return ops;
-            }
-
-            if (!isa<mlir_ts::PropertyRefOp, mlir_ts::AddressOfOp, mlir_ts::StoreOp, mlir_ts::ConstantOp>(&op))
-            {
-                return {};
-            }
-
-            // everything it reads is made in this run of ops: nothing to map from elsewhere
-            for (auto operand : op.getOperands())
-            {
-                if (!llvm::is_contained(ops, operand.getDefiningOp()))
-                {
-                    return {};
-                }
-            }
-
-            ops.push_back(&op);
-        }
-
-        return {};
-    }
-
     unsigned headerBytes()
     {
         return compileOptions.sizeBits() / 8;
@@ -753,46 +676,13 @@ class StackPromotionPass : public mlir::PassWrapper<StackPromotionPass, mlir::Op
         auto funcOp = site->getParentOfType<mlir_ts::FuncOp>();
         auto &entryBlock = funcOp.getBody().front();
 
-        mlir::Value stackObject;
+        auto newOp = cast<mlir_ts::NewOp>(site);
         mlir::OpBuilder builder(site);
-        if (auto newOp = dyn_cast<mlir_ts::NewOp>(site))
-        {
-            newOp.setStackAllocAttr(builder.getBoolAttr(true));
-            stackObject = newOp.getResult();
-            // where the allocation was - so before it moves, which the builder's position follows
-            builder.create<mlir_ts::ReleaseContentsOp>(site->getLoc(), stackObject, builder.getUnitAttr());
-            newOp->moveBefore(&entryBlock, entryBlock.begin());
-        }
-        else
-        {
-            // `new C()`: the body of `C..new` in place of the call, its `ts.New` on the stack
-            auto body = newMethodBody(newMethodOf(site));
-            mlir::IRMapping mapping;
-            for (auto *op : body)
-            {
-                if (auto newOp = dyn_cast<mlir_ts::NewOp>(op))
-                {
-                    mlir::OpBuilder entryBuilder(&entryBlock, entryBlock.begin());
-                    auto stackNewOp = entryBuilder.create<mlir_ts::NewOp>(site->getLoc(), newOp.getType(),
-                                                                          entryBuilder.getBoolAttr(true));
-                    // the call's ownership marks (`__owned_result`) describe the object it made
-                    for (auto attr : site->getDiscardableAttrs())
-                    {
-                        stackNewOp->setAttr(attr.getName(), attr.getValue());
-                    }
-
-                    stackObject = stackNewOp.getResult();
-                    mapping.map(newOp.getResult(), stackObject);
-                    builder.create<mlir_ts::ReleaseContentsOp>(site->getLoc(), stackObject, builder.getUnitAttr());
-                    continue;
-                }
-
-                builder.clone(*op, mapping)->setLoc(site->getLoc());
-            }
-
-            site->getResult(0).replaceAllUsesWith(stackObject);
-            site->erase();
-        }
+        newOp.setStackAllocAttr(builder.getBoolAttr(true));
+        mlir::Value stackObject = newOp.getResult();
+        // where the allocation was - so before it moves, which the builder's position follows
+        builder.create<mlir_ts::ReleaseContentsOp>(site->getLoc(), stackObject, builder.getUnitAttr());
+        newOp->moveBefore(&entryBlock, entryBlock.begin());
 
         if (!compileOptions.tracksOwnership())
         {