Rework Element.dependencies to be more efficient

2 files changed, 51 insertions, 39 deletions

diff --git a/buildstream/_pipeline.py b/buildstream/_pipeline.py
index 004776293..5dec9001d 100644
--- a/buildstream/_pipeline.py
+++ b/buildstream/_pipeline.py
@@ -24,6 +24,8 @@ import itertools
 from operator import itemgetter
 from collections import OrderedDict
 
+from pyroaring import BitMap  # pylint: disable=no-name-in-module
+
 from ._exceptions import PipelineError
 from ._message import Message, MessageType
 from ._profile import Topics, profile_start, profile_end
@@ -152,7 +154,7 @@ class Pipeline():
     def dependencies(self, targets, scope, *, recurse=True):
         # Keep track of 'visited' in this scope, so that all targets
         # share the same context.
-        visited = {}
+        visited = (BitMap(), BitMap())
 
         for target in targets:
             for element in target.dependencies(scope, recurse=recurse, visited=visited):
diff --git a/buildstream/element.py b/buildstream/element.py
index 438a77a13..f1f0273f6 100644
--- a/buildstream/element.py
+++ b/buildstream/element.py
@@ -81,9 +81,12 @@ from collections.abc import Mapping
 import contextlib
 from contextlib import contextmanager
 from functools import partial
+from itertools import chain
 import tempfile
 import string
 
+from pyroaring import BitMap  # pylint: disable=no-name-in-module
+
 from . import _yaml
 from ._variables import Variables
 from ._versions import BST_CORE_ARTIFACT_VERSION
@@ -388,7 +391,7 @@ class Element(Plugin):
         for source in self.__sources:
             yield source
 
-    def dependencies(self, scope, *, recurse=True, visited=None, recursed=False):
+    def dependencies(self, scope, *, recurse=True, visited=None):
         """dependencies(scope, *, recurse=True)
 
         A generator function which yields the dependencies of the given element.
@@ -406,47 +409,54 @@ class Element(Plugin):
         Yields:
            (:class:`.Element`): The dependencies in `scope`, in deterministic staging order
         """
-        if visited is None:
-            visited = {}
+        # The format of visited is (BitMap(), BitMap()), with the first BitMap
+        # containing element that have been visited for the `Scope.BUILD` case
+        # and the second one relating to the `Scope.RUN` case.
+        if not recurse:
+            if scope in (Scope.BUILD, Scope.ALL):
+                yield from self.__build_dependencies
+            if scope in (Scope.RUN, Scope.ALL):
+                yield from self.__runtime_dependencies
+        else:
+            def visit(element, scope, visited):
+                if scope == Scope.ALL:
+                    visited[0].add(element._unique_id)
+                    visited[1].add(element._unique_id)
 
-        full_name = self._get_full_name()
+                    for dep in chain(element.__build_dependencies, element.__runtime_dependencies):
+                        if dep._unique_id not in visited[0] and dep._unique_id not in visited[1]:
+                            yield from visit(dep, Scope.ALL, visited)
 
-        scope_set = set((Scope.BUILD, Scope.RUN)) if scope == Scope.ALL else set((scope,))
+                    yield element
+                elif scope == Scope.BUILD:
+                    visited[0].add(element._unique_id)
 
-        if full_name in visited and scope_set.issubset(visited[full_name]):
-            return
+                    for dep in element.__build_dependencies:
+                        if dep._unique_id not in visited[1]:
+                            yield from visit(dep, Scope.RUN, visited)
 
-        should_yield = False
-        if full_name not in visited:
-            visited[full_name] = scope_set
-            should_yield = True
-        else:
-            visited[full_name] |= scope_set
-
-        if recurse or not recursed:
-            if scope == Scope.ALL:
-                for dep in self.__build_dependencies:
-                    yield from dep.dependencies(Scope.ALL, recurse=recurse,
-                                                visited=visited, recursed=True)
-
-                for dep in self.__runtime_dependencies:
-                    if dep not in self.__build_dependencies:
-                        yield from dep.dependencies(Scope.ALL, recurse=recurse,
-                                                    visited=visited, recursed=True)
-
-            elif scope == Scope.BUILD:
-                for dep in self.__build_dependencies:
-                    yield from dep.dependencies(Scope.RUN, recurse=recurse,
-                                                visited=visited, recursed=True)
-
-            elif scope == Scope.RUN:
-                for dep in self.__runtime_dependencies:
-                    yield from dep.dependencies(Scope.RUN, recurse=recurse,
-                                                visited=visited, recursed=True)
-
-        # Yeild self only at the end, after anything needed has been traversed
-        if should_yield and (recurse or recursed) and scope != Scope.BUILD:
-            yield self
+                elif scope == Scope.RUN:
+                    visited[1].add(element._unique_id)
+
+                    for dep in element.__runtime_dependencies:
+                        if dep._unique_id not in visited[1]:
+                            yield from visit(dep, Scope.RUN, visited)
+
+                    yield element
+                else:
+                    yield element
+
+            if visited is None:
+                # Visited is of the form (Visited for Scope.BUILD, Visited for Scope.RUN)
+                visited = (BitMap(), BitMap())
+            else:
+                # We have already a visited set passed. we might be able to short-circuit
+                if scope in (Scope.BUILD, Scope.ALL) and self._unique_id in visited[0]:
+                    return
+                if scope in (Scope.RUN, Scope.ALL) and self._unique_id in visited[1]:
+                    return
+
+            yield from visit(self, scope, visited)
 
     def search(self, scope, name):
         """Search for a dependency by name