CycleInfo: Fix splitCriticalEdge

llvm/include/llvm/ADT/GenericCycleImpl.h

@@ -184,6 +184,24 @@ void GenericCycleInfo<ContextT>::moveTopLevelCycleToNewParent(CycleT *NewParent,
       It.second = NewParent;
 }
 
+template <typename ContextT>
+void GenericCycleInfo<ContextT>::addBlockToCycle(BlockT *Block, CycleT *Cycle) {
+  // FixMe: Appending NewBlock is fine as a set of blocks in a cycle. When
+  // printing, cycle NewBlock is at the end of list but it should be in the
+  // middle to represent actual traversal of a cycle.
+  Cycle->appendBlock(Block);
+  BlockMap.try_emplace(Block, Cycle);
+
+  CycleT *ParentCycle = Cycle->getParentCycle();
+  while (ParentCycle) {
+    Cycle = ParentCycle;
+    Cycle->appendBlock(Block);
+    ParentCycle = Cycle->getParentCycle();
+  }
+
+  BlockMapTopLevel.try_emplace(Block, Cycle);
+}
+
 /// \brief Main function of the cycle info computations.
 template <typename ContextT>
 void GenericCycleInfoCompute<ContextT>::run(BlockT *EntryBlock) {
@@ -368,17 +386,11 @@ void GenericCycleInfo<ContextT>::splitCriticalEdge(BlockT *Pred, BlockT *Succ,
                                                    BlockT *NewBlock) {
   // Edge Pred-Succ is replaced by edges Pred-NewBlock and NewBlock-Succ, all
   // cycles that had blocks Pred and Succ also get NewBlock.
-  CycleT *Cycle = this->getCycle(Pred);
-  if (Cycle && Cycle->contains(Succ)) {
-    while (Cycle) {
-      // FixMe: Appending NewBlock is fine as a set of blocks in a cycle. When
-      // printing cycle NewBlock is at the end of list but it should be in the
-      // middle to represent actual traversal of a cycle.
-      Cycle->appendBlock(NewBlock);
-      BlockMap.try_emplace(NewBlock, Cycle);
-      Cycle = Cycle->getParentCycle();
-    }
-  }
+  CycleT *Cycle = getSmallestCommonCycle(getCycle(Pred), getCycle(Succ));
+  if (!Cycle)
+    return;
+
+  addBlockToCycle(NewBlock, Cycle);
   assert(validateTree());
 }
 
@@ -392,6 +404,35 @@ auto GenericCycleInfo<ContextT>::getCycle(const BlockT *Block) const
   return BlockMap.lookup(Block);
 }
 
+/// \brief Find the innermost cycle containing both given cycles.
+///
+/// \returns the innermost cycle containing both \p A and \p B
+///          or nullptr if there is no such cycle.
+template <typename ContextT>
+auto GenericCycleInfo<ContextT>::getSmallestCommonCycle(CycleT *A,
+                                                        CycleT *B) const
+    -> CycleT * {
+  if (!A || !B)
+    return nullptr;
+
+  // If cycles A and B have different depth replace them with parent cycle
+  // until they have the same depth.
+  while (A->getDepth() > B->getDepth())
+    A = A->getParentCycle();
+  while (B->getDepth() > A->getDepth())
+    B = B->getParentCycle();
+
+  // Cycles A and B are at same depth but may be disjoint, replace them with
+  // parent cycles until we find cycle that contains both or we run out of
+  // parent cycles.
+  while (A != B) {
+    A = A->getParentCycle();
+    B = B->getParentCycle();
+  }
+
+  return A;
+}
+
 /// \brief get the depth for the cycle which containing a given block.
 ///
 /// \returns the depth for the innermost cycle containing \p Block or 0 if it is

llvm/include/llvm/ADT/GenericCycleInfo.h

@@ -248,6 +248,12 @@ template <typename ContextT> class GenericCycleInfo {
   /// the subtree.
   void moveTopLevelCycleToNewParent(CycleT *NewParent, CycleT *Child);
 
+  /// Assumes that \p Cycle is the innermost cycle containing \p Block.
+  /// \p Block will be appended to \p Cycle and all of its parent cycles.
+  /// \p Block will be added to BlockMap with \p Cycle and
+  /// BlockMapTopLevel with \p Cycle's top level parent cycle.
+  void addBlockToCycle(BlockT *Block, CycleT *Cycle);
+
 public:
   GenericCycleInfo() = default;
   GenericCycleInfo(GenericCycleInfo &&) = default;
@@ -261,6 +267,7 @@ template <typename ContextT> class GenericCycleInfo {
   const ContextT &getSSAContext() const { return Context; }
 
   CycleT *getCycle(const BlockT *Block) const;
+  CycleT *getSmallestCommonCycle(CycleT *A, CycleT *B) const;
   unsigned getCycleDepth(const BlockT *Block) const;
   CycleT *getTopLevelParentCycle(BlockT *Block);