Fix state-order traversal w.r.t. if branches without else

spcl · Oct 11, 2022 · 2f6f06d · 2f6f06d
1 parent 12732c5
commit 2f6f06d
Showing 1 changed file with 5 additions and 5 deletions.
diff --git a/dace/sdfg/analysis/cfg.py b/dace/sdfg/analysis/cfg.py
@@ -219,12 +219,10 @@ def _stateorder_topological_sort(sdfg: SDFG,
     """
     # Traverse states in custom order
     visited = visited or set()
-    if stop is not None:
-        visited.add(stop)
     stack = [start]
     while stack:
         node = stack.pop()
-        if node in visited:
+        if node in visited or node is stop:
             continue
         yield node
 
@@ -265,6 +263,9 @@ def _stateorder_topological_sort(sdfg: SDFG,
                 mergestate = stop
 
         for branch in oe:
+            if branch.dst is mergestate:
+                # If we hit the merge state (if without else), defer to end of branch traversal
+                continue
             for s in _stateorder_topological_sort(sdfg,
                                                   branch.dst,
                                                   ptree,
@@ -273,8 +274,7 @@ def _stateorder_topological_sort(sdfg: SDFG,
                                                   visited=visited):
                 yield s
                 visited.add(s)
-        if mergestate != stop:
-            stack.append(mergestate)
+        stack.append(mergestate)
 
 
 def stateorder_topological_sort(sdfg: SDFG) -> Iterator[SDFGState]: