Bläddra i källkod

Reroute Cache Feature:

This commit adds a Reroute Cache to Mantis' UI, which records the start/end
of reroute networks as the user builds the tree so that there is no need to trace
through the reroute network during tree parsing.

This is a significant optimization, saving as much as 0.2 seconds in my stego test
graph.
Joseph Brandenburg 2 veckor sedan
förälder
incheckning
d5dd06f335
7 ändrade filer med 293 tillägg och 79 borttagningar
  1. 22 1
      __init__.py
  2. 201 19
      base_definitions.py
  3. 24 15
      execution/readtree.py
  4. 8 1
      mantis_dataclasses.py
  5. 1 1
      menu_classes.py
  6. 2 2
      ops_nodegroup.py
  7. 35 40
      utilities/utilities.py

+ 22 - 1
__init__.py

@@ -389,6 +389,17 @@ def begin_selection_sync_timer(filename):
         SELECTION_POLL_RATE = bl_mantis_addon.preferences.SelectionSyncPollRate
         bpy.app.timers.register(timer_update_selection_sync)
 
+# Build the reroute cache on file load
+@persistent
+def on_load_build_reroute_cache(filename):
+    for tree in bpy.data.node_groups:
+        if tree.bl_idname in ['MantisTree', 'SchemaTree']:
+            # and I have to do this somewhere, this is the best place...
+            tree.is_quitting = False # reset the quitting flag
+            from .base_definitions import hash_tree, build_reroute_cache
+            tree.hash = str(hash_tree(tree, recurse=False))
+            build_reroute_cache(tree, tree.links)
+
 # I'll need to do some fiddling here when it comes time to try
 #   and make rig definitions animatable.
 @persistent
@@ -407,10 +418,18 @@ def on_undo_post_handler(scene): # the undo will trigger a depsgraph update
     for t in bpy.data.node_groups: # so we enable prevent_next_exec.
         if t.bl_idname in ['MantisTree', 'SchemaTree']:
             t.prevent_next_exec = True
-            t.hash=""
+            t.hash_recursive=""
             # set the tree to invalid to trigger a tree update
             # since the context data is wiped by an undo.
 
+@persistent
+def on_quit_handler(is_ui):
+    # we'll let the node trees know they're quitting
+    # this prevents them from spamming updates on file close.
+    for tree in bpy.data.node_groups:
+        if tree.bl_idname in ['MantisTree', 'SchemaTree']:
+            tree.is_quitting=True
+
 from .menu_classes import (node_context_menu_draw, node_add_menu_draw,
                            armature_add_menu_draw, import_menu_draw)
 
@@ -446,9 +465,11 @@ def register():
     bpy.app.handlers.load_post.insert(0, begin_selection_sync_timer)
     bpy.app.handlers.load_post.insert(0, autoload_components)
     bpy.app.handlers.load_post.insert(0, version_update_handler)
+    bpy.app.handlers.load_post.insert(0, on_load_build_reroute_cache)
     bpy.app.handlers.animation_playback_pre.insert(0, on_animation_playback_pre_handler)
     bpy.app.handlers.animation_playback_post.insert(0, on_animation_playback_post_handler)
     bpy.app.handlers.undo_post.insert(0, on_undo_post_handler)
+    bpy.app.handlers.exit_pre.insert(0, on_quit_handler)
     # I'm adding mine in first to ensure other addons don't mess up mine
     # but I am a good citizen! so my addon won't mess up yours! probably...
 

+ 201 - 19
base_definitions.py

