CodeGen: Handle missed case of block removal during BlockPlacement.

[iteratee]

Summary

There is a use after free bug in the existing code. Loop layout selects
a preferred exit block, and then lays out the loop. If this block is
removed during layout, it needs to be invalidated to prevent a use after
free.

[davidxl]

Can you explain the code path that leads to use after free?

[iteratee]

I added comments to the 3 lines that show the possible use after free.

lib/CodeGen/MachineBlockPlacement.cpp

1477

ExitingBB is selected here. This is a reference to a block, and it only exists on the call stack.

1493

buildChain is called here. buildChain may tail-duplicate and remove the block referred to by ExitingBB.

1498

ExitingBB is used here, after it may have been freed.

[davidxl]

lib/CodeGen/MachineBlockPlacement.cpp

1503

Should the right fix be to move the 'findBestLoopExit' after buildChain call?

[iteratee]

I found this by observing it in the debug logs while looking for something else, so I think the direct fix is worth committing, even if we do a better fix later.

lib/CodeGen/MachineBlockPlacement.cpp

1503

That would require re-writing findBestLoopExit to not assume that the blocks weren't all in a single chain already.

We should re-write it to do that, but I think that's out of scope for an existing use after free bug.

[davidxl]

lgtm.

As discussed, as this resolve the use after free error, it does point to possible missing loop rotation due to tail dup -- that needs to be investigated as follow up.

[iteratee]

Committed in rL285348

lib/CodeGen/MachineBlockPlacement.cpp

class MachineBlockPlacement : public MachineFunctionPass {
   const MachineBranchProbabilityInfo *MBPI;
 
   /// \brief A handle to the function-wide block frequency pass.
   std::unique_ptr<BranchFolder::MBFIWrapper> MBFI;
 
   /// \brief A handle to the loop info.
   MachineLoopInfo *MLI;
 
+  /// \brief Preferred loop exit.
+  /// Member variable for convenience. It may be removed by duplication deep
+  /// in the call stack.
+  MachineBasicBlock *PreferredLoopExit;
+
   /// \brief A handle to the target's instruction info.
   const TargetInstrInfo *TII;
 
   /// \brief A handle to the target's lowering info.
   const TargetLoweringBase *TLI;
 
   /// \brief A handle to the post dominator tree.
   MachineDominatorTree *MDT;
void MachineBlockPlacement::buildLoopChains(MachineLoop &L) {
   // fewer branches in the loop body.
   // When we use profile data to rotate the loop, this is unnecessary.
   MachineBasicBlock *LoopTop =
       RotateLoopWithProfile ? L.getHeader() : findBestLoopTop(L, LoopBlockSet);
 
   // If we selected just the header for the loop top, look for a potentially
   // profitable exit block in the event that rotating the loop can eliminate
   // branches by placing an exit edge at the bottom.
-  MachineBasicBlock *ExitingBB = nullptr;
+  PreferredLoopExit = nullptr;
   if (!RotateLoopWithProfile && LoopTop == L.getHeader())
-    ExitingBB = findBestLoopExit(L, LoopBlockSet);
+    PreferredLoopExit = findBestLoopExit(L, LoopBlockSet);
 
   BlockChain &LoopChain = *BlockToChain[LoopTop];
 
   // FIXME: This is a really lame way of walking the chains in the loop: we
   // walk the blocks, and use a set to prevent visiting a particular chain
   // twice.
   SmallPtrSet<BlockChain *, 4> UpdatedPreds;
   assert(LoopChain.UnscheduledPredecessors == 0);
   UpdatedPreds.insert(&LoopChain);
 
   for (MachineBasicBlock *LoopBB : LoopBlockSet)
     fillWorkLists(LoopBB, UpdatedPreds, &LoopBlockSet);
 
   buildChain(LoopTop, LoopChain, &LoopBlockSet);
 
   if (RotateLoopWithProfile)
     rotateLoopWithProfile(LoopChain, L, LoopBlockSet);
   else
-    rotateLoop(LoopChain, ExitingBB, LoopBlockSet);
+    rotateLoop(LoopChain, PreferredLoopExit, LoopBlockSet);
 
   DEBUG({
     // Crash at the end so we get all of the debugging output first.
     bool BadLoop = false;
     if (LoopChain.UnscheduledPredecessors) {
       BadLoop = true;
       dbgs() << "Loop chain contains a block without its preds placed!\n"
              << "  Loop header:  " << getBlockName(*L.block_begin()) << "\n"
auto RemovalCallback =
 
         // Handle the filter set
         if (BlockFilter) {
           BlockFilter->erase(RemBB);
         }
 
         // Remove the block from loop info.
         MLI->removeBlock(RemBB);
+        if (RemBB == PreferredLoopExit)
+          PreferredLoopExit = nullptr;
 
-        // TailDuplicator handles removing it from loops.
         DEBUG(dbgs() << "TailDuplicator deleted block: "
               << getBlockName(RemBB) << "\n");
       };
   auto RemovalCallbackRef =
       llvm::function_ref<void(MachineBasicBlock*)>(RemovalCallback);
 
   SmallVector<MachineBasicBlock *, 8> DuplicatedPreds;
   TailDup.tailDuplicateAndUpdate(IsSimple, BB, LPred,