[llvm] [Coroutines] fix coroutines + std::unique_ptr with async exceptions validation errors (PR #149691)

[Weibo He via llvm-commits]

================
@@ -973,6 +975,178 @@ static StructType *buildFrameType(Function &F, coro::Shape &Shape,
   return FrameTy;
 }
 
+// This function assumes that it is being called on basic block in reversed
+// post-order, meaning predecessors are visited before successors failing to do
+// so will cause VMap to be non-valid and will cause instructions to fail
+// mapping to their corresponding clones
+static void finalizeBasicBlockCloneAndTrackSuccessors(
+    BasicBlock *InitialBlock, BasicBlock *ClonedBlock, ValueToValueMapTy &VMap,
+    SmallSet<BasicBlock *, 20> &SuccessorBlocksSet) {
+  // This code will examine the basic block, fix issues caused by clones
+  //  for example - tailor cleanupret to the corresponding cleanuppad
+  //  it will use VMap to do so
+  // in addition, it will add the node successors to SuccessorBlocksSet
+
+  // Remap the instructions, VMap here aggregates instructions across
+  // multiple BasicBlocks, and we assume that traversal is reversed post-order,
+  // therefore successor blocks (for example instructions having funclet
+  // tags) will be mapped correctly to the new cloned cleanuppad
+  for (Instruction &ClonedBlockInstruction : *ClonedBlock) {
+    RemapInstruction(&ClonedBlockInstruction, VMap,
+                     RF_NoModuleLevelChanges | RF_IgnoreMissingLocals);
+  }
+
+  const auto &InitialBlockTerminator = InitialBlock->getTerminator();
+
+  // If it's cleanupret, find the correspondant cleanuppad (use the VMap to
+  // find it).
+  if (auto *InitialBlockTerminatorCleanupReturn =
+          dyn_cast<CleanupReturnInst>(InitialBlockTerminator)) {
+    // if none found do nothing
+    if (VMap.find(InitialBlockTerminatorCleanupReturn->getCleanupPad()) ==
+        VMap.end()) {
+      return;
+    }
+
+    auto *ClonedBlockTerminatorCleanupReturn =
+        cast<CleanupReturnInst>(ClonedBlock->getTerminator());
+
+    // Assuming reversed post-order traversal
+    llvm::Value *ClonedBlockCleanupPadValue =
+        VMap[InitialBlockTerminatorCleanupReturn->getCleanupPad()];
+    auto *ClonedBlockCleanupPad =
+        cast<CleanupPadInst>(ClonedBlockCleanupPadValue);
+    ClonedBlockTerminatorCleanupReturn->setCleanupPad(ClonedBlockCleanupPad);
+
+    // If it's a branch/invoke, keep track of its successors, we want to
+    // calculate dominance between CoroBegin and them also
+  } else if (auto *InitialBlockTerminatorBranch =
+                 dyn_cast<BranchInst>(InitialBlockTerminator)) {
+    for (unsigned int successorIdx = 0;
+         successorIdx < InitialBlockTerminatorBranch->getNumSuccessors();
+         ++successorIdx) {
+      SuccessorBlocksSet.insert(
+          InitialBlockTerminatorBranch->getSuccessor(successorIdx));
+    }
+  } else if (auto *InitialBlockTerminatorInvoke =
+                 dyn_cast<InvokeInst>(InitialBlockTerminator)) {
+    SuccessorBlocksSet.insert(InitialBlockTerminatorInvoke->getUnwindDest());
+  } else if (isa<ReturnInst>(InitialBlockTerminator)) {
+    // No action needed
+  } else {
+    InitialBlockTerminator->print(dbgs());
+    report_fatal_error("Terminator is not implemented in "
+                       "finalizeBasicBlockCloneAndTrackSuccessors");
+  }
+}
+
+// Dominance issue fixer for each predecessor satisfying predicate function
+void splitIfBasicBlockPredecessors(
+    BasicBlock *InitialBlock, BasicBlock *ReplacementBlock,
+    std::function<bool(BasicBlock *)> Predicate) {
+
+  SmallVector<BasicBlock *> InitialBlockPredecessors;
+
+  auto Predecessors = predecessors(InitialBlock);
+  std::copy_if(Predecessors.begin(), Predecessors.end(),
+               std::back_inserter(InitialBlockPredecessors), Predicate);
+
+  // Fixups on the predecessors terminator - point them to ReplacementBlock.
+  for (auto *InitialBlockPredecessor : InitialBlockPredecessors) {
+    auto *InitialBlockPredecessorTerminator =
+        InitialBlockPredecessor->getTerminator();
+    if (auto *InitialBlockPredecessorTerminatorInvoke =
+            dyn_cast<InvokeInst>(InitialBlockPredecessorTerminator)) {
+      if (InitialBlock ==
+          InitialBlockPredecessorTerminatorInvoke->getUnwindDest()) {
+        InitialBlockPredecessorTerminatorInvoke->setUnwindDest(
+            ReplacementBlock);
+      } else {
+        InitialBlockPredecessorTerminatorInvoke->setNormalDest(
+            ReplacementBlock);
+      }
+    } else if (auto *InitialBlockPredecessorTerminatorBranch =
+                   dyn_cast<BranchInst>(InitialBlockPredecessorTerminator)) {
+      for (unsigned int successorIdx = 0;
+           successorIdx <
+           InitialBlockPredecessorTerminatorBranch->getNumSuccessors();
+           ++successorIdx) {
+        if (InitialBlock ==
+            InitialBlockPredecessorTerminatorBranch->getSuccessor(
+                successorIdx)) {
+          InitialBlockPredecessorTerminatorBranch->setSuccessor(
+              successorIdx, ReplacementBlock);
+        }
+      }
+    } else if (auto *InitialBlockPredecessorTerminatorCleanupReturn =
+                   dyn_cast<CleanupReturnInst>(
+                       InitialBlockPredecessorTerminator)) {
+      InitialBlockPredecessorTerminatorCleanupReturn->setUnwindDest(
+          ReplacementBlock);
+    } else {
+      InitialBlockPredecessorTerminator->print(dbgs());
+      report_fatal_error(
+          "Terminator is not implemented in splitIfBasicBlockPredecessors");
+    }
+  }
+}
+
+// Fixer for the "Instruction does not dominate all uses!" bug
+// The fix consists of mapping problematic paths (where CoroBegin does not
+// dominate cleanup nodes) and duplicates them to 2 flows - the one that
+// insertSpills intended to create (using the spill) and another one, preserving
+// the logics of pre-splitting, which would be triggered if unwinding happened
+// before CoroBegin
+static void
+splitBasicBlocksNotDominatedByCoroBegin(const FrameDataInfo &FrameData,
+                                        const coro::Shape &Shape, Function *F,
+                                        DominatorTree &DT) {
+  ValueToValueMapTy VMap;
+  SmallSet<BasicBlock *, 20> SpillUserBlocksSet;
+
+  // Prepare the node set, logics will be run only on those nodes
+  for (const auto &E : FrameData.Spills) {
+    for (auto *U : E.second) {
+      auto CurrentBlock = U->getParent();
+      if (!DT.dominates(Shape.CoroBegin, CurrentBlock)) {
+        SpillUserBlocksSet.insert(CurrentBlock);
+      }
+    }
+  }
+
+  // Run is in reversed post order, to enforce visiting predecessors before
+  // successors
+  for (BasicBlock *CurrentBlock : ReversePostOrderTraversal<Function *>(F)) {
+    if (!SpillUserBlocksSet.contains(CurrentBlock)) {
+      continue;
+    }
+    SpillUserBlocksSet.erase(CurrentBlock);
+
+    // Preserve the current node. the duplicate will become the unspilled
+    // alternative
+    auto UnspilledAlternativeBlock =
+        CloneBasicBlock(CurrentBlock, VMap, ".unspilled_alternative", F);
----------------
NewSigma wrote:

Should we clear `VMap` before reuse it?

https://github.com/llvm/llvm-project/pull/149691