Skip to content

Commit 66742f0

Browse files
twoethsnazarhussain
twoeths
and
nazarhussain
authored
feat: improve type.hashTreeRoot() using batch (#409)
* feat: improve type.hashTreeRoot() using batch * feat: consume merkleizeBlockArray * fix: lint in ssz package * Fix old perf test usage * Fiz types * Fix lint errors * Fix lint errors * Rename 'getBlocksBytes' to 'getPaddedBytes64' * Revert "Rename 'getBlocksBytes' to 'getPaddedBytes64'" This reverts commit 15cf649. * chore: address PR's comment --------- Co-authored-by: Nazar Hussain <[email protected]> Co-authored-by: Tuyen Nguyen <[email protected]>
1 parent 089daed commit 66742f0

28 files changed

+684
-200
lines changed

‎packages/as-sha256/src/index.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ import type {HashObject} from "./hashObject.js";
44
import {byteArrayIntoHashObject, byteArrayToHashObject, hashObjectToByteArray} from "./hashObject.js";
55
import SHA256 from "./sha256.js";
66
export {HashObject, byteArrayToHashObject, hashObjectToByteArray, byteArrayIntoHashObject, SHA256};
7+
export {allocUnsafe};
78

89
let ctx: WasmContext;
910
export let simdEnabled: boolean;

‎packages/ssz/src/type/abstract.ts

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -145,6 +145,11 @@ export abstract class Type<V> {
145145
*/
146146
abstract hashTreeRoot(value: V): Uint8Array;
147147

148+
/**
149+
* Same to hashTreeRoot() but here we write result to output.
150+
*/
151+
abstract hashTreeRootInto(value: V, output: Uint8Array, offset: number): void;
152+
148153
// JSON support
149154

150155
/** Parse JSON representation of a type to value */

‎packages/ssz/src/type/arrayComposite.ts

Lines changed: 17 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -211,21 +211,29 @@ export function tree_deserializeFromBytesArrayComposite<ElementType extends Comp
211211
}
212212
}
213213

