Fix store_borrow generation in GenericCloner

Author: meg-gupta

include/swift/SILOptimizer/Utils/GenericCloner.h

@@ -41,6 +41,8 @@ class GenericCloner
       RemappedScopeCache;
 
   llvm::SmallVector<AllocStackInst *, 8> AllocStacks;
+  llvm::SmallVector<StoreBorrowInst *, 8> StoreBorrowsToCleanup;
+  llvm::SmallVector<TermInst *, 8> FunctionExits;
   AllocStackInst *ReturnValueAddr = nullptr;
 
 public:
@@ -92,6 +94,11 @@ class GenericCloner
     if (Callback)
       Callback(Orig, Cloned);
 
+    if (auto *termInst = dyn_cast<TermInst>(Cloned)) {
+      if (termInst->isFunctionExiting()) {
+        FunctionExits.push_back(termInst);
+      }
+    }
     SILClonerWithScopes<GenericCloner>::postProcess(Orig, Cloned);
   }
 

lib/SILOptimizer/Utils/GenericCloner.cpp

@@ -13,11 +13,13 @@
 #include "swift/SILOptimizer/Utils/GenericCloner.h"
 
 #include "swift/AST/Type.h"
+#include "swift/SIL/OwnershipUtils.h"
 #include "swift/SIL/SILBasicBlock.h"
 #include "swift/SIL/SILFunction.h"
 #include "swift/SIL/SILInstruction.h"
 #include "swift/SIL/SILModule.h"
 #include "swift/SIL/SILValue.h"
+#include "swift/SIL/ScopedAddressUtils.h"
 #include "swift/SILOptimizer/Utils/InstOptUtils.h"
 #include "swift/SILOptimizer/Utils/SILOptFunctionBuilder.h"
 #include "llvm/ADT/ArrayRef.h"
@@ -61,7 +63,6 @@ void GenericCloner::populateCloned() {
   assert(!ReturnValueAddr && "Stale cloner state.");
 
   SILFunction *Cloned = getCloned();
-
   // Create arguments for the entry block.
   SILBasicBlock *OrigEntryBB = &*Original.begin();
   SILBasicBlock *ClonedEntryBB = Cloned->createBasicBlock();
@@ -142,14 +143,19 @@ void GenericCloner::populateCloned() {
 
           // Store the new direct parameter to an alloc_stack.
           createAllocStack();
-          if (!NewArg->getArgumentConvention().isGuaranteedConvention()) {
+          SILValue addr;
+          if (NewArg->getArgumentConvention().isGuaranteedConvention() &&
+              NewArg->getFunction()->hasOwnership()) {
+            auto *sbi = getBuilder().createStoreBorrow(Loc, NewArg, ASI);
+            StoreBorrowsToCleanup.push_back(sbi);
+            addr = sbi;
+          } else {
             getBuilder().emitStoreValueOperation(Loc, NewArg, ASI,
                                                  StoreOwnershipQualifier::Init);
-          } else {
-            getBuilder().emitStoreBorrowOperation(Loc, NewArg, ASI);
+            addr = ASI;
           }
 
-          entryArgs.push_back(ASI);
+          entryArgs.push_back(addr);
           return true;
         }
       }
@@ -225,4 +231,25 @@ void GenericCloner::postFixUp(SILFunction *f) {
     getBuilder().setInsertionPoint(applyBlock);
     getBuilder().createUnreachable(apply->getLoc());
   }
+
+  SmallVector<SILBasicBlock *, 4> discoveredBlocks;
+  for (auto *sbi : StoreBorrowsToCleanup) {
+    ScopedAddressValue scopedAddress(sbi);
+    if (scopedAddress->getUses().empty()) {
+      scopedAddress.createScopeEnd(sbi->getNextInstruction()->getIterator(),
+                                   RegularLocation::getAutoGeneratedLocation());
+      continue;
+    }
+    discoveredBlocks.clear();
+    PrunedLiveness storeBorrowLiveness(&discoveredBlocks);
+    bool success = scopedAddress.computeLiveness(storeBorrowLiveness);
+    if (success) {
+      scopedAddress.endScopeAtLivenessBoundary(&storeBorrowLiveness);
+      continue;
+    }
+    for (auto *exit : FunctionExits) {
+      scopedAddress.createScopeEnd(exit->getIterator(),
+                                   RegularLocation::getAutoGeneratedLocation());
+    }
+  }
 }

test/SILOptimizer/specialize_default_witness_resilience_ossa.sil

@@ -38,9 +38,9 @@ public struct ConformingNonTrivialStruct : ResilientProtocol {
 // CHECK-LABEL: sil shared [ossa] @$s8defaultA4main26ConformingNonTrivialStructV_TB5
 // CHECK:       bb0(%0 : @guaranteed $ConformingNonTrivialStruct):
 // CHECK-NEXT:    [[TMP:%.*]] = alloc_stack $ConformingNonTrivialStruct
-// CHECK-NEXT:    store_borrow %0 to [[TMP]] : $*ConformingNonTrivialStruct
+// CHECK-NEXT:    [[SB:%.*]] = store_borrow %0 to [[TMP]] : $*ConformingNonTrivialStruct
 // CHECK:         [[FN:%.*]] = function_ref @$s8defaultB4main26ConformingNonTrivialStructV_TB5
-// CHECK-NEXT:    [[LOAD:%.*]] = load_borrow [[TMP]] : $*ConformingNonTrivialStruct
+// CHECK-NEXT:    [[LOAD:%.*]] = load_borrow [[SB]] : $*ConformingNonTrivialStruct
 // CHECK-NEXT:    [[RESULT:%.*]] = apply [[FN]]([[LOAD]])
 // CHECK:         dealloc_stack [[TMP]] : $*ConformingNonTrivialStruct
 // CHECK } // end sil function 's8defaultA4main16ConformingNonTrivialStructV_TB5'