Use nested set cache in key computation for CustomCommandLine.

RELNOTES: None
PiperOrigin-RevId: 183727976

8 files changed, 257 insertions, 18 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<>();
   }
 }
diff --git a/src/test/java/com/google/devtools/build/lib/actions/CustomCommandLineTest.java b/src/test/java/com/google/devtools/build/lib/actions/CustomCommandLineTest.java
index cf15941d9b..7e514bd37f 100644
--- a/src/test/java/com/google/devtools/build/lib/actions/CustomCommandLineTest.java
+++ b/src/test/java/com/google/devtools/build/lib/actions/CustomCommandLineTest.java
@@ -17,6 +17,7 @@ import static com.google.common.truth.Truth.assertThat;
 import static com.google.devtools.build.lib.analysis.actions.CustomCommandLine.builder;
 import static org.junit.Assert.fail;
 
+import com.google.common.base.Joiner;
 import com.google.common.collect.ImmutableList;
 import com.google.devtools.build.lib.actions.Artifact.SpecialArtifact;
 import com.google.devtools.build.lib.actions.Artifact.SpecialArtifactType;
@@ -29,10 +30,13 @@ 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.testutil.Scratch;
+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.devtools.build.lib.vfs.Root;
 import java.util.Arrays;
+import java.util.HashMap;
+import java.util.Map;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -900,6 +904,40 @@ public class CustomCommandLineTest {
   }
 
   @Test
+  public void testKeyComputation() {
+    NestedSet<String> values = NestedSetBuilder.<String>stableOrder().add("a").add("b").build();
+    ImmutableList<CustomCommandLine> commandLines =
+        ImmutableList.<CustomCommandLine>builder()
+            .add(builder().add("arg").build())
+            .add(builder().addFormatted("--foo=%s", "arg").build())
+            .add(builder().addPrefixed("--foo=%s", "arg").build())
+            .add(builder().addAll(values).build())
+            .add(builder().addAll(VectorArg.addBefore("--foo=%s").each(values)).build())
+            .add(builder().addAll(VectorArg.join("--foo=%s").each(values)).build())
+            .add(builder().addAll(VectorArg.format("--foo=%s").each(values)).build())
+            .add(builder().addAll(VectorArg.of(values).mapped(s -> s + "_mapped")).build())
+            .build();
+
+    // Ensure all these command lines have distinct keys
+    ActionKeyContext actionKeyContext = new ActionKeyContext();
+    Map<String, CustomCommandLine> digests = new HashMap<>();
+    for (CustomCommandLine commandLine : commandLines) {
+      Fingerprint fingerprint = new Fingerprint();
+      commandLine.addToFingerprint(actionKeyContext, fingerprint);
+      String digest = fingerprint.hexDigestAndReset();
+      CustomCommandLine previous = digests.putIfAbsent(digest, commandLine);
+      if (previous != null) {
+        fail(
+            String.format(
+                "Found two command lines with identical digest %s: '%s' and '%s'",
+                digest,
+                Joiner.on(' ').join(previous.arguments()),
+                Joiner.on(' ').join(commandLine.arguments())));
+      }
+    }
+  }
+
+  @Test
   public void testTreeFileArtifactArgThrowWithoutSubstitution() {
     Artifact treeArtifactOne = createTreeArtifact("myArtifact/treeArtifact1");
     Artifact treeArtifactTwo = createTreeArtifact("myArtifact/treeArtifact2");
diff --git a/src/test/java/com/google/devtools/build/lib/collect/nestedset/NestedSetFingerprintCacheTest.java b/src/test/java/com/google/devtools/build/lib/collect/nestedset/NestedSetFingerprintCacheTest.java
index 9fa980ce37..975f97c972 100644
--- a/src/test/java/com/google/devtools/build/lib/collect/nestedset/NestedSetFingerprintCacheTest.java
+++ b/src/test/java/com/google/devtools/build/lib/collect/nestedset/NestedSetFingerprintCacheTest.java
@@ -17,6 +17,8 @@ import static com.google.common.truth.Truth.assertThat;
 
 import com.google.common.collect.HashMultiset;
 import com.google.common.collect.Multiset;
+import com.google.devtools.build.lib.analysis.actions.CommandLineItem;
+import com.google.devtools.build.lib.analysis.actions.CommandLineItem.MapFn;
 import com.google.devtools.build.lib.util.Fingerprint;
 import org.junit.Before;
 import org.junit.Test;
@@ -28,11 +30,11 @@ import org.junit.runners.JUnit4;
 public class NestedSetFingerprintCacheTest {
 
   private class TestNestedSetFingerprintCache extends NestedSetFingerprintCache {
-    private final Multiset<Object> fingerprinted = HashMultiset.create();
+    private Multiset<Object> fingerprinted = HashMultiset.create();
 
     @Override
-    <T> void addToFingerprint(Fingerprint fingerprint, T object) {
-      super.addToFingerprint(fingerprint, object);
+    <T> void addToFingerprint(MapFn<? super T> mapFn, Fingerprint fingerprint, T object) {
+      super.addToFingerprint(mapFn, fingerprint, object);
       fingerprinted.add(object);
     }
   }
@@ -85,4 +87,29 @@ public class NestedSetFingerprintCacheTest {
       assertThat(entry.getCount()).isEqualTo(1);
     }
   }
+
+  @Test
+  public void testMapFn() {
+    // Make sure that the map function assigns completely different key spaces
+    NestedSet<String> a = NestedSetBuilder.<String>stableOrder().add("a0").add("a1").build();
+
+    Fingerprint defaultMapFnFingerprint = new Fingerprint();
+    cache.addNestedSetToFingerprint(defaultMapFnFingerprint, a);
+    Fingerprint explicitDefaultMapFnFingerprint = new Fingerprint();
+    cache.addNestedSetToFingerprint(
+        CommandLineItem.MapFn.DEFAULT, explicitDefaultMapFnFingerprint, a);
+    Fingerprint mappedFingerprint = new Fingerprint();
+    cache.addNestedSetToFingerprint(s -> s + "_mapped", mappedFingerprint, a);
+
+    String defaultMapFnDigest = defaultMapFnFingerprint.hexDigestAndReset();
+    String explicitDefaultMapFnDigest = explicitDefaultMapFnFingerprint.hexDigestAndReset();
+    String mappedDigest = mappedFingerprint.hexDigestAndReset();
+    assertThat(defaultMapFnDigest).isEqualTo(explicitDefaultMapFnDigest);
+    assertThat(mappedDigest).isNotEqualTo(defaultMapFnDigest);
+
+    assertThat(cache.fingerprinted.elementSet()).containsExactly("a0", "a1");
+    for (Multiset.Entry<Object> entry : cache.fingerprinted.entrySet()) {
+      assertThat(entry.getCount()).isEqualTo(2);
+    }
+  }
 }