From bbd7c5f5de569c489f996929f5590d1227f27e1e Mon Sep 17 00:00:00 2001
From: Tuyen Nguyen <vutuyen2636@gmail.com>
Date: Mon, 25 Mar 2024 13:14:15 +0700
Subject: [PATCH] fix: as-sha256 improve hashBlocks()

---
 packages/as-sha256/assembly/index.ts | 33 +++++++++++-----------------
 packages/as-sha256/src/index.ts      | 16 ++++----------
 2 files changed, 17 insertions(+), 32 deletions(-)

diff --git a/packages/as-sha256/assembly/index.ts b/packages/as-sha256/assembly/index.ts
index bb9142e7..c46874e0 100644
--- a/packages/as-sha256/assembly/index.ts
+++ b/packages/as-sha256/assembly/index.ts
@@ -43,11 +43,11 @@ const W64: u32[] = [
 const w64Ptr = W64.dataStart;
 
 // intermediate hash values stored in H0-H7
-var H0: u32, H1: u32, H2: u32, H3: u32, H4: u32, H5: u32, H6: u32, H7: u32;
+let H0: u32, H1: u32, H2: u32, H3: u32, H4: u32, H5: u32, H6: u32, H7: u32;
 
 // hash registers
 
-var a: u32, b: u32, c: u32, d: u32, e: u32, f: u32, g: u32, h: u32, i: u32, t1: u32, t2: u32;
+let a: u32, b: u32, c: u32, d: u32, e: u32, f: u32, g: u32, h: u32, i: u32, t1: u32, t2: u32;
 
 // 16 32bit message blocks
 const M = new ArrayBuffer(64);
@@ -66,10 +66,10 @@ export const output = new ArrayBuffer(DIGEST_LENGTH);
 const outputPtr = changetype<usize>(output);
 
 // number of bytes in M buffer
-var mLength = 0;
+let mLength = 0;
 
 // number of total bytes hashed
-var bytesHashed = 0;
+let bytesHashed = 0;
 
 @inline
 function load32(ptr: usize, offset: usize): u32 {
@@ -141,6 +141,7 @@ function SIG1(x: u32): u32 {
   return rotr(x, 17) ^ rotr(x, 19) ^ (x >>> 10);
 }
 
+let tmpW: u32;
 /**
  * Expand message blocks (16 32bit blocks), into extended message blocks (64 32bit blocks),
  * Apply SHA256 compression function on extended message blocks
@@ -158,25 +159,17 @@ function hashBlocks(wPtr: usize, mPtr: usize): void {
   g = H6;
   h = H7;
 
-  // Load message blocks into first 16 expanded message blocks
-  for (i = 0; i < 16; i++) {
-    store32(wPtr, i,
-      load32be(mPtr, i)
-    );
-  }
-  // Expand message blocks 17-64
-  for (i = 16; i < 64; i++) {
-    store32(wPtr, i,
-      SIG1(load32(wPtr, i - 2)) +
-      load32(wPtr, i - 7) +
-      SIG0(load32(wPtr, i - 15)) +
-      load32(wPtr, i - 16)
-    );
-  }
+  // 16 first u32 of expanded message block are same as message block
+  // rest of the 48 u32 are computed from the first 16 u32
+  // instead of computing expanded messsage blocks in a separate for loop,
+  // chain it to the below for loop to improve performance
 
   // Apply SHA256 compression function on expanded message blocks
   for (i = 0; i < 64; i++) {
-    t1 = h + EP1(e) + CH(e, f, g) + load32(kPtr, i) + load32(wPtr, i);
+    tmpW = i < 16 ? load32be(mPtr, i) : SIG1(load32(wPtr, i - 2)) + load32(wPtr, i - 7) +
+      SIG0(load32(wPtr, i - 15)) + load32(wPtr, i - 16);
+    store32(wPtr, i, tmpW);
+    t1 = h + EP1(e) + CH(e, f, g) + load32(kPtr, i) + tmpW;
     t2 = EP0(a) + MAJ(a, b, c);
     h = g;
     g = f;
diff --git a/packages/as-sha256/src/index.ts b/packages/as-sha256/src/index.ts
index 6fa451bd..fc6014d4 100644
--- a/packages/as-sha256/src/index.ts
+++ b/packages/as-sha256/src/index.ts
@@ -18,9 +18,7 @@ export function digest(data: Uint8Array): Uint8Array {
   if (data.length <= ctx.INPUT_LENGTH) {
     inputUint8Array.set(data);
     ctx.digest(data.length);
-    const output = new Uint8Array(32);
-    output.set(outputUint8Array);
-    return output;
+    return outputUint8Array.slice(0, 32);
   }
 
   ctx.init();
@@ -32,9 +30,7 @@ export function digest64(data: Uint8Array): Uint8Array {
   if (data.length === 64) {
     inputUint8Array.set(data);
     ctx.digest64(wasmInputValue, wasmOutputValue);
-    const output = new Uint8Array(32);
-    output.set(outputUint8Array);
-    return output;
+    return outputUint8Array.slice(0, 32);
   }
   throw new Error("InvalidLengthForDigest64");
 }
@@ -44,9 +40,7 @@ export function digest2Bytes32(bytes1: Uint8Array, bytes2: Uint8Array): Uint8Arr
     inputUint8Array.set(bytes1);
     inputUint8Array.set(bytes2, 32);
     ctx.digest64(wasmInputValue, wasmOutputValue);
-    const output = new Uint8Array(32);
-    output.set(outputUint8Array);
-    return output;
+    return outputUint8Array.slice(0, 32);
   }
   throw new Error("InvalidLengthForDigest64");
 }
@@ -98,7 +92,5 @@ function update(data: Uint8Array): void {
 
 function final(): Uint8Array {
   ctx.final(wasmOutputValue);
-  const output = new Uint8Array(32);
-  output.set(outputUint8Array);
-  return output;
+  return outputUint8Array.slice(0, 32);
 }