element: Add helper to get dependencies for a list of elements in a scopebschubert/optimize-deps

This allows us to remove the checks about 'visited' sets that were
sometimes used or not and allows us to open this for plugins that might
want this functionality

4 files changed, 75 insertions, 80 deletions

diff --git a/src/buildstream/_element.pyx b/src/buildstream/_element.pyx
index ca43dc304..cc0178a3c 100644
--- a/src/buildstream/_element.pyx
+++ b/src/buildstream/_element.pyx
@@ -40,7 +40,7 @@ def deps_visit_all(element, visited):
     yield element
 
 
-def dependencies(element, scope, *, recurse=True, visited=None):
+def dependencies(element, scope, *, recurse=True):
     # 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.
@@ -50,22 +50,38 @@ def dependencies(element, scope, *, recurse=True, visited=None):
         if scope in (SCOPE_RUN, SCOPE_ALL):
             yield from element._Element__runtime_dependencies
     else:
-        if visited is None:
-            # Visited is of the form (Visited for Scope.BUILD, Visited for Scope.RUN)
-            visited = (BitMap(), BitMap())
-
         if scope == SCOPE_ALL:
-            # We can use only one of the sets when checking for Scope.ALL, as we would get added to
-            # both anyways.
-            # This would break if we start reusing 'visited' and mixing scopes, but that is done
-            # nowhere in the codebase.
-            if element._unique_id not in visited[0]:
-                yield from deps_visit_all(element, visited[0])
+            yield from deps_visit_all(element, BitMap())
         elif scope == SCOPE_BUILD:
-            if element._unique_id not in visited[0]:
-                yield from deps_visit_build(element, visited[0], visited[1])
+            yield from deps_visit_build(element, BitMap(), BitMap())
         elif scope == SCOPE_RUN:
-            if element._unique_id not in visited[1]:
-                yield from deps_visit_run(element, visited[1])
+            yield from deps_visit_run(element, BitMap())
         else:
             yield element
+
+
+def dependencies_for_targets(elements, scope):
+    if scope == SCOPE_ALL:
+        visited = BitMap()
+
+        for element in elements:
+            if element._unique_id not in visited:
+                yield from deps_visit_all(element, visited)
+
+    elif scope == SCOPE_BUILD:
+        visited_build = BitMap()
+        visited_run = BitMap()
+
+        for element in elements:
+            if element._unique_id not in visited_build:
+                yield from deps_visit_build(element, visited_build, visited_run)
+
+    elif scope == SCOPE_RUN:
+        visited = BitMap()
+
+        for element in elements:
+            if element._unique_id not in visited:
+                yield from deps_visit_run(element, visited)
+
+    else:
+        yield from elements
diff --git a/src/buildstream/_pipeline.py b/src/buildstream/_pipeline.py
index 8de97bea6..727f41305 100644
--- a/src/buildstream/_pipeline.py
+++ b/src/buildstream/_pipeline.py
@@ -24,12 +24,10 @@ 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, PROFILER
-from . import Scope
+from . import Scope, Element
 from ._project import ProjectRefStorage
 from .types import _PipelineSelection
 
@@ -115,7 +113,7 @@ class Pipeline:
             # to happen, even for large projects (tested with the Debian stack). Although,
             # if it does become a problem we may have to set the recursion limit to a
             # greater value.
-            for element in self.dependencies(targets, Scope.ALL):
+            for element in Element.dependencies_for_targets(targets, Scope.ALL):
                 # Determine initial element state.
                 if not element._resolved_initial_state:
                     element._initialize_state()
@@ -125,44 +123,6 @@ class Pipeline:
                 # dependencies.
                 element._update_ready_for_runtime_and_cached()
 
-                if task:
-                    task.add_current_progress()
-
-    # check_remotes()
-    #
-    # Check if the target artifact is cached in any of the available remotes
-    #
-    # Args:
-    #    targets (list [Element]): The list of element targets
-    #
-    def check_remotes(self, targets):
-        with self._context.messenger.simple_task("Querying remotes for cached status", silent_nested=True) as task:
-            task.set_maximum_progress(len(targets))
-
-            for element in targets:
-                element._cached_remotely()
-
-                task.add_current_progress()
-
-    # dependencies()
-    #
-    # Generator function to iterate over elements and optionally
-    # also iterate over sources.
-    #
-    # Args:
-    #    targets (list of Element): The target Elements to loop over
-    #    scope (Scope): The scope to iterate over
-    #    recurse (bool): Whether to recurse into dependencies
-    #
-    def dependencies(self, targets, scope, *, recurse=True):
-        # Keep track of 'visited' in this scope, so that all targets
-        # share the same context.
-        visited = (BitMap(), BitMap())
-
-        for target in targets:
-            for element in target.dependencies(scope, recurse=recurse, visited=visited):
-                yield element
-
     # plan()
     #
     # Generator function to iterate over only the elements
@@ -197,7 +157,9 @@ class Pipeline:
     # the selected option.
     #
     def get_selection(self, targets, mode, *, silent=True):
