Merge pull request #41459 from drexin/wip-dealloc-recur

drexin · web-flow · commit 6d55ecd37dfc · 2022-03-14T08:12:48.000-07:00
[SILGen] Optimize generated dealloc for linearly recursive data struc…
diff --git a/lib/SILGen/SILGenDestructor.cpp b/lib/SILGen/SILGenDestructor.cpp
@@ -10,13 +10,15 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "ArgumentScope.h"
+#include "RValue.h"
 #include "SILGenFunction.h"
 #include "SILGenFunctionBuilder.h"
-#include "RValue.h"
-#include "ArgumentScope.h"
+#include "SwitchEnumBuilder.h"
 #include "swift/AST/GenericSignature.h"
 #include "swift/AST/SubstitutionMap.h"
 #include "swift/SIL/TypeLowering.h"
+#include "llvm/ADT/SmallSet.h"
 
 using namespace swift;
 using namespace Lowering;
@@ -219,6 +221,142 @@ void SILGenFunction::destroyClassMember(SILLocation cleanupLoc,
   }
 }
 
+/// Finds stored properties that have the same type as `cd` and thus form
+/// a recursive structure.
+///
+/// Example:
+///
+///   class Node<T> {
+///     let element: T
+///     let next: Node<T>?
+///   }
+///
+/// In the above example `next` is a recursive link and would be recognized
+/// by this function and added to the result set.
+static void findRecursiveLinks(ClassDecl *cd,
+                               llvm::SmallSetVector<VarDecl *, 4> &result) {
+  auto selfTy = cd->getDeclaredInterfaceType();
+
+  // Collect all stored properties that would form a recursive structure,
+  // so we can remove the recursion and prevent the call stack from
+  // overflowing.
+  for (VarDecl *vd : cd->getStoredProperties()) {
+    auto Ty = vd->getInterfaceType()->getOptionalObjectType();
+    if (Ty && Ty->getCanonicalType() == selfTy->getCanonicalType()) {
+      result.insert(vd);
+    }
+  }
+
+  // NOTE: Right now we only optimize linear recursion, so if there is more
+  // than one stored property of the same type, clear out the set and don't
+  // perform any recursion optimization.
+  if (result.size() > 1) {
+    result.clear();
+  }
+}
+
+void SILGenFunction::emitRecursiveChainDestruction(ManagedValue selfValue,
+                                                   ClassDecl *cd,
+                                                   VarDecl *recursiveLink,
+                                                   CleanupLocation cleanupLoc) {
+  auto selfTy = F.mapTypeIntoContext(cd->getDeclaredInterfaceType());
+
+  auto selfTyLowered = getTypeLowering(selfTy).getLoweredType();
+
+  SILBasicBlock *cleanBB = createBasicBlock();
+  SILBasicBlock *noneBB = createBasicBlock();
+  SILBasicBlock *notUniqueBB = createBasicBlock();
+  SILBasicBlock *uniqueBB = createBasicBlock();
+  SILBasicBlock *someBB = createBasicBlock();
+  SILBasicBlock *loopBB = createBasicBlock();
+
+  // var iter = self.link
+  // self.link = nil
+  auto Ty = getTypeLowering(F.mapTypeIntoContext(recursiveLink->getInterfaceType())).getLoweredType();
+  auto optionalNone = B.createOptionalNone(cleanupLoc, Ty);
+  SILValue varAddr =
+    B.createRefElementAddr(cleanupLoc, selfValue.getValue(), recursiveLink,
+                           Ty.getAddressType());
+  auto *iterAddr = B.createAllocStack(cleanupLoc, Ty);
+  SILValue addr = B.createBeginAccess(
+    cleanupLoc, varAddr, SILAccessKind::Modify, SILAccessEnforcement::Static,
+    true /*noNestedConflict*/, false /*fromBuiltin*/);
+  SILValue iter = B.createLoad(cleanupLoc, addr, LoadOwnershipQualifier::Take);
+  B.createStore(cleanupLoc, optionalNone, addr, StoreOwnershipQualifier::Init);
+  B.createEndAccess(cleanupLoc, addr, false /*is aborting*/);
+  B.createStore(cleanupLoc, iter, iterAddr, StoreOwnershipQualifier::Init);
+
+  B.createBranch(cleanupLoc, loopBB);
+
+  // while iter != nil {
+  {
+    B.emitBlock(loopBB);
+    auto iterBorrow =
+        ManagedValue::forUnmanaged(iterAddr).borrow(*this, cleanupLoc);
+    SwitchEnumBuilder switchBuilder(B, cleanupLoc, iterBorrow);
+    switchBuilder.addOptionalSomeCase(someBB);
+    switchBuilder.addOptionalNoneCase(noneBB);
+    std::move(switchBuilder).emit();
+  }
+
+  // if isKnownUniquelyReferenced(&iter) {
+  {
+    B.emitBlock(someBB);
+    auto isUnique = B.createIsUnique(cleanupLoc, iterAddr);
+    B.createCondBranch(cleanupLoc, isUnique, uniqueBB, notUniqueBB);
+  }
+
+  // we have a uniquely referenced link, so we need to deinit
+  {
+    B.emitBlock(uniqueBB);
+
+    // let tail = iter.unsafelyUnwrapped.next
+    // iter = tail
+    SILValue iterBorrow = B.createLoadBorrow(cleanupLoc, iterAddr);
+    auto *link = B.createUncheckedEnumData(
+        cleanupLoc, iterBorrow, getASTContext().getOptionalSomeDecl(),
+        selfTyLowered);
+
+    varAddr = B.createRefElementAddr(cleanupLoc, link, recursiveLink,
+                                     Ty.getAddressType());
+
+    addr = B.createBeginAccess(
+        cleanupLoc, varAddr, SILAccessKind::Read, SILAccessEnforcement::Static,
+        true /* noNestedConflict */, false /*fromBuiltin*/);
+
+    // The deinit of `iter` will decrement the ref count of the field
+    // containing the next element and potentially leading to its
+    // deinitialization, causing the recursion. The prevent that,
+    // we `load [copy]` here to ensure the object stays alive until
+    // we explicitly release it in the next step of the iteration.
+    iter = B.createLoad(cleanupLoc, addr, LoadOwnershipQualifier::Copy);
+    B.createEndAccess(cleanupLoc, addr, false /*is aborting*/);
+    B.createEndBorrow(cleanupLoc, iterBorrow);
+
+    B.createStore(cleanupLoc, iter, iterAddr, StoreOwnershipQualifier::Assign);
+
+    B.createBranch(cleanupLoc, loopBB);
+  }
+
+  // the next link in the chain is not unique, so we are done here
+  {
+    B.emitBlock(notUniqueBB);
+    B.createBranch(cleanupLoc, cleanBB);
+  }
+
+  // we reached the end of the chain
+  {
+    B.emitBlock(noneBB);
+    B.createBranch(cleanupLoc, cleanBB);
+  }
+
+  {
+    B.emitBlock(cleanBB);
+    B.createDestroyAddr(cleanupLoc, iterAddr);
+    B.createDeallocStack(cleanupLoc, iterAddr);
+  }
+}
+
 void SILGenFunction::emitClassMemberDestruction(ManagedValue selfValue,
                                                 ClassDecl *cd,
                                                 CleanupLocation cleanupLoc) {
@@ -231,10 +369,10 @@ void SILGenFunction::emitClassMemberDestruction(ManagedValue selfValue,
   /// For other cases, the basic blocks are not necessary and the destructor
   /// can just emit all the normal destruction code right into the current block.
   // If set, used as the basic block for the destroying of all members.
-  SILBasicBlock* normalMemberDestroyBB = nullptr;
+  SILBasicBlock *normalMemberDestroyBB = nullptr;
   // If set, used as the basic block after members have been destroyed,
   // and we're ready to perform final cleanups before returning.
-  SILBasicBlock* finishBB = nullptr;
+  SILBasicBlock *finishBB = nullptr;
 
   /// A distributed actor may be 'remote' in which case there is no need to
   /// destroy "all" members, because they never had storage to begin with.
@@ -247,13 +385,29 @@ void SILGenFunction::emitClassMemberDestruction(ManagedValue selfValue,
                                                finishBB);
   }
 
+  // Before we destroy all fields, we check if any of them are
+  // recursively the same type as `self`, so we can iteratively
+  // deinitialize them, to prevent deep recursion and potential
+  // stack overflows.
+
+  llvm::SmallSetVector<VarDecl *, 4> recursiveLinks;
+  findRecursiveLinks(cd, recursiveLinks);
+
   /// Destroy all members.
   {
     if (normalMemberDestroyBB)
       B.emitBlock(normalMemberDestroyBB);
 
-    for (VarDecl *vd : cd->getStoredProperties())
+    for (VarDecl *vd : cd->getStoredProperties()) {
+      if (recursiveLinks.contains(vd))
+        continue;
       destroyClassMember(cleanupLoc, selfValue, vd);
+    }
+
+    if (!recursiveLinks.empty()) {
+      assert(recursiveLinks.size() == 1 && "Only linear recursion supported.");
+      emitRecursiveChainDestruction(selfValue, cd, recursiveLinks[0], cleanupLoc);
+    }
 
     if (finishBB)
       B.createBranch(cleanupLoc, finishBB);
diff --git a/lib/SILGen/SILGenFunction.h b/lib/SILGen/SILGenFunction.h
@@ -680,6 +680,24 @@ class LLVM_LIBRARY_VISIBILITY SILGenFunction
   void emitClassMemberDestruction(ManagedValue selfValue, ClassDecl *cd,
                                   CleanupLocation cleanupLoc);
 
+  /// Generates code to destroy linearly recursive data structures, without
+  /// building up the call stack.
+  ///
+  /// E.x.: In the following we want to deinit next without recursing into next.
+  ///
+  /// class Node<A> {
+  ///   let value: A
+  ///   let next: Node<A>?
+  /// }
+  ///
+  /// \param selfValue The 'self' value.
+  /// \param cd The class declaration whose members are being destroyed.
+  /// \param recursiveLink The property that forms the recursive structure.
+  void emitRecursiveChainDestruction(ManagedValue selfValue,
+                                ClassDecl *cd,
+                                VarDecl* recursiveLink,
+                                CleanupLocation cleanupLoc);
+
   /// Generates a thunk from a foreign function to the native Swift convention.
   void emitForeignToNativeThunk(SILDeclRef thunk);
   /// Generates a thunk from a native function to foreign conventions.
diff --git a/test/Interpreter/deinit_recursive_no_overflow.swift b/test/Interpreter/deinit_recursive_no_overflow.swift
@@ -0,0 +1,20 @@
+// RUN: %target-run-simple-swift
+
+// REQUIRES: executable_test
+
+class Node {
+  var next: Node?
+}
+
+var first: Node? = nil
+for _ in 1...3_000_000 {
+    let next = Node()
+    next.next = first
+    first = next
+}
+
+@inline(never)
+func deallocList() {
+    first = nil
+}
+deallocList()
diff --git a/test/SILGen/deinit_recursive_branching.swift b/test/SILGen/deinit_recursive_branching.swift
@@ -0,0 +1,23 @@
+// RUN: %target-swift-emit-silgen %s | %FileCheck %s
+
+// Non-linearly recursive structures should not get optimized
+class Tree {
+  var left: Tree?
+  var right: Tree?
+}
+
+// CHECK: sil hidden [ossa] @$s26deinit_recursive_branching4TreeCfd : $@convention(method) (@guaranteed Tree) -> @owned Builtin.NativeObject {
+// CHECK: // [[SELF:%.*]] "self"
+// CHECK: bb0([[SELF]] : @guaranteed $Tree):
+// CHECK:   [[LEFT:%.*]] = ref_element_addr [[SELF]] : $Tree, #Tree.left
+// CHECK:   [[LEFT_ACCESS:%.*]] = begin_access [deinit] [static] [[LEFT]] : $*Optional<Tree>
+// CHECK:   destroy_addr [[LEFT_ACCESS]] : $*Optional<Tree>
+// CHECK:   end_access [[LEFT_ACCESS]] : $*Optional<Tree>
+// CHECK:   [[RIGHT:%.*]] = ref_element_addr [[SELF]] : $Tree, #Tree.right
+// CHECK:   [[RIGHT_ACCESS:%.*]] = begin_access [deinit] [static] [[RIGHT]] : $*Optional<Tree>
+// CHECK:   destroy_addr [[RIGHT_ACCESS]] : $*Optional<Tree>
+// CHECK:   end_access [[RIGHT_ACCESS]] : $*Optional<Tree>                // id: %9
+// CHECK:   [[SELF_NATIVE:%.*]] = unchecked_ref_cast [[SELF]] : $Tree to $Builtin.NativeObject
+// CHECK:   [[SELF_NATIVE_OWNED:%.*]] = unchecked_ownership_conversion [[SELF_NATIVE]] : $Builtin.NativeObject, @guaranteed to @owned
+// CHECK:   return [[SELF_NATIVE_OWNED]] : $Builtin.NativeObject
+// CHECK: } // end sil function '$s26deinit_recursive_branching4TreeCfd'
diff --git a/test/SILGen/deinit_recursive_linear.swift b/test/SILGen/deinit_recursive_linear.swift
@@ -0,0 +1,57 @@
+// RUN: %target-swift-emit-silgen %s | %FileCheck %s
+
+class Node {
+  var elem: [Int64] = []
+  var next: Node?
+}
+
+// CHECK: sil hidden [ossa] @$s23deinit_recursive_linear4NodeCfd : $@convention(method) (@guaranteed Node) -> @owned Builtin.NativeObject {
+// CHECK: [[SELF:%.*]] "self"
+// CHECK: bb0([[SELF]] : @guaranteed $Node):
+// CHECK:   [[ELEM:%.*]] = ref_element_addr [[SELF]] : $Node, #Node.elem
+// CHECK:   [[ELEM_ACCESS:%.*]] = begin_access [deinit] [static] [[ELEM]] : $*Array<Int64>
+// CHECK:   destroy_addr [[ELEM_ACCESS]] : $*Array<Int64>
+// CHECK:   end_access [[ELEM_ACCESS]] : $*Array<Int64>
+// CHECK:   [[NIL:%.*]] = enum $Optional<Node>, #Optional.none!enumelt
+// CHECK:   [[SELF_NEXT:%.*]] = ref_element_addr [[SELF]] : $Node, #Node.next
+// CHECK:   [[ITER:%.*]] = alloc_stack $Optional<Node>
+// CHECK:   [[SELF_NEXT_ACCESS:%.*]] = begin_access [modify] [static] [no_nested_conflict] [[SELF_NEXT]] : $*Optional<Node>
+// CHECK:   [[SELF_NEXT_COPY:%.*]] = load [take] [[SELF_NEXT_ACCESS]] : $*Optional<Node>
+// CHECK:   store [[NIL]] to [init] [[SELF_NEXT_ACCESS]] : $*Optional<Node>
+// CHECK:   end_access [[SELF_NEXT_ACCESS]] : $*Optional<Node>
+// CHECK:   store [[SELF_NEXT_COPY]] to [init] [[ITER]] : $*Optional<Node>
+// CHECK:   br [[LOOPBB:bb.*]] //
+
+// CHECK: [[LOOPBB]]:
+// CHECK:   [[ITER_COPY:%.*]] = load [copy] [[ITER]] : $*Optional<Node>
+// CHECK:   switch_enum [[ITER_COPY]] : $Optional<Node>, case #Optional.some!enumelt: [[IS_SOME_BB:bb.*]], case #Optional.none!enumelt: [[IS_NONE_BB:bb[0-9]+]]
+
+// CHECK: [[IS_SOME_BB]]([[NODE:%.*]] : @owned $Node):
+// CHECK:   destroy_value [[NODE]] : $Node
+// CHECK:   [[IS_UNIQUE:%.*]] = is_unique [[ITER]] : $*Optional<Node>
+// CHECK:   cond_br [[IS_UNIQUE]], [[IS_UNIQUE_BB:bb.*]], [[NOT_UNIQUE_BB:bb[0-9]*]]
+
+// CHECK: [[IS_UNIQUE_BB]]:
+// CHECK:   [[ITER_BORROW:%.*]] = load_borrow [[ITER]] : $*Optional<Node>
+// CHECK:   [[ITER_UNWRAPPED:%.*]] = unchecked_enum_data [[ITER_BORROW]] : $Optional<Node>, #Optional.some!enumelt
+// CHECK:   [[NEXT_ADDR:%.*]] = ref_element_addr [[ITER_UNWRAPPED]] : $Node, #Node.next
+// CHECK:   [[NEXT_ADDR_ACCESS:%.*]] = begin_access [read] [static] [no_nested_conflict] [[NEXT_ADDR]] : $*Optional<Node>
+// CHECK:   [[NEXT_COPY:%.*]] = load [copy] [[NEXT_ADDR_ACCESS]] : $*Optional<Node>
+// CHECK:   end_access [[NEXT_ADDR_ACCESS]] : $*Optional<Node>
+// CHECK:   end_borrow [[ITER_BORROW]] : $Optional<Node>
+// CHECK:   store [[NEXT_COPY]] to [assign] [[ITER]] : $*Optional<Node>
+// CHECK:   br [[LOOPBB]]
+
+// CHECK: [[NOT_UNIQUE_BB]]:
+// CHECK:   br bb6
+
+// CHECK: [[IS_NONE_BB]]:
+// CHECK:   br [[CLEAN_BB:bb[0-9]+]]
+
+// CHECK: [[CLEAN_BB]]:
+// CHECK:   destroy_addr [[ITER]] : $*Optional<Node>
+// CHECK:   dealloc_stack [[ITER]] : $*Optional<Node>
+// CHECK:   [[SELF_NATIVE:%.*]] = unchecked_ref_cast [[SELF]] : $Node to $Builtin.NativeObject
+// CHECK:   [[SELF_NATIVE_OWNED:%.*]] = unchecked_ownership_conversion [[SELF_NATIVE]] : $Builtin.NativeObject, @guaranteed to @owned
+// CHECK:   return [[SELF_NATIVE_OWNED]] : $Builtin.NativeObject
+// CHECK: } // end sil function '$s23deinit_recursive_linear4NodeCfd'
diff --git a/test/SILGen/deinit_recursive_linear_generic.swift b/test/SILGen/deinit_recursive_linear_generic.swift
