MachineBasicBlock::updateTerminator now requires an explicit layout successor.

[jyknight]

Summary

Previously, it tried to infer the correct destination block from the
successor list, but this is a rather tricky propspect, given the
existence of successors that occur mid-block (such as invoke, and
potentially in the future, callbr).

Instead, require the caller to pass in the expected fallthrough
successor explicitly. In most callers, the correct block is
immediately clear. But, in MachineBlockPlacement, we do need to record
the original ordering, before starting to reorder blocks.

[efriedma]

(such as invoke, and potentially in the future, callbr).

Is the endgame here that we relax the restriction that updateTerminator() requires an analyzable branch?

llvm/include/llvm/CodeGen/MachineBasicBlock.h

522

fallthrough.

llvm/lib/CodeGen/TailDuplicator.cpp

955

Is this related to the rest of the patch? I don't see any obvious interaction.

[jyknight]

Is the endgame here that we relax the restriction that updateTerminator() requires an analyzable branch?

That wasn't my goal. My goal is to remove the assumption that it's possible to reconstruct which branches should be present at the end of the block, based purely on CFG successors. Once upon a time, that may have been a simple prospect -- e.g. if there's a conditional branch, there should be 2 CFG successors, and at least one must have an actual branch instruction, so you just add/remove branches to the other one, depending on whether it's the fallthrough block. But then we added exception handling, and the EHPad successors mess up that simplicity. But, we could at least filter out EHPad blocks, because those can't also be the target of a normal jump.

However, I'm now proposing to make inlineasm_br not be a terminator instruction, which means the block can have extra successors in a new way. And unlike the EHPad blocks, the targets of an INLINEASM_BR may _also_ be used as normal jump/fallthroughs of another block.

I split this out as a separate review, because I think it's actually better that this is less magical and more explicit, regardless of whether we decide to make that INLINEASM_BR change.

llvm/lib/CodeGen/TailDuplicator.cpp

955

This code previously left the block in a bogus state, where the new CFG was not properly reflected in the instructions in the block (e.g. the machine instructions could fallthrough to the next block, while the CFG said it should actually jump somewhere else).

In the past, other future calls to updateTerminator managed to reconstruct the desired state of things, and insert the correct jump, based on the CFG successor. That no longer happens magically, so this code now must ensure that it has inserted a jump at the end of the block if it intends there to be one.

[efriedma]

Okay, that explanation makes sense.

llvm/lib/CodeGen/MachineBasicBlock.cpp

585

"or the end of the block is unreachable" makes me a little worried. If I'm understanding correctly, if do somehow end up in that situation, the layout successor won't actually be a successor of the BasicBlock, and assert(isSuccessor(PreviousLayoutSuccessor)) will fail.

llvm/lib/CodeGen/MachineBlockPlacement.cpp

2729

Commented out code?

[aheejin]

Are you planning to land this or D79793 is the alternative approach?

[jyknight]

Are you planning to land this or D79793 is the alternative approach?

They're not alternatives, I plan to land both.

llvm/lib/CodeGen/MachineBasicBlock.cpp

585

Um, yup. I think you've identified a serious issue.

Identifying whether the instruction stream represents a fallthrough is another place where we look at successor list to help figure out what the instructions mean. That's something else I should fix -- we currently don't actually represent the tail of the block being unreachable in the instruction stream. IMO, this is a mistake, and we ought to have an UNREACHABLE_MARKER pseudo-inst to represent this situation, and stop depending on querying the successors list to figure this out.

For the purposes of this patch and making incremental progress, I believe it's sufficient to just delete the assert and go with if(!PreviousLayoutSuccessor || !isSuccessor(PreviousLayoutSuccessor)) return;. That won't fully achieve my goal, but I'll work on the above idea coupled with a cleanup of all the fallthrough handling, here and everywhere else, as a separate change.

[efriedma]

llvm/lib/CodeGen/MachineBasicBlock.cpp

585

Yes, I think that works.

[aheejin]

llvm/lib/CodeGen/MachineBasicBlock.cpp

626

Does this new patch handle when a given layout successor is an EH pad? I might be mistaken, but I think that case should be excluded. But while this code is deleted, I can't find a similar counterpart in the new CL.

llvm/lib/Target/WebAssembly/WebAssemblyCFGSort.cpp

168

Here for example.. this seems to takes the next layout block as a fallthrough, regardless of whether the layout successor is an EH pad or not... No?

[efriedma]

