feat: improve type.hashTreeRoot() using batch

packages/as-sha256/src/index.ts

@@ -3,6 +3,7 @@ import {newInstance} from "./wasm";
 import {HashObject, byteArrayIntoHashObject, byteArrayToHashObject, hashObjectToByteArray} from "./hashObject";
 import SHA256 from "./sha256";
 export {HashObject, byteArrayToHashObject, hashObjectToByteArray, byteArrayIntoHashObject, SHA256};
+export {allocUnsafe};
 
 const ctx = newInstance();
 const wasmInputValue = ctx.input.value;

packages/ssz/src/type/abstract.ts

@@ -145,6 +145,11 @@ export abstract class Type<V> {
    */
   abstract hashTreeRoot(value: V): Uint8Array;
 
+  /**
+   * Same to hashTreeRoot() but here we write result to output.
+   */
+  abstract hashTreeRootInto(value: V, output: Uint8Array, offset: number): void;
+
   // JSON support
 
   /** Parse JSON representation of a type to value */

packages/ssz/src/type/arrayComposite.ts

@@ -211,21 +211,29 @@ export function tree_deserializeFromBytesArrayComposite<ElementType extends Comp
   }
 }
 
-/**
- * @param length In List length = value.length, Vector length = fixed value
- */
-export function value_getRootsArrayComposite<ElementType extends CompositeType<unknown, unknown, unknown>>(
+export function value_getBlocksBytesArrayComposite<ElementType extends CompositeType<unknown, unknown, unknown>>(
   elementType: ElementType,
   length: number,
-  value: ValueOf<ElementType>[]
-): Uint8Array[] {
-  const roots = new Array<Uint8Array>(length);
+  value: ValueOf<ElementType>[],
+  blocksBuffer: Uint8Array
+): Uint8Array {
+  const blockBytesLen = Math.ceil(length / 2) * 64;
+  if (blockBytesLen > blocksBuffer.length) {
+    throw new Error(`blocksBuffer is too small: ${blocksBuffer.length} < ${blockBytesLen}`);
+  }
+  const blocksBytes = blocksBuffer.subarray(0, blockBytesLen);
 
   for (let i = 0; i < length; i++) {
-    roots[i] = elementType.hashTreeRoot(value[i]);
+    elementType.hashTreeRootInto(value[i], blocksBytes, i * 32);
+  }
+
+  const isOddChunk = length % 2 === 1;
+  if (isOddChunk) {
+    // similar to append zeroHash(0)
+    blocksBytes.subarray(length * 32, blockBytesLen).fill(0);
   }
 
-  return roots;
+  return blocksBytes;
 }
 
 function readOffsetsArrayComposite(

packages/ssz/src/type/basic.ts

@@ -30,11 +30,18 @@ export abstract class BasicType<V> extends Type<V> {
   }
 
   hashTreeRoot(value: V): Uint8Array {
-    // TODO: Optimize
-    const uint8Array = new Uint8Array(32);
+    // cannot use allocUnsafe() here because hashTreeRootInto() may not fill the whole 32 bytes
+    const root = new Uint8Array(32);
+    this.hashTreeRootInto(value, root, 0);
+    return root;
+  }
+
+  hashTreeRootInto(value: V, output: Uint8Array, offset: number): void {
+    const uint8Array = output.subarray(offset, offset + 32);
+    // output could have preallocated data, some types may not fill the whole 32 bytes
+    uint8Array.fill(0);
     const dataView = new DataView(uint8Array.buffer, uint8Array.byteOffset, uint8Array.byteLength);
     this.value_serializeToBytes({uint8Array, dataView}, 0, value);
-    return uint8Array;
   }
 
   clone(value: V): V {

packages/ssz/src/type/bitArray.ts

@@ -1,10 +1,10 @@
 import {concatGindices, Gindex, Node, toGindex, Tree, HashComputationLevel} from "@chainsafe/persistent-merkle-tree";
 import {fromHexString, toHexString, byteArrayEquals} from "../util/byteArray";
-import {splitIntoRootChunks} from "../util/merkleize";
 import {CompositeType, LENGTH_GINDEX} from "./composite";
 import {BitArray} from "../value/bitArray";
 import {BitArrayTreeView} from "../view/bitArray";
 import {BitArrayTreeViewDU} from "../viewDU/bitArray";
+import {getBlocksBytes} from "./byteArray";
 
 /* eslint-disable @typescript-eslint/member-ordering */
 
@@ -40,8 +40,13 @@ export abstract class BitArrayType extends CompositeType<BitArray, BitArrayTreeV
 
   // Merkleization
 
-  protected getRoots(value: BitArray): Uint8Array[] {
-    return splitIntoRootChunks(value.uint8Array);
+  protected getBlocksBytes(value: BitArray): Uint8Array {
+    // reallocate this.blocksBuffer if needed
+    if (value.uint8Array.length > this.blocksBuffer.length) {
+      const chunkCount = Math.ceil(value.bitLen / 8 / 32);
+      this.blocksBuffer = new Uint8Array(Math.ceil(chunkCount / 2) * 64);
+    }
+    return getBlocksBytes(value.uint8Array, this.blocksBuffer);
   }
 
   // Proofs

packages/ssz/src/type/bitList.ts

@@ -1,5 +1,12 @@
-import {getNodesAtDepth, Node, packedNodeRootsToBytes, packedRootsBytesToNode} from "@chainsafe/persistent-merkle-tree";
-import {mixInLength, maxChunksToDepth} from "../util/merkleize";
+import {allocUnsafe} from "@chainsafe/as-sha256";
+import {
+  getNodesAtDepth,
+  merkleizeBlocksBytes,
+  Node,
+  packedNodeRootsToBytes,
+  packedRootsBytesToNode,
+} from "@chainsafe/persistent-merkle-tree";
+import {maxChunksToDepth} from "../util/merkleize";
 import {Require} from "../util/types";
 import {namedClass} from "../util/named";
 import {ByteViews} from "./composite";
@@ -29,6 +36,12 @@ export class BitListType extends BitArrayType {
   readonly maxSize: number;
   readonly maxChunkCount: number;
   readonly isList = true;
+  readonly mixInLengthBlockBytes = new Uint8Array(64);
+  readonly mixInLengthBuffer = Buffer.from(
+    this.mixInLengthBlockBytes.buffer,
+    this.mixInLengthBlockBytes.byteOffset,
+    this.mixInLengthBlockBytes.byteLength
+  );
 
   constructor(readonly limitBits: number, opts?: BitListOptions) {
     super();
@@ -101,7 +114,18 @@ export class BitListType extends BitArrayType {
   // Merkleization: inherited from BitArrayType
 
   hashTreeRoot(value: BitArray): Uint8Array {
-    return mixInLength(super.hashTreeRoot(value), value.bitLen);
+    const root = allocUnsafe(32);
+    this.hashTreeRootInto(value, root, 0);
+    return root;
+  }
+
+  hashTreeRootInto(value: BitArray, output: Uint8Array, offset: number): void {
+    super.hashTreeRootInto(value, this.mixInLengthBlockBytes, 0);
+    // mixInLength
+    this.mixInLengthBuffer.writeUIntLE(value.bitLen, 32, 6);
+    // one for hashTreeRoot(value), one for length
+    const chunkCount = 2;
+    merkleizeBlocksBytes(this.mixInLengthBlockBytes, chunkCount, output, offset);
   }
 
   // Proofs: inherited from BitArrayType

packages/ssz/src/type/byteArray.ts

@@ -8,7 +8,6 @@ import {
   getHashComputations,
 } from "@chainsafe/persistent-merkle-tree";
 import {fromHexString, toHexString, byteArrayEquals} from "../util/byteArray";
-import {splitIntoRootChunks} from "../util/merkleize";
 import {ByteViews} from "./abstract";
 import {CompositeType, LENGTH_GINDEX} from "./composite";
 
@@ -82,10 +81,21 @@ export abstract class ByteArrayType extends CompositeType<ByteArray, ByteArray,
     return Uint8Array.prototype.slice.call(data.uint8Array, start, end);
   }
 
+  value_toTree(value: ByteArray): Node {
+    // this saves 1 allocation of Uint8Array
+    const dataView = new DataView(value.buffer, value.byteOffset, value.byteLength);
+    return this.tree_deserializeFromBytes({uint8Array: value, dataView}, 0, value.length);
+  }
+
   // Merkleization
 
-  protected getRoots(value: ByteArray): Uint8Array[] {
-    return splitIntoRootChunks(value);
+  protected getBlocksBytes(value: ByteArray): Uint8Array {
+    // reallocate this.blocksBuffer if needed
+    if (value.length > this.blocksBuffer.length) {
+      const chunkCount = Math.ceil(value.length / 32);
+      this.blocksBuffer = new Uint8Array(Math.ceil(chunkCount / 2) * 64);
+    }
+    return getBlocksBytes(value, this.blocksBuffer);
   }
 
   // Proofs
@@ -149,3 +159,16 @@ export abstract class ByteArrayType extends CompositeType<ByteArray, ByteArray,
 
   protected abstract assertValidSize(size: number): void;
 }
+
+export function getBlocksBytes(value: Uint8Array, blocksBuffer: Uint8Array): Uint8Array {
+  if (value.length > blocksBuffer.length) {
+    throw new Error(`data length ${value.length} exceeds blocksBuffer length ${blocksBuffer.length}`);
+  }
+
+  blocksBuffer.set(value);
+  const valueLen = value.length;
+  const blockByteLen = Math.ceil(valueLen / 64) * 64;
+  // all padding bytes must be zero, this is similar to set zeroHash(0)
+  blocksBuffer.subarray(valueLen, blockByteLen).fill(0);
+  return blocksBuffer.subarray(0, blockByteLen);
+}

packages/ssz/src/type/byteList.ts

@@ -1,11 +1,18 @@
-import {getNodesAtDepth, Node, packedNodeRootsToBytes, packedRootsBytesToNode} from "@chainsafe/persistent-merkle-tree";
-import {mixInLength, maxChunksToDepth} from "../util/merkleize";
+import {allocUnsafe} from "@chainsafe/as-sha256";
+import {
+  getNodesAtDepth,
+  Node,
+  packedNodeRootsToBytes,
+  packedRootsBytesToNode,
+  merkleizeBlocksBytes,
+  merkleizeBlockArray,
+} from "@chainsafe/persistent-merkle-tree";
+import {maxChunksToDepth} from "../util/merkleize";
 import {Require} from "../util/types";
 import {namedClass} from "../util/named";
 import {addLengthNode, getChunksNodeFromRootNode, getLengthFromRootNode} from "./arrayBasic";
 import {ByteViews} from "./composite";
 import {ByteArrayType, ByteArray} from "./byteArray";
-
 /* eslint-disable @typescript-eslint/member-ordering */
 
 export interface ByteListOptions {
@@ -34,6 +41,14 @@ export class ByteListType extends ByteArrayType {
   readonly maxSize: number;
   readonly maxChunkCount: number;
   readonly isList = true;
+  readonly blockArray: Uint8Array[] = [];
+  private blockBytesLen = 0;
+  readonly mixInLengthBlockBytes = new Uint8Array(64);
+  readonly mixInLengthBuffer = Buffer.from(
+    this.mixInLengthBlockBytes.buffer,
+    this.mixInLengthBlockBytes.byteOffset,
+    this.mixInLengthBlockBytes.byteLength
+  );
 
   constructor(readonly limitBytes: number, opts?: ByteListOptions) {
     super();
@@ -89,7 +104,49 @@ export class ByteListType extends ByteArrayType {
   // Merkleization: inherited from ByteArrayType
 
   hashTreeRoot(value: ByteArray): Uint8Array {
-    return mixInLength(super.hashTreeRoot(value), value.length);
+    const root = allocUnsafe(32);
+    this.hashTreeRootInto(value, root, 0);
+    return root;
+  }
+
+  /**
+   * Use  merkleizeBlockArray() instead of merkleizeBlocksBytes() to avoid big memory allocation
+   */
+  hashTreeRootInto(value: Uint8Array, output: Uint8Array, offset: number): void {
+    // should not call super.hashTreeRoot() here
+    // use  merkleizeBlockArray() instead of merkleizeBlocksBytes() to avoid big memory allocation
+    // reallocate this.blockArray if needed
+    if (value.length > this.blockBytesLen) {
+      const newBlockCount = Math.ceil(value.length / 64);
+      // this.blockBytesLen should be a multiple of 64
+      const oldBlockCount = Math.ceil(this.blockBytesLen / 64);
+      const blockDiff = newBlockCount - oldBlockCount;
+      const newBlocksBytes = new Uint8Array(blockDiff * 64);
+      for (let i = 0; i < blockDiff; i++) {
+        this.blockArray.push(newBlocksBytes.subarray(i * 64, (i + 1) * 64));
+        this.blockBytesLen += 64;
+      }
+    }
+
+    // populate this.blockArray
+    for (let i = 0; i < value.length; i += 64) {
+      const block = this.blockArray[i / 64];
+      // zero out the last block if it's over value.length
+      if (i + 64 > value.length) {
+        block.fill(0);
+      }
+      block.set(value.subarray(i, Math.min(i + 64, value.length)));
+    }
+
+    // compute hashTreeRoot
+    const blockLimit = Math.ceil(value.length / 64);
+    merkleizeBlockArray(this.blockArray, blockLimit, this.maxChunkCount, this.mixInLengthBlockBytes, 0);
+
+    // mixInLength
+    this.mixInLengthBuffer.writeUIntLE(value.length, 32, 6);
+    // one for hashTreeRoot(value), one for length
+    const chunkCount = 2;
+    merkleizeBlocksBytes(this.mixInLengthBlockBytes, chunkCount, output, offset);
   }
 
   // Proofs: inherited from BitArrayType

packages/ssz/src/type/composite.ts

@@ -1,3 +1,4 @@
+import {allocUnsafe} from "@chainsafe/as-sha256";
 import {
   concatGindices,
   createProof,
@@ -7,10 +8,11 @@ import {
   Proof,
   ProofType,
   Tree,
+  merkleizeBlocksBytes,
   HashComputationLevel,
 } from "@chainsafe/persistent-merkle-tree";
 import {byteArrayEquals} from "../util/byteArray";
-import {merkleize, symbolCachedPermanentRoot, ValueWithCachedPermanentRoot} from "../util/merkleize";
+import {cacheRoot, symbolCachedPermanentRoot, ValueWithCachedPermanentRoot} from "../util/merkleize";
 import {treePostProcessFromProofNode} from "../util/proof/treePostProcessFromProofNode";
 import {Type, ByteViews, JsonPath, JsonPathProp} from "./abstract";
 export {ByteViews};
@@ -59,6 +61,7 @@ export abstract class CompositeType<V, TV, TVDU> extends Type<V> {
    * Required for ContainerNodeStruct to ensure no dangerous types are constructed.
    */
   abstract readonly isViewMutable: boolean;
+  protected blocksBuffer = new Uint8Array(0);
 
   constructor(
     /**
@@ -216,13 +219,30 @@ export abstract class CompositeType<V, TV, TVDU> extends Type<V> {
       }
     }
 
-    const root = merkleize(this.getRoots(value), this.maxChunkCount);
+    const root = allocUnsafe(32);
+    const safeCache = true;
+    this.hashTreeRootInto(value, root, 0, safeCache);
 
+    // hashTreeRootInto will cache the root if cachePermanentRootStruct is true
+
+    return root;
+  }
+
+  hashTreeRootInto(value: V, output: Uint8Array, offset: number, safeCache = false): void {
+    // Return cached mutable root if any
     if (this.cachePermanentRootStruct) {
-      (value as ValueWithCachedPermanentRoot)[symbolCachedPermanentRoot] = root;
+      const cachedRoot = (value as ValueWithCachedPermanentRoot)[symbolCachedPermanentRoot];
+      if (cachedRoot) {
+        output.set(cachedRoot, offset);
+        return;
+      }
     }
 
-    return root;
+    const blocksBuffer = this.getBlocksBytes(value);
+    merkleizeBlocksBytes(blocksBuffer, this.maxChunkCount, output, offset);
+    if (this.cachePermanentRootStruct) {
+      cacheRoot(value as ValueWithCachedPermanentRoot, output, offset, safeCache);
+    }
   }
 
   // For debugging and testing this feature
@@ -236,7 +256,12 @@ export abstract class CompositeType<V, TV, TVDU> extends Type<V> {
   //   and feed those numbers directly to the hasher input with a DataView
   // - The return of the hasher should be customizable too, to reduce conversions from Uint8Array
   //   to hashObject and back.
-  protected abstract getRoots(value: V): Uint8Array[];
+
+  /**
+   * Get multiple SHA256 blocks, each is 64 bytes long.
+   * If chunk count is not even, need to append zeroHash(0)
+   */
+  protected abstract getBlocksBytes(value: V): Uint8Array;
 
   // Proofs API
 

packages/ssz/src/type/container.ts

@@ -130,6 +130,8 @@ export class ContainerType<Fields extends Record<string, Type<unknown>>> extends
     // Refactor this constructor to allow customization without pollutin the options
     this.TreeView = opts?.getContainerTreeViewClass?.(this) ?? getContainerTreeViewClass(this);
     this.TreeViewDU = opts?.getContainerTreeViewDUClass?.(this) ?? getContainerTreeViewDUClass(this);
+    const fieldBytes = this.fieldsEntries.length * 32;
+    this.blocksBuffer = new Uint8Array(Math.ceil(fieldBytes / 64) * 64);
   }
 
   static named<Fields extends Record<string, Type<unknown>>>(
@@ -272,15 +274,13 @@ export class ContainerType<Fields extends Record<string, Type<unknown>>> extends
 
   // Merkleization
 
-  protected getRoots(struct: ValueOfFields<Fields>): Uint8Array[] {
-    const roots = new Array<Uint8Array>(this.fieldsEntries.length);
-
+  protected getBlocksBytes(struct: ValueOfFields<Fields>): Uint8Array {
     for (let i = 0; i < this.fieldsEntries.length; i++) {
       const {fieldName, fieldType} = this.fieldsEntries[i];
-      roots[i] = fieldType.hashTreeRoot(struct[fieldName]);
+      fieldType.hashTreeRootInto(struct[fieldName], this.blocksBuffer, i * 32);
     }
-
-    return roots;
+    // remaining bytes are zeroed as we never write them
+    return this.blocksBuffer;
   }
 
   // Proofs