1 files changed, 23 insertions, 2 deletions
diff --git a/source/slang/slang-ir-dominators.cpp b/source/slang/slang-ir-dominators.cpp
index e57099321..8527fbf36 100644
--- a/source/slang/slang-ir-dominators.cpp
+++ b/source/slang/slang-ir-dominators.cpp
@@ -188,7 +188,7 @@ Int IRDominatorTree::getBlockIndex(IRBlock* block)
 
 bool IRDominatorTree::isUnreachable(IRBlock* block)
 {
-    return !mapBlockToIndex.containsKey(block);
+    return !reachableSet.contains(block);
 }
 
 
@@ -333,9 +333,24 @@ struct PostorderComputationContext : public DepthFirstSearchContext
     }
 };
 
+void computeReachableSet(IRGlobalValueWithCode* code, HashSet<IRBlock*>& outSet)
+{
+    DepthFirstSearchContext context;
+    if (code->getFirstBlock())
+        context.walk(code->getFirstBlock(), [](IRBlock* block) {return block->getSuccessors(); });
+    outSet = _Move(context.visited);
+}
+
 /// Compute a postorder traversal of the blocks in `code`, writing the resulting order to `outOrder`.
 void computePostorder(IRGlobalValueWithCode* code, List<IRBlock*>& outOrder)
 {
+    HashSet<IRBlock*> reachableSet;
+    computePostorder(code, outOrder, reachableSet);
+}
+
+/// Compute a postorder traversal of the blocks in `code`, writing the resulting order to `outOrder`.
+void computePostorder(IRGlobalValueWithCode* code, List<IRBlock*>& outOrder, HashSet<IRBlock*>& outReachableSet)
+{
     PostorderComputationContext context;
     context.order = &outOrder;
     if (code->getFirstBlock())
@@ -352,6 +367,7 @@ void computePostorder(IRGlobalValueWithCode* code, List<IRBlock*>& outOrder)
     }
     prefix.addRange(outOrder);
     outOrder = _Move(prefix);
+    outReachableSet = _Move(context.visited);
 }
 
 void computePostorderOnReverseCFG(IRGlobalValueWithCode* code, List<IRBlock*>& outOrder)
@@ -397,6 +413,10 @@ struct DominatorTreeComputationContext
     // traversal, so that we can look up a block based on its "name"
     //
     List<IRBlock*> postorder;
+    // 
+    // Also maintain a set of reachable blocks.
+    //
+    HashSet<IRBlock*> reachableSet;
 
     //
     // We need a way to map our actual IR blocks to their names for
@@ -426,7 +446,7 @@ struct DominatorTreeComputationContext
     void iterativelyComputeImmediateDominators(IRGlobalValueWithCode* code)
     {
         // First we compute the postorder traversal order for the blocks in the CFG.
-        computePostorder(code, postorder);
+        computePostorder(code, postorder, reachableSet);
 
         // We will initialize our map from the block objects to their "name"
         // (index in the traversal order), before moving on.
@@ -746,6 +766,7 @@ struct DominatorTreeComputationContext
         RefPtr<IRDominatorTree> dominatorTree = new IRDominatorTree();
         dominatorTree->code = code;
         dominatorTree->nodes.setCount(blockCount);
+        dominatorTree->reachableSet = _Move(reachableSet);
 
         // We will iterate over all of the blocks, and fill in the corresponding
         // dominator tree node for each.