From 9de8564c77fce967a4ed5ae8f65c74940dc1201c Mon Sep 17 00:00:00 2001
From: Fredrik Medley <fredrik.medley@gmail.com>
Date: Thu, 19 Aug 2021 19:49:33 +0200
Subject: [PATCH] Reimplement getInputMapping() using walkInputs()

Keep a single way of retreiving the inputs, instead of keeping two
implementations in sync. Using walkInputs() will be slightly slower,
but as getInputMapping() is only called on cache miss, I think the
maintenance problem is more important to solve.
---
 .../build/lib/exec/SpawnInputExpander.java    | 32 ++++++++++++++-----
 1 file changed, 24 insertions(+), 8 deletions(-)

diff --git a/src/main/java/com/google/devtools/build/lib/exec/SpawnInputExpander.java b/src/main/java/com/google/devtools/build/lib/exec/SpawnInputExpander.java
index 5cc3ebefa2141a..0084e6051684a5 100644
--- a/src/main/java/com/google/devtools/build/lib/exec/SpawnInputExpander.java
+++ b/src/main/java/com/google/devtools/build/lib/exec/SpawnInputExpander.java
@@ -17,6 +17,7 @@
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.Sets;
 import com.google.devtools.build.lib.actions.ActionInput;
 import com.google.devtools.build.lib.actions.ActionInputHelper;
 import com.google.devtools.build.lib.actions.Artifact;
@@ -44,6 +45,7 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 import java.util.SortedMap;
 import java.util.TreeMap;
 
@@ -245,18 +247,32 @@ public SortedMap<PathFragment, ActionInput> getInputMapping(
       PathFragment baseDirectory,
       MetadataProvider actionInputFileCache)
       throws IOException, ForbiddenActionInputException {
+    Set<Object> visitedNodes = Sets.newConcurrentHashSet();
     TreeMap<PathFragment, ActionInput> inputMap = new TreeMap<>();
-    addInputs(inputMap, spawn.getInputFiles(), artifactExpander, baseDirectory);
-    addRunfilesToInputs(
-        inputMap,
-        spawn.getRunfilesSupplier(),
-        actionInputFileCache,
-        artifactExpander,
-        baseDirectory);
-    addFilesetManifests(spawn.getFilesetMappings(), inputMap, baseDirectory);
+    walkInputs(
+      spawn, artifactExpander, baseDirectory, actionInputFileCache,
+      (Object nodeKey, SpawnExecutionContext.InputWalker walker) -> {
+        getInputMappingVisitor(visitedNodes, inputMap, walker);
+      }
+    );
     return inputMap;
   }
 
+  private void getInputMappingVisitor(
+      Set<Object> visitedNodes,
+      TreeMap<PathFragment, ActionInput> inputMap,
+      SpawnExecutionContext.InputWalker walker)
+      throws IOException, ForbiddenActionInputException {
+    inputMap.putAll(walker.getLeavesInputMapping());
+    walker.visitNonLeaves(
+      (Object nodeKey, SpawnExecutionContext.InputWalker subWalker) -> {
+        if (visitedNodes.add(nodeKey)) {
+          getInputMappingVisitor(visitedNodes, inputMap, subWalker);
+        }
+      }
+    );
+  }
+
   public void walkInputs(
       Spawn spawn,
       ArtifactExpander artifactExpander,