Use nested set cache in key computation for CustomCommandLine.

RELNOTES: None
PiperOrigin-RevId: 183727976

6 files changed, 189 insertions, 15 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/actions/ActionKeyContext.java b/src/main/java/com/google/devtools/build/lib/actions/ActionKeyContext.java
index 6378675147..053c12024a 100644
--- a/src/main/java/com/google/devtools/build/lib/actions/ActionKeyContext.java
+++ b/src/main/java/com/google/devtools/build/lib/actions/ActionKeyContext.java
@@ -14,6 +14,7 @@
 
 package com.google.devtools.build.lib.actions;
 
+import com.google.devtools.build.lib.analysis.actions.CommandLineItem;
 import com.google.devtools.build.lib.collect.nestedset.NestedSet;
 import com.google.devtools.build.lib.collect.nestedset.NestedSetFingerprintCache;
 import com.google.devtools.build.lib.util.Fingerprint;
@@ -28,6 +29,11 @@ public class ActionKeyContext {
     nestedSetFingerprintCache.addNestedSetToFingerprint(fingerprint, nestedSet);
   }
 
+  public <T> void addNestedSetToFingerprint(
+      CommandLineItem.MapFn<? super T> mapFn, Fingerprint fingerprint, NestedSet<T> nestedSet) {
+    nestedSetFingerprintCache.addNestedSetToFingerprint(mapFn, fingerprint, nestedSet);
+  }
+
   public void clear() {
     nestedSetFingerprintCache.clear();
   }
diff --git a/src/main/java/com/google/devtools/build/lib/analysis/actions/CommandLine.java b/src/main/java/com/google/devtools/build/lib/analysis/actions/CommandLine.java
index 1ca9df707a..0e5617c042 100644
--- a/src/main/java/com/google/devtools/build/lib/analysis/actions/CommandLine.java
+++ b/src/main/java/com/google/devtools/build/lib/analysis/actions/CommandLine.java
@@ -17,9 +17,11 @@ package com.google.devtools.build.lib.analysis.actions;
 import com.google.common.base.Joiner;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Iterables;
+import com.google.devtools.build.lib.actions.ActionKeyContext;
 import com.google.devtools.build.lib.actions.Artifact.ArtifactExpander;
 import com.google.devtools.build.lib.actions.CommandLineExpansionException;
 import com.google.devtools.build.lib.collect.CollectionUtils;
