Have DigestMap support multiple hash functions.

RELNOTES: None
PiperOrigin-RevId: 208837641

6 files changed, 93 insertions, 29 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/collect/nestedset/BUILD b/src/main/java/com/google/devtools/build/lib/collect/nestedset/BUILD
index 134eaca919..853c35f9e9 100644
--- a/src/main/java/com/google/devtools/build/lib/collect/nestedset/BUILD
+++ b/src/main/java/com/google/devtools/build/lib/collect/nestedset/BUILD
@@ -55,6 +55,7 @@ java_library(
         ":nestedset",
         "//src/main/java/com/google/devtools/build/lib:util",
         "//src/main/java/com/google/devtools/build/lib/actions:commandline_item",
+        "//src/main/java/com/google/devtools/build/lib/vfs",
         "//third_party:guava",
     ],
 )
diff --git a/src/main/java/com/google/devtools/build/lib/collect/nestedset/DigestMap.java b/src/main/java/com/google/devtools/build/lib/collect/nestedset/DigestMap.java
index da6f6a47d5..30d53f0a14 100644
--- a/src/main/java/com/google/devtools/build/lib/collect/nestedset/DigestMap.java
+++ b/src/main/java/com/google/devtools/build/lib/collect/nestedset/DigestMap.java
@@ -15,6 +15,8 @@ package com.google.devtools.build.lib.collect.nestedset;
 
 import com.google.common.base.Preconditions;
 import com.google.devtools.build.lib.util.Fingerprint;
+import com.google.devtools.build.lib.vfs.DigestHashFunction;
+import com.google.devtools.build.lib.vfs.DigestHashFunction.DigestLength;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicReferenceArray;
 import java.util.concurrent.locks.StampedLock;
