From 17b8da37ff4f70b82e0bed563d9ce234f64b5987 Mon Sep 17 00:00:00 2001
From: Tuyen Nguyen <vutuyen2636@gmail.com>
Date: Tue, 28 May 2024 14:34:56 +0900
Subject: [PATCH] fix: improve for loop in as-sha256 executeHashComputations

---
 .../src/hasher/as-sha256.ts                   | 125 ++++++++++++++----
 packages/persistent-merkle-tree/src/node.ts   |   2 +-
 2 files changed, 97 insertions(+), 30 deletions(-)

diff --git a/packages/persistent-merkle-tree/src/hasher/as-sha256.ts b/packages/persistent-merkle-tree/src/hasher/as-sha256.ts
index ec484ddb..b92db603 100644
--- a/packages/persistent-merkle-tree/src/hasher/as-sha256.ts
+++ b/packages/persistent-merkle-tree/src/hasher/as-sha256.ts
@@ -1,6 +1,6 @@
 import {digest2Bytes32, digest64HashObjects, HashObject, batchHash4HashObjectInputs} from "@chainsafe/as-sha256";
 import type {Hasher} from "./types";
-import {HashComputation} from "../node";
+import {HashComputation, Node} from "../node";
 
 export const hasher: Hasher = {
   digest64: digest2Bytes32,
@@ -36,39 +36,106 @@ export const hasher: Hasher = {
       }
 
       // HashComputations of the same level are safe to batch
-      const batch = Math.floor(hcArr.length / 4);
-      for (let i = 0; i < batch; i++) {
-        const index = i * 4;
-        // access array once
-        const {src0: src0_0, src1: src1_0, dest: dest_0} = hcArr[index];
-        const {src0: src0_1, src1: src1_1, dest: dest_1} = hcArr[index + 1];
-        const {src0: src0_2, src1: src1_2, dest: dest_2} = hcArr[index + 2];
-        const {src0: src0_3, src1: src1_3, dest: dest_3} = hcArr[index + 3];
+      let src0_0: Node | null = null;
+      let src1_0: Node | null = null;
+      let dest_0: Node | null = null;
+      let src0_1: Node | null = null;
+      let src1_1: Node | null = null;
+      let dest_1: Node | null = null;
+      let src0_2: Node | null = null;
+      let src1_2: Node | null = null;
+      let dest_2: Node | null = null;
+      let src0_3: Node | null = null;
+      let src1_3: Node | null = null;
+      let dest_3: Node | null = null;
 
-        const [o0, o1, o2, o3] = batchHash4HashObjectInputs([
-          src0_0,
-          src1_0,
-          src0_1,
-          src1_1,
-          src0_2,
-          src1_2,
-          src0_3,
-          src1_3,
-        ]);
-        if (o0 == null || o1 == null || o2 == null || o3 == null) {
-          throw Error(`batchHash4HashObjectInputs return null at batch ${i} level ${level}`);
+      for (const [i, hc] of hcArr.entries()) {
+        const indexInBatch = i % 4;
+        switch (indexInBatch) {
+          case 0:
+            src0_0 = hc.src0;
+            src1_0 = hc.src1;
+            dest_0 = hc.dest;
+            break;
+          case 1:
+            src0_1 = hc.src0;
+            src1_1 = hc.src1;
+            dest_1 = hc.dest;
+            break;
+          case 2:
+            src0_2 = hc.src0;
+            src1_2 = hc.src1;
+            dest_2 = hc.dest;
+            break;
+          case 3:
+            src0_3 = hc.src0;
+            src1_3 = hc.src1;
+            dest_3 = hc.dest;
+            break;
+          default:
+            throw Error(`Unexpected indexInBatch ${indexInBatch}`);
+        }
+
+        if (
+          indexInBatch === 3 &&
+          src0_0 !== null &&
+          src1_0 !== null &&
+          dest_0 !== null &&
+          src0_1 !== null &&
+          src1_1 !== null &&
+          dest_1 !== null &&
+          src0_2 !== null &&
+          src1_2 !== null &&
+          dest_2 !== null &&
+          src0_3 !== null &&
+          src1_3 !== null &&
+          dest_3 !== null
+        ) {
+          const [o0, o1, o2, o3] = batchHash4HashObjectInputs([
+            src0_0,
+            src1_0,
+            src0_1,
+            src1_1,
+            src0_2,
+            src1_2,
+            src0_3,
+            src1_3,
+          ]);
+          if (o0 == null || o1 == null || o2 == null || o3 == null) {
+            throw Error(`batchHash4HashObjectInputs return null at batch ${i} level ${level}`);
+          }
+          dest_0.applyHash(o0);
+          dest_1.applyHash(o1);
+          dest_2.applyHash(o2);
+          dest_3.applyHash(o3);
+
+          src0_0 = null;
+          src1_0 = null;
+          dest_0 = null;
+          src0_1 = null;
+          src1_1 = null;
+          dest_1 = null;
+          src0_2 = null;
+          src1_2 = null;
+          dest_2 = null;
+          src0_3 = null;
+          src1_3 = null;
+          dest_3 = null;
         }
-        dest_0.applyHash(o0);
-        dest_1.applyHash(o1);
-        dest_2.applyHash(o2);
-        dest_3.applyHash(o3);
       }
 
       // remaining
-      for (let i = batch * 4; i < hcArr.length; i++) {
-        const {src0, src1, dest} = hcArr[i];
-        const output = digest64HashObjects(src0, src1);
-        dest.applyHash(output);
+      if (src0_0 !== null && src1_0 !== null && dest_0 !== null) {
+        dest_0.applyHash(digest64HashObjects(src0_0, src1_0));
+      }
+      if (src0_1 !== null && src1_1 !== null && dest_1 !== null) {
+        dest_1.applyHash(digest64HashObjects(src0_1, src1_1));
+      }
+      if (src0_2 !== null && src1_2 !== null && dest_2 !== null) {
+        dest_2.applyHash(digest64HashObjects(src0_2, src1_2));
+      }
+      if (src0_3 !== null && src1_3 !== null && dest_3 !== null) {
+        dest_3.applyHash(digest64HashObjects(src0_3, src1_3));
       }
     }
   },
diff --git a/packages/persistent-merkle-tree/src/node.ts b/packages/persistent-merkle-tree/src/node.ts
index 8b1e5695..e2d69c1b 100644
--- a/packages/persistent-merkle-tree/src/node.ts
+++ b/packages/persistent-merkle-tree/src/node.ts
@@ -86,7 +86,7 @@ export class BranchNode extends Node {
   batchHash(): Uint8Array {
     const hashComputations: HashComputation[][] = [];
     getHashComputations(this, 0, hashComputations);
-    hasher.executeHashComputations(hashComputations);
+    executeHashComputations(hashComputations);
 
     if (this.h0 === null) {
       throw Error("Root is not computed by batch");