Simplify blake3 hasher and improve performance

bazelbuild · Jul 25, 2023 · 4ee7857 · 4ee7857
1 parent f48c0f6
commit 4ee7857
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 34 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
@@ -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;
@@ -38,45 +36,23 @@ public final class Blake3MessageDigest extends MessageDigest {
     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];
+  private byte[] oneByteArray = new byte[1];
 
   public Blake3MessageDigest() {
     super("BLAKE3");
     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);
+    oneByteArray[0] = b;
+    engineUpdate(oneByteArray, 0, 1);
   }
 
   @Override
@@ -85,8 +61,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 +75,6 @@ public Object clone() throws CloneNotSupportedException {
 
   @Override
   public void engineReset() {
-    buffer.clear();
     System.arraycopy(INITIAL_STATE, 0, hasher, 0, STATE_SIZE);
   }
 
@@ -133,7 +106,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
@@ -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);
   }