Merge pull request #34930 from gottesmm/pr-d981ca4c510954e4ca19818a4303454d4bda80e3

[semantic-arc] Prevent future pointer invalidation issues in the OwnedToGuaranteedPhiOp transform.

Author: gottesmm

lib/SILOptimizer/SemanticARC/Context.h

@@ -64,7 +64,44 @@ struct LLVM_LIBRARY_VISIBILITY Context {
   /// our LiveRange can not see through joined live ranges, we know that we
   /// should only be able to have a single owned value introducer for each
   /// consumed operand.
-  FrozenMultiMap<SILValue, Operand *> joinedOwnedIntroducerToConsumedOperands;
+  ///
+  /// NOTE: To work around potential invalidation of our consuming operands when
+  /// adding values to edges on the CFG, we store our Operands as a
+  /// SILBasicBlock and an operand number. We only add values to edges and never
+  /// remove/modify edges so the operand number should be safe.
+  struct ConsumingOperandState {
+    PointerUnion<SILBasicBlock *, SILInstruction *> parent;
+    unsigned operandNumber;
+
+    ConsumingOperandState() : parent(nullptr), operandNumber(UINT_MAX) {}
+
+    ConsumingOperandState(Operand *op)
+        : parent(), operandNumber(op->getOperandNumber()) {
+      if (auto *ti = dyn_cast<TermInst>(op->getUser())) {
+        parent = ti->getParent();
+      } else {
+        parent = op->getUser();
+      }
+    }
+
+    ConsumingOperandState(const ConsumingOperandState &other) :
+        parent(other.parent), operandNumber(other.operandNumber) {}
+
+    ConsumingOperandState &operator=(const ConsumingOperandState &other) {
+      parent = other.parent;
+      operandNumber = other.operandNumber;
+      return *this;
+    }
+
+    ~ConsumingOperandState() = default;
+
+    operator bool() const {
+      return bool(parent) && operandNumber != UINT_MAX;
+    }
+  };
+
+  FrozenMultiMap<SILValue, ConsumingOperandState>
+      joinedOwnedIntroducerToConsumedOperands;
 
   /// If set to true, then we should only run cheap optimizations that do not
   /// build up data structures or analyze code in depth.

lib/SILOptimizer/SemanticARC/CopyValueOpts.cpp

@@ -224,9 +224,9 @@ bool SemanticARCOptVisitor::performGuaranteedCopyValueOptimization(
     });
 
     if (canOptimizePhi) {
+      Context::ConsumingOperandState state(opPhi);
       opPhi.visitResults([&](SILValue value) {
-        ctx.joinedOwnedIntroducerToConsumedOperands.insert(value,
-                                                           opPhi.getOperand());
+        ctx.joinedOwnedIntroducerToConsumedOperands.insert(value, state);
         return true;
       });
     }

lib/SILOptimizer/SemanticARC/OwnedToGuaranteedPhiOpt.cpp

@@ -19,12 +19,18 @@
 #include "Context.h"
 #include "OwnershipPhiOperand.h"
 #include "Transforms.h"
+#include "swift/Basic/STLExtras.h"
 
 using namespace swift;
 using namespace swift::semanticarc;
 
+namespace {
+using ConsumingOperandState = Context::ConsumingOperandState;
+} // anonymous namespace
+
+template <typename OperandRangeTy>
 static bool canEliminatePhi(
-    Context::FrozenMultiMapRange optimizableIntroducerRange,
+    OperandRangeTy optimizableIntroducerRange,
     ArrayRef<OwnershipPhiOperand> incomingValueOperandList,
     SmallVectorImpl<OwnedValueIntroducer> &ownedValueIntroducerAccumulator) {
   for (auto incomingValueOperand : incomingValueOperandList) {
@@ -161,9 +167,19 @@ bool swift::semanticarc::tryConvertOwnedPhisToGuaranteedPhis(Context &ctx) {
     // eliminated if it was not for the given phi. If all of them are, we can
     // optimize!
     {
-      auto rawFoundOptimizableIntroducerArray = pair.second;
-      if (!canEliminatePhi(rawFoundOptimizableIntroducerArray,
-                           incomingValueOperandList, ownedValueIntroducers)) {
+      std::function<Operand *(const Context::ConsumingOperandState)> lambda =
+        [&](const Context::ConsumingOperandState &state) -> Operand * {
+            unsigned opNum = state.operandNumber;
+            if (state.parent.is<SILBasicBlock *>()) {
+              SILBasicBlock *block = state.parent.get<SILBasicBlock *>();
+              return &block->getTerminator()->getAllOperands()[opNum];
+            }
+            SILInstruction *inst = state.parent.get<SILInstruction *>();
+            return &inst->getAllOperands()[opNum];
+        };
+      auto operandsTransformed = makeTransformRange(pair.second, lambda);
+      if (!canEliminatePhi(operandsTransformed, incomingValueOperandList,
+                           ownedValueIntroducers)) {
         continue;
       }
     }

lib/SILOptimizer/SemanticARC/OwnershipPhiOperand.h

@@ -64,6 +64,9 @@ class LLVM_LIBRARY_VISIBILITY OwnershipPhiOperand {
     }
   }
 
+  operator const Operand *() const { return op; }
+  operator Operand *() { return op; }
+
   Operand *getOperand() const { return op; }
   SILValue getValue() const { return op->get(); }
   SILType getType() const { return op->get()->getType(); }