From 3c81c52c9d5776235965085ade27b4fb1c6273d5 Mon Sep 17 00:00:00 2001
From: Tyler Williams <williams.tyler@gmail.com>
Date: Mon, 24 Jul 2023 16:41:08 -0700
Subject: [PATCH] Simplify blake3 hasher and improve performance

---
 .../lib/vfs/bazel/Blake3MessageDigest.java    | 37 ++-----------------
 src/main/native/blake3_jni.cc                 |  5 ++-
 2 files changed, 7 insertions(+), 35 deletions(-)

diff --git a/src/main/java/com/google/devtools/build/lib/vfs/bazel/Blake3MessageDigest.java b/src/main/java/com/google/devtools/build/lib/vfs/bazel/Blake3MessageDigest.java
index 50d2ace80f51a8..0ab2916bddb115 100644
--- a/src/main/java/com/google/devtools/build/lib/vfs/bazel/Blake3MessageDigest.java
+++ b/src/main/java/com/google/devtools/build/lib/vfs/bazel/Blake3MessageDigest.java
@@ -13,8 +13,6 @@
 // limitations under the License.
 package com.google.devtools.build.lib.vfs.bazel;
 
-import static java.lang.Math.min;
-
 import com.google.devtools.build.lib.jni.JniLoader;
 import java.nio.ByteBuffer;
 import java.security.DigestException;
@@ -37,12 +35,6 @@ public final class Blake3MessageDigest extends MessageDigest {
   static {
     initialize_hasher(INITIAL_STATE);
   }
-
-  // To reduce the number of calls made via JNI, buffer up to this many bytes
-  // before updating the hasher.
-  public static final int ONESHOT_THRESHOLD = 8 * 1024;
-
-  private final ByteBuffer buffer = ByteBuffer.allocate(ONESHOT_THRESHOLD);
   private final byte[] hasher = new byte[STATE_SIZE];
 
   public Blake3MessageDigest() {
@@ -50,33 +42,14 @@ public Blake3MessageDigest() {
     System.arraycopy(INITIAL_STATE, 0, hasher, 0, STATE_SIZE);
   }
 
-  private void flush() {
-    if (buffer.position() > 0) {
-      blake3_hasher_update(hasher, buffer.array(), buffer.position());
-      buffer.clear();
-    }
-  }
-
   @Override
   public void engineUpdate(byte[] data, int offset, int length) {
-    while (length > 0) {
-      int numToCopy = min(length, buffer.remaining());
-      buffer.put(data, offset, numToCopy);
-      length -= numToCopy;
-      offset += numToCopy;
-
-      if (buffer.remaining() == 0) {
-        flush();
-      }
-    }
+    blake3_hasher_update(hasher, data, offset, length);
   }
 
   @Override
   public void engineUpdate(byte b) {
-    if (buffer.remaining() == 0) {
-      flush();
-    }
-    buffer.put(b);
+    engineUpdate(new byte[] {b}, 0, 1);
   }
 
   @Override
@@ -85,8 +58,6 @@ public void engineUpdate(ByteBuffer input) {
   }
 
   private byte[] getOutput(int outputLength) {
-    flush();
-
     byte[] retByteArray = new byte[outputLength];
     blake3_hasher_finalize(hasher, retByteArray, outputLength);
 
@@ -101,7 +72,6 @@ public Object clone() throws CloneNotSupportedException {
 
   @Override
   public void engineReset() {
-    buffer.clear();
     System.arraycopy(INITIAL_STATE, 0, hasher, 0, STATE_SIZE);
   }
 
@@ -133,7 +103,8 @@ public int engineDigest(byte[] buf, int off, int len) throws DigestException {
 
   public static final native void initialize_hasher(byte[] hasher);
 
-  public static final native void blake3_hasher_update(byte[] hasher, byte[] input, int inputLen);
+  public static final native void blake3_hasher_update(
+      byte[] hasher, byte[] input, int offset, int inputLen);
 
   public static final native void blake3_hasher_finalize(byte[] hasher, byte[] out, int outLen);
 }
diff --git a/src/main/native/blake3_jni.cc b/src/main/native/blake3_jni.cc
index 56201888e91a56..3a0bdf14aae213 100644
--- a/src/main/native/blake3_jni.cc
+++ b/src/main/native/blake3_jni.cc
@@ -46,12 +46,13 @@ Java_com_google_devtools_build_lib_vfs_bazel_Blake3MessageDigest_initialize_1has
 
 extern "C" JNIEXPORT void JNICALL
 Java_com_google_devtools_build_lib_vfs_bazel_Blake3MessageDigest_blake3_1hasher_1update(
-    JNIEnv *env, jobject obj, jbyteArray jhasher, jbyteArray input,
+    JNIEnv *env, jobject obj, jbyteArray jhasher, jbyteArray input, jint offset,
     jint input_len) {
   blake3_hasher *hasher = (blake3_hasher *)get_byte_array(env, jhasher);
   if (hasher) {
     jbyte *input_addr = get_byte_array(env, input);
-    blake3_hasher_update(hasher, input_addr, input_len);
+    blake3_hasher_update(hasher, input_addr + (offset * sizeof(jbyte)),
+                         input_len);
     release_byte_array(env, input, input_addr);
     release_byte_array(env, jhasher, (jbyte *)hasher);
   }