-        def redirect_and_log():
+        if mode == _PipelineSelection.NONE:
+            elements = targets
+        elif mode == _PipelineSelection.REDIRECT:
             # Redirect and log if permitted
             elements = []
             for t in targets:
@@ -206,21 +168,19 @@ class Pipeline:
                     self._message(MessageType.INFO, "Element '{}' redirected to '{}'".format(t.name, new_elm.name))
                 if new_elm not in elements:
                     elements.append(new_elm)
-            return elements
+        elif mode == _PipelineSelection.PLAN:
+            elements = self.plan(targets)
+        else:
+            if mode == _PipelineSelection.ALL:
+                scope = Scope.ALL
+            elif mode == _PipelineSelection.BUILD:
+                scope = Scope.BUILD
+            elif mode == _PipelineSelection.RUN:
+                scope = Scope.RUN
+
+            elements = list(Element.dependencies_for_targets(targets, scope))
 
-        # Work around python not having a switch statement; this is
-        # much clearer than the if/elif/else block we used to have.
-        #
-        # Note that the lambda is necessary so that we don't evaluate
-        # all possible values at run time; that would be slow.
-        return {
-            _PipelineSelection.NONE: lambda: targets,
-            _PipelineSelection.REDIRECT: redirect_and_log,
-            _PipelineSelection.PLAN: lambda: self.plan(targets),
-            _PipelineSelection.ALL: lambda: list(self.dependencies(targets, Scope.ALL)),
-            _PipelineSelection.BUILD: lambda: list(self.dependencies(targets, Scope.BUILD)),
-            _PipelineSelection.RUN: lambda: list(self.dependencies(targets, Scope.RUN)),
-        }[mode]()
+        return elements
 
     # except_elements():
     #
@@ -241,7 +201,7 @@ class Pipeline:
         if not except_targets:
             return elements
 
-        targeted = list(self.dependencies(targets, Scope.ALL))
+        targeted = list(Element.dependencies_for_targets(targets, Scope.ALL))
         visited = []
 
         def find_intersection(element):
diff --git a/src/buildstream/_stream.py b/src/buildstream/_stream.py
index 18297c2e4..189efe835 100644
--- a/src/buildstream/_stream.py
+++ b/src/buildstream/_stream.py
@@ -54,7 +54,7 @@ from ._state import State
 from .types import _KeyStrength, _PipelineSelection, _SchedulerErrorAction
 from .plugin import Plugin
 from . import utils, _yaml, _site
-from . import Scope
+from . import Scope, Element
 
 
 # Stream()
@@ -218,7 +218,10 @@ class Stream:
                     "Element must be fetched.".format(element._get_full_name())
                 )
 
-            missing_deps = [dep for dep in self._pipeline.dependencies([element], scope) if not dep._cached()]
+            missing_deps = [
+                dep for dep in element.dependencies(scope)
+                if not dep._cached()
+            ]
             if missing_deps:
                 if not pull_dependencies:
                     raise StreamError(
@@ -1356,7 +1359,7 @@ class Stream:
         # Inform the frontend of the full list of elements
         # and the list of elements which will be processed in this run
         #
-        self.total_elements = list(self._pipeline.dependencies(self.targets, Scope.ALL))
+        self.total_elements = list(Element.dependencies_for_targets(self.targets, Scope.ALL))
 
         if self._session_start_callback is not None:
             self._session_start_callback()
diff --git a/src/buildstream/element.py b/src/buildstream/element.py
index bd6700d6b..0a02cd655 100644
--- a/src/buildstream/element.py
+++ b/src/buildstream/element.py
@@ -439,10 +439,26 @@ class Element(Plugin):
         for source in self.__sources:
             yield source
 
-    def dependencies(self, scope: Scope, *, recurse: bool = True, visited=None) -> Iterator["Element"]:
-        """dependencies(scope, *, recurse=True)
+    @classmethod
+    def dependencies_for_targets(cls, targets, scope):
+        """A generator function which yields the dependencies for the given elements, recursively.
+
+        This ensures that no elements is handed twice and that the order is correct, with the basemost elements
+        in the given `scope` coming first, and the targets last.
+
+        If `scope` is `None`, will simply return all the elements.
+
+        Args:
+            targets (list): list List of elements for which to get the dependencies
+            scope: (:class:`.Scope`): The scope to iterate in
+
+        Yields:
+            (:class:`.Element`): The dependencies in `scope`, in deterministic staging order
+        """
+        return _element.dependencies_for_targets(targets, scope)
 
-        A generator function which yields the dependencies of the given element.
+    def dependencies(self, scope, *, recurse=True) -> Iterator["Element"]:
+        """A generator function which yields the dependencies of the given element.
 
         If `recurse` is specified (the default), the full dependencies will be listed
         in deterministic staging order, starting with the basemost elements in the
@@ -457,7 +473,7 @@ class Element(Plugin):
         Yields:
            The dependencies in `scope`, in deterministic staging order
         """
-        return _element.dependencies(self, scope, recurse=recurse, visited=visited)
+        return _element.dependencies(self, scope, recurse=recurse)
 
     def search(self, scope: Scope, name: str) -> Optional["Element"]:
         """Search for a dependency by name