Adding build tree extraction functionality

tests/completions/completions.py: Adding caching build tree flag to list of recognised flags in test
_frontend/cli.py: Adding caching build tree flags to cli
_artifactcache/artifactcache.py: Adding abstract get_buildtree_dir function
_artifactcache/cascache.py: Adding get_buildtree_dir function
element.py: Adding function to return build tree path

Adding functionality to allow the use of a cached build tree in a workspace

_stream.py: Adding staging of cached build tree

5 files changed, 116 insertions, 26 deletions

diff --git a/buildstream/_artifactcache/artifactcache.py b/buildstream/_artifactcache/artifactcache.py
index 5feae93f4..dd3699ac3 100644
--- a/buildstream/_artifactcache/artifactcache.py
+++ b/buildstream/_artifactcache/artifactcache.py
@@ -366,7 +366,7 @@ class ArtifactCache():
     #
     # Returns: path to extracted artifact
     #
-    def extract(self, element, key):
+    def extract(self, element, key, dest=None):
         raise ImplError("Cache '{kind}' does not implement extract()"
                         .format(kind=type(self).__name__))
 
@@ -483,6 +483,21 @@ class ArtifactCache():
         raise ImplError("Cache '{kind}' does not implement calculate_cache_size()"
                         .format(kind=type(self).__name__))
 
+    # does get_buildtree_dir():
+    #
+    # Returns build tree from cache if exists
+    #
+    # Args:
+    #    element (Element): The Element whose cache is being checked
+    #    key: the related cache key
+    #    directory: the directory to return if exists
+    # Returns:
+    # string: directory path or None
+    #
+    def get_buildtree_dir(self, element, key):
+        raise ImplError("Cache '{kind}' does not implement get_buildtree_dir()"
+                        .format(kind=type(self).__name__))
+
     ################################################
     #               Local Private Methods          #
     ################################################
diff --git a/buildstream/_artifactcache/cascache.py b/buildstream/_artifactcache/cascache.py
index 1b2dc198f..e30a12214 100644
--- a/buildstream/_artifactcache/cascache.py
+++ b/buildstream/_artifactcache/cascache.py
@@ -75,12 +75,14 @@ class CASCache(ArtifactCache):
         # This assumes that the repository doesn't have any dangling pointers
         return os.path.exists(refpath)
 
-    def extract(self, element, key):
+    def extract(self, element, key, dest=None):
         ref = self.get_artifact_fullname(element, key)
 
         tree = self.resolve_ref(ref, update_mtime=True)
 
-        dest = os.path.join(self.extractdir, element._get_project().name, element.normal_name, tree.hash)
+        if dest is None:
+            dest = os.path.join(self.extractdir, element._get_project().name, element.normal_name, tree.hash)
+
         if os.path.isdir(dest):
             # artifact has already been extracted
             return dest
@@ -106,6 +108,9 @@ class CASCache(ArtifactCache):
 
         return dest
 
+    def get_buildtree_dir(self, element, key):
+        return os.path.join(self.extract(element, key), "buildtree")
+
     def commit(self, element, content, keys):
         refs = [self.get_artifact_fullname(element, key) for key in keys]
 
