element.py: Migrate dependencies() to a cython filedanielsilverstone-ct/cythonize-bits

This migrates Element.dependencies() core visitor to a Cython
extension which roughly doubles its speed.  There may be ways
to further improve this, but Cython can't switch to a fully C
function when `yield` and `yield from` are involved.

Signed-off-by: Daniel Silverstone <daniel.silverstone@codethink.co.uk>

3 files changed, 95 insertions, 29 deletions

diff --git a/setup.py b/setup.py
index ed211dd09..a9e2295f0 100755
--- a/setup.py
+++ b/setup.py
@@ -400,6 +400,7 @@ BUILD_EXTENSIONS = []
 
 register_cython_module("buildstream._yaml")
 register_cython_module("buildstream._variables", dependencies=["buildstream._yaml"])
+register_cython_module("buildstream._element")
 
 #####################################################
 #             Main setup() Invocation               #
diff --git a/src/buildstream/_element.pyx b/src/buildstream/_element.pyx
new file mode 100644
index 000000000..372026241
--- /dev/null
+++ b/src/buildstream/_element.pyx
@@ -0,0 +1,83 @@
+#
+#  Copyright (C) 2016-2019 Codethink Limited
+#
+#  This program is free software; you can redistribute it and/or
+#  modify it under the terms of the GNU Lesser General Public
+#  License as published by the Free Software Foundation; either
+#  version 2 of the License, or (at your option) any later version.
+#
+#  This library is distributed in the hope that it will be useful,
+#  but WITHOUT ANY WARRANTY; without even the implied warranty of
+#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.	 See the GNU
+#  Lesser General Public License for more details.
+#
+#  You should have received a copy of the GNU Lesser General Public
+#  License along with this library. If not, see <http://www.gnu.org/licenses/>.
+#
+#  Authors:
+#        Tristan Van Berkom <tristan.vanberkom@codethink.co.uk>
+#        Daniel Silverstone <daniel.silverstone@codethink.co.uk>
+
+"""
+Element - Cython helpers
+========================
+
+Functions in here are helpers for Element and so are permitted to access
+the __members of it.
+"""
+
+# Mirror of the types.Scope
+#
+# If you add something to that, mirror the change here and in
+# Element.dependencies().  Alternatively if types.py inherited
+# this enumeration then life might be better.
+cpdef enum DepScope:
+    ALL
+    BUILD
+    RUN
+    NONE
+
+# deps_visit()
+#
+# Visit dependencies of an element in the given scope in staging order.
+#
+# See Element.dependencies() for more details
+#
+# Args:
+#   (Element) element: The element whose dependencies need visiting
+#   (DepScope) scope:  The scope of the dependencies to visit
+#   ((BitMap, BitMap)) visited: A tuple of bitmaps for already-visited elements
+#
+# Yields:
+#   (Element): Elements, in staging order
+def deps_visit(object element, DepScope scope, object visited):
+    if scope == DepScope.ALL:
+        visited[0].add(element._unique_id)
+        visited[1].add(element._unique_id)
+
+        for dep in element._Element__build_dependencies:
+            if dep._unique_id not in visited[0] and dep._unique_id not in visited[1]:
+                yield from deps_visit(dep, DepScope.ALL, visited)
+
+        for dep in element._Element__runtime_dependencies:
+            if dep._unique_id not in visited[0] and dep._unique_id not in visited[1]:
+                yield from deps_visit(dep, DepScope.ALL, visited)
+
+        yield element
+    elif scope == DepScope.BUILD:
+        visited[0].add(element._unique_id)
+
+        for dep in element._Element__build_dependencies:
+            if dep._unique_id not in visited[1]:
+                yield from deps_visit(dep, DepScope.RUN, visited)
+
+    elif scope == DepScope.RUN:
+        visited[1].add(element._unique_id)
+
+        for dep in element._Element__runtime_dependencies:
+            if dep._unique_id not in visited[1]:
+                yield from deps_visit(dep, DepScope.RUN, visited)
+
+        yield element
+    else:
+        yield element
diff --git a/src/buildstream/element.py b/src/buildstream/element.py
index 8e006ea6b..af8ee80a4 100644
--- a/src/buildstream/element.py
+++ b/src/buildstream/element.py
@@ -104,6 +104,7 @@ from .sandbox._config import SandboxConfig
 from .sandbox._sandboxremote import SandboxRemote
 from .types import Consistency, CoreWarnings, Scope, _KeyStrength, _UniquePriorityQueue
 from ._artifact import Artifact
+from ._element import deps_visit, DepScope
 
 from .storage.directory import Directory
 from .storage._filebaseddirectory import FileBasedDirectory
@@ -419,34 +420,6 @@ class Element(Plugin):
             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)
-
-                    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)
-
-                    yield element
-                elif scope == Scope.BUILD:
-                    visited[0].add(element._unique_id)
-
-                    for dep in element.__build_dependencies:
-                        if dep._unique_id not in visited[1]:
-                            yield from visit(dep, Scope.RUN, visited)
-
-                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())
@@ -457,7 +430,16 @@ class Element(Plugin):
                 if scope in (Scope.RUN, Scope.ALL) and self._unique_id in visited[1]:
                     return
 
-            yield from visit(self, scope, visited)
+            if scope == Scope.ALL:
+                scope = DepScope.ALL
+            elif scope == Scope.BUILD:
+                scope = DepScope.BUILD
+            elif scope == Scope.RUN:
+                scope = DepScope.RUN
+            else:
+                scope = DepScope.NONE
+
+            yield from deps_visit(self, scope, visited)
 
     def search(self, scope, name):
         """Search for a dependency by name