1 files changed, 9 insertions, 17 deletions

diff --git a/source/slang/slang-ir-loop-unroll.cpp b/source/slang/slang-ir-loop-unroll.cpp
index 79b00f60a..2f689ebde 100644
--- a/source/slang/slang-ir-loop-unroll.cpp
+++ b/source/slang/slang-ir-loop-unroll.cpp
@@ -47,7 +47,7 @@ static bool _eliminateDeadBlocks(List<IRBlock*>& blocks, IRBlock* unreachableBlo
     return changed;
 }
 
-List<IRBlock*> _collectBlocksInLoop(Dictionary<IRBlock*, int>& blockOrdering, IRLoop* loopInst)
+List<IRBlock*> _collectBlocksInLoop(IRDominatorTree* dom, IRLoop* loopInst)
 {
     List<IRBlock*> loopBlocks;
     HashSet<IRBlock*> loopBlocksSet;
@@ -58,7 +58,6 @@ List<IRBlock*> _collectBlocksInLoop(Dictionary<IRBlock*, int>& blockOrdering, IR
     };
     auto firstBlock = as<IRBlock>(loopInst->block.get());
     auto breakBlock = as<IRBlock>(loopInst->breakBlock.get());
-    auto breakBlockOrdering = blockOrdering[breakBlock].GetValue();
 
     addBlock(firstBlock);
     for (Index i = 0; i < loopBlocks.getCount(); i++)
@@ -68,18 +67,19 @@ List<IRBlock*> _collectBlocksInLoop(Dictionary<IRBlock*, int>& blockOrdering, IR
         {
             if (succ == breakBlock)
                 continue;
-            auto successorOrdering = blockOrdering[block].GetValue();
-            // The target must be post-dominated by the break block in order to be considered
-            // the body of the loop.
-            // Since we don't support arbitrary goto or multi-level continue, the simple
-            // ordering comparison is sufficient to serve as a post-dominance check.
-            if (successorOrdering < breakBlockOrdering)
+            if (dom->dominates(firstBlock, succ) && !dom->dominates(breakBlock, succ))
                 addBlock(succ);
         }
     }
     return loopBlocks;
 }
 
+List<IRBlock*> collectBlocksInLoop(IRGlobalValueWithCode* func,  IRLoop* loopInst)
+{
+    auto dom = computeDominatorTree(func);
+    return _collectBlocksInLoop(dom, loopInst);
+}
+
 static int _getLoopMaxIterationsToUnroll(IRLoop* loopInst)
 {
     static constexpr int kMaxIterationsToAttempt = 100;
@@ -483,15 +483,7 @@ bool unrollLoopsInFunc(
         // Remove any continue jumps from the loop.
         eliminateContinueBlocks(module, loop);
 
-        auto postOrderReverseCFG = getPostorderOnReverseCFG(func);
-        Dictionary<IRBlock*, int> blockOrdering;
-        
-        for (Index i = 0; i < postOrderReverseCFG.getCount(); i++)
-        {
-            blockOrdering[postOrderReverseCFG[i]] = (int)i;
-        }
-
-        auto blocks = _collectBlocksInLoop(blockOrdering, loop);
+        auto blocks = collectBlocksInLoop(func, loop);
         auto loopLoc = loop->sourceLoc;
         if (!_unrollLoop(module, loop, blocks))
         {