SIL: automatically set the reborrow flags for phi-arguments when creating an end_borrow or setting the operand of an end_borrow

This is done for all transitively incoming phi-arguments of an end_borrow operand.

Author: eeckstein

include/swift/SIL/OwnershipUtils.h

@@ -1404,6 +1404,10 @@ bool isNestedLexicalBeginBorrow(BeginBorrowInst *bbi);
 /// then the move_value is redundant.
 bool isRedundantMoveValue(MoveValueInst *mvi);
 
+/// Sets the reborrow flags for all transitively incoming phi-arguments of
+/// `forEndBorrowValue`, which is the operand value of an `end_borrow`.
+void updateReborrowFlags(SILValue forEndBorrowValue);
+
 } // namespace swift
 
 #endif

include/swift/SIL/SILBuilder.h

@@ -941,14 +941,7 @@ class SILBuilder {
     return lowering.emitStore(*this, Loc, Src, DestAddr, Qualifier);
   }
 
-  EndBorrowInst *createEndBorrow(SILLocation loc, SILValue borrowedValue) {
-    ASSERT(!SILArgument::isTerminatorResult(borrowedValue) &&
-               "terminator results do not have end_borrow");
-    ASSERT(!isa<SILFunctionArgument>(borrowedValue) &&
-           "Function arguments should never have an end_borrow");
-    return insert(new (getModule())
-                      EndBorrowInst(getSILDebugLocation(loc), borrowedValue));
-  }
+  EndBorrowInst *createEndBorrow(SILLocation loc, SILValue borrowedValue);
 
   BeginAccessInst *createBeginAccess(SILLocation loc, SILValue address,
                                      SILAccessKind accessKind,

include/swift/SIL/SILValue.h

@@ -1073,9 +1073,11 @@ class Operand {
     removeFromCurrent();
     TheValue = newValue;
     insertIntoCurrent();
+    updateReborrowFlags();
     verify();
   }
 
+  void updateReborrowFlags();
   void verify() const;
 
   /// Swap the given operand with the current one.

lib/SIL/IR/SILBuilder.cpp

@@ -13,6 +13,7 @@
 #include "swift/SIL/SILBuilder.h"
 #include "swift/AST/Expr.h"
 #include "swift/Basic/Assertions.h"
+#include "swift/SIL/OwnershipUtils.h"
 #include "swift/SIL/Projection.h"
 #include "swift/SIL/SILGlobalVariable.h"
 
@@ -681,6 +682,17 @@ void SILBuilder::emitScopedBorrowOperation(SILLocation loc, SILValue original,
     createEndBorrow(loc, value);
 }
 
+EndBorrowInst *SILBuilder::createEndBorrow(SILLocation loc, SILValue borrowedValue) {
+  ASSERT(!SILArgument::isTerminatorResult(borrowedValue) &&
+             "terminator results do not have end_borrow");
+  ASSERT(!isa<SILFunctionArgument>(borrowedValue) &&
+         "Function arguments should never have an end_borrow");
+  updateReborrowFlags(borrowedValue);
+  return insert(new (getModule())
+                    EndBorrowInst(getSILDebugLocation(loc), borrowedValue));
+}
+
+
 SILPhiArgument *SILBuilder::createSwitchOptional(
                                 SILLocation loc, SILValue operand,
                                 SILBasicBlock *someBB, SILBasicBlock *noneBB,

lib/SIL/IR/SILValue.cpp

@@ -441,6 +441,12 @@ llvm::raw_ostream &swift::operator<<(llvm::raw_ostream &os,
 //                                  Operand
 //===----------------------------------------------------------------------===//
 
+void Operand::updateReborrowFlags() {
+  if (isa<EndBorrowInst>(getUser())) {
+    swift::updateReborrowFlags(get());
+  }
+}
+
 void Operand::verify() const {
   if (isa<BorrowedFromInst>(getUser()) && getOperandNumber() == 0) {
     assert(isa<SILArgument>(get()) || isa<SILUndef>(get()));

lib/SIL/Utils/OwnershipUtils.cpp

@@ -2422,3 +2422,20 @@ bool swift::isRedundantMoveValue(MoveValueInst *mvi) {
   auto moveHasEscape = findPointerEscape(mvi);
   return moveHasEscape == originalHasEscape;
 }
+
+void swift::updateReborrowFlags(SILValue forEndBorrowValue) {
+  ValueWorklist worklist(forEndBorrowValue);
+  while (SILValue v = worklist.pop()) {
+    if (auto *bfi = dyn_cast<BorrowedFromInst>(v)) {
+      v = bfi->getBorrowedValue();
+    }
+    if (auto *arg = dyn_cast<SILPhiArgument>(v)) {
+      if (arg->isPhi() && !arg->isReborrow()) {
+        arg->setReborrow(true);
+        for (auto *predBlock : arg->getParent()->getPredecessorBlocks()) {
+          worklist.pushIfNotVisited(arg->getIncomingPhiValue(predBlock));
+        }
+      }
+    }
+  }
+}