1 files changed, 49 insertions, 4 deletions
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 b48d4a2298..083ef06269 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
@@ -15,9 +15,14 @@
 package com.google.devtools.build.lib.collect.nestedset;
 
 import com.google.common.annotations.VisibleForTesting;
+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 java.util.HashSet;
 import java.util.Map;
+import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
 
 /** Computes fingerprints for nested sets, reusing sub-computations from children. */
@@ -26,7 +31,10 @@ public class NestedSetFingerprintCache {
   private static final int EMPTY_SET_DIGEST = 104_395_303;
 
   /** Memoize the subresults. We have to have one cache per type of command item map function. */
-  private Map<CommandLineItem.MapFn<?>, DigestMap> mapFnToFingerprints = createMap();
+  private Map<CommandLineItem.MapFn<?>, DigestMap> mapFnToDigestMap = createMap();
+
+  private final Set<Class<?>> seenMapFns = new HashSet<>();
+  private final Multiset<Class<?>> seenParametrizedMapFns = HashMultiset.create();
 
   public <T> void addNestedSetToFingerprint(Fingerprint fingerprint, NestedSet<T> nestedSet) {
     addNestedSetToFingerprint(CommandLineItem.MapFn.DEFAULT, fingerprint, nestedSet);
@@ -34,20 +42,32 @@ public class NestedSetFingerprintCache {
 
   public <T> void addNestedSetToFingerprint(
       CommandLineItem.MapFn<? super T> mapFn, Fingerprint fingerprint, NestedSet<T> nestedSet) {
+    if (mapFn instanceof CommandLineItem.CapturingMapFn) {
+      addNestedSetToFingerprintSlow(mapFn, fingerprint, nestedSet);
+      return;
+    }
     // Only top-level nested sets can be empty, so we can bail here
     if (nestedSet.isEmpty()) {
       fingerprint.addInt(EMPTY_SET_DIGEST);
       return;
     }
-    DigestMap digestMap =
-        mapFnToFingerprints.computeIfAbsent(mapFn, k -> new DigestMap(DIGEST_SIZE, 1024));
+    DigestMap digestMap = mapFnToDigestMap.computeIfAbsent(mapFn, this::newDigestMap);
     fingerprint.addInt(nestedSet.getOrder().ordinal());
     Object children = nestedSet.rawChildren();
     addToFingerprint(mapFn, fingerprint, digestMap, children);
   }
 
+  private <T> void addNestedSetToFingerprintSlow(
+      MapFn<? super T> mapFn, Fingerprint fingerprint, NestedSet<T> nestedSet) {
+    for (T object : nestedSet) {
+      fingerprint.addString(mapFn.expandToCommandLine(object));
+    }
+  }
+
   public void clear() {
-    mapFnToFingerprints = createMap();
+    mapFnToDigestMap = createMap();
+    seenMapFns.clear();
+    seenParametrizedMapFns.clear();
   }
 
   @SuppressWarnings("unchecked")
@@ -78,4 +98,29 @@ public class NestedSetFingerprintCache {
   private static Map<CommandLineItem.MapFn<?>, DigestMap> createMap() {
     return new ConcurrentHashMap<>();
   }
+
+  private DigestMap newDigestMap(CommandLineItem.MapFn<?> mapFn) {
+    Class<?> mapFnClass = mapFn.getClass();
+    if (mapFn instanceof CommandLineItem.ParametrizedMapFn) {
+      int occurrences = seenParametrizedMapFns.add(mapFnClass, 1) + 1;
+      if (occurrences > ((CommandLineItem.ParametrizedMapFn) mapFn).maxInstancesAllowed()) {
+        throw new IllegalArgumentException(
+            String.format(
+                "Too many instances of CommandLineItem.ParametrizedMapFn '%s' detected. "
+                    + "Please construct fewer instances or use CommandLineItem.CapturingMapFn.",
+                mapFnClass.getName()));
+      }
+    } else {
+      if (!seenMapFns.add(mapFnClass)) {
+        throw new IllegalArgumentException(
+            String.format(
+                "Illegal mapFn implementation: '%s'. "
+                    + "CommandLineItem.MapFn instances must be singletons. "
+                    + "Please see CommandLineItem.ParametrizedMapFn or "
+                    + "CommandLineItem.CapturingMapFn for alternatives.",
+                mapFnClass.getName()));
+      }
+    }
+    return new DigestMap(DIGEST_SIZE, 1024);
+  }
 }