rewrite topological sort, avoiding calls to the expensive 'links' property of the NodeSocket class

api/arrange.py

@@ -1,34 +1,40 @@
 import bpy
 import typing
+from collections import deque
 
 def _arrange(node_tree, padding: typing.Tuple[float, float] = (50, 25)):
     # Organize the nodes into columns based on their links.
     columns: typing.List[typing.List[typing.Any]] = []
-    def contains_link(node, column):
+    
-        return any(
+    def topo_sort(graph):
-            any(
+        in_degree = {u: 0 for u in graph}
-                any(link.from_node == node for link in input.links)
+        for u in graph:
-                for input in n.inputs
+            for v in graph[u]:
-            )
+                in_degree[v] += 1
-            for n in column
+        queue = deque([u for u in in_degree if in_degree[u] == 0])
-        )
+        topo_order = []
-    for node in reversed(node_tree.nodes):
+        while queue:
-        if (x := next(
+            u = queue.popleft()
-            filter(
+            topo_order.append(u)
-                lambda x: contains_link(node, x[1]),
+            for v in graph[u]:
-                enumerate(columns)
+                in_degree[v] -= 1
-            ),
+                if in_degree[v] == 0:
-            None
+                    queue.append(v)
-        )) is not None:
+        return topo_order
-            if x[0] > 0:
+    
-                columns[x[0] - 1].append(node)
+    graph = { node:set() for node in node_tree.nodes }
-            else:
+    for link in node_tree.links:
-                columns.insert(x[0], [node])
+        graph[link.from_node].add(link.to_node)
+    sorted_nodes = topo_sort(graph)
+
+    column_index = {}
+    for node in reversed(sorted_nodes):
+        column_index[node] = max([ column_index[node] for node in graph[node] ], default=-1) + 1
+        if column_index[node] == len(columns):
+            columns.append([node])
         else:
-            if len(columns) == 0:
+            columns[column_index[node]].append(node)
-                columns.append([node])
+    columns = reversed(columns)
-            else:
-                columns[len(columns) - 1].append(node)
     
     # Arrange the columns, computing the size of the node manually so arrangement can be done without UI being visible.
     UI_SCALE = bpy.context.preferences.view.ui_scale