improve level compute algo

paraffin/utils.py

@@ -136,26 +136,26 @@ def dag_to_levels(graph) -> HirachicalStages:
         >>> dag_to_levels(G)
         {0: [1], 1: [2, 3], 2: [4]}
     """
-    nodes = []
+    # Initialize the levels dictionary
     levels = {}
-    for start_node in graph.nodes():
-        if len(list(graph.predecessors(start_node))) == 0:
-            if start_node not in nodes:
-                for node in nx.bfs_tree(graph, start_node):
-                    if node not in nodes:
-                        nodes.append(node)
-                        # find the longest path from the start_node to the current node
-                        # to determine the level of the current node
-                        level = 0
-                        for path in nx.all_simple_paths(graph, start_node, node):
-                            level = max(level, len(path) - 1)
-                        try:
-                            levels[level].append(node)
-                        except KeyError:
-                            levels[level] = [node]
-                    else:
-                        # this part has already been added
-                        break
+    node_to_level = {}
+
+    # Perform topological sort
+    for node in nx.topological_sort(graph):
+        # Determine the level of the current node
+        if len(list(graph.predecessors(node))) == 0:
+            # No predecessors, this is a root node
+            level = 0
+        else:
+            # Level is one more than the maximum level of all predecessors
+            level = max(node_to_level[pred] for pred in graph.predecessors(node)) + 1
+
+        # Assign the node to its level
+        node_to_level[node] = level
+        if level not in levels:
+            levels[level] = []
+        levels[level].append(node)
+
     return levels
 
 

tests/test_utils.py

@@ -62,3 +62,36 @@ def test_dag_to_levles_4():
     levels = dag_to_levels(digraph)
 
     assert levels == {0: ["A", "B", "C"], 1: ["D", "E"]}
+
+def test_dag_to_levels_5():
+    """
+    ```mermaid
+    flowchart TD
+        A --> D
+        B --> D
+        B --> E
+        C --> E
+
+        D --> F
+        E --> F
+
+        B --> G
+        E --> G
+
+        F --> H
+        C --> H
+
+        G --> I
+    ```
+    """
+    digraph = nx.DiGraph()
+    digraph.add_edges_from([
+        ("A", "D"), ("B", "D"), ("B", "E"), ("C", "E"),
+        ("D", "F"), ("E", "F"),
+        ("B", "G"), ("E", "G"),
+        ("F", "H"), ("C", "H"),
+        ("G", "I")
+    ])
+    levels = dag_to_levels(digraph)
+
+    assert levels == {0: ["A", "B", "C"], 1: ["D", "E"], 2: ["F", "G"], 3: ["H", "I"]}