Fix possible infinite loop in shrink wrapping when searching for save/restore

points.

There is an infinite loop that can occur in Shrink Wrapping while searching 
for the Save/Restore points. 

Part of this search checks whether the save/restore points are located in
different loop nests and if so, uses the (post) dominator trees to find the
immediate (post) dominator blocks. However, if the current block does not have
any immediate (post) dominators then this search will result in an infinite
loop. This can occur in code containing an infinite loop.

The modification checks whether the immediate (post) dominator is different from
the current save/restore block. If it is not, then the search terminates and the
current location is not considered as a valid save/restore point for shrink wrapping.

Phabricator: http://reviews.llvm.org/D11607
llvm-svn: 244247

1 files changed, 24 insertions, 6 deletions

diff --git a/llvm/lib/CodeGen/ShrinkWrap.cpp b/llvm/lib/CodeGen/ShrinkWrap.cpp
index f8ec1f358c4..261b8402bc0 100644
--- a/llvm/lib/CodeGen/ShrinkWrap.cpp
+++ b/llvm/lib/CodeGen/ShrinkWrap.cpp
@@ -309,12 +309,30 @@ void ShrinkWrap::updateSaveRestorePoints(MachineBasicBlock &MBB) {
     // Fix (C).
     if (Save && Restore && Save != Restore &&
         MLI->getLoopFor(Save) != MLI->getLoopFor(Restore)) {
-      if (MLI->getLoopDepth(Save) > MLI->getLoopDepth(Restore))
-        // Push Save outside of this loop.
-        Save = FindIDom<>(*Save, Save->predecessors(), *MDT);
-      else
-        // Push Restore outside of this loop.
-        Restore = FindIDom<>(*Restore, Restore->successors(), *MPDT);
+      if (MLI->getLoopDepth(Save) > MLI->getLoopDepth(Restore)) {
+        // Push Save outside of this loop if immediate dominator is different
+        // from save block. If immediate dominator is not different, bail out. 
+        MachineBasicBlock *IDom = FindIDom<>(*Save, Save->predecessors(), *MDT);
+        if (IDom != Save)
+          Save = IDom;
+        else {
+          Save = nullptr;
+          break;
+        }
+      }
+      else {
+        // Push Restore outside of this loop if immediate post-dominator is
+        // different from restore block. If immediate post-dominator is not
+        // different, bail out. 
+        MachineBasicBlock *IPdom =
+          FindIDom<>(*Restore, Restore->successors(), *MPDT);
+        if (IPdom != Restore)
+          Restore = IPdom; 
+        else {
+          Restore = nullptr;
+          break;
+        }
+      }      
     }
   }
 }