Fix store_borrow generation in AddressLowering

Author: meg-gupta

lib/SILOptimizer/Mandatory/AddressLowering.cpp

@@ -1774,25 +1774,32 @@ void CallArgRewriter::rewriteIndirectArgument(Operand *operand) {
   // Allocate temporary storage for a loadable operand.
   AllocStackInst *allocInst =
       argBuilder.createAllocStack(callLoc, argValue->getType());
-
-  operand->set(allocInst);
-
   if (apply.getArgumentConvention(*operand).isOwnedConvention()) {
     argBuilder.createTrivialStoreOr(apply.getLoc(), argValue, allocInst,
                                     StoreOwnershipQualifier::Init);
     apply.insertAfterFullEvaluation([&](SILBuilder &callBuilder) {
       callBuilder.createDeallocStack(callLoc, allocInst);
     });
+    operand->set(allocInst);
   } else {
     auto borrow = argBuilder.emitBeginBorrowOperation(callLoc, argValue);
-    argBuilder.emitStoreBorrowOperation(callLoc, borrow, allocInst);
-
+    auto *store =
+        argBuilder.emitStoreBorrowOperation(callLoc, borrow, allocInst);
+    auto *storeBorrow = dyn_cast<StoreBorrowInst>(store);
     apply.insertAfterFullEvaluation([&](SILBuilder &callBuilder) {
+      if (storeBorrow) {
+        callBuilder.emitEndBorrowOperation(callLoc, storeBorrow);
+      }
       if (borrow != argValue) {
         callBuilder.emitEndBorrowOperation(callLoc, borrow);
       }
       callBuilder.createDeallocStack(callLoc, allocInst);
     });
+    if (storeBorrow) {
+      operand->set(storeBorrow);
+    } else {
+      operand->set(allocInst);
+    }
   }
 }
 

test/SILOptimizer/address_lowering.sil

@@ -262,7 +262,8 @@ bb0(%0 : @owned $T):
 // CHECK:   [[TMP:%.*]] = alloc_stack $C
 // CHECK:   [[B:%.*]] = begin_borrow [[LD]] : $C
 // CHECK:   [[SB:%.*]] = store_borrow [[B]] to [[TMP]] : $*C
-// CHECK:   apply %{{.*}}([[TMP]]) : $@convention(thin) (@in_guaranteed C) -> ()
+// CHECK:   apply %{{.*}}([[SB]]) : $@convention(thin) (@in_guaranteed C) -> ()
+// CHECK:   end_borrow [[SB]] : $*C
 // CHECK:   end_borrow [[B]] : $C
 // CHECK:   dealloc_stack [[TMP]] : $*C
 // CHECK:   destroy_value [[LD]] : $C