[SILGen] Optimize generated dealloc for linearly recursive data structures

Adds detection of linearly recursive data structures by finding stored properties that share the type of the class the dealloc is being generated for. Each link will then be deallocated in a loop, while ensuring to keep the next link alive to prevent the recursion. This prevents stack overflows for long chains while also improving performance.

rdar://89162954

Author: drexin

lib/SILGen/SILGenDestructor.cpp

@@ -14,6 +14,7 @@
 #include "SILGenFunctionBuilder.h"
 #include "RValue.h"
 #include "ArgumentScope.h"
+#include "llvm/ADT/SmallSet.h"
 #include "swift/AST/GenericSignature.h"
 #include "swift/AST/SubstitutionMap.h"
 #include "swift/SIL/TypeLowering.h"
@@ -219,6 +220,140 @@ void SILGenFunction::destroyClassMember(SILLocation cleanupLoc,
   }
 }
 
+llvm::SmallSetVector<VarDecl*, 4> findRecursiveLinks(DeclContext* DC, ClassDecl *cd) {
+  auto SelfTy = DC->mapTypeIntoContext(cd->getDeclaredType());
+
+  // Collect all stored properties that would form a recursive structure,
+  // so we can remove the recursion and prevent the call stack from
+  // overflowing.
+  llvm::SmallSetVector<VarDecl*, 4> recursiveLinks;
+  for (VarDecl *vd : cd->getStoredProperties()) {
+    auto Ty = vd->getType()->getOptionalObjectType();
+    if (Ty && Ty->getCanonicalType() == SelfTy->getCanonicalType()) {
+      recursiveLinks.insert(vd);
+    }
+  }
+
+  // NOTE: Right now we only optimize linear recursion, so if there is more than one link,
+  // clear out the set and don't perform any recursion optimization.
+  if (recursiveLinks.size() < 1) {
+    recursiveLinks.clear();
+  }
+
+  return recursiveLinks;
+}
+
+
+void SILGenFunction::emitRecursiveChainDestruction(ManagedValue selfValue,
+                                                   ClassDecl *cd,
+                                                   SmallSetVector<VarDecl*, 4> recursiveLinks,
+                                                   CleanupLocation cleanupLoc) {
+  auto SelfTy = F.mapTypeIntoContext(cd->getDeclaredType());
+
+  assert(recursiveLinks.size() <= 1 && "Only linear recursion supported.");
+
+  auto SelfTyLowered = getTypeLowering(SelfTy).getLoweredType();
+  for (VarDecl* vd : recursiveLinks) {
+    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(vd->getType()).getLoweredType();
+    auto optionalNone = B.createOptionalNone(cleanupLoc, Ty);
+    SILValue varAddr =
+      B.createRefElementAddr(cleanupLoc, selfValue.getValue(), vd,
+                             Ty.getAddressType());
+    auto iterAddr = B.createAllocStack(cleanupLoc, Ty);
+    SILValue addr = B.createBeginAccess(
+      cleanupLoc, varAddr, SILAccessKind::Modify, SILAccessEnforcement::Static,
+      false /*noNestedConflict*/, false /*fromBuiltin*/);
+    SILValue iter = B.createLoad(cleanupLoc, addr, LoadOwnershipQualifier::Copy);
+    B.createStore(cleanupLoc, optionalNone, addr, StoreOwnershipQualifier::Assign);
+    B.createEndAccess(cleanupLoc, addr, false /*is aborting*/);
+    B.createStore(cleanupLoc, iter, iterAddr, StoreOwnershipQualifier::Init);
+
+    B.createBranch(cleanupLoc, loopBB);
+
+    // while iter != nil {
+    B.emitBlock(loopBB);
+    SILValue operand = B.createLoad(cleanupLoc, iterAddr, LoadOwnershipQualifier::Copy);
+    auto operandCopy = B.createCopyValue(cleanupLoc, operand);
+    auto operandAddr = B.createAllocStack(cleanupLoc, Ty);
+    B.createStore(cleanupLoc, operandCopy, operandAddr, StoreOwnershipQualifier::Init);
+    B.createDestroyValue(cleanupLoc, operand);
+    B.createSwitchEnumAddr(
+      cleanupLoc, operandAddr, nullptr,
+      {{getASTContext().getOptionalSomeDecl(), someBB},
+       {std::make_pair(getASTContext().getOptionalNoneDecl(), noneBB)}});
+
+    // if isKnownUniquelyReferenced(&iter) {
+    B.emitBlock(someBB);
+    B.createDestroyAddr(cleanupLoc, operandAddr);
+    B.createDeallocStack(cleanupLoc, operandAddr);
+    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);
+
+    // NOTE: We increment the ref count of the tail instead of unlinking it,
+    //       because custom deinit implementations of subclasses may access
+    //       it and it would be semantically wrong to unset it before that.
+    //       Making the tail non-uniquely referenced prevents the recursion.
+
+    // let tail = iter.unsafelyUnwrapped.next
+    // iter = tail
+    SILValue _iter = B.createLoad(cleanupLoc, iterAddr, LoadOwnershipQualifier::Copy);
+    auto iterBorrow = B.createBeginBorrow(cleanupLoc, _iter);
+    auto iterBorrowAddr = B.createAllocStack(cleanupLoc, Ty);
+    B.createStoreBorrow(cleanupLoc, iterBorrow, iterBorrowAddr);
+    auto xx = B.createLoadBorrow(cleanupLoc, iterBorrowAddr);
+    auto* link = B.createUncheckedEnumData(cleanupLoc,
+                                           xx,
+                                           getASTContext().getOptionalSomeDecl(),
+                                           SelfTyLowered);
+    varAddr = B.createRefElementAddr(cleanupLoc,
+                                     link,
+                                     vd,
+                                     Ty.getAddressType());
+
+    addr = B.createBeginAccess(
+      cleanupLoc, varAddr, SILAccessKind::Read, SILAccessEnforcement::Static,
+      false /* noNestedConflict */, false /*fromBuiltin*/);
+    iter = B.createLoad(cleanupLoc, addr, LoadOwnershipQualifier::Copy);
+    B.createEndAccess(cleanupLoc, addr, false /*is aborting*/);
+    B.createStore(cleanupLoc, iter, iterAddr, StoreOwnershipQualifier::Assign);
+
+    B.createEndBorrow(cleanupLoc, xx);
+    B.createEndBorrow(cleanupLoc, iterBorrow);
+
+    B.createDestroyValue(cleanupLoc, _iter);
+    B.createDeallocStack(cleanupLoc, iterBorrowAddr);
+
+    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.createDeallocStack(cleanupLoc, operandAddr);
+    B.createBranch(cleanupLoc, cleanBB);
+
+
+    B.emitBlock(cleanBB);
+    B.createDestroyAddr(cleanupLoc, iterAddr);
+    B.createDeallocStack(cleanupLoc, iterAddr);
+  }
+}
+
 void SILGenFunction::emitClassMemberDestruction(ManagedValue selfValue,
                                                 ClassDecl *cd,
                                                 CleanupLocation cleanupLoc) {
@@ -239,21 +374,28 @@ void SILGenFunction::emitClassMemberDestruction(ManagedValue selfValue,
   /// 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.
   if (cd->isDistributedActor()) {
-    finishBB = createBasicBlock();
     normalMemberDestroyBB = createBasicBlock();
-
+    finishBB = createBasicBlock();
     emitDistributedActorClassMemberDestruction(cleanupLoc, selfValue, cd,
                                                normalMemberDestroyBB,
                                                finishBB);
   }
 
+  auto recursiveLinks = findRecursiveLinks(F.getDeclContext(), cd);
+
   /// 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())
+      emitRecursiveChainDestruction(selfValue, cd, recursiveLinks, cleanupLoc);
 
     if (finishBB)
       B.createBranch(cleanupLoc, finishBB);

lib/SILGen/SILGenFunction.h

@@ -680,6 +680,18 @@ class LLVM_LIBRARY_VISIBILITY SILGenFunction
   void emitClassMemberDestruction(ManagedValue selfValue, ClassDecl *cd,
                                   CleanupLocation cleanupLoc);
 
+  /// Generates code to destroy recursive data structures, without building
+  /// up the call stack.
+  ///
+  /// \param selfValue The 'self' value.
+  /// \param cd The class declaration whose members are being destroyed.
+  /// \param recursiveLinks The set of stored properties that form the
+  ///        recursive data structure.
+  void emitRecursiveChainDestruction(ManagedValue selfValue,
+                                ClassDecl *cd,
+                                llvm::SmallSetVector<VarDecl*, 4> recursiveLinks,
+                                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.