llvm/lib/CodeGen/MachineBasicBlock.cpp

626

There are three possibilities here:

1. The end of the block is unreachable. This is case we discussed in the other comment.
2. We have fallthrough. The original layout successor can't be an EHPad because it's illegal to fall through into an EHPad.
3. We don't have fallthrough. In this case, the original layout successor is irrelevant.

[efriedma]

llvm/lib/CodeGen/MachineBasicBlock.cpp

585

Hmm. Actually, isSuccessor() isn't sufficient, I think. Consider the case where the end of the block is unreachable, and the layout successor is an EHPad.

[aheejin]

There are many places in this patch that just gives the layout successor as an argument to updateTerminator, which does not consider EH pad possibility. I think the main reason we enumerated the successor iterator is to filter out EH pads, and I'm not sure if deleting that and delegating the task of giving the right layout successor to the caller side is a good idea, because each caller has to evaluate that separately.
Also I'm not very sure PreviousLayoutSuccessor stands for; updateTerminator does not change layout (= order of BBs), so the layout successor should stay the same before and after calling updateTerminator. Depending on branches, after calling updateTerminator, we may not end up falling through to the PreviousLayoutSuccessor variable.

llvm/include/llvm/CodeGen/MachineBasicBlock.h

526

I'm not very sure PreviousLayoutSuccessor stands for. updateTerminator does not change layout (= order of BBs), so the layout successor should stay the same before and after calling updateTerminator. Depending on branches created, after calling updateTerminator, we may not end up falling through to the PreviousLayoutSuccessor variable, but in that case it still succeeds MBB in layout. PreviousLayoutSuccessor, to me, sounds like the layout successor will change after calling this function. How about just LayoutSuccessor or something..?

llvm/lib/CodeGen/MachineBasicBlock.cpp

626

I'm not sure about case 2; the reason is here: https://reviews.llvm.org/D79605#2035699

[jyknight]

llvm/include/llvm/CodeGen/MachineBasicBlock.h

526

In some (but not all) of the callers, the current actual block ordering of the function has been modified _before_ calling this function. Therefore, the argument passed as "PreviousLayoutSuccessor" is no longer the actual, current, layout successor, at the time of this call.

If there is currently a fallthrough at the end of the block, a branch may need to be inserted, and the target of that branch is PreviousLayoutSuccessor.

llvm/lib/CodeGen/MachineBasicBlock.cpp

585

Agreed.

626

The problem right now (discussed in the other comment) is that we cannot easily distinguish between actual fallthrough and unreachable block-end.

So, yes, right now we may end up in a situation that appears as though a fallthrough occurs into an EHPad, but, that's not actually possible. It's only because in that situation, the block-end must in fact have been unreachable.

[jyknight]

Handle the case where the apparent fallthrough is not actually a
fallthrough, because the end of the block is unreachable.

[efriedma]

LGTM with a couple minor comments.

llvm/lib/CodeGen/MachineBasicBlock.cpp

593

nit: please clang-format

622

Might as well assert(!PreviousLayoutSuccessor->isEHPad()).

[aheejin]

llvm/include/llvm/CodeGen/MachineBasicBlock.h

526

Then the change in WebAssemblyCFGSort.cpp can be incorrect. What you gave as an argument there is the 'current' layout successor. The order change happens in the caller, SortBlocks(), and you have to track all changes there to figure out what the previous layout successor had been.

I'm still not sure if this approach, that pushing all the responsibilities of figuring out "previous layout successor" to the callers, is feasible or scalable. The order change can happen anywhere, sometimes not right before you call MachineBasicBlock::updateTerminator().

[void]

Is this and D79793 ready for submission? :-)

[void]

In D79605#2065825, @void wrote:

Is this and D79793 ready for submission? :-)

[Sorry to be a pest.] If there's anything I can do to help this and the related patch along, please let me know. It's fixes a major bug in an important feature. :-)

[jyknight]

llvm/include/llvm/CodeGen/MachineBasicBlock.h

526

The change in WebAssemblyCFGSort.cpp is passing the original/correct successor, because it's looking it up based on the block id created in RenumberBlocks (at the top of SortBlocks), which is before the blocks are reordered.

llvm/include/llvm/CodeGen/MachineBasicBlock.h

#endif
   void moveBefore(MachineBasicBlock *NewAfter);
   void moveAfter(MachineBasicBlock *NewBefore);
 
   /// Returns true if this and MBB belong to the same section.
   bool sameSection(const MachineBasicBlock *MBB) const {
     return getSectionID() == MBB->getSectionID();
   }
 