@@ -33,7 +35,7 @@ import java.util.concurrent.locks.StampedLock;
  */
 final class DigestMap {
   private static final Object INSERTION_IN_PROGRESS = new Object();
-  private final int digestSize;
+  private final DigestLength digestLength;
   private final StampedLock readWriteLock = new StampedLock();
 
   static class Table {
@@ -42,23 +44,23 @@ final class DigestMap {
     final AtomicReferenceArray<Object> keys;
     final byte[] bytes;
 
-    Table(int tableSize, int digestSize) {
+    Table(int tableSize, int digestLength) {
       this.tableSize = tableSize;
       this.nextResize = getNextResize(tableSize);
       this.keys = new AtomicReferenceArray<>(tableSize);
-      this.bytes = new byte[tableSize * digestSize];
+      this.bytes = new byte[tableSize * digestLength];
     }
   }
 
   private volatile Table table;
   private final AtomicInteger allocatedSlots;
 
-  DigestMap(int digestSize, int initialSize) {
+  DigestMap(DigestHashFunction digestHashFunction, int initialSize) {
     Preconditions.checkArgument(
         initialSize > 0 && (initialSize & (initialSize - 1)) == 0,
         "initialSize must be a power of 2 greater than 0");
-    this.digestSize = digestSize;
-    this.table = new Table(initialSize, digestSize);
+    this.digestLength = digestHashFunction.getDigestLength();
+    this.table = new Table(initialSize, digestLength.getDigestMaximumLength());
     this.allocatedSlots = new AtomicInteger();
   }
 
@@ -67,7 +69,9 @@ final class DigestMap {
     Table table = this.table; // Read once for duration of method
     int index = findKey(table, key);
     if (index >= 0) {
-      fingerprint.addBytes(table.bytes, index * digestSize, digestSize);
+      int offset = index * this.digestLength.getDigestMaximumLength();
+      int digestLength = this.digestLength.getDigestLength(table.bytes, offset);
+      fingerprint.addBytes(table.bytes, offset, digestLength);
       return true;
     }
     return false;
@@ -124,7 +128,9 @@ final class DigestMap {
       readWriteLock.unlockRead(stamp);
     }
     // This can be done outside of the read lock since the slot is immutable once inserted
-    readTo.addBytes(table.bytes, index * digestSize, digestSize);
+    int offset = index * this.digestLength.getDigestMaximumLength();
+    int digestLength = this.digestLength.getDigestLength(table.bytes, offset);
+    readTo.addBytes(table.bytes, offset, digestLength);
   }
 
   // Inserts a key into the passed table and returns the index.
@@ -140,7 +146,10 @@ final class DigestMap {
           // Failure to do so could lead to a double insertion.
           continue;
         }
-        digest.digestAndReset(table.bytes, index * digestSize, digestSize);
+        digest.digestAndReset(
+            table.bytes,
+            index * digestLength.getDigestMaximumLength(),
+            digestLength.getDigestMaximumLength());
         table.keys.set(index, key);
         return index;
       } else if (currentKey == key) {
@@ -160,7 +169,7 @@ final class DigestMap {
   }
 
   private void resizeTableWriteLocked() {
-    int digestSize = this.digestSize;
+    int digestSize = this.digestLength.getDigestMaximumLength();
     Table oldTable = this.table;
     Table newTable = new Table(oldTable.tableSize * 2, digestSize);
     for (int i = 0; i < oldTable.tableSize; ++i) {
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 8eb4131793..e0481b98f6 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
@@ -20,6 +20,7 @@ import com.google.common.collect.Multiset;
 import com.google.devtools.build.lib.actions.CommandLineItem;
 import com.google.devtools.build.lib.actions.CommandLineItem.MapFn;
 import com.google.devtools.build.lib.util.Fingerprint;
+import com.google.devtools.build.lib.vfs.DigestHashFunction;
 import java.util.HashSet;
 import java.util.Map;
 import java.util.Set;
@@ -27,7 +28,6 @@ import java.util.concurrent.ConcurrentHashMap;
 
 /** Computes fingerprints for nested sets, reusing sub-computations from children. */
 public class NestedSetFingerprintCache {
-  private static final int DIGEST_SIZE = 16;
   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. */
@@ -121,6 +121,8 @@ public class NestedSetFingerprintCache {
                 mapFnClass.getName()));
       }
     }
-    return new DigestMap(DIGEST_SIZE, 1024);
+    // TODO(b/112460990): Use the value from DigestHashFunction.getDefault(), but check for
+    // contention.
+    return new DigestMap(DigestHashFunction.MD5, 1024);
   }
 }
diff --git a/src/main/java/com/google/devtools/build/lib/vfs/DigestHashFunction.java b/src/main/java/com/google/devtools/build/lib/vfs/DigestHashFunction.java
index 123e4b5a9c..e059ba832d 100644
--- a/src/main/java/com/google/devtools/build/lib/vfs/DigestHashFunction.java
+++ b/src/main/java/com/google/devtools/build/lib/vfs/DigestHashFunction.java
@@ -18,6 +18,7 @@ import com.google.common.annotations.VisibleForTesting;
 import com.google.common.collect.ImmutableList;
 import com.google.common.hash.HashFunction;
 import com.google.common.hash.Hashing;
+import com.google.devtools.build.lib.vfs.DigestHashFunction.DigestLength.DigestLengthImpl;
 import com.google.devtools.common.options.Converter;
 import com.google.devtools.common.options.OptionsParsingException;
 import java.security.MessageDigest;
@@ -39,6 +40,31 @@ public class DigestHashFunction {
   // This map must be declared first to make sure that calls to register() have it ready.
   private static final HashMap<String, DigestHashFunction> hashFunctionRegistry = new HashMap<>();
 
+  /** Describes the length of a digest. */
+  public interface DigestLength {
+    /** Returns the length of a digest by inspecting its bytes. Used for variable-length digests. */
+    default int getDigestLength(byte[] bytes, int offset) {
+      return getDigestMaximumLength();
+    }
+
+    /** Returns the maximum length a digest can turn into. */
+    int getDigestMaximumLength();
+
+    /** Default implementation that simply returns a fixed length. */
+    class DigestLengthImpl implements DigestLength {
+      private final int length;
+
+      DigestLengthImpl(HashFunction hashFunction) {
+        this.length = hashFunction.bits() / 8;
+      }
+
+      @Override
+      public int getDigestMaximumLength() {
+        return length;
+      }
+    }
+  }
+
   public static final DigestHashFunction MD5 = register(Hashing.md5(), "MD5");
   public static final DigestHashFunction SHA1 = register(Hashing.sha1(), "SHA-1", "SHA1");
   public static final DigestHashFunction SHA256 = register(Hashing.sha256(), "SHA-256", "SHA256");
@@ -47,17 +73,24 @@ public class DigestHashFunction {
   private static boolean defaultHasBeenSet = false;
 
   private final HashFunction hashFunction;
+  private final DigestLength digestLength;
   private final String name;
   private final MessageDigest messageDigestPrototype;
   private final boolean messageDigestPrototypeSupportsClone;
 
-  private DigestHashFunction(HashFunction hashFunction, String name) {
+  private DigestHashFunction(HashFunction hashFunction, DigestLength digestLength, String name) {
     this.hashFunction = hashFunction;
+    this.digestLength = digestLength;
     this.name = name;
     this.messageDigestPrototype = getMessageDigestInstance();
     this.messageDigestPrototypeSupportsClone = supportsClone(messageDigestPrototype);
   }
 
+  public static DigestHashFunction register(
+      HashFunction hash, String hashName, String... altNames) {
+    return register(hash, new DigestLengthImpl(hash), hashName, altNames);
+  }
+
   /**
    * Creates a new DigestHashFunction that is registered to be recognized by its name in {@link
    * DigestFunctionConverter}.
@@ -70,7 +103,7 @@ public class DigestHashFunction {
    * @throws IllegalArgumentException if the name is already registered.
    */
   public static DigestHashFunction register(
-      HashFunction hash, String hashName, String... altNames) {
+      HashFunction hash, DigestLength digestLength, String hashName, String... altNames) {
     try {
       MessageDigest.getInstance(hashName);
     } catch (NoSuchAlgorithmException e) {
@@ -80,7 +113,7 @@ public class DigestHashFunction {
           e);
     }
 
-    DigestHashFunction hashFunction = new DigestHashFunction(hash, hashName);
+    DigestHashFunction hashFunction = new DigestHashFunction(hash, digestLength, hashName);
     List<String> names = ImmutableList.<String>builder().add(hashName).add(altNames).build();
     synchronized (hashFunctionRegistry) {
       for (String name : names) {
@@ -177,6 +210,10 @@ public class DigestHashFunction {
     }
   }
 
+  public DigestLength getDigestLength() {
+    return digestLength;
+  }
+
   public boolean isValidDigest(byte[] digest) {
     // TODO(b/109764197): Remove this check to accept variable-length hashes.
     return digest != null && digest.length * 8 == hashFunction.bits();
diff --git a/src/test/java/com/google/devtools/build/lib/BUILD b/src/test/java/com/google/devtools/build/lib/BUILD
index 50d954e62e..a4941a77e7 100644
--- a/src/test/java/com/google/devtools/build/lib/BUILD
+++ b/src/test/java/com/google/devtools/build/lib/BUILD
@@ -232,6 +232,7 @@ java_test(
         "//src/main/java/com/google/devtools/build/lib/skyframe/serialization",
         "//src/main/java/com/google/devtools/build/lib/skyframe/serialization:constants",
         "//src/main/java/com/google/devtools/build/lib/skyframe/serialization/testutils",
+        "//src/main/java/com/google/devtools/build/lib/vfs",
         "//third_party:mockito",
         "//third_party/protobuf:protobuf_java",
     ],
diff --git a/src/test/java/com/google/devtools/build/lib/collect/nestedset/DigestMapTest.java b/src/test/java/com/google/devtools/build/lib/collect/nestedset/DigestMapTest.java
index 35fb7ed3da..67d9b86009 100644
--- a/src/test/java/com/google/devtools/build/lib/collect/nestedset/DigestMapTest.java
+++ b/src/test/java/com/google/devtools/build/lib/collect/nestedset/DigestMapTest.java
@@ -15,7 +15,9 @@ package com.google.devtools.build.lib.collect.nestedset;
 
 import static com.google.common.truth.Truth.assertThat;
 
+import com.google.common.collect.ImmutableList;
 import com.google.devtools.build.lib.util.Fingerprint;
+import com.google.devtools.build.lib.vfs.DigestHashFunction;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Random;
@@ -23,12 +25,26 @@ import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicReference;
 import org.junit.Test;
 import org.junit.runner.RunWith;
-import org.junit.runners.JUnit4;
+import org.junit.runners.Parameterized;
+import org.junit.runners.Parameterized.Parameter;
+import org.junit.runners.Parameterized.Parameters;
 
 /** Tests for {@link DigestMap}. */
-@RunWith(JUnit4.class)
+@RunWith(Parameterized.class)
 public class DigestMapTest {
 
+  @Parameters(name = "Hash: {0}")
+  public static Iterable<Object[]> hashFunction() {
+    return ImmutableList.of(
+        new Object[] {DigestHashFunction.MD5}, new Object[] {DigestHashFunction.SHA256});
+  }
+
+  @Parameter public DigestHashFunction digestHashFunction;
+
+  private Fingerprint fingerprint() {
+    return new Fingerprint(digestHashFunction);
+  }
+
   @Test
   public void simpleTest() {
     int count = 128; // Must be smaller than byte for this test or we'll overflow
@@ -37,20 +53,18 @@ public class DigestMapTest {
       keys[i] = new Object();
     }
 
-    DigestMap digestMap = new DigestMap(16, 4);
+    DigestMap digestMap = new DigestMap(digestHashFunction, 4);
     for (int i = 0; i < count; ++i) {
-      Fingerprint digest = new Fingerprint().addInt(i);
-      Fingerprint fingerprint = new Fingerprint();
+      Fingerprint digest = fingerprint().addInt(i);
+      Fingerprint fingerprint = fingerprint();
       digestMap.insertAndReadDigest(keys[i], digest, fingerprint);
-      Fingerprint reference =
-          new Fingerprint().addBytes(new Fingerprint().addInt(i).digestAndReset());
+      Fingerprint reference = fingerprint().addBytes(fingerprint().addInt(i).digestAndReset());
       assertThat(fingerprint.hexDigestAndReset()).isEqualTo(reference.hexDigestAndReset());
     }
     for (int i = 0; i < count; ++i) {
-      Fingerprint fingerprint = new Fingerprint();
+      Fingerprint fingerprint = fingerprint();
       assertThat(digestMap.readDigest(keys[i], fingerprint)).isTrue();
-      Fingerprint reference =
-          new Fingerprint().addBytes(new Fingerprint().addInt(i).digestAndReset());
+      Fingerprint reference = fingerprint().addBytes(fingerprint().addInt(i).digestAndReset());
       assertThat(fingerprint.hexDigestAndReset()).isEqualTo(reference.hexDigestAndReset());
     }
   }
@@ -62,7 +76,7 @@ public class DigestMapTest {
     for (int i = 0; i < count; ++i) {
       keys[i] = new Object();
     }
-    DigestMap digestMap = new DigestMap(16, 4);
+    DigestMap digestMap = new DigestMap(digestHashFunction, 4);
 
     AtomicBoolean done = new AtomicBoolean();
     AtomicReference<Exception> exception = new AtomicReference<>();
@@ -76,13 +90,13 @@ public class DigestMapTest {
                 while (!done.get()) {
                   int index = random.nextInt(count);
                   Object key = keys[index];
-                  Fingerprint fingerprint = new Fingerprint();
+                  Fingerprint fingerprint = fingerprint();
                   if (!digestMap.readDigest(key, fingerprint)) {
-                    Fingerprint digest = new Fingerprint().addInt(index);
+                    Fingerprint digest = fingerprint().addInt(index);
                     digestMap.insertAndReadDigest(key, digest, fingerprint);
                   }
                   Fingerprint reference =
-                      new Fingerprint().addBytes(new Fingerprint().addInt(index).digestAndReset());
+                      fingerprint().addBytes(fingerprint().addInt(index).digestAndReset());
                   String hexDigest = fingerprint.hexDigestAndReset();
                   String referenceDigest = reference.hexDigestAndReset();
                   if (!hexDigest.equals(referenceDigest)) {