[ownership] Only allow BranchPropagatedUser to be constructed from Operands.

include/swift/SIL/ApplySite.h

@@ -130,8 +130,13 @@ class ApplySite {
     llvm_unreachable("covered switch");                                        \
   } while (0)
 
+  /// Return the callee operand as a value.
+  SILValue getCallee() const { return getCalleeOperand()->get(); }
+
   /// Return the callee operand.
-  SILValue getCallee() const { FOREACH_IMPL_RETURN(getCallee()); }
+  const Operand *getCalleeOperand() const {
+    FOREACH_IMPL_RETURN(getCalleeOperand());
+  }
 
   /// Return the callee value by looking through function conversions until we
   /// find a function_ref, partial_apply, or unrecognized callee value.

include/swift/SIL/BranchPropagatedUser.h

@@ -31,17 +31,28 @@ class BranchPropagatedUser {
   InnerTy user;
 
 public:
-  BranchPropagatedUser(SILInstruction *inst) : user(inst) {
-    assert(!isa<CondBranchInst>(inst));
-  }
-
-  BranchPropagatedUser(CondBranchInst *cbi) : user(cbi) {}
-
-  BranchPropagatedUser(CondBranchInst *cbi, unsigned successorIndex)
-      : user(cbi, successorIndex) {
-    assert(successorIndex == CondBranchInst::TrueIdx ||
-           successorIndex == CondBranchInst::FalseIdx);
+  BranchPropagatedUser(Operand *op) : user() {
+    auto *opUser = op->getUser();
+    auto *cbi = dyn_cast<CondBranchInst>(opUser);
+    if (!cbi) {
+      user = InnerTy(opUser, 0);
+      return;
+    }
+    unsigned operandIndex = op->getOperandNumber();
+    if (cbi->isConditionOperandIndex(operandIndex)) {
+      // TODO: Is this correct?
+      user = InnerTy(cbi, CondBranchInst::TrueIdx);
+      return;
+    }
+    bool isTrueOperand = cbi->isTrueOperandIndex(operandIndex);
+    if (isTrueOperand) {
+      user = InnerTy(cbi, CondBranchInst::TrueIdx);
+    } else {
+      user = InnerTy(cbi, CondBranchInst::FalseIdx);
+    }
   }
+  BranchPropagatedUser(const Operand *op)
+      : BranchPropagatedUser(const_cast<Operand *>(op)) {}
 
   BranchPropagatedUser(const BranchPropagatedUser &other) : user(other.user) {}
   BranchPropagatedUser &operator=(const BranchPropagatedUser &other) {
@@ -96,6 +107,19 @@ class BranchPropagatedUser {
     NumLowBitsAvailable =
         llvm::PointerLikeTypeTraits<InnerTy>::NumLowBitsAvailable
   };
+
+private:
+  BranchPropagatedUser(SILInstruction *inst) : user(inst) {
+    assert(!isa<CondBranchInst>(inst));
+  }
+
+  BranchPropagatedUser(CondBranchInst *cbi) : user(cbi) {}
+
+  BranchPropagatedUser(CondBranchInst *cbi, unsigned successorIndex)
+      : user(cbi, successorIndex) {
+    assert(successorIndex == CondBranchInst::TrueIdx ||
+           successorIndex == CondBranchInst::FalseIdx);
+  }
 };
 
 } // namespace swift

include/swift/SIL/OwnershipUtils.h

@@ -310,14 +310,13 @@ struct BorrowScopeIntroducingValue {
   /// called with a scope that is not local.
   ///
   /// The intention is that this method can be used instead of
-  /// BorrowScopeIntroducingValue::getLocalScopeEndingInstructions() to avoid
+  /// BorrowScopeIntroducingValue::getLocalScopeEndingUses() to avoid
   /// introducing an intermediate array when one needs to transform the
   /// instructions before storing them.
   ///
   /// NOTE: To determine if a scope is a local scope, call
   /// BorrowScopeIntoducingValue::isLocalScope().
-  void visitLocalScopeEndingInstructions(
-      function_ref<void(SILInstruction *)> visitor) const;
+  void visitLocalScopeEndingUses(function_ref<void(Operand *)> visitor) const;
 
   bool isLocalScope() const { return kind.isLocalScope(); }
 

include/swift/SIL/SILInstruction.h

@@ -1861,7 +1861,8 @@ class ApplyInstBase<Impl, Base, false> : public Base {
   /// The operand number of the first argument.
   static unsigned getArgumentOperandNumber() { return NumStaticOperands; }
 
-  SILValue getCallee() const { return getAllOperands()[Callee].get(); }
+  const Operand *getCalleeOperand() const { return &getAllOperands()[Callee]; }
+  SILValue getCallee() const { return getCalleeOperand()->get(); }
 
   /// Gets the origin of the callee by looking through function type conversions
   /// until we find a function_ref, partial_apply, or unrecognized value.
@@ -7212,7 +7213,10 @@ class CondBranchInst final
          ProfileCounter FalseBBCount, SILFunction &F);
 
 public:
-  SILValue getCondition() const { return getAllOperands()[ConditionIdx].get(); }
+  const Operand *getConditionOperand() const {
+    return &getAllOperands()[ConditionIdx];
+  }
+  SILValue getCondition() const { return getConditionOperand()->get(); }
   void setCondition(SILValue newCondition) {
     getAllOperands()[ConditionIdx].set(newCondition);
   }
@@ -7258,6 +7262,11 @@ class CondBranchInst final
     return getAllOperands().slice(NumFixedOpers + getNumTrueArgs());
   }
 
+  /// Returns true if \p op is mapped to the condition operand of the cond_br.
+  bool isConditionOperand(Operand *op) const {
+    return getConditionOperand() == op;
+  }
+
   bool isConditionOperandIndex(unsigned OpIndex) const {
     assert(OpIndex < getNumOperands() &&
            "OpIndex must be an index for an actual operand");

lib/SIL/OwnershipUtils.cpp

@@ -121,20 +121,24 @@ void BorrowScopeIntroducingValue::getLocalScopeEndingInstructions(
   llvm_unreachable("Covered switch isn't covered?!");
 }
 
-void BorrowScopeIntroducingValue::visitLocalScopeEndingInstructions(
-    function_ref<void(SILInstruction *)> visitor) const {
+void BorrowScopeIntroducingValue::visitLocalScopeEndingUses(
+    function_ref<void(Operand *)> visitor) const {
   assert(isLocalScope() && "Should only call this given a local scope");
   switch (kind) {
   case BorrowScopeIntroducerKind::SILFunctionArgument:
     llvm_unreachable("Should only call this with a local scope");
   case BorrowScopeIntroducerKind::BeginBorrow:
-    for (auto *inst : cast<BeginBorrowInst>(value)->getEndBorrows()) {
-      visitor(inst);
+    for (auto *use : cast<BeginBorrowInst>(value)->getUses()) {
+      if (isa<EndBorrowInst>(use->getUser())) {
+        visitor(use);
+      }
     }
     return;
   case BorrowScopeIntroducerKind::LoadBorrow:
-    for (auto *inst : cast<LoadBorrowInst>(value)->getEndBorrows()) {
-      visitor(inst);
+    for (auto *use : cast<LoadBorrowInst>(value)->getUses()) {
+      if (isa<EndBorrowInst>(use->getUser())) {
+        visitor(use);
+      }
     }
     return;
   }
@@ -203,8 +207,8 @@ bool BorrowScopeIntroducingValue::areInstructionsWithinScope(
     return true;
 
   // Otherwise, gather up our local scope ending instructions.
-  visitLocalScopeEndingInstructions(
-      [&scratchSpace](SILInstruction *i) { scratchSpace.emplace_back(i); });
+  visitLocalScopeEndingUses(
+      [&scratchSpace](Operand *op) { scratchSpace.emplace_back(op); });
 
   LinearLifetimeChecker checker(visitedBlocks, deadEndBlocks);
   return checker.validateLifetime(value, scratchSpace, instructions);

lib/SIL/SILOwnershipVerifier.cpp

@@ -93,19 +93,19 @@ class SILValueOwnershipChecker {
 
   /// The list of lifetime ending users that we found. Only valid if check is
   /// successful.
-  SmallVector<BranchPropagatedUser, 16> lifetimeEndingUsers;
+  SmallVector<Operand *, 16> lifetimeEndingUsers;
 
   /// The list of non lifetime ending users that we found. Only valid if check
   /// is successful.
-  SmallVector<BranchPropagatedUser, 16> regularUsers;
+  SmallVector<Operand *, 16> regularUsers;
 
   /// The list of implicit non lifetime ending users that we found. This
   /// consists of instructions like end_borrow that end a scoped lifetime. We
   /// must treat those as regular uses and ensure that our value is not
   /// destroyed while that sub-scope is valid.
   ///
   /// TODO: Rename to SubBorrowScopeUsers?
-  SmallVector<BranchPropagatedUser, 4> implicitRegularUsers;
+  SmallVector<Operand *, 4> implicitRegularUsers;
 
   /// The set of blocks that we have visited.
   SmallPtrSetImpl<SILBasicBlock *> &visitedBlocks;
@@ -133,22 +133,25 @@ class SILValueOwnershipChecker {
     if (!result.getValue())
       return false;
 
+    SmallVector<BranchPropagatedUser, 32> allLifetimeEndingUsers;
+    llvm::copy(lifetimeEndingUsers, std::back_inserter(allLifetimeEndingUsers));
     SmallVector<BranchPropagatedUser, 32> allRegularUsers;
     llvm::copy(regularUsers, std::back_inserter(allRegularUsers));
     llvm::copy(implicitRegularUsers, std::back_inserter(allRegularUsers));
+
     LinearLifetimeChecker checker(visitedBlocks, deadEndBlocks);
     auto linearLifetimeResult = checker.checkValue(
-        value, lifetimeEndingUsers, allRegularUsers, errorBehavior);
+        value, allLifetimeEndingUsers, allRegularUsers, errorBehavior);
     result = !linearLifetimeResult.getFoundError();
 
     return result.getValue();
   }
 
 private:
   bool checkUses();
-  bool gatherUsers(SmallVectorImpl<BranchPropagatedUser> &lifetimeEndingUsers,
-                   SmallVectorImpl<BranchPropagatedUser> &regularUsers,
-                   SmallVectorImpl<BranchPropagatedUser> &implicitRegularUsers);
+  bool gatherUsers(SmallVectorImpl<Operand *> &lifetimeEndingUsers,
+                   SmallVectorImpl<Operand *> &regularUsers,
+                   SmallVectorImpl<Operand *> &implicitRegularUsers);
 
   bool checkValueWithoutLifetimeEndingUses();
 
@@ -157,10 +160,10 @@ class SILValueOwnershipChecker {
 
   bool isGuaranteedFunctionArgWithLifetimeEndingUses(
       SILFunctionArgument *arg,
-      const SmallVectorImpl<BranchPropagatedUser> &lifetimeEndingUsers) const;
+      const SmallVectorImpl<Operand *> &lifetimeEndingUsers) const;
   bool isSubobjectProjectionWithLifetimeEndingUses(
       SILValue value,
-      const SmallVectorImpl<BranchPropagatedUser> &lifetimeEndingUsers) const;
+      const SmallVectorImpl<Operand *> &lifetimeEndingUsers) const;
 
   /// Depending on our initialization, either return false or call Func and
   /// throw an error.
@@ -181,9 +184,9 @@ class SILValueOwnershipChecker {
 } // end anonymous namespace
 
 bool SILValueOwnershipChecker::gatherUsers(
-    SmallVectorImpl<BranchPropagatedUser> &lifetimeEndingUsers,
-    SmallVectorImpl<BranchPropagatedUser> &nonLifetimeEndingUsers,
-    SmallVectorImpl<BranchPropagatedUser> &implicitRegularUsers) {
+    SmallVectorImpl<Operand *> &lifetimeEndingUsers,
+    SmallVectorImpl<Operand *> &nonLifetimeEndingUsers,
+    SmallVectorImpl<Operand *> &implicitRegularUsers) {
 
   // See if Value is guaranteed. If we are guaranteed and not forwarding, then
   // we need to look through subobject uses for more uses. Otherwise, if we are
@@ -198,19 +201,7 @@ bool SILValueOwnershipChecker::gatherUsers(
 
   // Then gather up our initial list of users.
   SmallVector<Operand *, 8> users;
-  std::copy(value->use_begin(), value->use_end(), std::back_inserter(users));
-
-  auto addCondBranchToList = [](SmallVectorImpl<BranchPropagatedUser> &list,
-                                CondBranchInst *cbi, unsigned operandIndex) {
-    if (cbi->isConditionOperandIndex(operandIndex)) {
-      list.emplace_back(cbi);
-      return;
-    }
-
-    bool isTrueOperand = cbi->isTrueOperandIndex(operandIndex);
-    list.emplace_back(cbi, isTrueOperand ? CondBranchInst::TrueIdx
-                                         : CondBranchInst::FalseIdx);
-  };
+  llvm::copy(value->getUses(), std::back_inserter(users));
 
   bool foundError = false;
   while (!users.empty()) {
@@ -267,19 +258,10 @@ bool SILValueOwnershipChecker::gatherUsers(
         opOwnershipKindMap.getLifetimeConstraint(ownershipKind);
     if (lifetimeConstraint == UseLifetimeConstraint::MustBeInvalidated) {
       LLVM_DEBUG(llvm::dbgs() << "        Lifetime Ending User: " << *user);
-      if (auto *cbi = dyn_cast<CondBranchInst>(user)) {
-        addCondBranchToList(lifetimeEndingUsers, cbi, op->getOperandNumber());
-      } else {
-        lifetimeEndingUsers.emplace_back(user);
-      }
+      lifetimeEndingUsers.push_back(op);
     } else {
       LLVM_DEBUG(llvm::dbgs() << "        Regular User: " << *user);
-      if (auto *cbi = dyn_cast<CondBranchInst>(user)) {
-        addCondBranchToList(nonLifetimeEndingUsers, cbi,
-                            op->getOperandNumber());
-      } else {
-        nonLifetimeEndingUsers.emplace_back(user);
-      }
+      nonLifetimeEndingUsers.push_back(op);
     }
 
     // If our base value is not guaranteed, we do not to try to visit
@@ -296,12 +278,14 @@ bool SILValueOwnershipChecker::gatherUsers(
         // For correctness reasons we use indices to make sure that we can
         // append to NonLifetimeEndingUsers without needing to deal with
         // iterator invalidation.
-        SmallVector<SILInstruction *, 4> endBorrowInsts;
         for (unsigned i : indices(nonLifetimeEndingUsers)) {
           if (auto *bbi = dyn_cast<BeginBorrowInst>(
-                  nonLifetimeEndingUsers[i].getInst())) {
-            llvm::copy(bbi->getEndBorrows(),
-                       std::back_inserter(implicitRegularUsers));
+                  nonLifetimeEndingUsers[i]->getUser())) {
+            for (auto *use : bbi->getUses()) {
+              if (isa<EndBorrowInst>(use->getUser())) {
+                implicitRegularUsers.push_back(use);
+              }
+            }
           }
         }
       }
@@ -381,8 +365,8 @@ bool SILValueOwnershipChecker::gatherUsers(
         // them to ensure that all of BBArg's uses are completely
         // enclosed within the end_borrow of this argument.
         for (auto *op : succArg->getUses()) {
-          if (auto *ebi = dyn_cast<EndBorrowInst>(op->getUser())) {
-            implicitRegularUsers.push_back(ebi);
+          if (isa<EndBorrowInst>(op->getUser())) {
+            implicitRegularUsers.push_back(op);
           }
         }
       }
@@ -493,33 +477,31 @@ bool SILValueOwnershipChecker::checkValueWithoutLifetimeEndingUses() {
 
 bool SILValueOwnershipChecker::isGuaranteedFunctionArgWithLifetimeEndingUses(
     SILFunctionArgument *arg,
-    const llvm::SmallVectorImpl<BranchPropagatedUser> &lifetimeEndingUsers)
-    const {
+    const llvm::SmallVectorImpl<Operand *> &lifetimeEndingUsers) const {
   if (arg->getOwnershipKind() != ValueOwnershipKind::Guaranteed)
     return true;
 
   return handleError([&] {
     llvm::errs() << "    Function: '" << arg->getFunction()->getName() << "'\n"
                  << "    Guaranteed function parameter with life ending uses!\n"
                  << "    Value: " << *arg;
-    for (const auto &user : lifetimeEndingUsers) {
-      llvm::errs() << "    Lifetime Ending User: " << *user;
+    for (const auto *use : lifetimeEndingUsers) {
+      llvm::errs() << "    Lifetime Ending User: " << *use->getUser();
     }
     llvm::errs() << '\n';
   });
 }
 
 bool SILValueOwnershipChecker::isSubobjectProjectionWithLifetimeEndingUses(
     SILValue value,
-    const llvm::SmallVectorImpl<BranchPropagatedUser> &lifetimeEndingUsers)
-    const {
+    const llvm::SmallVectorImpl<Operand *> &lifetimeEndingUsers) const {
   return handleError([&] {
     llvm::errs() << "    Function: '" << value->getFunction()->getName()
                  << "'\n"
                  << "    Subobject projection with life ending uses!\n"
                  << "    Value: " << *value;
-    for (const auto &user : lifetimeEndingUsers) {
-      llvm::errs() << "    Lifetime Ending User: " << *user;
+    for (const auto *use : lifetimeEndingUsers) {
+      llvm::errs() << "    Lifetime Ending User: " << *use->getUser();
     }
     llvm::errs() << '\n';
   });