-  /// Update the terminator instructions in block to account for changes to the
-  /// layout. If the block previously used a fallthrough, it may now need a
-  /// branch, and if it previously used branching it may now be able to use a
-  /// fallthrough.
-  void updateTerminator();
+  /// Update the terminator instructions in block to account for changes to
+  /// block layout which may have been made. PreviousLayoutSuccessor should be
+  /// set to the block which may have been used as fallthrough before the block
+  /// layout was modified.  If the block previously fell through to that block,
+  /// it may now need a branch. If it previously branched to another block, it
+  /// may now be able to fallthrough to the current layout successor.
+  void updateTerminator(MachineBasicBlock *PreviousLayoutSuccessor);
 
   // Machine-CFG mutators
 
   /// Add Succ as a successor of this MachineBasicBlock.  The Predecessors list
   /// of Succ is automatically updated. PROB parameter is stored in
   /// Probabilities list. The default probability is set as unknown. Mixing
   /// known and unknown probabilities in successor list is not allowed. When all
   /// successors have unknown probabilities, 1 / N is returned as the

llvm/lib/CodeGen/BBSectionsPrepare.cpp

if (MBB.isEndSection())
       continue;
 
     // It might be possible to optimize branches by flipping the branch
     // condition.
     Cond.clear();
     MachineBasicBlock *TBB = nullptr, *FBB = nullptr; // For analyzeBranch.
     if (TII->analyzeBranch(MBB, TBB, FBB, Cond))
       continue;
