[llvm] r285348 - CodeGen: Handle missed case of block removal during BlockPlacement.

[Kyle Butt via llvm-commits]

Author: iteratee
Date: Thu Oct 27 16:37:20 2016
New Revision: 285348

URL: http://llvm.org/viewvc/llvm-project?rev=285348&view=rev
Log:
CodeGen: Handle missed case of block removal during BlockPlacement.

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.

Modified:
    llvm/trunk/lib/CodeGen/MachineBlockPlacement.cpp

Modified: llvm/trunk/lib/CodeGen/MachineBlockPlacement.cpp
URL: http://llvm.org/viewvc/llvm-project/llvm/trunk/lib/CodeGen/MachineBlockPlacement.cpp?rev=285348&r1=285347&r2=285348&view=diff
==============================================================================
--- llvm/trunk/lib/CodeGen/MachineBlockPlacement.cpp (original)
+++ llvm/trunk/lib/CodeGen/MachineBlockPlacement.cpp Thu Oct 27 16:37:20 2016
@@ -282,6 +282,11 @@ class MachineBlockPlacement : public Mac
   /// \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;
 
@@ -1474,9 +1479,9 @@ void MachineBlockPlacement::buildLoopCha
   // 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];
 
@@ -1495,7 +1500,7 @@ void MachineBlockPlacement::buildLoopCha
   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.
@@ -1928,8 +1933,9 @@ bool MachineBlockPlacement::maybeTailDup
 
         // 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");
       };