Add BLAKE3 hasher to vfs

This PR adds the Blake3Hasher and Blake3HashFunction classes to vfs and
makes them available under the flag --digest_function=BLAKE3.

PiperOrigin-RevId: 550525978
Change-Id: Iedc0886c51755585d56b4d8f47676d3be5bbedba

(cherry picked from commit cc49d68)

BUILD

@@ -93,6 +93,7 @@ genrule(
 pkg_tar(
     name = "bootstrap-jars",
     srcs = [
+        "@blake3",
         "@com_google_protobuf//:protobuf_java",
         "@com_google_protobuf//:protobuf_java_util",
         "@com_google_protobuf//:protobuf_javalite",

MODULE.bazel

@@ -16,6 +16,7 @@ bazel_dep(name = "platforms", version = "0.0.5")
 bazel_dep(name = "rules_pkg", version = "0.7.0")
 bazel_dep(name = "stardoc", version = "0.5.0", repo_name = "io_bazel_skydoc")
 bazel_dep(name = "zstd-jni", version = "1.5.2-3")
+bazel_dep(name = "blake3", version = "1.3.3")
 bazel_dep(name = "zlib", version = "1.2.13")
 
 # The following are required when building without WORKSPACE SUFFIX

distdir_deps.bzl

@@ -257,7 +257,7 @@ DIST_DEPS = {
         ],
     },
     "blake3": {
-        "archive": "v1.3.3.zip",
+        "archive": "1.3.3.zip",
         "sha256": "bb529ba133c0256df49139bd403c17835edbf60d2ecd6463549c6a5fe279364d",
         "urls": [
             "https://github.com/BLAKE3-team/BLAKE3/archive/refs/tags/1.3.3.zip",

src/main/java/com/google/devtools/build/lib/BUILD

@@ -102,6 +102,7 @@ filegroup(
         "//src/main/java/com/google/devtools/build/lib/versioning:srcs",
         "//src/main/java/com/google/devtools/build/lib/vfs/inmemoryfs:srcs",
         "//src/main/java/com/google/devtools/build/lib/vfs:srcs",
+        "//src/main/java/com/google/devtools/build/lib/vfs/bazel:srcs",
         "//src/main/java/com/google/devtools/build/lib/windows:srcs",
         "//src/main/java/com/google/devtools/build/lib/worker:srcs",
         "//src/main/java/com/google/devtools/build/skyframe:srcs",
@@ -403,6 +404,7 @@ java_library(
         "//src/main/java/com/google/devtools/build/lib/vfs",
         "//src/main/java/com/google/devtools/build/lib/vfs:output_service",
         "//src/main/java/com/google/devtools/build/lib/vfs:pathfragment",
+        "//src/main/java/com/google/devtools/build/lib/vfs/bazel",
         "//src/main/java/com/google/devtools/build/lib/windows",
         "//src/main/java/com/google/devtools/build/lib/worker:worker_metric",
         "//src/main/java/com/google/devtools/build/skyframe",

src/main/java/com/google/devtools/build/lib/bazel/BUILD

@@ -143,6 +143,7 @@ java_library(
         "//src/main/java/com/google/devtools/build/lib/util:os",
         "//src/main/java/com/google/devtools/build/lib/vfs",
         "//src/main/java/com/google/devtools/build/lib/vfs:pathfragment",
+        "//src/main/java/com/google/devtools/build/lib/vfs/bazel",
         "//src/main/java/com/google/devtools/build/lib/windows",
         "//src/main/java/com/google/devtools/common/options",
         "//src/main/protobuf:failure_details_java_proto",

src/main/java/com/google/devtools/build/lib/bazel/BazelFileSystemModule.java

@@ -31,6 +31,7 @@
 import com.google.devtools.build.lib.vfs.FileSystem;
 import com.google.devtools.build.lib.vfs.JavaIoFileSystem;
 import com.google.devtools.build.lib.vfs.PathFragment;
+import com.google.devtools.build.lib.vfs.bazel.BazelHashFunctions;
 import com.google.devtools.build.lib.windows.WindowsFileSystem;
 import com.google.devtools.common.options.OptionsParsingException;
 import com.google.devtools.common.options.OptionsParsingResult;
@@ -44,6 +45,10 @@
  * com.google.devtools.build.lib.vfs.FileSystem} class use {@code SHA256} by default.
  */
 public class BazelFileSystemModule extends BlazeModule {
+  static {
+    BazelHashFunctions.ensureRegistered();
+  }
+
   @Override
   public ModuleFileSystem getFileSystem(
       OptionsParsingResult startupOptions, PathFragment realExecRootBase)

src/main/java/com/google/devtools/build/lib/vfs/bazel/BUILD

@@ -0,0 +1,28 @@
+load("@rules_java//java:defs.bzl", "java_library")
+
+package(
+    default_applicable_licenses = ["//:license"],
+    default_visibility = ["//src:__subpackages__"],
+)
+
+filegroup(
+    name = "srcs",
+    srcs = glob(["**"]),
+    visibility = ["//src:__subpackages__"],
+)
+
+java_library(
+    name = "bazel",
+    srcs = glob(
+        [
+            "*.java",
+        ],
+    ),
+    deps = [
+        "//src/main/java/com/google/devtools/build/lib/jni",
+        "//src/main/java/com/google/devtools/build/lib/vfs",
+        "//third_party:error_prone_annotations",
+        "//third_party:guava",
+        "//third_party:jsr305",
+    ],
+)

src/main/java/com/google/devtools/build/lib/vfs/bazel/BazelHashFunctions.java

@@ -0,0 +1,45 @@
+// Copyright 2023 The Bazel Authors. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+package com.google.devtools.build.lib.vfs.bazel;
+
+import com.google.devtools.build.lib.jni.JniLoader;
+import com.google.devtools.build.lib.vfs.DigestHashFunction;
+import java.security.Security;
+import javax.annotation.Nullable;
+
+/** Bazel specific {@link DigestHashFunction}s. */
+public final class BazelHashFunctions {
+  @Nullable public static final DigestHashFunction BLAKE3;
+
+  static {
+    DigestHashFunction hashFunction = null;
+
+    if (JniLoader.isJniAvailable()) {
+      try {
+        Security.addProvider(new Blake3Provider());
+        hashFunction = DigestHashFunction.register(new Blake3HashFunction(), "BLAKE3");
+      } catch (UnsatisfiedLinkError ignored) {
+        // This can happen if bazel was compiled manually (with compile.sh),
+        // on windows. In that case jni is available, but missing the blake3
+        // symbols necessary to register the hasher.
+      }
+    }
+
+    BLAKE3 = hashFunction;
+  }
+
+  public static void ensureRegistered() {}
+
+  private BazelHashFunctions() {}
+}

src/main/java/com/google/devtools/build/lib/vfs/bazel/Blake3HashFunction.java

@@ -0,0 +1,88 @@
+// Copyright 2023 The Bazel Authors. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+package com.google.devtools.build.lib.vfs.bazel;
+
+import static com.google.common.base.Preconditions.checkArgument;
+import static com.google.common.base.Preconditions.checkPositionIndexes;
+
+import com.google.common.hash.Funnel;
+import com.google.common.hash.HashCode;
+import com.google.common.hash.HashFunction;
+import com.google.common.hash.Hasher;
+import java.nio.ByteBuffer;
+import java.nio.charset.Charset;
+
+/** A {@link HashFunction} for BLAKE3. */
+public final class Blake3HashFunction implements HashFunction {
+  @Override
+  public int bits() {
+    return 256;
+  }
+
+  @Override
+  public Hasher newHasher() {
+    return new Blake3Hasher(new Blake3MessageDigest());
+  }
+
+  @Override
+  public Hasher newHasher(int expectedInputSize) {
+    checkArgument(
+        expectedInputSize >= 0, "expectedInputSize must be >= 0 but was %s", expectedInputSize);
+    return newHasher();
+  }
+
+  /* The following methods implement the {HashFunction} interface. */
+
+  @Override
+  public <T> HashCode hashObject(T instance, Funnel<? super T> funnel) {
+    return newHasher().putObject(instance, funnel).hash();
+  }
+
+  @Override
+  public HashCode hashUnencodedChars(CharSequence input) {
+    int len = input.length();
+    return newHasher(len * 2).putUnencodedChars(input).hash();
+  }
+
+  @Override
+  public HashCode hashString(CharSequence input, Charset charset) {
+    return newHasher().putString(input, charset).hash();
+  }
+
+  @Override
+  public HashCode hashInt(int input) {
+    return newHasher(4).putInt(input).hash();
+  }
+
+  @Override
+  public HashCode hashLong(long input) {
+    return newHasher(8).putLong(input).hash();
+  }
+
+  @Override
+  public HashCode hashBytes(byte[] input) {
+    return hashBytes(input, 0, input.length);
+  }
+
+  @Override
+  public HashCode hashBytes(byte[] input, int off, int len) {
+    checkPositionIndexes(off, off + len, input.length);
+    return newHasher(len).putBytes(input, off, len).hash();
+  }
+
+  @Override
+  public HashCode hashBytes(ByteBuffer input) {
+    return newHasher(input.remaining()).putBytes(input).hash();
+  }
+}

src/main/java/com/google/devtools/build/lib/vfs/bazel/Blake3Hasher.java

@@ -0,0 +1,146 @@
+// Copyright 2023 The Bazel Authors. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+package com.google.devtools.build.lib.vfs.bazel;
+
+import static com.google.common.base.Preconditions.checkState;
+
+import com.google.common.hash.Funnel;
+import com.google.common.hash.HashCode;
+import com.google.common.hash.Hasher;
+import com.google.errorprone.annotations.CanIgnoreReturnValue;
+import java.nio.ByteBuffer;
+import java.nio.charset.Charset;
+
+/** A {@link Hasher} for BLAKE3. */
+public final class Blake3Hasher implements Hasher {
+  private final Blake3MessageDigest messageDigest;
+  private boolean isDone = false;
+
+  public Blake3Hasher(Blake3MessageDigest blake3MessageDigest) {
+    messageDigest = blake3MessageDigest;
+  }
+
+  /* The following methods implement the {Hasher} interface. */
+
+  @Override
+  @CanIgnoreReturnValue
+  public Hasher putBytes(ByteBuffer b) {
+    messageDigest.engineUpdate(b);
+    return this;
+  }
+
+  @Override
+  @CanIgnoreReturnValue
+  public Hasher putBytes(byte[] bytes, int off, int len) {
+    messageDigest.engineUpdate(bytes, off, len);
+    return this;
+  }
+
+  @Override
+  @CanIgnoreReturnValue
+  public Hasher putBytes(byte[] bytes) {
+    messageDigest.engineUpdate(bytes, 0, bytes.length);
+    return this;
+  }
+
+  @Override
+  @CanIgnoreReturnValue
+  public Hasher putByte(byte b) {
+    messageDigest.engineUpdate(b);
+    return this;
+  }
+
+  @Override
+  public HashCode hash() {
+    checkState(!isDone);
+    isDone = true;
+
+    return HashCode.fromBytes(messageDigest.engineDigest());
+  }
+
+  @Override
+  @CanIgnoreReturnValue
+  public final Hasher putBoolean(boolean b) {
+    return putByte(b ? (byte) 1 : (byte) 0);
+  }
+
+  @Override
+  @CanIgnoreReturnValue
+  public final Hasher putDouble(double d) {
+    return putLong(Double.doubleToRawLongBits(d));
+  }
+
+  @Override
+  @CanIgnoreReturnValue
+  public final Hasher putFloat(float f) {
+    return putInt(Float.floatToRawIntBits(f));
+  }
+
+  @Override
+  @CanIgnoreReturnValue
+  public Hasher putUnencodedChars(CharSequence charSequence) {
+    for (int i = 0, len = charSequence.length(); i < len; i++) {
+      putChar(charSequence.charAt(i));
+    }
+    return this;
+  }
+
+  @Override
+  @CanIgnoreReturnValue
+  public Hasher putString(CharSequence charSequence, Charset charset) {
+    return putBytes(charSequence.toString().getBytes(charset));
+  }
+
+  @Override
+  @CanIgnoreReturnValue
+  public Hasher putShort(short s) {
+    putByte((byte) s);
+    putByte((byte) (s >>> 8));
+    return this;
+  }
+
+  @Override
+  @CanIgnoreReturnValue
+  public Hasher putInt(int i) {
+    putByte((byte) i);
+    putByte((byte) (i >>> 8));
+    putByte((byte) (i >>> 16));
+    putByte((byte) (i >>> 24));
+    return this;
+  }
+
+  @Override
+  @CanIgnoreReturnValue
+  public Hasher putLong(long l) {
+    for (int i = 0; i < 64; i += 8) {
+      putByte((byte) (l >>> i));
+    }
+    return this;
+  }
+
+  @Override
+  @CanIgnoreReturnValue
+  public Hasher putChar(char c) {
+    putByte((byte) c);
+    putByte((byte) (c >>> 8));
+    return this;
+  }
+
+  @Override
+  @CanIgnoreReturnValue
+  public <T> Hasher putObject(T instance, Funnel<? super T> funnel) {
+    funnel.funnel(instance, this);
+    return this;
+  }
+}