Added getArgForOperand for BranchInst and CondBranchInst.

Allows a SIL pass to follow a def-use chain through phis.

Other terminators can also propagate values through block arguments, but they
always need special handling.

Author: atrick

include/swift/SIL/SILInstruction.h

@@ -62,6 +62,7 @@ class SILInstructionResultArray;
 class SILOpenedArchetypesState;
 class SILType;
 class SILArgument;
+class SILPHIArgument;
 class SILUndef;
 class Stmt;
 class StringLiteralExpr;
@@ -6835,6 +6836,11 @@ class BranchInst final
 
   unsigned getNumArgs() const { return getAllOperands().size(); }
   SILValue getArg(unsigned i) const { return getAllOperands()[i].get(); }
+
+  /// Return the SILPHIArgument for the given operand.
+  ///
+  /// See SILArgument.cpp.
+  const SILPHIArgument *getArgForOperand(const Operand *oper) const;
 };
 
 /// A conditional branch.
@@ -6976,6 +6982,15 @@ class CondBranchInst final
   SILValue getArgForDestBB(const SILBasicBlock *DestBB,
                            unsigned ArgIndex) const;
 
+  /// Return the SILPHIArgument from either the true or false destination for
+  /// the given operand.
+  ///
+  /// Returns nullptr for an operand with no block argument
+  /// (i.e the branch condition).
+  ///
+  /// See SILArgument.cpp.
+  const SILPHIArgument *getArgForOperand(const Operand *oper) const;
+
   void swapSuccessors();
 };
 

lib/SIL/SILArgument.cpp

@@ -188,6 +188,28 @@ SILValue SILPHIArgument::getIncomingValue(SILBasicBlock *BB) {
   return getIncomingValueForPred(Parent, BB, Index);
 }
 
+const SILPHIArgument *BranchInst::getArgForOperand(const Operand *oper) const {
+  assert(oper->getUser() == this);
+  return cast<SILPHIArgument>(
+      getDestBB()->getArgument(oper->getOperandNumber()));
+}
+
+const SILPHIArgument *
+CondBranchInst::getArgForOperand(const Operand *oper) const {
+  assert(oper->getUser() == this);
+
+  unsigned operIdx = oper->getOperandNumber();
+  if (isTrueOperandIndex(operIdx)) {
+    return cast<SILPHIArgument>(getTrueBB()->getArgument(
+        operIdx - getTrueOperands().front().getOperandNumber()));
+  }
+  if (isFalseOperandIndex(operIdx)) {
+    return cast<SILPHIArgument>(getFalseBB()->getArgument(
+        operIdx - getFalseOperands().front().getOperandNumber()));
+  }
+  return nullptr;
+}
+
 //===----------------------------------------------------------------------===//
 //                            SILFunctionArgument
 //===----------------------------------------------------------------------===//