Add BLAKE3 hasher to vfs

BUILD

@@ -100,6 +100,7 @@ pkg_tar(
         "@com_google_protobuf//:protobuf_java_util",
         "@com_google_protobuf//:protobuf_javalite",
         "@zstd-jni//:zstd-jni",
+        "@blake3//:blake3",
     ],
     package_dir = "derived/jars",
     strip_prefix = "external",

MODULE.bazel

@@ -14,6 +14,7 @@ bazel_dep(name = "platforms", version = "0.0.6")
 bazel_dep(name = "rules_pkg", version = "0.7.0")
 bazel_dep(name = "stardoc", version = "0.5.3", 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")
 bazel_dep(name = "rules_cc", version = "0.0.6")
 bazel_dep(name = "rules_java", version = "6.1.1")

distdir_deps.bzl

@@ -274,7 +274,7 @@ DIST_DEPS = {
         "package_version": "1.5.2-3",
     },
     "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

@@ -96,6 +96,7 @@ filegroup(
         "//src/main/java/com/google/devtools/build/lib/util:srcs",
         "//src/main/java/com/google/devtools/build/lib/versioning: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/vfs/inmemoryfs:srcs",
         "//src/main/java/com/google/devtools/build/lib/windows:srcs",
         "//src/main/java/com/google/devtools/build/lib/worker:srcs",
@@ -443,6 +444,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/runtime/BazelFileSystemModule.java

@@ -29,6 +29,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;
@@ -42,6 +43,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/util:abrupt_exit_exception",
+        "//src/main/java/com/google/devtools/build/lib/jni",
+        "//src/main/java/com/google/devtools/build/lib/vfs",
+	"//src/main/java/com/google/devtools/common/options",
+        "//third_party:guava",
+    ],
+)

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

@@ -0,0 +1,16 @@
+package com.google.devtools.build.lib.vfs.bazel;
+
+import com.google.devtools.build.lib.vfs.DigestHashFunction;
+import java.security.Security;
+
+public final class BazelHashFunctions {
+  static {
+    Security.addProvider(new Blake3Provider());
+  }
+
+  public static final DigestHashFunction BLAKE3 =
+      DigestHashFunction.register(new Blake3HashFunction(), "BLAKE3");
+
+  public static void ensureRegistered() {}
+  ;
+}

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

@@ -0,0 +1,63 @@
+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;
+
+public final class Blake3HashFunction implements HashFunction {
+  public int bits() {
+    return 256;
+  }
+
+  public Hasher newHasher() {
+    return new Blake3Hasher(new Blake3MessageDigest());
+  }
+
+  /* The following methods implement the {HashFunction} interface. */
+
+  public <T extends Object> HashCode hashObject(T instance, Funnel<? super T> funnel) {
+    return newHasher().putObject(instance, funnel).hash();
+  }
+
+  public HashCode hashUnencodedChars(CharSequence input) {
+    int len = input.length();
+    return newHasher(len * 2).putUnencodedChars(input).hash();
+  }
+
+  public HashCode hashString(CharSequence input, Charset charset) {
+    return newHasher().putString(input, charset).hash();
+  }
+
+  public HashCode hashInt(int input) {
+    return newHasher(4).putInt(input).hash();
+  }
+
+  public HashCode hashLong(long input) {
+    return newHasher(8).putLong(input).hash();
+  }
+
+  public HashCode hashBytes(byte[] input) {
+    return hashBytes(input, 0, input.length);
+  }
+
+  public HashCode hashBytes(byte[] input, int off, int len) {
+    checkPositionIndexes(off, off + len, input.length);
+    return newHasher(len).putBytes(input, off, len).hash();
+  }
+
+  public HashCode hashBytes(ByteBuffer input) {
+    return newHasher(input.remaining()).putBytes(input).hash();
+  }
+
+  public Hasher newHasher(int expectedInputSize) {
+    checkArgument(
+        expectedInputSize >= 0, "expectedInputSize must be >= 0 but was %s", expectedInputSize);
+    return newHasher();
+  }
+}

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

