Analysis/depgraph: improve get_modifiers complexity

1 files changed, 46 insertions, 34 deletions

diff --git a/miasm2/analysis/depgraph.py b/miasm2/analysis/depgraph.py
index 72941e1b..8d6da3b2 100644
--- a/miasm2/analysis/depgraph.py
+++ b/miasm2/analysis/depgraph.py
@@ -246,47 +246,59 @@ class DependencyDict(object):
         resolution"""
         return self._pending
 
-    def _get_modifiers_in_cache(self, depnode, force=False):
-        """Recursively find nodes in the path of @depnode which are modifiers.
-        Update the internal cache If @depnode is already managed (ie. in
-        @depnode_queued), abort
-        @depnode: DependencyNode
-        @force (optionnal): boolean, compute modifiers even if @depnode is a
-        modifier node.
-        """
-
-        # Base case
-        if depnode not in self._cache:
-            # Constant does not have any dependencies
-            return [depnode] if depnode.modifier else []
-
-        if depnode.modifier and not force:
-            return [depnode]
-
-        # Recursion
-        dependencies = self._cache[depnode]
+    def _get_modifiers_in_cache(self, nodes_heads):
+        """Find modifier nodes in cache starting from @nodes_heads.
+        Returns new cache"""
+        # 'worklist_depnode' order is needed (depth first)
+        worklist_depnodes = list(nodes_heads)
+        # Temporary cache
+        cache = {}
+        # Partially resolved 'cache' elements
+        worklist = []
+
+        # Build worklist and cache for non modifiers
+        while worklist_depnodes:
+            depnode = worklist_depnodes.pop()
+            # Resolve node dependencies
+            if depnode in cache:
+                # Depnode previously resolved
+                continue
 
-        out = set()
-        # Launch on each depnodes
-        parallels = []
-        for depnode in dependencies:
-            parallels.append(self._get_modifiers_in_cache(depnode))
+            if depnode not in self._cache:
+                # Final node
+                if not depnode.modifier:
+                    cache[depnode] = []
+                continue
 
-        if parallels:
-            for parallel in itertools.product(*(p for p in parallels if p)):
+            # Propagate to son
+            dependencies = self._cache[depnode]
+            for son in dependencies:
+                worklist_depnodes.append(son)
+            # Save partially resolved dependency
+            worklist.append((depnode, dependencies))
+
+        # Convert worklist to cache
+        while worklist:
+            depnode, dependencies = worklist.pop()
+            parallels = []
+            for node in dependencies:
+                if node.modifier:
+                    parallels.append([node])
+                else:
+                    parallels.append(cache[node])
+            out = set()
+            for parallel in itertools.product(*[p for p in parallels if p]):
                 out.update(parallel)
+            cache[depnode] = out
 
-        return out
+        return cache
 
-    def clean_modifiers_in_cache(self):
+    def clean_modifiers_in_cache(self, node_heads):
         """Remove intermediary states (non modifier depnodes) in the internal
         cache values"""
 
-        cache_out = CacheWrapper()
-        for depnode in self._cache:
-            cache_out[depnode] = self._get_modifiers_in_cache(depnode,
-                                                              force=True)
-        self._cache = cache_out
+        self._cache = CacheWrapper(self._get_modifiers_in_cache(node_heads))
+
 
     def _build_depgraph(self, depnode):
         """Recursively build the final list of DiGraph, and clean up unmodifier
@@ -972,7 +984,7 @@ class DependencyGraph(object):
 
         for final_depdict in depdicts:
             # Keep only relevant nodes
-            final_depdict.clean_modifiers_in_cache()
+            final_depdict.clean_modifiers_in_cache(input_depnodes)
             final_depdict.filter_used_nodes(input_depnodes)
 
             # Remove duplicate solutions