+import com.google.devtools.build.lib.util.Fingerprint;
 
 /** A representation of a list of arguments, often a command executed by {@link SpawnAction}. */
 public abstract class CommandLine {
@@ -47,6 +49,13 @@ public abstract class CommandLine {
     return arguments();
   }
 
+  public void addToFingerprint(ActionKeyContext actionKeyContext, Fingerprint fingerprint)
+      throws CommandLineExpansionException {
+    for (String s : arguments()) {
+      fingerprint.addString(s);
+    }
+  }
+
   /** Returns a {@link CommandLine} backed by a copy of the given list of arguments. */
   public static CommandLine of(Iterable<String> arguments) {
     final Iterable<String> immutableArguments = CollectionUtils.makeImmutable(arguments);
diff --git a/src/main/java/com/google/devtools/build/lib/analysis/actions/CustomCommandLine.java b/src/main/java/com/google/devtools/build/lib/analysis/actions/CustomCommandLine.java
index a987c99917..eb0c1c72c4 100644
--- a/src/main/java/com/google/devtools/build/lib/analysis/actions/CustomCommandLine.java
+++ b/src/main/java/com/google/devtools/build/lib/analysis/actions/CustomCommandLine.java
@@ -23,6 +23,7 @@ import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.Interner;
 import com.google.common.collect.Iterables;
 import com.google.common.collect.Lists;
+import com.google.devtools.build.lib.actions.ActionKeyContext;
 import com.google.devtools.build.lib.actions.Artifact;
 import com.google.devtools.build.lib.actions.Artifact.ArtifactExpander;
 import com.google.devtools.build.lib.actions.Artifact.TreeFileArtifact;
@@ -30,6 +31,7 @@ import com.google.devtools.build.lib.cmdline.Label;
 import com.google.devtools.build.lib.collect.nestedset.NestedSet;
 import com.google.devtools.build.lib.concurrent.BlazeInterners;
 import com.google.devtools.build.lib.concurrent.ThreadSafety.Immutable;
+import com.google.devtools.build.lib.util.Fingerprint;
 import com.google.devtools.build.lib.util.LazyString;
 import com.google.devtools.build.lib.vfs.PathFragment;
 import com.google.errorprone.annotations.CompileTimeConstant;
@@ -42,6 +44,7 @@ import java.util.List;
 import java.util.Map;
 import java.util.Set;
 import java.util.TreeSet;
+import java.util.UUID;
 import javax.annotation.Nullable;
 
 /** A customizable, serializable class for building memory efficient command lines. */
@@ -59,6 +62,12 @@ public final class CustomCommandLine extends CommandLine {
      *     ArgvFragment doesn't have any args, it should return {@code argi} unmodified.
      */
     int eval(List<Object> arguments, int argi, ImmutableList.Builder<String> builder);
+
+    int addToFingerprint(
+        List<Object> arguments,
+        int argi,
+        ActionKeyContext actionKeyContext,
+        Fingerprint fingerprint);
   }
 
   /**
@@ -73,7 +82,19 @@ public final class CustomCommandLine extends CommandLine {
       return argi; // Doesn't consume any arguments, so return argi unmodified
     }
 
+    @Override
+    public int addToFingerprint(
+        List<Object> arguments,
+        int argi,
+        ActionKeyContext actionKeyContext,
+        Fingerprint fingerprint) {
+      addToFingerprint(actionKeyContext, fingerprint);
+      return argi; // Doesn't consume any arguments, so return argi unmodified
+    }
+
     abstract void eval(ImmutableList.Builder<String> builder);
+
+    abstract void addToFingerprint(ActionKeyContext actionKeyContext, Fingerprint fingerprint);
   }
 
   /** Deprecated. Do not use. TODO(b/64841073): Remove this */
@@ -86,6 +107,13 @@ public final class CustomCommandLine extends CommandLine {
     }
 
     public abstract Iterable<String> argv();
+
+    @Override
+    final void addToFingerprint(ActionKeyContext actionKeyContext, Fingerprint fingerprint) {
+      for (String arg : argv()) {
+        fingerprint.addString(arg);
+      }
+    }
   }
 
   /**
@@ -303,6 +331,12 @@ public final class CustomCommandLine extends CommandLine {
 
     private static final class VectorArgFragment implements ArgvFragment {
       private static Interner<VectorArgFragment> interner = BlazeInterners.newStrongInterner();
+      private static final UUID FORMAT_EACH_UUID =
+          UUID.fromString("f830781f-2e0d-4e3b-9b99-ece7f249e0f3");
+      private static final UUID BEFORE_EACH_UUID =
+          UUID.fromString("07d22a0d-2691-4f1c-9f47-5294de1f94e4");
+      private static final UUID JOIN_WITH_UUID =
+          UUID.fromString("c96ed6f0-9220-40f6-9e0c-1c0c5e0b47e4");
 
       private final boolean isNestedSet;
       private final boolean hasMapEach;
@@ -369,6 +403,48 @@ public final class CustomCommandLine extends CommandLine {
         return argi;
       }
 
+      @SuppressWarnings("unchecked")
+      @Override
+      public int addToFingerprint(
+          List<Object> arguments,
+          int argi,
+          ActionKeyContext actionKeyContext,
+          Fingerprint fingerprint) {
+        if (isNestedSet) {
+          NestedSet<Object> values = (NestedSet<Object>) arguments.get(argi++);
+          CommandLineItem.MapFn<Object> mapFn =
+              hasMapEach
+                  ? (CommandLineItem.MapFn<Object>) arguments.get(argi++)
+                  : CommandLineItem.MapFn.DEFAULT;
+          actionKeyContext.addNestedSetToFingerprint(mapFn, fingerprint, values);
+        } else {
+          int count = (Integer) arguments.get(argi++);
+          CommandLineItem.MapFn<Object> mapFn =
+              hasMapEach
+                  ? (CommandLineItem.MapFn<Object>)
+                      arguments.get(argi + count) // Peek ahead to mapFn
+                  : CommandLineItem.MapFn.DEFAULT;
+          for (int i = 0; i < count; ++i) {
+            fingerprint.addString(mapFn.expandToCommandLine(arguments.get(argi++)));
+          }
+          if (hasMapEach) {
+            ++argi; // Consume mapFn
+          }
+        }
+        if (hasFormatEach) {
+          fingerprint.addUUID(FORMAT_EACH_UUID);
+          fingerprint.addString((String) arguments.get(argi++));
+        }
+        if (hasBeforeEach) {
+          fingerprint.addUUID(BEFORE_EACH_UUID);
+          fingerprint.addString((String) arguments.get(argi++));
+        } else if (hasJoinWith) {
+          fingerprint.addUUID(JOIN_WITH_UUID);
+          fingerprint.addString((String) arguments.get(argi++));
+        }
+        return argi;
+      }
+
       @Override
       public boolean equals(Object o) {
         if (this == o) {
@@ -394,6 +470,7 @@ public final class CustomCommandLine extends CommandLine {
 
   private static class FormatArg implements ArgvFragment {
     private static final FormatArg INSTANCE = new FormatArg();
+    private static final UUID FORMAT_UUID = UUID.fromString("377cee34-e947-49e0-94a2-6ab95b396ec4");
 
     private static void push(List<Object> arguments, String formatStr, Object... args) {
       arguments.add(INSTANCE);
@@ -413,10 +490,26 @@ public final class CustomCommandLine extends CommandLine {
       builder.add(String.format(formatStr, args));
       return argi;
     }
+
+    @Override
+    public int addToFingerprint(
+        List<Object> arguments,
+        int argi,
+        ActionKeyContext actionKeyContext,
+        Fingerprint fingerprint) {
+      int argCount = (Integer) arguments.get(argi++);
+      fingerprint.addUUID(FORMAT_UUID);
+      fingerprint.addString((String) arguments.get(argi++));
+      for (int i = 0; i < argCount; ++i) {
+        fingerprint.addString(CommandLineItem.expandToCommandLine(arguments.get(argi++)));
+      }
+      return argi;
+    }
   }
 
   private static class PrefixArg implements ArgvFragment {
     private static final PrefixArg INSTANCE = new PrefixArg();
+    private static final UUID PREFIX_UUID = UUID.fromString("a95eccdf-4f54-46fc-b925-c8c7e1f50c95");
 
     private static void push(List<Object> arguments, String before, Object arg) {
       arguments.add(INSTANCE);
@@ -431,6 +524,18 @@ public final class CustomCommandLine extends CommandLine {
       builder.add(before + CommandLineItem.expandToCommandLine(arg));
       return argi;
     }
+
+    @Override
+    public int addToFingerprint(
+        List<Object> arguments,
+        int argi,
+        ActionKeyContext actionKeyContext,
+        Fingerprint fingerprint) {
+      fingerprint.addUUID(PREFIX_UUID);
+      fingerprint.addString((String) arguments.get(argi++));
+      fingerprint.addString(CommandLineItem.expandToCommandLine(arguments.get(argi++)));
+      return argi;
+    }
   }
 
   /**
@@ -491,6 +596,7 @@ public final class CustomCommandLine extends CommandLine {
   private static final class ExpandedTreeArtifactExecPathsArg
       extends TreeArtifactExpansionArgvFragment {
     private final Artifact treeArtifact;
+    private static final UUID TREE_UUID = UUID.fromString("13b7626b-c77d-4a30-ad56-ff08c06b1cee");
 
     private ExpandedTreeArtifactExecPathsArg(Artifact treeArtifact) {
       Preconditions.checkArgument(
@@ -513,6 +619,12 @@ public final class CustomCommandLine extends CommandLine {
       return String.format(
           "ExpandedTreeArtifactExecPathsArg{ treeArtifact: %s}", treeArtifact.getExecPathString());
     }
+
+    @Override
+    void addToFingerprint(ActionKeyContext actionKeyContext, Fingerprint fingerprint) {
+      fingerprint.addUUID(TREE_UUID);
+      fingerprint.addPath(treeArtifact.getExecPath());
+    }
   }
 
   /**
@@ -1096,4 +1208,35 @@ public final class CustomCommandLine extends CommandLine {
       return arg;
     }
   }
+
+  @Override
+  public void addToFingerprint(ActionKeyContext actionKeyContext, Fingerprint fingerprint) {
+    int count = arguments.size();
+    for (int i = 0; i < count; ) {
+      Object arg = arguments.get(i++);
+      Object substitutedArg = substituteTreeFileArtifactArgvFragment(arg);
+      if (substitutedArg instanceof Iterable) {
+        addSimpleVectorArgToFingerprint(
+            (Iterable<?>) substitutedArg, actionKeyContext, fingerprint);
+      } else if (substitutedArg instanceof ArgvFragment) {
+        i =
+            ((ArgvFragment) substitutedArg)
+                .addToFingerprint(arguments, i, actionKeyContext, fingerprint);
+      } else {
+        fingerprint.addString(CommandLineItem.expandToCommandLine(substitutedArg));
+      }
+    }
+  }
+
+  @SuppressWarnings("unchecked")
+  private void addSimpleVectorArgToFingerprint(
+      Iterable<?> arg, ActionKeyContext actionKeyContext, Fingerprint fingerprint) {
+    if (arg instanceof NestedSet) {
+      actionKeyContext.addNestedSetToFingerprint(fingerprint, (NestedSet<Object>) arg);
+    } else {
+      for (Object value : arg) {
+        fingerprint.addString(CommandLineItem.expandToCommandLine(value));
+      }
+    }
+  }
 }
diff --git a/src/main/java/com/google/devtools/build/lib/analysis/actions/ParameterFileWriteAction.java b/src/main/java/com/google/devtools/build/lib/analysis/actions/ParameterFileWriteAction.java
index df2b8894b2..6cbc02607e 100644
--- a/src/main/java/com/google/devtools/build/lib/analysis/actions/ParameterFileWriteAction.java
+++ b/src/main/java/com/google/devtools/build/lib/analysis/actions/ParameterFileWriteAction.java
@@ -67,14 +67,19 @@ public final class ParameterFileWriteAction extends AbstractFileWriteAction {
    *
    * @param owner the action owner
    * @param inputs the list of TreeArtifacts that must be resolved and expanded before evaluating
-   *     the contents of {@link commandLine}.
+   *     the contents of {@link CommandLine}.
    * @param output the Artifact that will be created by executing this Action
    * @param commandLine the contents to be written to the file
    * @param type the type of the file
    * @param charset the charset of the file
    */
-  public ParameterFileWriteAction(ActionOwner owner, Iterable<Artifact> inputs, Artifact output,
-      CommandLine commandLine, ParameterFileType type, Charset charset) {
+  public ParameterFileWriteAction(
+      ActionOwner owner,
+      Iterable<Artifact> inputs,
+      Artifact output,
+      CommandLine commandLine,
+      ParameterFileType type,
+      Charset charset) {
     super(owner, ImmutableList.copyOf(inputs), output, false);
     this.commandLine = commandLine;
     this.type = type;
@@ -172,7 +177,7 @@ public final class ParameterFileWriteAction extends AbstractFileWriteAction {
     f.addString(String.valueOf(makeExecutable));
     f.addString(type.toString());
     f.addString(charset.toString());
-    f.addStrings(commandLine.arguments());
+    commandLine.addToFingerprint(actionKeyContext, f);
     return f.hexDigestAndReset();
   }
 }
diff --git a/src/main/java/com/google/devtools/build/lib/analysis/actions/SpawnAction.java b/src/main/java/com/google/devtools/build/lib/analysis/actions/SpawnAction.java
index 74bd33a0d6..adf0dbeeef 100644
--- a/src/main/java/com/google/devtools/build/lib/analysis/actions/SpawnAction.java
+++ b/src/main/java/com/google/devtools/build/lib/analysis/actions/SpawnAction.java
@@ -337,7 +337,7 @@ public class SpawnAction extends AbstractAction implements ExecutionInfoSpecifie
       throws CommandLineExpansionException {
     Fingerprint f = new Fingerprint();
     f.addString(GUID);
-    f.addStrings(argv.arguments());
+    argv.addToFingerprint(actionKeyContext, f);
     f.addString(getMnemonic());
     // We don't need the toolManifests here, because they are a subset of the inputManifests by
     // definition and the output of an action shouldn't change whether something is considered a
diff --git a/src/main/java/com/google/devtools/build/lib/collect/nestedset/NestedSetFingerprintCache.java b/src/main/java/com/google/devtools/build/lib/collect/nestedset/NestedSetFingerprintCache.java
index c19b67e0cd..b4d341143c 100644
--- a/src/main/java/com/google/devtools/build/lib/collect/nestedset/NestedSetFingerprintCache.java
+++ b/src/main/java/com/google/devtools/build/lib/collect/nestedset/NestedSetFingerprintCache.java
@@ -23,30 +23,40 @@ import java.util.concurrent.ConcurrentHashMap;
 /** Computes fingerprints for nested sets, reusing sub-computations from children. */
 public class NestedSetFingerprintCache {
   private static final byte[] EMPTY_SET_BYTES = new byte[] {};
-  private Map<Object, byte[]> fingerprints = createMap();
+
+  /** Memoize the subresults. We have to have one cache per type of command item map function. */
+  private Map<CommandLineItem.MapFn<?>, Map<Object, byte[]>> mapFnToFingerprints = createMap();
 
   public <T> void addNestedSetToFingerprint(Fingerprint fingerprint, NestedSet<T> nestedSet) {
+    addNestedSetToFingerprint(CommandLineItem.MapFn.DEFAULT, fingerprint, nestedSet);
+  }
+
+  public <T> void addNestedSetToFingerprint(
+      CommandLineItem.MapFn<? super T> mapFn, Fingerprint fingerprint, NestedSet<T> nestedSet) {
+    Map<Object, byte[]> fingerprints =
+        mapFnToFingerprints.computeIfAbsent(mapFn, k -> new ConcurrentHashMap<>());
     fingerprint.addInt(nestedSet.getOrder().ordinal());
     Object children = nestedSet.rawChildren();
-    byte[] bytes = getBytes(children);
+    byte[] bytes = getBytes(mapFn, fingerprints, children);
     fingerprint.addBytes(bytes);
   }
 
   public void clear() {
-    fingerprints = createMap();
+    mapFnToFingerprints = createMap();
   }
 
   @SuppressWarnings("unchecked")
-  private byte[] getBytes(Object children) {
+  private <T> byte[] getBytes(
+      CommandLineItem.MapFn<? super T> mapFn, Map<Object, byte[]> fingerprints, Object children) {
     byte[] bytes = fingerprints.get(children);
     if (bytes == null) {
       if (children instanceof Object[]) {
         Fingerprint fingerprint = new Fingerprint();
         for (Object child : (Object[]) children) {
           if (child instanceof Object[]) {
-            fingerprint.addBytes(getBytes(child));
+            fingerprint.addBytes(getBytes(mapFn, fingerprints, child));
           } else {
-            addToFingerprint(fingerprint, child);
+            addToFingerprint(mapFn, fingerprint, (T) child);
           }
         }
         bytes = fingerprint.digestAndReset();
@@ -58,7 +68,7 @@ public class NestedSetFingerprintCache {
       } else if (children != NestedSet.EMPTY_CHILDREN) {
         // Single item
         Fingerprint fingerprint = new Fingerprint();
-        addToFingerprint(fingerprint, children);
+        addToFingerprint(mapFn, fingerprint, (T) children);
         bytes = fingerprint.digestAndReset();
       } else {
         // Empty nested set
@@ -69,11 +79,12 @@ public class NestedSetFingerprintCache {
   }
 
   @VisibleForTesting
-  <T> void addToFingerprint(Fingerprint fingerprint, T object) {
-    fingerprint.addString(CommandLineItem.expandToCommandLine(object));
+  <T> void addToFingerprint(
+      CommandLineItem.MapFn<? super T> mapFn, Fingerprint fingerprint, T object) {
+    fingerprint.addString(mapFn.expandToCommandLine(object));
   }
 
-  private static Map<Object, byte[]> createMap() {
+  private static Map<CommandLineItem.MapFn<?>, Map<Object, byte[]>> createMap() {
     return new ConcurrentHashMap<>();
   }
 }