@@ -0,0 +1,117 @@
+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;
+
+public final class Blake3Hasher implements Hasher {
+  private Blake3MessageDigest messageDigest;
+  private boolean isDone = false;
+
+  public Blake3Hasher(Blake3MessageDigest blake3MessageDigest) {
+    messageDigest = blake3MessageDigest;
+  }
+
+  /* The following methods implement the {Hasher} interface. */
+
+  @CanIgnoreReturnValue
+  public Hasher putBytes(ByteBuffer b) {
+    messageDigest.engineUpdate(b);
+    return this;
+  }
+
+  @CanIgnoreReturnValue
+  public Hasher putBytes(byte[] bytes, int off, int len) {
+    messageDigest.engineUpdate(bytes, off, len);
+    return this;
+  }
+
+  @CanIgnoreReturnValue
+  public Hasher putBytes(byte[] bytes) {
+    messageDigest.engineUpdate(bytes, 0, bytes.length);
+    return this;
+  }
+
+  @CanIgnoreReturnValue
+  public Hasher putByte(byte b) {
+    messageDigest.engineUpdate(new byte[] {b});
+    return this;
+  }
+
+  public HashCode hash() throws IllegalStateException {
+    checkState(!isDone);
+    isDone = true;
+
+    return HashCode.fromBytes(messageDigest.engineDigest());
+  }
+
+  @CanIgnoreReturnValue
+  public final Hasher putBoolean(boolean b) {
+    return putByte(b ? (byte) 1 : (byte) 0);
+  }
+
+  @CanIgnoreReturnValue
+  public final Hasher putDouble(double d) {
+    return putLong(Double.doubleToRawLongBits(d));
+  }
+
+  @CanIgnoreReturnValue
+  public final Hasher putFloat(float f) {
+    return putInt(Float.floatToRawIntBits(f));
+  }
+
+  @CanIgnoreReturnValue
+  public Hasher putUnencodedChars(CharSequence charSequence) {
+    for (int i = 0, len = charSequence.length(); i < len; i++) {
+      putChar(charSequence.charAt(i));
+    }
+    return this;
+  }
+
+  @CanIgnoreReturnValue
+  public Hasher putString(CharSequence charSequence, Charset charset) {
+    return putBytes(charSequence.toString().getBytes(charset));
+  }
+
+  @CanIgnoreReturnValue
+  public Hasher putShort(short s) {
+    putByte((byte) s);
+    putByte((byte) (s >>> 8));
+    return this;
+  }
+
+  @CanIgnoreReturnValue
+  public Hasher putInt(int i) {
+    putByte((byte) i);
+    putByte((byte) (i >>> 8));
+    putByte((byte) (i >>> 16));
+    putByte((byte) (i >>> 24));
+    return this;
+  }
+
+  @CanIgnoreReturnValue
+  public Hasher putLong(long l) {
+    for (int i = 0; i < 64; i += 8) {
+      putByte((byte) (l >>> i));
+    }
+    return this;
+  }
+
+  @CanIgnoreReturnValue
+  public Hasher putChar(char c) {
+    putByte((byte) c);
+    putByte((byte) (c >>> 8));
+    return this;
+  }
+
+  @CanIgnoreReturnValue
+  public <T extends Object> Hasher putObject(T instance, Funnel<? super T> funnel) {
+    funnel.funnel(instance, this);
+    return this;
+  }
+}

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

@@ -0,0 +1,38 @@
+// 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;
+
+final class Blake3JNI {
+  private Blake3JNI() {}
+
+  static {
+    JniLoader.loadJni();
+  }
+
+  public static final native long allocate_and_initialize_hasher();
+
+  public static final native void blake3_hasher_reset(long self);
+
+  public static final native void blake3_hasher_close(long self);
+
+  public static final native void blake3_hasher_update(long self, byte[] input, int input_len);
+
+  public static final native void blake3_hasher_finalize_and_reset(
+      long self, byte[] out, int out_len);
+
+  public static final native void oneshot(byte[] input, int input_len, byte[] out, int out_len);
+}