element: optimize getting dependencies in cython

This makes the dependency code roughly 8x faster

4 files changed, 75 insertions, 59 deletions

diff --git a/.pylintrc b/.pylintrc
index 25d5647b0..39ab2ac6a 100644
--- a/.pylintrc
+++ b/.pylintrc
@@ -5,6 +5,7 @@
 # run arbitrary code
 extension-pkg-whitelist=
     buildstream.node,
+    buildstream._element,
     buildstream._loader._loader,
     buildstream._loader.loadelement,
     buildstream._loader.types,
diff --git a/setup.py b/setup.py
index 2d24fd533..f5ba0215f 100755
--- a/setup.py
+++ b/setup.py
@@ -309,6 +309,7 @@ def register_cython_module(module_name, dependencies=None):
 BUILD_EXTENSIONS = []
 
 register_cython_module("buildstream.node")
+register_cython_module("buildstream._element")
 register_cython_module("buildstream._loader._loader")
 register_cython_module("buildstream._loader.loadelement")
 register_cython_module("buildstream._loader.types", dependencies=["buildstream.node"])
diff --git a/src/buildstream/_element.pyx b/src/buildstream/_element.pyx
new file mode 100644
index 000000000..ca43dc304
--- /dev/null
+++ b/src/buildstream/_element.pyx
@@ -0,0 +1,71 @@
+from pyroaring import BitMap
+
+from .types import Scope
+
+# FIXME: we should make those as enums consumable from Cython
+cdef SCOPE_ALL = Scope.ALL
+cdef SCOPE_BUILD = Scope.BUILD
+cdef SCOPE_RUN = Scope.RUN
+
+
+def deps_visit_run(element, visited):
+    visited.add(element._unique_id)
+
+    for dep in element._Element__runtime_dependencies:
+        if dep._unique_id not in visited:
+            yield from deps_visit_run(dep, visited)
+
+    yield element
+
+
+def deps_visit_build(element, visited_build, visited_run):
+    visited_build.add(element._unique_id)
+
+    for dep in element._Element__build_dependencies:
+        if dep._unique_id not in visited_run:
+            yield from deps_visit_run(dep, visited_run)
+
+
+def deps_visit_all(element, visited):
+    visited.add(element._unique_id)
+
+    for dep in element._Element__build_dependencies:
+        if dep._unique_id not in visited:
+            yield from deps_visit_all(dep, visited)
+
+    for dep in element._Element__runtime_dependencies:
+        if dep._unique_id not in visited:
+            yield from deps_visit_all(dep, visited)
+
+    yield element
+
+
+def dependencies(element, scope, *, recurse=True, visited=None):
+    # 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 element._Element__build_dependencies
+        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])
+        elif scope == SCOPE_BUILD:
+            if element._unique_id not in visited[0]:
+                yield from deps_visit_build(element, visited[0], visited[1])
+        elif scope == SCOPE_RUN:
+            if element._unique_id not in visited[1]:
+                yield from deps_visit_run(element, visited[1])
+        else:
+            yield element
diff --git a/src/buildstream/element.py b/src/buildstream/element.py
index 9239a2c6d..bd6700d6b 100644
--- a/src/buildstream/element.py
+++ b/src/buildstream/element.py
@@ -83,8 +83,6 @@ from functools import partial
 import string
 from typing import cast, TYPE_CHECKING, Any, Dict, Iterator, List, Optional
 
-from pyroaring import BitMap  # pylint: disable=no-name-in-module
-
 from . import _yaml
 from ._variables import Variables
 from ._versions import BST_CORE_ARTIFACT_VERSION
@@ -93,6 +91,7 @@ from .exceptions import ErrorDomain, LoadErrorReason
 from .utils import FileListResult
 from . import utils
 from . import _cachekey
+from . import _element
 from . import _site
 from ._platform import Platform
 from .node import Node
@@ -458,63 +457,7 @@ class Element(Plugin):
         Yields:
            The dependencies in `scope`, in deterministic staging order
         """
-        # 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_run(element, visited):
-                visited.add(element._unique_id)
-
-                for dep in element.__runtime_dependencies:
-                    if dep._unique_id not in visited:
-                        yield from visit_run(dep, visited)
-
-                yield element
-
-            def visit_build(element, visited_build, visited_run):
-                visited_build.add(element._unique_id)
-
-                for dep in element.__build_dependencies:
-                    if dep._unique_id not in visited_run:
-                        yield from visit_run(dep, visited_run)
-
-            def visit_all(element, visited):
-                visited.add(element._unique_id)
-
-                for dep in element.__build_dependencies:
-                    if dep._unique_id not in visited:
-                        yield from visit_all(dep, visited)
-
-                for dep in element.__runtime_dependencies:
-                    if dep._unique_id not in visited:
-                        yield from visit_all(dep, visited)
-
-                yield element
-
-            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 self._unique_id not in visited[0]:
-                    yield from visit_all(self, visited[0])
-            elif scope == Scope.BUILD:
-                if self._unique_id not in visited[0]:
-                    yield from visit_build(self, visited[0], visited[1])
-            elif scope == Scope.RUN:
-                if self._unique_id not in visited[1]:
-                    yield from visit_run(self, visited[1])
-            else:
-                yield self
+        return _element.dependencies(self, scope, recurse=recurse, visited=visited)
 
     def search(self, scope: Scope, name: str) -> Optional["Element"]:
         """Search for a dependency by name