diff --git a/buildstream/_frontend/cli.py b/buildstream/_frontend/cli.py
index 20624e2ac..9295f991e 100644
--- a/buildstream/_frontend/cli.py
+++ b/buildstream/_frontend/cli.py
@@ -446,7 +446,9 @@ def pull(app, elements, deps, remote):
         all:   All dependencies
     """
     with app.initialized(session_name="Pull"):
-        app.stream.pull(elements, selection=deps, remote=remote)
+        app.stream.pull(elements,
+                        selection=deps,
+                        remote=remote)
 
 
 ##################################################################
@@ -681,11 +683,14 @@ def workspace():
               help="Overwrite files existing in checkout directory")
 @click.option('--track', 'track_', default=False, is_flag=True,
               help="Track and fetch new source references before checking out the workspace")
+@click.option('--use-cached-buildtree', 'use_cached_buildtree', default='when-local',
+              type=click.Choice(['use-cached', 'ignore-cached', 'when-local']),
+              help="Using cached build trees")
 @click.argument('element',
                 type=click.Path(readable=False))
 @click.argument('directory', type=click.Path(file_okay=False))
 @click.pass_obj
-def workspace_open(app, no_checkout, force, track_, element, directory):
+def workspace_open(app, no_checkout, force, track_, element, directory, use_cached_buildtree):
     """Open a workspace for manual source modification"""
 
     if os.path.exists(directory):
@@ -702,7 +707,8 @@ def workspace_open(app, no_checkout, force, track_, element, directory):
         app.stream.workspace_open(element, directory,
                                   no_checkout=no_checkout,
                                   track_first=track_,
-                                  force=force)
+                                  force=force,
+                                  use_cached_buildtree=use_cached_buildtree)
 
 
 ##################################################################
@@ -762,10 +768,13 @@ def workspace_close(app, remove_dir, all_, elements):
               help="Track and fetch the latest source before resetting")
 @click.option('--all', '-a', 'all_', default=False, is_flag=True,
               help="Reset all open workspaces")
+@click.option('--use-cached-buildtree', 'use_cached_buildtree', default='when-local',
+              type=click.Choice(['use-cached', 'ignore-cached', 'when-local']),
+              help="Using cached build trees")
 @click.argument('elements', nargs=-1,
                 type=click.Path(readable=False))
 @click.pass_obj
-def workspace_reset(app, soft, track_, all_, elements):
+def workspace_reset(app, soft, track_, all_, elements, use_cached_buildtree):
     """Reset a workspace to its original state"""
 
     # Check that the workspaces in question exist
@@ -785,7 +794,10 @@ def workspace_reset(app, soft, track_, all_, elements):
         if all_:
             elements = tuple(element_name for element_name, _ in app.context.get_workspaces().list())
 
-        app.stream.workspace_reset(elements, soft=soft, track_first=track_)
+        app.stream.workspace_reset(elements,
+                                   soft=soft,
+                                   track_first=track_,
+                                   use_cached_buildtree=use_cached_buildtree)
 
 
 ##################################################################
diff --git a/buildstream/_stream.py b/buildstream/_stream.py
index f17d641de..7fc4e0522 100644
--- a/buildstream/_stream.py
+++ b/buildstream/_stream.py
@@ -435,6 +435,29 @@ class Stream():
             raise StreamError("Error while staging dependencies into a sandbox"
                               ": '{}'".format(e), detail=e.detail, reason=e.reason) from e
 
+    # __stage_cached_buildtree
+    #
+    # Stage a cached build tree if it exists
+    #
+    # Args:
+    #     use_cached_buildtree (bool): Whether or not to use cached buildtrees
+    #     element: element in use
+    #     cache_activity, stage_activity, message: (str)
+    #
+    # Returns:
+    #    (bool): True if the build tree was staged
+    #
+    def __stage_cached_buildtree(self, element, target_dir):
+        buildtree_path = None
+        if element._cached():
+            with element.timed_activity("Extracting cached build tree"):
+                buildtree_path = element._buildtree_path(element._get_cache_key())
+        if buildtree_path is not None:
+            with element.timed_activity("Staging cached build tree to {}".format(target_dir)):
+                shutil.copytree(buildtree_path, target_dir)
+                return True
+        return False
+
     # workspace_open
     #
     # Open a project workspace
@@ -445,11 +468,25 @@ class Stream():
     #    no_checkout (bool): Whether to skip checking out the source
     #    track_first (bool): Whether to track and fetch first
     #    force (bool): Whether to ignore contents in an existing directory
+    #    use_cached_buildtree(str): Whether or not to use cached buildtrees
     #
     def workspace_open(self, target, directory, *,
                        no_checkout,
                        track_first,
-                       force):
+                       force,
+                       use_cached_buildtree):
+
+        workspaces = self._context.get_workspaces()
+        assert use_cached_buildtree in ('use-cached', 'when-local', 'ignore-cached')
+
+        # Make cached_buildtree a boolean based on the flag assigned to it
+        # If flag was `when-local`, assigning value based on whether or not the project uses a remote cache
+        if use_cached_buildtree == 'use-cached':
+            use_cached_buildtree = True
+        elif use_cached_buildtree == 'when-local' and not self._artifacts.has_fetch_remotes():
+            use_cached_buildtree = True
+        else:
+            use_cached_buildtree = False
 
         if track_first:
             track_targets = (target,)
@@ -462,22 +499,6 @@ class Stream():
         target = elements[0]
         workdir = os.path.abspath(directory)
 
-        if not list(target.sources()):
-            build_depends = [x.name for x in target.dependencies(Scope.BUILD, recurse=False)]
-            if not build_depends:
-                raise StreamError("The given element has no sources")
-            detail = "Try opening a workspace on one of its dependencies instead:\n"
-            detail += "  \n".join(build_depends)
-            raise StreamError("The given element has no sources", detail=detail)
-
-        workspaces = self._context.get_workspaces()
-
-        # Check for workspace config
-        workspace = workspaces.get_workspace(target._get_full_name())
-        if workspace and not force:
-            raise StreamError("Workspace '{}' is already defined at: {}"
-                              .format(target.name, workspace.path))
-
         # If we're going to checkout, we need at least a fetch,
         # if we were asked to track first, we're going to fetch anyway.
         #
@@ -487,6 +508,26 @@ class Stream():
                 track_elements = elements
             self._fetch(elements, track_elements=track_elements)
 
+        # Check for workspace config
+        workspace = workspaces.get_workspace(target._get_full_name())
+        if workspace and not force:
+            raise StreamError("Workspace '{}' is already defined at: {}"
+                              .format(target.name, workspace.path))
+
+        if use_cached_buildtree:
+            if self.__stage_cached_buildtree(target, workdir):
+                workspaces.save_config()
+                self._message(MessageType.INFO, "Saved workspace configuration")
+                return
+
+        if not list(target.sources()):
+            build_depends = [x.name for x in target.dependencies(Scope.BUILD, recurse=False)]
+            if not build_depends:
+                raise StreamError("The given element has no sources")
+            detail = "Try opening a workspace on one of its dependencies instead:\n"
+            detail += "  \n".join(build_depends)
+            raise StreamError("The given element has no sources", detail=detail)
+
         if not no_checkout and target._get_consistency() != Consistency.CACHED:
             raise StreamError("Could not stage uncached source. " +
                               "Use `--track` to track and " +
@@ -547,8 +588,9 @@ class Stream():
     #    targets (list of str): The target elements to reset the workspace for
     #    soft (bool): Only reset workspace state
     #    track_first (bool): Whether to also track the sources first
+    #    use_cached_buildtree(str): Whether or not to use cached buildtrees
     #
-    def workspace_reset(self, targets, *, soft, track_first):
+    def workspace_reset(self, targets, *, soft, track_first, use_cached_buildtree):
 
         if track_first:
             track_targets = targets
@@ -592,6 +634,15 @@ class Stream():
             workspaces.delete_workspace(element._get_full_name())
             workspaces.create_workspace(element._get_full_name(), workspace.path)
 
+            if use_cached_buildtree == 'when-local':
+                use_cached_buildtree = os.path.isdir(os.path.join(workspace.path, 'buildtree'))
+
+            if use_cached_buildtree and self.__stage_cached_buildtree(element, workspace.path):
+                self._message(MessageType.INFO, "Reset workspace state for {} at: {}"
+                              .format(element.name, workspace.path))
+                return
+
+            workspaces.create_workspace(element.name, workspace.path)
             with element.timed_activity("Staging sources to {}".format(workspace.path)):
                 element._open_workspace()
 
diff --git a/buildstream/element.py b/buildstream/element.py
index e2a032197..ef71ccd6e 100644
--- a/buildstream/element.py
+++ b/buildstream/element.py
@@ -893,6 +893,13 @@ class Element(Plugin):
     #            Private Methods used in BuildStream            #
     #############################################################
 
+    # _buildtree_path():
+    #
+    # Returns the path of the cached build tree if it exists
+    #
+    def _buildtree_path(self, key):
+        return self.__artifacts.get_buildtree_dir(self, key)
+
     # _new_from_meta():
     #
     # Recursively instantiate a new Element instance, it's sources