@@ -13,12 +13,21 @@ from .utilities.utilities import (prRed, prGreen, prPurple, prWhite,
 from .utilities.utilities import get_socket_maps, relink_socket_map, do_relink
 from .mantis_dataclasses import MantisSocketTemplate
 
+MANTIS_VERSION_MAJOR=0
+MANTIS_VERSION_MINOR=13
+MANTIS_VERSION_SUB=5
+
 FLOAT_EPSILON=0.0001 # used to check against floating point inaccuracy
 links_sort_key= lambda a : (-a.multi_input_sort_id, -a.sub_sort_id)
 
 selection_sync_node_to_object = {} # what objects a node spawns
 selection_sync_object_to_node = {} # what node an object comes from
 
+# The reroute cache is used to accelerate reading reroute nodes
+# it uses the tree name and hash to ensure reroutes are only updated when needed
+# reroutes recursively find their source and remember it
+reroute_cache = {}
+
 def TellClasses():
     #Why use a function to do this? Because I don't need every class to register.
     return [ MantisTree,
@@ -44,15 +53,13 @@ def valid_interface_types(cls : NodeTree, socket_idname : str):
         return ['NodeSocketGeometry']
     else: # once versioning is finished this will be unnecesary.
         return socket_idname in tell_valid_bl_idnames()
-
 def fix_reroute_colors(tree):
     context = bpy.context
     if any((tree.is_executing, tree.is_exporting, tree.do_live_update==False, context.space_data is None) ):
         return
-    from collections import deque
     from .utilities.utilities import find_reroute_start_socket
     from .socket_definitions import MantisSocket
-    reroutes_without_color = deque()
+    reroutes_without_color = []
     for n in tree.nodes:
         if n.bl_idname=='NodeReroute' and n.inputs[0].bl_idname == "NodeSocketColor":
             reroutes_without_color.append(n)
@@ -69,11 +76,11 @@ def fix_reroute_colors(tree):
 
 #functions to identify the state of the system using hashes
 # this function runs a lot so it should be optimized as well as possible.
-def hash_tree(tree):
+def hash_tree(tree, recurse=True):
     trees=set(); links=[]; hash_data=""
     for node in tree.nodes:
         hash_data+=str(node.name)
-        if hasattr(node, 'node_tree') and node.node_tree:
+        if recurse and hasattr(node, 'node_tree') and node.node_tree:
             trees.add(node.node_tree)
     for other_tree in trees:
         hash_data+=str(hash_tree(other_tree))
@@ -84,9 +91,157 @@ def hash_tree(tree):
     links.sort(); hash_data+=''.join(links)
     return hash(hash_data)
 
-MANTIS_VERSION_MAJOR=0
-MANTIS_VERSION_MINOR=13
-MANTIS_VERSION_SUB=5
+# TODO: I should see if I can make these work without every using NodeSocket.links
+#  as that getter is very slow at scale
+# however with the cache in place it is probably negligible
+def update_reroute_cache_item_source_recursively(reroute, tree_cache, source):
+    # find all the inputs to this reroute and update their destination sets
+    # every reroute ancestor pushes to this destination... but not every recursive-child
+    #   of the root reroute necessarily does, so pushing back like this is right
+    # since each reroute node only has one ancestor, this is easy
+    from .mantis_dataclasses import RerouteCacheItem
+    tree= reroute.id_data
+    reroutes = {reroute}
+    while(reroutes):
+        reroute = reroutes.pop()
+        # for making a new cache item
+        if (cache_object := tree_cache.get(reroute.name)) is None:
+            cache_object = RerouteCacheItem(
+                    reroute_name=reroute.name,
+                    reroute_tree_name=tree.name,
+                )
+            tree_cache[reroute.name] = cache_object
+        cache_object.source = source
+        if reroute.outputs[0].is_linked: # reroutes are always single-output
+            for l in reroute.outputs[0].links: # with potentially many links
+                if l.to_node.bl_idname == 'NodeReroute':
+                    reroutes.add(l.to_node)
+
+def update_reroute_cache_item_destination_recursively(reroute, tree_cache, destination):
+    # find all the inputs to this reroute and update their destination sets
+    # every reroute ancestor pushes to this destination... but not every recursive-child
+    #   of the root reroute necessarily does, so pushing back like this is right
+    # since each reroute node only has one ancestor, this is easy
+    from .mantis_dataclasses import RerouteCacheItem
+    tree= reroute.id_data
+    while(reroute):
+        # for making a new cache item
+        if (cache_object := tree_cache.get(reroute.name)) is None:
+            cache_object = RerouteCacheItem(
+                    reroute_name=reroute.name,
+                    reroute_tree_name=tree.name,
+                    destinations= {destination},
+                )
+            tree_cache[reroute.name] = cache_object
+        else:
+            cache_object.destinations.add(destination)
+        # Odd: if I check reroute.inputs[0].is_linked with live updates, CRASH.
+        if rr_links := reroute.inputs[0].links:
+            # reroutes are always single-input with at most 1 link
+            if (l := rr_links[0]).from_node == 'NodeReroute':
+                reroute = l.from_node
+                continue
+        reroute = None
+
+# Update the reroute cache from the link
+# NOTE this is updated as the user adds links by the node.insert_link() callback
+def reroute_cache_handle_link(tree, tree_cache, link):
+    from .mantis_dataclasses import RerouteCacheItem
+    # if it is a link to or from a reroute, add or update the cache item
+    # first check if the cache item is stale
+    # then try and fill the source and destination - make sure to recursively fill them
+    # if it isn't connected to a source or destination, it's part of a network
+    # We only need to check the ends, and we can do both the destination and source
+    #  passes simultaneously since they don't depend on each other.
+    # NOTE that we use socket index to make things explicit and fast
+    if link.from_node.bl_idname == 'NodeReroute' and link.to_node.bl_idname != 'NodeReroute':
+        other_node = link.to_node
+        other_socket = link.to_socket
+        for socket_index, check_socket in enumerate(other_node.inputs):
+            if check_socket.identifier == other_socket.identifier: break
+        else:
+            raise RuntimeError("Error, cannot make reroute cache for link: "
+                               f"{link.from_node.name}:{link.from_socket.name} ==>"
+                               f"{link.to_node.name}:{link.to_socket.name}")
+        reroute = link.from_node
+        destination = (other_node.name, socket_index)
+        if (cache_object := tree_cache.get(reroute.name)) is None:
+            cache_object = RerouteCacheItem(
+                reroute_name=reroute.name,
+                reroute_tree_name=tree.name,
+                destinations= {destination},
+            )
+            tree_cache[reroute.name]=cache_object
+        update_reroute_cache_item_destination_recursively(reroute, tree_cache,
+            destination=destination)
+    elif link.from_node.bl_idname != 'NodeReroute' and link.to_node.bl_idname == 'NodeReroute':
+        other_node = link.from_node
+        other_socket = link.from_socket
+        for socket_index, check_socket in enumerate(other_node.outputs):
+            if check_socket.identifier == other_socket.identifier: break
+        else:
+            raise RuntimeError("Error, cannot make reroute cache for link: "
+                               f"{link.from_node.name}:{link.from_socket.name} ==>"
+                               f"{link.to_node.name}:{link.to_socket.name}")
+        reroute = link.to_node
+        source = (link.from_node.name, socket_index)
+        if (cache_object := tree_cache.get(reroute.name)) is None:
+            cache_object = RerouteCacheItem(
+                reroute_name=reroute.name,
+                reroute_tree_name=tree.name,
+                source = source,
+            )
+            tree_cache[reroute.name]=cache_object
+        update_reroute_cache_item_source_recursively(reroute, tree_cache,
+            source=source)
+
+def build_reroute_cache(tree, links):
+    # we'll track this to know if a tree update should invalidate the hash on incremental update
+    tree.number_of_reroute_links = 0
+    # make sure to clear the previous cache or confusing things will happen
+    tree_cache = {}; reroute_cache[tree.name] = tree_cache
+    tree_cache['_MANTIS_RESERVED_TREE_HASH'] = tree.hash
+    for link in links:
+        if link.from_node.bl_idname == 'NodeReroute' or link.to_node.bl_idname == 'NodeReroute':
+            tree.number_of_reroute_links += 1
+        reroute_cache_handle_link(tree, tree_cache, link)
+
+def tree_update_common(tree): # set the reroute colors
+    if tree.is_quitting: return
+    if (bpy.app.version == (4,4,0)): # CTD bug in that version
+        return
+    fix_reroute_colors(tree)
+    # intialize it if e.g the tree is new.
+    if (tree_cache := reroute_cache.get(tree.name)) is None:
+        tree_cache = {}; reroute_cache[tree.name] = tree_cache
+    # events fire in this order:
+    #   - node.insert_link (one for each node)
+    #   - node.update (one for each node)
+    #   - tree.update
+    # so if the node's insert_link handled things, we don't need to check
+    tree.hash = str(hash_tree(tree, recurse=False))
+    if tree.reroute_link_handled: # we can update the hash, it's still good.
+        if tree_cache := reroute_cache.get(tree.name):
+            tree_cache["_MANTIS_RESERVED_TREE_HASH"]=tree.hash
+        tree.reroute_link_handled = False # always make sure to reset the flag
+    else:
+        reroute_links = []
+        # set the UI tree hash for e.g. reroute cache
+        for link in tree.links:
+            if link.from_node.bl_idname == 'NodeReroute' or link.to_node.bl_idname == 'NodeReroute':
+                reroute_links.append(link)
+        if len(reroute_links) == tree.number_of_reroute_links:
+            if tree_cache := reroute_cache.get(tree.name): # we can update the hash, it's still good.
+                tree_cache["_MANTIS_RESERVED_TREE_HASH"] = tree.hash
+            else: # there is no cache after all. How did that happen?
+                build_reroute_cache(tree, reroute_links)
+            # NOTE I put a lot of effort into catching deleted links but it just isn't practical
+            #      to catch this and it is almost free to rebuild the cache.
+        else:
+            build_reroute_cache(tree, reroute_links)
+            # reset the count but do not update the hash.
+        tree.number_of_reroute_links = len(reroute_links)
+
 
 class MantisTree(NodeTree):
     '''A custom node tree type that will show up in the editor type list'''
@@ -95,7 +250,8 @@ class MantisTree(NodeTree):
     bl_icon = 'OUTLINER_OB_ARMATURE'
 
     tree_valid:BoolProperty(default=False)
-    hash:StringProperty(default='')
+    hash:StringProperty(default='') # use for UI updates, reroute cache
+    hash_recursive:StringProperty(default='') # use for checking Mantis state
     do_live_update:BoolProperty(default=True) # use this to disable updates for e.g. scripts
     # operator settings for re-exporting the tree.
     filepath:StringProperty(default="", subtype='FILE_PATH')
@@ -103,6 +259,7 @@ class MantisTree(NodeTree):
     #
     is_executing:BoolProperty(default=False)
     is_exporting:BoolProperty(default=False)
+    is_quitting:BoolProperty(default=False)
     execution_id:StringProperty(default='')
     # prev_execution_id:StringProperty(default='')
     mantis_version:IntVectorProperty(default=[
@@ -114,6 +271,11 @@ class MantisTree(NodeTree):
     #added to work around a bug in 4.5.0 LTS
     interface_helper : StringProperty(default='')
 
+    # for reroute cache, to check if it is still good
+    number_of_reroute_links : IntProperty(default=0)
+    # flag to indicate a link was already handled this update
+    reroute_link_handled : BoolProperty(default=False)
+
     parsed_tree={}
 
     # NOTE that these are class properties that should be serialized
@@ -129,16 +291,19 @@ class MantisTree(NodeTree):
         def valid_socket_type(cls : NodeTree, socket_idname: str):
             return valid_interface_types(cls, socket_idname)
 
-    def update(self): # set the reroute colors
-        if (bpy.app.version >= (4,4,0)):
-            fix_reroute_colors(self)
+    def update(self):
+        tree_update_common(self)
 
     def update_tree(self, context = None, force=False, error_popups=False):
         if self.is_exporting:
             return
         my_hash = str( hash_tree(self) )
-        if (my_hash != self.hash) or force:
-            self.hash = my_hash
+        my_ui_hash = str( hash_tree(self, recurse=False))
+        tree_cache = reroute_cache.get(self.name)
+        if tree_cache['_MANTIS_RESERVED_TREE_HASH'] != my_ui_hash:
+            build_reroute_cache(self, self.links) # somehow it became invalid, rebuild it.
+        if (my_hash != self.hash_recursive) or force:
+            self.hash_recursive = my_hash
             self.is_executing = True
             from .execution import readtree
             prGreen("Validating Tree: %s" % self.name)
@@ -154,7 +319,7 @@ class MantisTree(NodeTree):
             except Exception as e:
                 prRed("Failed to update node tree due to error.")
                 self.tree_valid = False
-                self.hash='' # unset the hash to mark the tree as un-parsed.
+                self.hash_recursive='' # unset the hash to mark the tree as un-parsed.
                 raise e
             finally:
                 scene.render.use_lock_interface = False
@@ -185,7 +350,7 @@ class MantisTree(NodeTree):
 
     def execute_tree(self,context, error_popups = False, profile=False):
         self.prevent_next_exec = False
-        if not self.hash:
+        if not self.hash_recursive:
             return
         if self.is_exporting or self.is_executing:
             return
@@ -211,20 +376,25 @@ class SchemaTree(NodeTree):
     do_live_update:BoolProperty(default=True) # default to true so that updates work
     is_executing:BoolProperty(default=False)
     is_exporting:BoolProperty(default=False)
+    is_quitting:BoolProperty(default=False)
 
     mantis_version:IntVectorProperty(default=[
         MANTIS_VERSION_MAJOR, MANTIS_VERSION_MINOR, MANTIS_VERSION_SUB])
     # see the note in MantisTree
     interface_helper : StringProperty(default='')
 
+    # for reroute cache, to check if it is still good
+    number_of_reroute_links : IntProperty(default=0)
+    # flag to indicate a link was already handled this update
+    reroute_link_handled : BoolProperty(default=False)
+
     if (bpy.app.version < (4, 4, 0) or bpy.app.version >= (4,5,0)):  # in 4.4 this leads to a crash
         @classmethod
         def valid_socket_type(cls : NodeTree, socket_idname: str):
             return valid_interface_types(cls, socket_idname)
 
-    def update(self): # set the reroute colors
-        if (bpy.app.version >= (4,4,0)):
-            fix_reroute_colors(self)
+    def update(self):
+        tree_update_common(self)
 
 
 #TODO: do a better job explaining how MantisNode and MantisUINode relate.
@@ -290,6 +460,18 @@ class MantisUINode:
             if template.use_multi_input: # this is an array
                 socket.display_shape = 'SQUARE_DOT'
 
+    def insert_link(self, link):
+        if (bpy.app.version == (4, 4, 0)):
+            return # this causes a crash due to a bug.
+        tree = self.id_data
+        if tree.reroute_link_handled: return
+        if ( link.from_node.bl_idname == 'NodeReroute' or
+             link.to_node.bl_idname == 'NodeReroute' ):
+                tree.number_of_reroute_links+=1
+                reroute_cache_handle_link(tree=tree,
+                    tree_cache=reroute_cache.get(tree.name), link=link)
+                tree.reroute_link_handled = True
+        
 # NOTE
 # MantisUINode has no proper mantis libraryclass on its own, so:
 # register()'s are in SUBclasses; they have to run AFTER library is defined

+ 24 - 15
execution/readtree.py

@@ -277,7 +277,7 @@ def solve_schema_to_tree(mantis_node, all_mantis_nodes, roots=[], error_popups=F
     except Exception as e:
     #     # the schema will run the error cleanup code, we just need to raise or not
         solved_nodes = {}
-        mantis_node.base_tree.hash=''
+        mantis_node.base_tree.hash_recursive=''
         raise execution_error_cleanup(mantis_node, e, show_error=error_popups)
     for k,v in solved_nodes.items():
         all_mantis_nodes[k]=v
@@ -412,6 +412,7 @@ def parse_tree(base_tree, error_popups=False):
     base_tree.execution_id = uuid4().__str__() # set the unique id of this execution
 
     from ..base_definitions import MantisExecutionContext
+    from ..utilities.node_common import GraphError
     mContext = MantisExecutionContext(base_tree=base_tree)
 
     import time
@@ -420,15 +421,25 @@ def parse_tree(base_tree, error_popups=False):
     #  then they stick around because the function definition inits them once and keeps a reference
     # so instead I have to supply them to avoid ugly code or bugs elsewhere
     # it's REALLY confusing when you run into this sort of problem. So it warrants four entire lines of comments!      
-    dummy_nodes, all_mantis_nodes, all_schema =  data_from_tree(
-        base_tree,
-        tree_path = [None], 
-        mContext=mContext,
-        dummy_nodes = {},
-        all_mantis_nodes = {},
-        all_schema={},
-        in_schema=False
-        )
+    try:
+        dummy_nodes, all_mantis_nodes, all_schema =  data_from_tree(
+            base_tree,
+            tree_path = [None], 
+            mContext=mContext,
+            dummy_nodes = {},
+            all_mantis_nodes = {},
+            all_schema={},
+            in_schema=False
+            )
+    except GraphError as e:
+        base_tree.tree_valid = False; base_tree.hash_recursive=""
+        if error_popups:
+            from bpy import context
+            def error_popup_draw(self, context):
+                error_popup_draw_inner(self.layout, e, ())
+            context.window_manager.popup_menu(error_popup_draw, title="Error", icon='ERROR')
+        else: raise e
+        return # it is not possible to parse the tree until GraphErrors are resolved.
     for dummy in dummy_nodes.values():    # reroute the links in the group nodes
         dummy.reroute_links(all_mantis_nodes)
     prGreen(f"Pulling data from tree took {time.time() - data_start_time} seconds")
@@ -575,11 +586,9 @@ def parse_tree(base_tree, error_popups=False):
     return kept_mantis_node
 
 from ..utilities.utilities import switch_mode
-
-
-def error_popup_draw_inner(layout, exception, ui_sig):
-    message = f"Error: {exception}\n" \
-              f"See node: {ui_sig[1:]}."
+def error_popup_draw_inner(layout, exception, ui_sig = ()):
+    node_string = f"See node: {ui_sig[1:]}." if ui_sig else ""
+    message = f"Error: {exception}\n" + node_string  
     lines = message.splitlines()
     for line in lines:
         layout.label(text=line)

+ 8 - 1
mantis_dataclasses.py

@@ -99,4 +99,11 @@ class MantisSocketFiller():
                f"property_name: {self.property_name}, " \
                f"default_value: {self.default_value}, " \
                f"array_index: {self.array_index or 0}, " \
-               f"value_index: {self.value_index or None}" 
+               f"value_index: {self.value_index or None}"
+
+@dataclass(slots=True)
+class RerouteCacheItem:
+    reroute_name      : str = field(default="") # the reroute node name in the tree 
+    reroute_tree_name : str = field(default="") # the tree to find the reroute in
+    source            : Tuple[str, str] = field(default_factory=tuple) # the beginning of the reroute network
+    destinations      : set[Tuple[str, str]] = field(default_factory=set) # (node_name, socket_name)

+ 1 - 1
menu_classes.py

@@ -119,7 +119,7 @@ class MantisActiveTreePanel(Panel):
         spc = area.spaces[0]
         nt = spc.node_tree
         layout = self.layout
-        layout.label(text=f"Tree Hash: {nt.hash}")
+        layout.label(text=f"Tree Hash: {nt.hash_recursive}")
         layout.prop(nt, "do_live_update", text='Live Updates',)
         layout.operator("mantis.invalidate_node_tree")
         layout.operator("mantis.execute_node_tree")

+ 2 - 2
ops_nodegroup.py

@@ -27,7 +27,7 @@ def tag_mantis_node_for_update(context, node):
         mantis_node = node_tree.parsed_tree.get(
                     get_signature_from_edited_tree(node, context))
         if mantis_node is None:
-            node_tree.hash = ''
+            node_tree.hash_recursive = ''
         else:
             mantis_node.reset_execution_recursive() # tag this as a change.
 #########################################################################3
@@ -256,7 +256,7 @@ class InvalidateNodeTree(Operator):
     def execute(self, context):
         print("Clearing Active Node Tree cache")
         tree=context.space_data.path[0].node_tree
-        tree.execution_id=''; tree.hash=''
+        tree.execution_id=''; tree.hash_recursive=''
         return {"FINISHED"}
 
 class ExecuteNodeTree(Operator):

+ 35 - 40
utilities/utilities.py

@@ -26,32 +26,26 @@ def prOrange(*args): print (*[wrapOrange(arg) for arg in args])
 def float_lerp(a : float, b : float, factor : float)  -> float:
     return (a * (1.0 - factor)) + (b * factor)
 
-# TODO: BUG may live here
-#I need to stress-test this to make sure it doesn't hang on:
-# - empty reroute nodes with no links
-# - links starting with reroute nodes with no inputs
-# - links ending with reroute nodes with no outputs
 def find_reroute_start_socket(reroute, track='BACK'):
-    # "BACK" traces back through the tree
-    # "FORWARD" traces forward through the tree
-    socket = None
-    while (reroute and track == 'BACK'):
-        if len(reroute.inputs[0].links) == 1:
-            link = reroute.inputs[0].links[0]
-            socket = link.from_socket
-            if link.from_node.bl_idname == 'NodeReroute':
-                reroute =  link.from_node
-            else:
-                link, reroute = None, None
-    while (reroute and track == 'FORWARD'):
-        if len(reroute.outputs[0].links) == 1:
-            link = reroute.outputs[0].links[0]
-            socket = link.to_socket
-            if link.to_node.bl_idname == 'NodeReroute':
-                reroute =  link.to_node
-            else:
-                link, reroute = None, None
-    return socket
+    # "BACK" traces back through the tree and returns one socket
+    # "FORWARD" traces forward through the tree and returns a list of sockets
+    from ..base_definitions import reroute_cache
+    from .node_common import GraphError
+    # We assume the cache is good, the tree is responsible for this
+    tree = reroute.id_data
+    tree_cache = reroute_cache.get(tree.name)
+    if track == 'FORWARD':
+        cache_item = tree_cache.get(reroute.name); sockets = []
+        if not cache_item: # I think this is OK, they don't always go somewhere
+            return # And this FORWARD trace is always a background task
+        for node_name, socket_index in cache_item.destinations:
+            sockets.append(tree.nodes.get(node_name).inputs[socket_index])
+        return sockets
+    if track == 'BACK':
+        cache_item = tree_cache.get(reroute.name)
+        if not cache_item or not cache_item.source:
+            raise GraphError(f"There is a missing link in the tree at reroute node: {reroute.name}")
+        return tree.nodes.get(cache_item.source[0]).outputs[cache_item.source[1]]
 
 # this creates fake links that have the same interface as Blender's
 # so that I can bypass Reroutes
@@ -208,25 +202,26 @@ def do_relink(node, socket, map, in_out='INPUT', parent_name = ''):
         for sub_val in val:
             # this will only happen once because it assigns socket, so it is safe to do in the for loop.
             if socket is None:
-                socket = sub_val
+                socket = sub_val; sockets = [socket]
                 if sub_val.bl_idname == "NodeReroute":
                     # we have to trace the reroute node...
                     if in_out == 'INPUT':
-                        socket = find_reroute_start_socket(sub_val)
+                        socket = find_reroute_start_socket(sub_val); sockets = [socket]
                     else:
-                        socket = find_reroute_start_socket(sub_val, track="FORWARD")
-                sock_type = socket.interface_type
-                name = unique_socket_name(node, socket, tree)
-                if parent_name:
-                    interface_socket = update_interface(tree.interface, name, interface_in_out, sock_type, parent_name)
-                if in_out =='INPUT':
-                    socket = node.inputs.new(sock_type, name, identifier=interface_socket.identifier)
-                else:
-                    socket = node.outputs.new(sock_type, name, identifier=interface_socket.identifier)
-                if parent_name == 'Array': socket.display_shape='SQUARE_DOT'
-                if parent_name == 'Constant': socket.display_shape='CIRCLE_DOT'
-                # then move it up and delete the other link.
-                # this also needs to modify the interface of the node tree.
+                        sockets = find_reroute_start_socket(sub_val, track="FORWARD")
+                for socket in sockets:
+                    sock_type = socket.interface_type
+                    name = unique_socket_name(node, socket, tree)
+                    if parent_name:
+                        interface_socket = update_interface(tree.interface, name, interface_in_out, sock_type, parent_name)
+                    if in_out =='INPUT':
+                        socket = node.inputs.new(sock_type, name, identifier=interface_socket.identifier)
+                    else:
+                        socket = node.outputs.new(sock_type, name, identifier=interface_socket.identifier)
+                    if parent_name == 'Array': socket.display_shape='SQUARE_DOT'
+                    if parent_name == 'Constant': socket.display_shape='CIRCLE_DOT'
+                    # then move it up and delete the other link.
+                    # this also needs to modify the interface of the node tree.
             if isinstance(sub_val, NodeSocket):
                 counter_id=sub_val.node.name+"::"+sub_val.identifier
                 l = None