Implement SpawnInputExpander.walkInputs

src/main/java/com/google/devtools/build/lib/exec/SpawnInputExpander.java

@@ -33,11 +33,13 @@
 import com.google.devtools.build.lib.actions.RunfilesSupplier;
 import com.google.devtools.build.lib.actions.Spawn;
 import com.google.devtools.build.lib.actions.cache.VirtualActionInput;
+import com.google.devtools.build.lib.collect.nestedset.NestedSet;
 import com.google.devtools.build.lib.collect.nestedset.NestedSetBuilder;
 import com.google.devtools.build.lib.collect.nestedset.Order;
 import com.google.devtools.build.lib.vfs.Path;
 import com.google.devtools.build.lib.vfs.PathFragment;
 import java.io.IOException;
+import java.util.Arrays;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -95,7 +97,7 @@ public SpawnInputExpander(
     this.relSymlinkBehavior = relSymlinkBehavior;
   }
 
-  private void addMapping(
+  private static void addMapping(
       Map<PathFragment, ActionInput> inputMappings,
       PathFragment targetLocation,
       ActionInput input,
@@ -215,13 +217,13 @@ void addFilesetManifest(
       }
   }
 
-  private void addInputs(
+  private static void addInputs(
       Map<PathFragment, ActionInput> inputMap,
-      Spawn spawn,
+      NestedSet<? extends ActionInput> inputFiles,
       ArtifactExpander artifactExpander,
       PathFragment baseDirectory) {
     List<ActionInput> inputs =
-        ActionInputHelper.expandArtifacts(spawn.getInputFiles(), artifactExpander);
+        ActionInputHelper.expandArtifacts(inputFiles, artifactExpander);
     for (ActionInput input : inputs) {
       addMapping(inputMap, input.getExecPath(), input, baseDirectory);
     }
@@ -243,7 +245,7 @@ public SortedMap<PathFragment, ActionInput> getInputMapping(
       MetadataProvider actionInputFileCache)
       throws IOException, ForbiddenActionInputException {
     TreeMap<PathFragment, ActionInput> inputMap = new TreeMap<>();
-    addInputs(inputMap, spawn, artifactExpander, baseDirectory);
+    addInputs(inputMap, spawn.getInputFiles(), artifactExpander, baseDirectory);
     addRunfilesToInputs(
         inputMap,
         spawn.getRunfilesSupplier(),
@@ -254,6 +256,120 @@ public SortedMap<PathFragment, ActionInput> getInputMapping(
     return inputMap;
   }
 
+  public interface InputWalker {
+    SortedMap<PathFragment, ActionInput> getLeavesInputMapping()
+        throws IOException, ForbiddenActionInputException;
+    void visitNonLeaves(InputVisitor visitor)
+        throws IOException, ForbiddenActionInputException;
+  }
+
+  public interface InputVisitor {
+    /** Visits a part of the input hierarchy.
+      *
+      * <p>{@code nodeKey} can be used as key when memoizing visited parts of the hierarchy.
+      */
+    void visit(Object nodeKey, InputWalker walker)
+        throws IOException, ForbiddenActionInputException;
+  }
+
+  /**
+    * Visits the input files hierarchy in a depth first manner.
+    *
+    * <p>Similar to {@link getInputMapping} but allows for early exit, by not visiting children,
+    * when walking through the input hierarchy. By applying memoization, the retreival process
+    * of the inputs can be speeded up.
+    *
+    * <p>{@code baseDirectory} is prepended to every path in the input key. This is useful if the
+    * mapping is used in a context where the directory relative to which the keys are interpreted
+    * is not the same as the execroot.
+    */
+  public void walkInputs(
+      Spawn spawn,
+      ArtifactExpander artifactExpander,
+      PathFragment baseDirectory,
+      MetadataProvider actionInputFileCache,
+      InputVisitor visitor)
+      throws IOException, ForbiddenActionInputException {
+    walkNestedSetInputs(baseDirectory, spawn.getInputFiles(), artifactExpander, visitor);
+
+    RunfilesSupplier runfilesSupplier = spawn.getRunfilesSupplier();
+    visitor.visit(
+        // The list of variables affecting the functional expressions below.
+        Arrays.asList(
+            // Assuming that artifactExpander and actionInputFileCache, different for each spawn,
+            // always expand the same way.
+            this,  // For accessing addRunfilesToInputs.
+            runfilesSupplier,
+            baseDirectory),
+        new InputWalker() {
+          public SortedMap<PathFragment, ActionInput> getLeavesInputMapping()
+              throws IOException, ForbiddenActionInputException {
+            TreeMap<PathFragment, ActionInput> inputMap = new TreeMap<>();
+            addRunfilesToInputs(
+                inputMap,
+                runfilesSupplier, actionInputFileCache, artifactExpander, baseDirectory);
+            return inputMap;
+          }
+
+          public void visitNonLeaves(InputVisitor childVisitor) {}
+        }
+    );
+
+    Map<Artifact, ImmutableList<FilesetOutputSymlink>> filesetMappings =
+        spawn.getFilesetMappings();
+    // filesetMappings is assumed to be very small, so no need to implement visitNonLeaves() for
+    // improved runtime.
+    visitor.visit(
+        // The list of variables affecting the functional expressions below.
+        Arrays.asList(
+            this,  // For accessing addFilesetManifests.
+            filesetMappings, baseDirectory),
+        new InputWalker() {
+          public SortedMap<PathFragment, ActionInput> getLeavesInputMapping()
+              throws ForbiddenRelativeSymlinkException {
+            TreeMap<PathFragment, ActionInput> inputMap = new TreeMap<>();
+            addFilesetManifests(filesetMappings, inputMap, baseDirectory);
+            return inputMap;
+          }
+
+          public void visitNonLeaves(InputVisitor childVisitor) {}
+        }
+    );
+  }
+
+  /** Walks through one level of a {@link NestedSet} of {@link ActionInput}s. */
+  private void walkNestedSetInputs(
+      PathFragment baseDirectory,
+      NestedSet<? extends ActionInput> someInputFiles,
+      ArtifactExpander artifactExpander,
+      InputVisitor visitor)
+      throws IOException, ForbiddenActionInputException {
+    visitor.visit(
+        // addInputs is static so no need to add 'this' as dependent key.
+        Arrays.asList(
+          // Assuming that artifactExpander, different for each spawn, always expands the same way.
+          someInputFiles.toNode(), baseDirectory),
+        new InputWalker() {
+          public SortedMap<PathFragment, ActionInput> getLeavesInputMapping() {
+            TreeMap<PathFragment, ActionInput> inputMap = new TreeMap<>();
+            addInputs(
+                inputMap,
+                NestedSetBuilder.wrap(someInputFiles.getOrder(), someInputFiles.getLeaves()),
+                artifactExpander,
+                baseDirectory);
+            return inputMap;
+          }
+
+          public void visitNonLeaves(InputVisitor childVisitor)
+              throws IOException, ForbiddenActionInputException {
+            for (NestedSet<? extends ActionInput> subInputs : someInputFiles.getNonLeaves()) {
+              walkNestedSetInputs(baseDirectory, subInputs, artifactExpander, childVisitor);
+            }
+          }
+        }
+    );
+  }
+
   /**
    * Exception signaling that an input was not a regular file: most likely a directory. This
    * exception is currently never thrown in practice since we do not enforce "strict" mode.