Disable stack promotion for classes with isolated deinit

SwiftCompilerSources/Sources/AST/Declarations.swift

@@ -57,6 +57,10 @@ final public class StructDecl: NominalTypeDecl {
 
 final public class ClassDecl: NominalTypeDecl {
   public var superClass: Type? { Type(bridgedOrNil: bridged.Class_getSuperclass()) }
+
+  final public var destructor: DestructorDecl {
+    bridged.Class_getDestructor().getAs(DestructorDecl.self)
+  }
 }
 
 final public class ProtocolDecl: NominalTypeDecl {}
@@ -85,7 +89,9 @@ public class AbstractFunctionDecl: ValueDecl {}
 
 final public class ConstructorDecl: AbstractFunctionDecl {}
 
-final public class DestructorDecl: AbstractFunctionDecl {}
+final public class DestructorDecl: AbstractFunctionDecl {
+  final public var isIsolated: Bool { bridged.Destructor_isIsolated() }
+}
 
 public class FuncDecl: AbstractFunctionDecl {}
 

SwiftCompilerSources/Sources/Optimizer/FunctionPasses/StackPromotion.swift

@@ -10,6 +10,7 @@
 //
 //===----------------------------------------------------------------------===//
 
+import AST
 import SIL
 
 /// Promotes heap allocated objects to the stack.
@@ -74,6 +75,17 @@ private func tryPromoteAlloc(_ allocRef: AllocRefInstBase,
     return false
   }
 
+  if let dtor = (allocRef.type.nominal as? ClassDecl)?.destructor {
+    if dtor.isIsolated {
+      // Classes (including actors) with isolated deinit can escape implicitly.
+      //
+      // We could optimize this further and allow promotion if we can prove that
+      // deinit will take fast path (i.e. it will not schedule a job).
+      // But for now, let's keep things simple and disable promotion conservatively.
+      return false
+    }
+  }
+
   // The most important check: does the object escape the current function?
   if allocRef.isEscaping(context) {
     return false

include/swift/AST/ASTBridging.h

@@ -313,6 +313,8 @@ struct BridgedDeclObj {
   SWIFT_IMPORT_UNSAFE BRIDGED_INLINE OptionalBridgedDeclObj NominalType_getValueTypeDestructor() const;
   BRIDGED_INLINE bool Struct_hasUnreferenceableStorage() const;
   SWIFT_IMPORT_UNSAFE BRIDGED_INLINE BridgedASTType Class_getSuperclass() const;
+  SWIFT_IMPORT_UNSAFE BRIDGED_INLINE BridgedDeclObj Class_getDestructor() const;
+  BRIDGED_INLINE bool Destructor_isIsolated() const;
 };
 
 struct BridgedASTNode {

include/swift/AST/ASTBridgingImpl.h

@@ -142,6 +142,16 @@ BridgedASTType BridgedDeclObj::Class_getSuperclass() const {
   return {getAs<swift::ClassDecl>()->getSuperclass().getPointer()};
 }
 
+BridgedDeclObj BridgedDeclObj::Class_getDestructor() const {
+  return {getAs<swift::ClassDecl>()->getDestructor()};
+}
+
+bool BridgedDeclObj::Destructor_isIsolated() const {
+  auto dd = getAs<swift::DestructorDecl>();
+  auto ai = swift::getActorIsolation(dd);
+  return ai.isActorIsolated();
+}
+
 //===----------------------------------------------------------------------===//
 // MARK: BridgedASTNode
 //===----------------------------------------------------------------------===//

test/Concurrency/Runtime/async_task_locals_isolated_deinit.swift

@@ -126,14 +126,6 @@ class ClassNoOp: Probe {
 
 let tests = TestSuite("Isolated Deinit")
 
-// Dummy global variable to suppress stack propagation
-// TODO: Remove it after disabling allocation on stack for classes with isolated deinit
-var x: AnyObject? = nil
-func preventAllocationOnStack(_ object: AnyObject) {
-  x = object
-  x = nil
-}
-
 if #available(SwiftStdlib 5.1, *) {
   tests.test("class sync fast path") {
     let group = DispatchGroup()
@@ -142,7 +134,7 @@ if #available(SwiftStdlib 5.1, *) {
       // FIXME: isolated deinit should be clearing task locals
       await TL.$number.withValue(42) {
         await AnotherActor.shared.performTesting {
-          preventAllocationOnStack(ClassNoOp(expectedNumber: 0, group: group))
+          _ = ClassNoOp(expectedNumber: 0, group: group)
         }
       }
     }
@@ -154,7 +146,7 @@ if #available(SwiftStdlib 5.1, *) {
     group.enter(1)
     Task {
       TL.$number.withValue(99) {
-        preventAllocationOnStack(ClassNoOp(expectedNumber: 0, group: group))
+        _ = ClassNoOp(expectedNumber: 0, group: group)
       }
     }
     group.wait()
@@ -168,7 +160,7 @@ if #available(SwiftStdlib 5.1, *) {
       TL.$number.withValue(99) {
         // Despite last release happening not on the actor itself,
         // this is still a fast path due to optimisation for deallocating actors.
-        preventAllocationOnStack(ActorNoOp(expectedNumber: 0, group: group))
+        _ = ActorNoOp(expectedNumber: 0, group: group)
       }
     }
     group.wait()
@@ -180,7 +172,7 @@ if #available(SwiftStdlib 5.1, *) {
     Task {
       TL.$number.withValue(99) {
         // Using ProxyActor breaks optimization
-        preventAllocationOnStack(ProxyActor(expectedNumber: 0, group: group))
+        _ = ProxyActor(expectedNumber: 0, group: group)
       }
     }
     group.wait()
@@ -190,8 +182,8 @@ if #available(SwiftStdlib 5.1, *) {
     let group = DispatchGroup()
     group.enter(2)
     Task {
-      preventAllocationOnStack(ActorNoOp(expectedNumber: 0, group: group))
-      preventAllocationOnStack(ClassNoOp(expectedNumber: 0, group: group))
+      _ = ActorNoOp(expectedNumber: 0, group: group)
+      _ = ClassNoOp(expectedNumber: 0, group: group)
     }
     group.wait()
   }

test/Concurrency/voucher_propagation.swift

@@ -266,14 +266,6 @@ func adopt(voucher: voucher_t?) {
   os_release(voucher_adopt(os_retain(voucher)))
 }
 
-// Dummy global variable to suppress stack propagation
-// TODO: Remove it after disabling allocation on stack for classes with isolated deinit
-var x: AnyObject? = nil
-func preventAllocationOnStack(_ object: AnyObject) {
-  x = object
-  x = nil
-}
-
 let tests = TestSuite("Voucher Propagation")
 
 if #available(SwiftStdlib 5.1, *) {
@@ -444,15 +436,15 @@ if #available(SwiftStdlib 5.1, *) {
       Task {
         await AnotherActor.shared.performTesting {
           adopt(voucher: v1)
-          preventAllocationOnStack(ClassWithIsolatedDeinit(expectedVoucher: v1, group: group))
+          _ = ClassWithIsolatedDeinit(expectedVoucher: v1, group: group)
         }
         await AnotherActor.shared.performTesting {
           adopt(voucher: v2)
-          preventAllocationOnStack(ActorWithSelfIsolatedDeinit(expectedVoucher: v2, group: group))
+          _ = ActorWithSelfIsolatedDeinit(expectedVoucher: v2, group: group)
         }
         await AnotherActor.shared.performTesting {
           adopt(voucher: v3)
-          preventAllocationOnStack(ActorWithDeinitIsolatedOnAnother(expectedVoucher: v3, group: group))
+          _ = ActorWithDeinitIsolatedOnAnother(expectedVoucher: v3, group: group)
         }
       }
       group.wait()
@@ -467,11 +459,11 @@ if #available(SwiftStdlib 5.1, *) {
       Task {
         do {
           adopt(voucher: v1)
-          preventAllocationOnStack(ActorWithDeinitIsolatedOnAnother(expectedVoucher: v1, group: group))
+          _ = ActorWithDeinitIsolatedOnAnother(expectedVoucher: v1, group: group)
         }
         do {
           adopt(voucher: v2)
-          preventAllocationOnStack(ClassWithIsolatedDeinit(expectedVoucher: v2, group: group))
+          _ = ClassWithIsolatedDeinit(expectedVoucher: v2, group: group)
         }
       }
       group.wait()

test/SILOptimizer/stack_promotion_isolated_deinit.swift

@@ -0,0 +1,71 @@
+// RUN: %target-swift-frontend -parse-as-library -O -module-name=test %s -emit-sil -enable-experimental-feature IsolatedDeinit | %FileCheck %s
+// REQUIRES: swift_in_compiler
+
+@globalActor actor AnotherActor: GlobalActor {
+  static let shared = AnotherActor()
+}
+
+final class Inner {}
+
+final class IsolatedDeinit {
+    var inner: Inner?
+    @AnotherActor deinit {}
+}
+
+final class Container {
+    var ref: IsolatedDeinit?
+}
+
+// CHECK-LABEL: sil [noinline] {{.*}}@$s4test0A16ContainerOutsideyyF : $@convention(thin) () -> () {
+// CHECK: [[C:%.*]] = alloc_ref [bare] [stack] $Container
+// CHECK: [[ID:%.*]] = alloc_ref $IsolatedDeinit
+// CHECK: dealloc_stack_ref [[C]] : $Container
+// CHECK: return
+@inline(never)
+public func testContainerOutside() {
+    // container can be promoted
+    let container = Container()
+    let obj = IsolatedDeinit()
+    container.ref = obj
+}
+
+// CHECK-LABEL: sil [noinline] @$s4test0A15ContainerInsideyyF : $@convention(thin) () -> () {
+// CHECK: [[D:%.*]] = alloc_ref $IsolatedDeinit
+// CHECK: [[C:%.*]] = alloc_ref [bare] [stack] $Container
+// CHECK: dealloc_stack_ref [[C]] : $Container
+// CHECK: return
+@inline(never)
+public func testContainerInside() {
+    let obj = IsolatedDeinit()
+    // container can be promoted
+    let container = Container()
+    container.ref = obj
+}
+
+// CHECK-LABEL: sil [noinline] @$s4test0A12InnerOutsideyyF : $@convention(thin) () -> () {
+// CHECK: [[I:%.*]] = alloc_ref $Inner
+// CHECK: [[D:%.*]] = alloc_ref $IsolatedDeinit
+// CHECK: [[DI:%.*]] = end_init_let_ref [[D]] : $IsolatedDeinit
+// CHECK: strong_release [[DI]] : $IsolatedDeinit
+// CHECK: return
+@inline(never)
+public func testInnerOutside() {
+    // inner cannot be promoted, because it escapes to isolated deinit
+    let inner = Inner()
+    let obj = IsolatedDeinit()
+    obj.inner = inner
+}
+
+// CHECK-LABEL: sil [noinline] @$s4test0A11InnerInsideyyF : $@convention(thin) () -> () {
+// CHECK: [[D:%.*]] = alloc_ref $IsolatedDeinit
+// CHECK: [[DI:%.*]] = end_init_let_ref [[D]] : $IsolatedDeinit
+// CHECK: [[I:%.*]] = alloc_ref $Inner
+// CHECK: strong_release [[DI]] : $IsolatedDeinit
+// CHECK: return
+@inline(never)
+public func testInnerInside() {
+    let obj = IsolatedDeinit()
+    // inner cannot be promoted, because it escapes to isolated deinit
+    let inner = Inner()
+    obj.inner = inner
+}