214-
/**
215-
* @param length In List length = value.length, Vector length = fixed value
216-
*/
217-
export function value_getRootsArrayComposite<ElementType extends CompositeType<unknown, unknown, unknown>>(
214+
export function value_getBlocksBytesArrayComposite<ElementType extends CompositeType<unknown, unknown, unknown>>(
218215
elementType: ElementType,
219216
length: number,
220-
value: ValueOf<ElementType>[]
221-
): Uint8Array[] {
222-
const roots = new Array<Uint8Array>(length);
217+
value: ValueOf<ElementType>[],
218+
blocksBuffer: Uint8Array
219+
): Uint8Array {
220+
const blockBytesLen = Math.ceil(length / 2) * 64;
221+
if (blockBytesLen > blocksBuffer.length) {
222+
throw new Error(`blocksBuffer is too small: ${blocksBuffer.length} < ${blockBytesLen}`);
223+
}
224+
const blocksBytes = blocksBuffer.subarray(0, blockBytesLen);
223225

224226
for (let i = 0; i < length; i++) {
225-
roots[i] = elementType.hashTreeRoot(value[i]);
227+
elementType.hashTreeRootInto(value[i], blocksBytes, i * 32);
228+
}
229+
230+
const isOddChunk = length % 2 === 1;
231+
if (isOddChunk) {
232+
// similar to append zeroHash(0)
233+
blocksBytes.subarray(length * 32, blockBytesLen).fill(0);
226234
}
227235

228-
return roots;
236+
return blocksBytes;
229237
}
230238

231239
function readOffsetsArrayComposite(

‎packages/ssz/src/type/basic.ts

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -30,11 +30,17 @@ export abstract class BasicType<V> extends Type<V> {
3030
}
3131

3232
hashTreeRoot(value: V): Uint8Array {
33-
// TODO: Optimize
34-
const uint8Array = new Uint8Array(32);
33+
const root = new Uint8Array(32);
34+
this.hashTreeRootInto(value, root, 0);
35+
return root;
36+
}
37+
38+
hashTreeRootInto(value: V, output: Uint8Array, offset: number): void {
39+
const uint8Array = output.subarray(offset, offset + 32);
40+
// output could have preallocated data, some types may not fill the whole 32 bytes
41+
uint8Array.fill(0);
3542
const dataView = new DataView(uint8Array.buffer, uint8Array.byteOffset, uint8Array.byteLength);
3643
this.value_serializeToBytes({uint8Array, dataView}, 0, value);
37-
return uint8Array;
3844
}
3945

4046
clone(value: V): V {

‎packages/ssz/src/type/bitArray.ts

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,10 @@
11
import {concatGindices, Gindex, Node, toGindex, Tree, HashComputationLevel} from "@chainsafe/persistent-merkle-tree";
22
import {fromHexString, toHexString, byteArrayEquals} from "../util/byteArray.js";
3-
import {splitIntoRootChunks} from "../util/merkleize.js";
43
import {CompositeType, LENGTH_GINDEX} from "./composite.js";
54
import {BitArray} from "../value/bitArray.js";
65
import {BitArrayTreeView} from "../view/bitArray.js";
76
import {BitArrayTreeViewDU} from "../viewDU/bitArray.js";
7+
import {getBlocksBytes} from "./byteArray.js";
88

99
/* eslint-disable @typescript-eslint/member-ordering */
1010

@@ -40,8 +40,13 @@ export abstract class BitArrayType extends CompositeType<BitArray, BitArrayTreeV
4040

4141
// Merkleization
4242

43-
protected getRoots(value: BitArray): Uint8Array[] {
44-
return splitIntoRootChunks(value.uint8Array);
43+
protected getBlocksBytes(value: BitArray): Uint8Array {
44+
// reallocate this.blocksBuffer if needed
45+
if (value.uint8Array.length > this.blocksBuffer.length) {
46+
const chunkCount = Math.ceil(value.bitLen / 8 / 32);
47+
this.blocksBuffer = new Uint8Array(Math.ceil(chunkCount / 2) * 64);
48+
}
49+
return getBlocksBytes(value.uint8Array, this.blocksBuffer);
4550
}
4651

4752
// Proofs

‎packages/ssz/src/type/bitList.ts

Lines changed: 27 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,12 @@
1-
import {getNodesAtDepth, Node, packedNodeRootsToBytes, packedRootsBytesToNode} from "@chainsafe/persistent-merkle-tree";
2-
import {mixInLength, maxChunksToDepth} from "../util/merkleize.js";
1+
import {
2+
getNodesAtDepth,
3+
Node,
4+
packedNodeRootsToBytes,
5+
packedRootsBytesToNode,
6+
merkleizeBlocksBytes,
7+
} from "@chainsafe/persistent-merkle-tree";
8+
import {allocUnsafe} from "@chainsafe/as-sha256";
9+
import {maxChunksToDepth} from "../util/merkleize.js";
310
import {Require} from "../util/types.js";
411
import {namedClass} from "../util/named.js";
512
import {ByteViews} from "./composite.js";
@@ -29,6 +36,12 @@ export class BitListType extends BitArrayType {
2936
readonly maxSize: number;
3037
readonly maxChunkCount: number;
3138
readonly isList = true;
39+
readonly mixInLengthBlockBytes = new Uint8Array(64);
40+
readonly mixInLengthBuffer = Buffer.from(
41+
this.mixInLengthBlockBytes.buffer,
42+
this.mixInLengthBlockBytes.byteOffset,
43+
this.mixInLengthBlockBytes.byteLength
44+
);
3245

3346
constructor(readonly limitBits: number, opts?: BitListOptions) {
3447
super();
@@ -101,7 +114,18 @@ export class BitListType extends BitArrayType {
101114
// Merkleization: inherited from BitArrayType
102115

103116
hashTreeRoot(value: BitArray): Uint8Array {
104-
return mixInLength(super.hashTreeRoot(value), value.bitLen);
117+
const root = allocUnsafe(32);
118+
this.hashTreeRootInto(value, root, 0);
119+
return root;
120+
}
121+
122+
hashTreeRootInto(value: BitArray, output: Uint8Array, offset: number): void {
123+
super.hashTreeRootInto(value, this.mixInLengthBlockBytes, 0);
124+
// mixInLength
125+
this.mixInLengthBuffer.writeUIntLE(value.bitLen, 32, 6);
126+
// one for hashTreeRoot(value), one for length
127+
const chunkCount = 2;
128+
merkleizeBlocksBytes(this.mixInLengthBlockBytes, chunkCount, output, offset);
105129
}
106130

107131
// Proofs: inherited from BitArrayType

‎packages/ssz/src/type/byteArray.ts

Lines changed: 26 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,6 @@ import {
88
getHashComputations,
99
} from "@chainsafe/persistent-merkle-tree";
1010
import {fromHexString, toHexString, byteArrayEquals} from "../util/byteArray.js";
11-
import {splitIntoRootChunks} from "../util/merkleize.js";
1211
import {ByteViews} from "./abstract.js";
1312
import {CompositeType, LENGTH_GINDEX} from "./composite.js";
1413

@@ -82,10 +81,21 @@ export abstract class ByteArrayType extends CompositeType<ByteArray, ByteArray,
8281
return Uint8Array.prototype.slice.call(data.uint8Array, start, end);
8382
}
8483

84+
value_toTree(value: ByteArray): Node {
85+
// this saves 1 allocation of Uint8Array
86+
const dataView = new DataView(value.buffer, value.byteOffset, value.byteLength);
87+
return this.tree_deserializeFromBytes({uint8Array: value, dataView}, 0, value.length);
88+
}
89+
8590
// Merkleization
8691

87-
protected getRoots(value: ByteArray): Uint8Array[] {
88-
return splitIntoRootChunks(value);
92+
protected getBlocksBytes(value: ByteArray): Uint8Array {
93+
// reallocate this.blocksBuffer if needed
94+
if (value.length > this.blocksBuffer.length) {
95+
const chunkCount = Math.ceil(value.length / 32);
96+
this.blocksBuffer = new Uint8Array(Math.ceil(chunkCount / 2) * 64);
97+
}
98+
return getBlocksBytes(value, this.blocksBuffer);
8999
}
90100

91101
// Proofs
@@ -149,3 +159,16 @@ export abstract class ByteArrayType extends CompositeType<ByteArray, ByteArray,
149159

150160
protected abstract assertValidSize(size: number): void;
151161
}
162+
163+
export function getBlocksBytes(value: Uint8Array, blocksBuffer: Uint8Array): Uint8Array {
164+
if (value.length > blocksBuffer.length) {
165+
throw new Error(`data length ${value.length} exceeds blocksBuffer length ${blocksBuffer.length}`);
166+
}
167+
168+
blocksBuffer.set(value);
169+
const valueLen = value.length;
170+
const blockByteLen = Math.ceil(valueLen / 64) * 64;
171+
// all padding bytes must be zero, this is similar to set zeroHash(0)
172+
blocksBuffer.subarray(valueLen, blockByteLen).fill(0);
173+
return blocksBuffer.subarray(0, blockByteLen);
174+
}

‎packages/ssz/src/type/byteList.ts

Lines changed: 61 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,18 @@
1-
import {getNodesAtDepth, Node, packedNodeRootsToBytes, packedRootsBytesToNode} from "@chainsafe/persistent-merkle-tree";
2-
import {mixInLength, maxChunksToDepth} from "../util/merkleize.js";
1+
import {
2+
getNodesAtDepth,
3+
Node,
4+
packedNodeRootsToBytes,
5+
packedRootsBytesToNode,
6+
merkleizeBlocksBytes,
7+
merkleizeBlockArray,
8+
} from "@chainsafe/persistent-merkle-tree";
9+
import {allocUnsafe} from "@chainsafe/as-sha256";
10+
import {maxChunksToDepth} from "../util/merkleize.js";
311
import {Require} from "../util/types.js";
412
import {namedClass} from "../util/named.js";
513
import {addLengthNode, getChunksNodeFromRootNode, getLengthFromRootNode} from "./arrayBasic.js";
614
import {ByteViews} from "./composite.js";
715
import {ByteArrayType, ByteArray} from "./byteArray.js";
8-
916
/* eslint-disable @typescript-eslint/member-ordering */
1017

1118
export interface ByteListOptions {
@@ -34,6 +41,14 @@ export class ByteListType extends ByteArrayType {
3441
readonly maxSize: number;
3542
readonly maxChunkCount: number;
3643
readonly isList = true;
44+
readonly blockArray: Uint8Array[] = [];
45+
private blockBytesLen = 0;
46+
readonly mixInLengthBlockBytes = new Uint8Array(64);
47+
readonly mixInLengthBuffer = Buffer.from(
48+
this.mixInLengthBlockBytes.buffer,
49+
this.mixInLengthBlockBytes.byteOffset,
50+
this.mixInLengthBlockBytes.byteLength
51+
);
3752

3853
constructor(readonly limitBytes: number, opts?: ByteListOptions) {
3954
super();
@@ -89,7 +104,49 @@ export class ByteListType extends ByteArrayType {
89104
// Merkleization: inherited from ByteArrayType
90105

91106
hashTreeRoot(value: ByteArray): Uint8Array {
92-
return mixInLength(super.hashTreeRoot(value), value.length);
107+
const root = allocUnsafe(32);
108+
this.hashTreeRootInto(value, root, 0);
109+
return root;
110+
}
111+
112+
/**
113+
* Use merkleizeBlockArray() instead of merkleizeBlocksBytes() to avoid big memory allocation
114+
*/
115+
hashTreeRootInto(value: Uint8Array, output: Uint8Array, offset: number): void {
116+
// should not call super.hashTreeRoot() here
117+
// use merkleizeBlockArray() instead of merkleizeBlocksBytes() to avoid big memory allocation
118+
// reallocate this.blockArray if needed
119+
if (value.length > this.blockBytesLen) {
120+
const newBlockCount = Math.ceil(value.length / 64);
121+
// this.blockBytesLen should be a multiple of 64
122+
const oldBlockCount = Math.ceil(this.blockBytesLen / 64);
123+
const blockDiff = newBlockCount - oldBlockCount;
124+
const newBlocksBytes = new Uint8Array(blockDiff * 64);
125+
for (let i = 0; i < blockDiff; i++) {
126+
this.blockArray.push(newBlocksBytes.subarray(i * 64, (i + 1) * 64));
127+
this.blockBytesLen += 64;
128+
}
129+
}
130+
131+
// populate this.blockArray
132+
for (let i = 0; i < value.length; i += 64) {
133+
const block = this.blockArray[i / 64];
134+
// zero out the last block if it's over value.length
135+
if (i + 64 > value.length) {
136+
block.fill(0);
137+
}
138+
block.set(value.subarray(i, Math.min(i + 64, value.length)));
139+
}
140+
141+
// compute hashTreeRoot
142+
const blockLimit = Math.ceil(value.length / 64);
143+
merkleizeBlockArray(this.blockArray, blockLimit, this.maxChunkCount, this.mixInLengthBlockBytes, 0);
144+
145+
// mixInLength
146+
this.mixInLengthBuffer.writeUIntLE(value.length, 32, 6);
147+
// one for hashTreeRoot(value), one for length
148+
const chunkCount = 2;
149+
merkleizeBlocksBytes(this.mixInLengthBlockBytes, chunkCount, output, offset);
93150
}
94151

95152
// Proofs: inherited from BitArrayType

‎packages/ssz/src/type/composite.ts

Lines changed: 30 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
import {allocUnsafe} from "@chainsafe/as-sha256";
12
import {
23
concatGindices,
34
createProof,
@@ -7,10 +8,11 @@ import {
78
Proof,
89
ProofType,
910
Tree,
11+
merkleizeBlocksBytes,
1012
HashComputationLevel,
1113
} from "@chainsafe/persistent-merkle-tree";
1214
import {byteArrayEquals} from "../util/byteArray.js";
13-
import {merkleize, symbolCachedPermanentRoot, ValueWithCachedPermanentRoot} from "../util/merkleize.js";
15+
import {cacheRoot, symbolCachedPermanentRoot, ValueWithCachedPermanentRoot} from "../util/merkleize.js";
1416
import {treePostProcessFromProofNode} from "../util/proof/treePostProcessFromProofNode.js";
1517
import {Type, ByteViews, JsonPath, JsonPathProp} from "./abstract.js";
1618
export {ByteViews};
@@ -59,6 +61,7 @@ export abstract class CompositeType<V, TV, TVDU> extends Type<V> {
5961
* Required for ContainerNodeStruct to ensure no dangerous types are constructed.
6062
*/
6163
abstract readonly isViewMutable: boolean;
64+
protected blocksBuffer = new Uint8Array(0);
6265

6366
constructor(
6467
/**
@@ -216,13 +219,30 @@ export abstract class CompositeType<V, TV, TVDU> extends Type<V> {
216219
}
217220
}
218221

219-
const root = merkleize(this.getRoots(value), this.maxChunkCount);
222+
const root = allocUnsafe(32);
223+
const safeCache = true;
224+
this.hashTreeRootInto(value, root, 0, safeCache);
220225

226+
// hashTreeRootInto will cache the root if cachePermanentRootStruct is true
227+
228+
return root;
229+
}
230+
231+
hashTreeRootInto(value: V, output: Uint8Array, offset: number, safeCache = false): void {
232+
// Return cached mutable root if any
221233
if (this.cachePermanentRootStruct) {
222-
(value as ValueWithCachedPermanentRoot)[symbolCachedPermanentRoot] = root;
234+
const cachedRoot = (value as ValueWithCachedPermanentRoot)[symbolCachedPermanentRoot];
235+
if (cachedRoot) {
236+
output.set(cachedRoot, offset);
237+
return;
238+
}
223239
}
224240

225-
return root;
241+
const blocksBuffer = this.getBlocksBytes(value);
242+
merkleizeBlocksBytes(blocksBuffer, this.maxChunkCount, output, offset);
243+
if (this.cachePermanentRootStruct) {
244+
cacheRoot(value as ValueWithCachedPermanentRoot, output, offset, safeCache);
245+
}
226246
}
227247

228248
// For debugging and testing this feature
@@ -236,7 +256,12 @@ export abstract class CompositeType<V, TV, TVDU> extends Type<V> {
236256
// and feed those numbers directly to the hasher input with a DataView
237257
// - The return of the hasher should be customizable too, to reduce conversions from Uint8Array
238258
// to hashObject and back.
239-
protected abstract getRoots(value: V): Uint8Array[];
259+
260+
/**
261+
* Get multiple SHA256 blocks, each is 64 bytes long.
262+
* If chunk count is not even, need to append zeroHash(0)
263+
*/
264+
protected abstract getBlocksBytes(value: V): Uint8Array;
240265

241266
// Proofs API
242267

‎packages/ssz/src/type/container.ts

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -130,6 +130,8 @@ export class ContainerType<Fields extends Record<string, Type<unknown>>> extends
130130
// Refactor this constructor to allow customization without pollutin the options
131131
this.TreeView = opts?.getContainerTreeViewClass?.(this) ?? getContainerTreeViewClass(this);
132132
this.TreeViewDU = opts?.getContainerTreeViewDUClass?.(this) ?? getContainerTreeViewDUClass(this);
133+
const fieldBytes = this.fieldsEntries.length * 32;
134+
this.blocksBuffer = new Uint8Array(Math.ceil(fieldBytes / 64) * 64);
133135
}
134136

135137
static named<Fields extends Record<string, Type<unknown>>>(
@@ -272,15 +274,13 @@ export class ContainerType<Fields extends Record<string, Type<unknown>>> extends
272274

273275
// Merkleization
274276

275-
protected getRoots(struct: ValueOfFields<Fields>): Uint8Array[] {
276-
const roots = new Array<Uint8Array>(this.fieldsEntries.length);
277-
277+
protected getBlocksBytes(struct: ValueOfFields<Fields>): Uint8Array {
278278
for (let i = 0; i < this.fieldsEntries.length; i++) {
279279
const {fieldName, fieldType} = this.fieldsEntries[i];
280-
roots[i] = fieldType.hashTreeRoot(struct[fieldName]);
280+
fieldType.hashTreeRootInto(struct[fieldName], this.blocksBuffer, i * 32);
281281
}
282-
283-
return roots;
282+
// remaining bytes are zeroed as we never write them
283+
return this.blocksBuffer;
284284
}
285285

286286
// Proofs

0 commit comments

Comments
 (0)