-    MBB.updateTerminator();
+    MBB.updateTerminator(FTMBB);
   }
 }
 
 // This function provides the BBCluster information associated with a function.
 // Returns true if a valid association exists and false otherwise.
 static bool getBBClusterInfoForFunction(
     const MachineFunction &MF, const StringMap<StringRef> FuncAliasMap,
     const ProgramBBClusterInfoMapTy &ProgramBBClusterInfo,

llvm/lib/CodeGen/BranchRelaxation.cpp

MachineBasicBlock *BranchRelaxation::splitBlockBeforeInstr(MachineInstr &MI,
   BlockInfo.insert(BlockInfo.begin() + NewBB->getNumber(), BasicBlockInfo());
 
   NewBB->transferSuccessors(OrigBB);
   OrigBB->addSuccessor(NewBB);
   OrigBB->addSuccessor(DestBB);
 
   // Cleanup potential unconditional branch to successor block.
   // Note that updateTerminator may change the size of the blocks.
-  NewBB->updateTerminator();
-  OrigBB->updateTerminator();
+  OrigBB->updateTerminator(NewBB);
 
   // Figure out how large the OrigBB is.  As the first half of the original
   // block, it cannot contain a tablejump.  The size includes
   // the new jump we added.  (It should be possible to do this without
   // recounting everything, but it's very confusing, and this is rarely
   // executed.)
   BlockInfo[OrigBB->getNumber()].Size = computeBlockSize(*OrigBB);
 

llvm/lib/CodeGen/MachineBasicBlock.cpp

 void MachineBasicBlock::moveBefore(MachineBasicBlock *NewAfter) {
   getParent()->splice(NewAfter->getIterator(), getIterator());
 }
 
 void MachineBasicBlock::moveAfter(MachineBasicBlock *NewBefore) {
   getParent()->splice(++NewBefore->getIterator(), getIterator());
 }
 
-void MachineBasicBlock::updateTerminator() {
+void MachineBasicBlock::updateTerminator(
+    MachineBasicBlock *PreviousLayoutSuccessor) {
+  LLVM_DEBUG(dbgs() << "Updating terminators on " << printMBBReference(*this)
+                    << "\n");
+
   const TargetInstrInfo *TII = getParent()->getSubtarget().getInstrInfo();
   // A block with no successors has no concerns with fall-through edges.
   if (this->succ_empty())
     return;
 
   MachineBasicBlock *TBB = nullptr, *FBB = nullptr;
   SmallVector<MachineOperand, 4> Cond;
   DebugLoc DL = findBranchDebugLoc();
   bool B = TII->analyzeBranch(*this, TBB, FBB, Cond);
   (void) B;
   assert(!B && "UpdateTerminators requires analyzable predecessors!");
   if (Cond.empty()) {
     if (TBB) {
       // The block has an unconditional branch. If its successor is now its
       // layout successor, delete the branch.
       if (isLayoutSuccessor(TBB))
         TII->removeBranch(*this);
     } else {
-      // The block has an unconditional fallthrough. If its successor is not its
-      // layout successor, insert a branch. First we have to locate the only
-      // non-landing-pad successor, as that is the fallthrough block.
-      for (succ_iterator SI = succ_begin(), SE = succ_end(); SI != SE; ++SI) {
-        if ((*SI)->isEHPad())
-          continue;
-        assert(!TBB && "Found more than one non-landing-pad successor!");
-        TBB = *SI;
-      }
+      // The block has an unconditional fallthrough, or the end of the block is
+      // unreachable.
 
-      // If there is no non-landing-pad successor, the block has no fall-through
-      // edges to be concerned with.
-      if (!TBB)
+      // Unfortunately, whether the end of the block is unreachable is not
+      // immediately obvious; we must fall back to checking the successor list,
+      // and assuming that if the passed in block is in the succesor list and
+      // not an EHPad, it must be the intended target.
+      if (!PreviousLayoutSuccessor || !isSuccessor(PreviousLayoutSuccessor)
+          || PreviousLayoutSuccessor->isEHPad())
         return;
 
-      // Finally update the unconditional successor to be reached via a branch
-      // if it would not be reached by fallthrough.
-      if (!isLayoutSuccessor(TBB))
-        TII->insertBranch(*this, TBB, nullptr, Cond, DL);
+      // If the unconditional successor block is not the current layout
+      // successor, insert a branch to jump to it.
+      if (!isLayoutSuccessor(PreviousLayoutSuccessor))
+        TII->insertBranch(*this, PreviousLayoutSuccessor, nullptr, Cond, DL);
     }
     return;
   }
 
   if (FBB) {
     // The block has a non-fallthrough conditional branch. If one of its
     // successors is its layout successor, rewrite it to a fallthrough
     // conditional branch.
     if (isLayoutSuccessor(TBB)) {
       if (TII->reverseBranchCondition(Cond))
         return;
       TII->removeBranch(*this);
       TII->insertBranch(*this, FBB, nullptr, Cond, DL);
     } else if (isLayoutSuccessor(FBB)) {
       TII->removeBranch(*this);
       TII->insertBranch(*this, TBB, nullptr, Cond, DL);
     }
     return;
   }
 
-  // Walk through the successors and find the successor which is not a landing
-  // pad and is not the conditional branch destination (in TBB) as the
-  // fallthrough successor.
-  MachineBasicBlock *FallthroughBB = nullptr;
-  for (succ_iterator SI = succ_begin(), SE = succ_end(); SI != SE; ++SI) {
-    if ((*SI)->isEHPad() || *SI == TBB)
-      continue;
-    assert(!FallthroughBB && "Found more than one fallthrough successor.");
-    FallthroughBB = *SI;
-  }
-
-  if (!FallthroughBB) {
-    if (canFallThrough()) {
-      // We fallthrough to the same basic block as the conditional jump targets.
-      // Remove the conditional jump, leaving unconditional fallthrough.
-      // FIXME: This does not seem like a reasonable pattern to support, but it
-      // has been seen in the wild coming out of degenerate ARM test cases.
-      TII->removeBranch(*this);
-
-      // Finally update the unconditional successor to be reached via a branch if
-      // it would not be reached by fallthrough.
-      if (!isLayoutSuccessor(TBB))
-        TII->insertBranch(*this, TBB, nullptr, Cond, DL);
-      return;
-    }
-
-    // We enter here iff exactly one successor is TBB which cannot fallthrough
-    // and the rest successors if any are EHPads.  In this case, we need to
-    // change the conditional branch into unconditional branch.
+  // We now know we're going to fallthrough to PreviousLayoutSuccessor.
+  assert(PreviousLayoutSuccessor);
+  assert(isSuccessor(PreviousLayoutSuccessor));
+
+  if (PreviousLayoutSuccessor == TBB) {
+    // We had a fallthrough to the same basic block as the conditional jump
+    // targets.  Remove the conditional jump, leaving an unconditional
+    // fallthrough or an unconditional jump.
     TII->removeBranch(*this);
+    if (!isLayoutSuccessor(TBB)) {
       Cond.clear();
       TII->insertBranch(*this, TBB, nullptr, Cond, DL);
+    }
     return;
   }
 
   // The block has a fallthrough conditional branch.
   if (isLayoutSuccessor(TBB)) {
     if (TII->reverseBranchCondition(Cond)) {
       // We can't reverse the condition, add an unconditional branch.
       Cond.clear();
-      TII->insertBranch(*this, FallthroughBB, nullptr, Cond, DL);
+      TII->insertBranch(*this, PreviousLayoutSuccessor, nullptr, Cond, DL);
       return;
     }
     TII->removeBranch(*this);
-    TII->insertBranch(*this, FallthroughBB, nullptr, Cond, DL);
-  } else if (!isLayoutSuccessor(FallthroughBB)) {
+    TII->insertBranch(*this, PreviousLayoutSuccessor, nullptr, Cond, DL);
+  } else if (!isLayoutSuccessor(PreviousLayoutSuccessor)) {
     TII->removeBranch(*this);
-    TII->insertBranch(*this, TBB, FallthroughBB, Cond, DL);
+    TII->insertBranch(*this, TBB, PreviousLayoutSuccessor, Cond, DL);
   }
 }
 
 void MachineBasicBlock::validateSuccProbs() const {
 #ifndef NDEBUG
   int64_t Sum = 0;
   for (auto Prob : Probs)
     Sum += Prob.getNumerator();

 
 MachineBasicBlock *MachineBasicBlock::SplitCriticalEdge(
     MachineBasicBlock *Succ, Pass &P,
     std::vector<SparseBitVector<>> *LiveInSets) {
   if (!canSplitCriticalEdge(Succ))
     return nullptr;
 
   MachineFunction *MF = getParent();
+  MachineBasicBlock *PrevFallthrough = getNextNode();
   DebugLoc DL;  // FIXME: this is nowhere
 
   MachineBasicBlock *NMBB = MF->CreateMachineBasicBlock();
   MF->insert(std::next(MachineFunction::iterator(this)), NMBB);
   LLVM_DEBUG(dbgs() << "Splitting critical edge: " << printMBBReference(*this)
                     << " -- " << printMBBReference(*NMBB) << " -- "
                     << printMBBReference(*Succ) << '\n');
 
MachineBasicBlock *MachineBasicBlock::SplitCriticalEdge(
   // SlotIndexes.
   SmallVector<MachineInstr*, 4> Terminators;
   if (Indexes) {
     for (instr_iterator I = getFirstInstrTerminator(), E = instr_end();
          I != E; ++I)
       Terminators.push_back(&*I);
   }
 
-  updateTerminator();
+  // Since we replaced all uses of Succ with NMBB, that should also be treated
+  // as the fallthrough successor
+  if (Succ == PrevFallthrough)
+    PrevFallthrough = NMBB;
+  updateTerminator(PrevFallthrough);
 
   if (Indexes) {
     SmallVector<MachineInstr*, 4> NewTerminators;
     for (instr_iterator I = getFirstInstrTerminator(), E = instr_end();
          I != E; ++I)
       NewTerminators.push_back(&*I);
 
     for (SmallVectorImpl<MachineInstr*>::iterator I = Terminators.begin(),

llvm/lib/CodeGen/MachineBlockPlacement.cpp

if (!FunctionBlockSet.empty()) {
       BadFunc = true;
       for (MachineBasicBlock *RemainingBB : FunctionBlockSet)
         dbgs() << "Function contains blocks never placed into a chain!\n"
                << "  Bad block:    " << getBlockName(RemainingBB) << "\n";
     }
     assert(!BadFunc && "Detected problems with the block placement.");
   });
 
+  // Remember original layout ordering, so we can update terminators after
+  // reordering to point to the original layout successor.
+  SmallVector<MachineBasicBlock *, 4> OriginalLayoutSuccessors(
+      F->getNumBlockIDs());
+  {
+    MachineBasicBlock *LastMBB = nullptr;
+    for (auto &MBB : *F) {
+      if (LastMBB != nullptr)
+        OriginalLayoutSuccessors[LastMBB->getNumber()] = &MBB;
+      LastMBB = &MBB;
+    }
+    OriginalLayoutSuccessors[F->back().getNumber()] = nullptr;
+  }
+
   // Splice the blocks into place.
   MachineFunction::iterator InsertPos = F->begin();
   LLVM_DEBUG(dbgs() << "[MBP] Function: " << F->getName() << "\n");
   for (MachineBasicBlock *ChainBB : FunctionChain) {
     LLVM_DEBUG(dbgs() << (ChainBB == *FunctionChain.begin() ? "Placing chain "
                                                             : "          ... ")
                       << getBlockName(ChainBB) << "\n");
     if (InsertPos != MachineFunction::iterator(ChainBB))
       F->splice(InsertPos, ChainBB);
     else
       ++InsertPos;
 
     // Update the terminator of the previous block.
     if (ChainBB == *FunctionChain.begin())
       continue;
     MachineBasicBlock *PrevBB = &*std::prev(MachineFunction::iterator(ChainBB));
 
     // FIXME: It would be awesome of updateTerminator would just return rather
#endif
     //   PrevBB->updateTerminator();
     //   Cond.clear();
     //   TBB = FBB = nullptr;
     //   if (TII->analyzeBranch(*PrevBB, TBB, FBB, Cond)) {
     //     // FIXME: This should never take place.
     //     TBB = FBB = nullptr;
     //   }
     // }
-    if (!TII->analyzeBranch(*PrevBB, TBB, FBB, Cond))
-      PrevBB->updateTerminator();
+    if (!TII->analyzeBranch(*PrevBB, TBB, FBB, Cond)) {
+      PrevBB->updateTerminator(OriginalLayoutSuccessors[PrevBB->getNumber()]);
+    }
   }
 
   // Fixup the last block.
   Cond.clear();
   MachineBasicBlock *TBB = nullptr, *FBB = nullptr; // For analyzeBranch.
-  if (!TII->analyzeBranch(F->back(), TBB, FBB, Cond))
-    F->back().updateTerminator();
+  if (!TII->analyzeBranch(F->back(), TBB, FBB, Cond)) {
+    MachineBasicBlock *PrevBB = &F->back();
+    PrevBB->updateTerminator(OriginalLayoutSuccessors[PrevBB->getNumber()]);
+  }
 
   BlockWorkList.clear();
   EHPadWorkList.clear();
 }
 
 void MachineBlockPlacement::optimizeBranches() {
   BlockChain &FunctionChain = *BlockToChain[&F->front()];
   SmallVector<MachineOperand, 4> Cond; // For analyzeBranch.
if (!TII->analyzeBranch(*ChainBB, TBB, FBB, Cond, /*AllowModify*/ true)) {
         LLVM_DEBUG(dbgs() << "Reverse order of the two branches: "
                           << getBlockName(ChainBB) << "\n");
         LLVM_DEBUG(dbgs() << "    Edge probability: "
                           << MBPI->getEdgeProbability(ChainBB, FBB) << " vs "
                           << MBPI->getEdgeProbability(ChainBB, TBB) << "\n");
         DebugLoc dl; // FIXME: this is nowhere
         TII->removeBranch(*ChainBB);
         TII->insertBranch(*ChainBB, FBB, TBB, Cond, dl);
-        ChainBB->updateTerminator();
       }
     }
   }
 }
 
 void MachineBlockPlacement::alignBlocks() {
   // Walk through the backedges of the function now that we have fully laid out
   // the basic blocks and align the destination of each backedge. We don't rely

llvm/lib/CodeGen/TailDuplicator.cpp

 bool TailDuplicator::tailDuplicate(bool IsSimple, MachineBasicBlock *TailBB,
                           MachineBasicBlock *ForcedLayoutPred,
                           SmallVectorImpl<MachineBasicBlock *> &TDBBs,
                           SmallVectorImpl<MachineInstr *> &Copies,
                           SmallVectorImpl<MachineBasicBlock *> *CandidatePtr) {
   LLVM_DEBUG(dbgs() << "\n*** Tail-duplicating " << printMBBReference(*TailBB)
                     << '\n');
 
+  bool ShouldUpdateTerminators = TailBB->canFallThrough();
+
   DenseSet<unsigned> UsedByPhi;
   getRegsUsedByPHIs(*TailBB, &UsedByPhi);
 
   if (IsSimple)
     return duplicateSimpleBB(TailBB, TDBBs, UsedByPhi, Copies);
 
   // Iterate through all the unique predecessors and tail-duplicate this
   // block into them, if possible. Copying the list ahead of time also
for (MachineBasicBlock *PredBB : Preds) {
 
     // Update the CFG.
     PredBB->removeSuccessor(PredBB->succ_begin());
     assert(PredBB->succ_empty() &&
            "TailDuplicate called on block with multiple successors!");
     for (MachineBasicBlock *Succ : TailBB->successors())
       PredBB->addSuccessor(Succ, MBPI->getEdgeProbability(TailBB, Succ));
 
+    // Update branches in pred to jump to tail's layout successor if needed.
+    if (ShouldUpdateTerminators)
+      PredBB->updateTerminator(TailBB->getNextNode());
+
     Changed = true;
     ++NumTailDups;
   }
 
   // If TailBB was duplicated into all its predecessors except for the prior
   // block, which falls through unconditionally, move the contents of this
   // block into the prior block.
   MachineBasicBlock *PrevBB = ForcedLayoutPred;
if (PrevBB->succ_size() == 1 &&
     } else {
       TII->removeBranch(*PrevBB);
       // No PHIs to worry about, just splice the instructions over.
       PrevBB->splice(PrevBB->end(), TailBB, TailBB->begin(), TailBB->end());
     }
     PrevBB->removeSuccessor(PrevBB->succ_begin());
     assert(PrevBB->succ_empty());
     PrevBB->transferSuccessors(TailBB);
+
+    // Update branches in PrevBB based on Tail's layout successor.
+    if (ShouldUpdateTerminators)
+      PrevBB->updateTerminator(TailBB->getNextNode());
+
     TDBBs.push_back(PrevBB);
     Changed = true;
   }
 
   // If this is after register allocation, there are no phis to fix.
   if (!PreRegAlloc)
     return Changed;
 

llvm/lib/Target/AArch64/AArch64ConditionOptimizer.cpp

void AArch64ConditionOptimizer::modifyCmp(MachineInstr *CmpMI,
   MachineInstr &BrMI = *MBB->getFirstTerminator();
 
   // Change condition in branch instruction.
   BuildMI(*MBB, BrMI, BrMI.getDebugLoc(), TII->get(AArch64::Bcc))
       .addImm(Cmp)
       .add(BrMI.getOperand(1));
   BrMI.eraseFromParent();
 
-  MBB->updateTerminator();
-
   ++NumConditionsAdjusted;
 }
 
 // Parse a condition code returned by analyzeBranch, and compute the CondCode
 // corresponding to TBB.
 // Returns true if parsing was successful, otherwise false is returned.
 static bool parseCond(ArrayRef<MachineOperand> Cond, AArch64CC::CondCode &CC) {
   // A normal br.cond simply has the condition code.

llvm/lib/Target/AArch64/AArch64ConditionalCompares.cpp

void SSACCmpConv::convert(SmallVectorImpl<MachineBasicBlock *> &RemovedBlocks) {
   if (isZBranch) {
     bool isNZ = CmpMI->getOpcode() == AArch64::CBNZW ||
                 CmpMI->getOpcode() == AArch64::CBNZX;
     BuildMI(*Head, CmpMI, CmpMI->getDebugLoc(), TII->get(AArch64::Bcc))
         .addImm(isNZ ? AArch64CC::NE : AArch64CC::EQ)
         .add(CmpMI->getOperand(1)); // Branch target.
   }
   CmpMI->eraseFromParent();
-  Head->updateTerminator();
+  Head->updateTerminator(CmpBB->getNextNode());
 
   RemovedBlocks.push_back(CmpBB);
   CmpBB->eraseFromParent();
   LLVM_DEBUG(dbgs() << "Result:\n" << *Head);
   ++NumConverted;
 }
 
 int SSACCmpConv::expectedCodeSizeDelta() const {

llvm/lib/Target/ARM/ARMConstantIslandPass.cpp

adjustJTTargetBlockForward(MachineBasicBlock *BB, MachineBasicBlock *JTBB) {
   // try to move it; otherwise, create a new block following the jump
   // table that branches back to the actual target. This is a very simple
   // heuristic. FIXME: We can definitely improve it.
   MachineBasicBlock *TBB = nullptr, *FBB = nullptr;
   SmallVector<MachineOperand, 4> Cond;
   SmallVector<MachineOperand, 4> CondPrior;
   MachineFunction::iterator BBi = BB->getIterator();
   MachineFunction::iterator OldPrior = std::prev(BBi);
+  MachineFunction::iterator OldNext = std::next(BBi);
 
   // If the block terminator isn't analyzable, don't try to move the block
   bool B = TII->analyzeBranch(*BB, TBB, FBB, Cond);
 
   // If the block ends in an unconditional branch, move it. The prior block
   // has to have an analyzable terminator for us to move this one. Be paranoid
   // and make sure we're not trying to move the entry block of the function.
   if (!B && Cond.empty() && BB != &MF->front() &&
       !TII->analyzeBranch(*OldPrior, TBB, FBB, CondPrior)) {
     BB->moveAfter(JTBB);
-    OldPrior->updateTerminator();
-    BB->updateTerminator();
+    OldPrior->updateTerminator(BB);
+    BB->updateTerminator(OldNext != MF->end() ? &*OldNext : nullptr);
     // Update numbering to account for the block being moved.
     MF->RenumberBlocks();
     ++NumJTMoved;
     return nullptr;
   }
 
   // Create a new MBB for the code after the jump BB.
   MachineBasicBlock *NewBB =

llvm/lib/Target/Hexagon/HexagonEarlyIfConv.cpp

void HexagonEarlyIfConversion::mergeBlocks(MachineBasicBlock *PredB,
   LLVM_DEBUG(dbgs() << "Merging blocks " << PrintMB(PredB) << " and "
                     << PrintMB(SuccB) << "\n");
   bool TermOk = hasUncondBranch(SuccB);
   eliminatePhis(SuccB);
   HII->removeBranch(*PredB);
   PredB->removeSuccessor(SuccB);
   PredB->splice(PredB->end(), SuccB, SuccB->begin(), SuccB->end());
   PredB->transferSuccessorsAndUpdatePHIs(SuccB);
+  MachineBasicBlock *OldLayoutSuccessor = SuccB->getNextNode();
   removeBlock(SuccB);
   if (!TermOk)
-    PredB->updateTerminator();
+    PredB->updateTerminator(OldLayoutSuccessor);
 }
 
 void HexagonEarlyIfConversion::simplifyFlowGraph(const FlowPattern &FP) {
+  MachineBasicBlock *OldLayoutSuccessor = FP.SplitB->getNextNode();
   if (FP.TrueB)
     removeBlock(FP.TrueB);
   if (FP.FalseB)
     removeBlock(FP.FalseB);
 
-  FP.SplitB->updateTerminator();
+  FP.SplitB->updateTerminator(OldLayoutSuccessor);
   if (FP.SplitB->succ_size() != 1)
     return;
 
   MachineBasicBlock *SB = *FP.SplitB->succ_begin();
   if (SB->pred_size() != 1)
     return;
 
   // By now, the split block has only one successor (SB), and SB has only

llvm/lib/Target/WebAssembly/WebAssemblyCFGSort.cpp

#endif
   for (const MachineInstr &Term : MBB->terminators()) {
 #ifndef NDEBUG
     AnyBarrier |= Term.isBarrier();
 #endif
     AllAnalyzable &= Term.isBranch() && !Term.isIndirectBranch();
   }
   assert((AnyBarrier || AllAnalyzable) &&
          "analyzeBranch needs to analyze any block with a fallthrough");
+
+  // Find the layout successor from the original block order.
+  MachineFunction *MF = MBB->getParent();
+  MachineBasicBlock *OriginalSuccessor =
+      unsigned(MBB->getNumber() + 1) < MF->getNumBlockIDs()
+          ? MF->getBlockNumbered(MBB->getNumber() + 1)
+          : nullptr;
+
   if (AllAnalyzable)
-    MBB->updateTerminator();
+    MBB->updateTerminator(OriginalSuccessor);
 }
 
 namespace {
 // EH pads are selected first regardless of the block comparison order.
 // When only one of the BBs is an EH pad, we give a higher priority to it, to
 // prevent common mismatches between possibly throwing calls and ehpads they
 // unwind to, as in the example below:
 //

 
 /// Sort the blocks, taking special care to make sure that regions are not
 /// interrupted by blocks not dominated by their header.
 /// TODO: There are many opportunities for improving the heuristics here.
 /// Explore them.
 static void sortBlocks(MachineFunction &MF, const MachineLoopInfo &MLI,
                        const WebAssemblyExceptionInfo &WEI,
                        const MachineDominatorTree &MDT) {
+  // Remember original layout ordering, so we can update terminators after
+  // reordering to point to the original layout successor.
+  MF.RenumberBlocks();
+
   // Prepare for a topological sort: Record the number of predecessors each
   // block has, ignoring loop backedges.
-  MF.RenumberBlocks();
   SmallVector<unsigned, 16> NumPredsLeft(MF.getNumBlockIDs(), 0);
   for (MachineBasicBlock &MBB : MF) {
     unsigned N = MBB.pred_size();
     if (MachineLoop *L = MLI.getLoopFor(&MBB))
       if (L->getHeader() == &MBB)
         for (const MachineBasicBlock *Pred : MBB.predecessors())
           if (L->contains(Pred))
             --N;