feat: new type for list of uint64 (#352)

* feat: implement ListUintNum64Type * feat: use ListUintNum64Type for lodestar Balances type
ChainSafe · Mar 12, 2024 · e131b5a · e131b5a
1 parent 4ba45d3
commit e131b5a
Show file tree

Hide file tree

Showing 5 changed files with 288 additions and 77 deletions.
diff --git a/packages/persistent-merkle-tree/src/packedNode.ts b/packages/persistent-merkle-tree/src/packedNode.ts
@@ -1,11 +1,45 @@
 import {subtreeFillToContents} from "./subtree";
 import {Node, LeafNode, getNodeH, setNodeH} from "./node";
 
+const NUMBER_2_POW_32 = 2 ** 32;
+
 export function packedRootsBytesToNode(depth: number, dataView: DataView, start: number, end: number): Node {
   const leafNodes = packedRootsBytesToLeafNodes(dataView, start, end);
   return subtreeFillToContents(leafNodes, depth);
 }
 
+/**
+ * Pack a list of uint64 numbers into a list of LeafNodes.
+ * Each value is UintNum64, which is 8 bytes long, which is 2 h values.
+ * Each 4 of them forms a LeafNode.
+ *
+ *      v0            v1           v2          v3
+ * |-------------|-------------|-------------|-------------|
+ *
+ *    h0     h1     h2     h3     h4     h5     h6     h7
+ * |------|------|------|------|------|------|------|------|
+ */
+export function packedUintNum64sToLeafNodes(values: number[]): LeafNode[] {
+  const leafNodes = new Array<LeafNode>(Math.ceil(values.length / 4));
+  for (let i = 0; i < values.length; i++) {
+    const nodeIndex = Math.floor(i / 4);
+    const leafNode = leafNodes[nodeIndex] ?? new LeafNode(0, 0, 0, 0, 0, 0, 0, 0);
+    const vIndex = i % 4;
+    const hIndex = 2 * vIndex;
+    const value = values[i];
+    // same logic to UintNumberType.value_serializeToBytes() for 8 bytes
+    if (value === Infinity) {
+      setNodeH(leafNode, hIndex, 0xffffffff);
+      setNodeH(leafNode, hIndex + 1, 0xffffffff);
+    } else {
+      setNodeH(leafNode, hIndex, value & 0xffffffff);
+      setNodeH(leafNode, hIndex + 1, (value / NUMBER_2_POW_32) & 0xffffffff);
+    }
+    leafNodes[nodeIndex] = leafNode;
+  }
+  return leafNodes;
+}
+
 /**
  * Optimized deserialization of linear bytes to consecutive leaf nodes
  */

diff --git a/packages/persistent-merkle-tree/test/unit/packedNode.test.ts b/packages/persistent-merkle-tree/test/unit/packedNode.test.ts
@@ -1,14 +1,15 @@
 import {HashObject} from "@chainsafe/as-sha256";
 import {expect} from "chai";
 import {LeafNode, Node} from "../../src";
-import {packedNodeRootsToBytes, packedRootsBytesToLeafNodes} from "../../src/packedNode";
+import {packedNodeRootsToBytes, packedRootsBytesToLeafNodes, packedUintNum64sToLeafNodes} from "../../src/packedNode";
 
 describe("subtree / packedNode single node", () => {
   const testCases: {
     id: string;
     size: number;
     nodes: Node[];
     outStr: string;
+    testPackedNumbers?: boolean;
   }[] = [
     {
       id: "One byte",
@@ -48,11 +49,37 @@ describe("subtree / packedNode single node", () => {
       nodes: [LeafNode.fromHashObject({h0: 0x0708090a, h1: 0x01020304, h2: 0, h3: 0, h4: 0, h5: 0, h6: 0, h7: 0})],
       outStr: "0x0a09080704030201",
     },
+    {
+      id: "2 h values fits uint64 number",
+      size: 8,
+      nodes: [LeafNode.fromHashObject({h0: 0x0708090a, h1: 0x0102030, h2: 0, h3: 0, h4: 0, h5: 0, h6: 0, h7: 0})],
+      outStr: "0x0a09080730201000",
+      testPackedNumbers: true,
+    },
     {
       id: "32 bytes zero",
       size: 32,
       nodes: [LeafNode.fromHashObject({h0: 0, h1: 0, h2: 0, h3: 0, h4: 0, h5: 0, h6: 0, h7: 0})],
       outStr: "0x0000000000000000000000000000000000000000000000000000000000000000",
+      testPackedNumbers: true,
+    },
+    {
+      id: "32 bytes max",
+      size: 32,
+      nodes: [
+        LeafNode.fromHashObject({
+          h0: 0xffffffff,
+          h1: 0xffffffff,
+          h2: 0xffffffff,
+          h3: 0xffffffff,
+          h4: 0xffffffff,
+          h5: 0xffffffff,
+          h6: 0xffffffff,
+          h7: 0xffffffff,
+        }),
+      ],
+      outStr: "0xffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff",
+      testPackedNumbers: true,
     },
     {
       id: "32 bytes same",
@@ -77,9 +104,82 @@ describe("subtree / packedNode single node", () => {
       ],
       outStr: "0x18735c375d8e7b2922ef64f165d10b9ace103f326627f0a21d0fe8a0a573f83f",
     },
+    // same to random tests but trim h1, h3, h5, h7 to make h{2*i} + h{2*i + 1} fits uint64
+    {
+      id: "8 bytes random fits unit64 number",
+      size: 32,
+      nodes: [
+        LeafNode.fromHashObject({
+          h0: 928805656,
+          h1: 6959632,
+          h2: 0,
+          h3: 0,
+          h4: 0,
+          h5: 0,
+          h6: 0,
+          h7: 0,
+        }),
+      ],
+      outStr: "0x18735c3710326a00000000000000000000000000000000000000000000000000",
+      testPackedNumbers: true,
+    },
+    {
+      id: "16 bytes random fits unit64 number",
+      size: 32,
+      nodes: [
+        LeafNode.fromHashObject({
+          h0: 928805656,
+          h1: 6959632,
+          h2: 4049923874,
+          h3: 258446,
+          h4: 0,
+          h5: 0,
+          h6: 0,
+          h7: 0,
+        }),
+      ],
+      outStr: "0x18735c3710326a0022ef64f18ef1030000000000000000000000000000000000",
+      testPackedNumbers: true,
+    },
+    {
+      id: "24 bytes random fits unit64 number",
+      size: 32,
+      nodes: [
+        LeafNode.fromHashObject({
+          h0: 928805656,
+          h1: 6959632,
+          h2: 4049923874,
+          h3: 258446,
+          h4: 842993870,
+          h5: 273364,
+          h6: 0,
+          h7: 0,
+        }),
+      ],
+      outStr: "0x18735c3710326a0022ef64f18ef10300ce103f32d42b04000000000000000000",
+      testPackedNumbers: true,
+    },
+    {
+      id: "32 bytes random fits unit64 number",
+      size: 32,
+      nodes: [
+        LeafNode.fromHashObject({
+          h0: 928805656,
+          h1: 6959632,
+          h2: 4049923874,
+          h3: 258446,
+          h4: 842993870,
+          h5: 273364,
+          h6: 2699562781,
+          h7: 107324,
+        }),
+      ],
+      outStr: "0x18735c3710326a0022ef64f18ef10300ce103f32d42b04001d0fe8a03ca30100",
+      testPackedNumbers: true,
+    },
   ];
 
-  for (const {id, size, nodes, outStr} of testCases) {
+  for (const {id, size, nodes, outStr, testPackedNumbers} of testCases) {
     it(`${id} - packedNodeRootsToBytes`, () => {
       const uint8Array = new Uint8Array(size);
       const dataView = new DataView(uint8Array.buffer, uint8Array.byteOffset, uint8Array.byteLength);
@@ -93,6 +193,28 @@ describe("subtree / packedNode single node", () => {
       const nodesRes = packedRootsBytesToLeafNodes(dataView, 0, size);
       expect(onlyHashObject(nodesRes[0].rootHashObject)).to.deep.equal(onlyHashObject(nodes[0].rootHashObject));
     });
+
+    // 1 UintNum64 = 8 bytes
+    if (testPackedNumbers) {
+      const NUMBER_2_POW_32 = 2 ** 32;
+      it(`${id} - packedUintNum64sToLeafNodes, value size=${Math.floor(size / 8)}`, () => {
+        const values: number[] = [];
+        const uint8Array = new Uint8Array(Buffer.from(outStr.replace("0x", ""), "hex"));
+        const dataView = new DataView(uint8Array.buffer, uint8Array.byteOffset, uint8Array.byteLength);
+        for (let i = 0; i < size; i += 8) {
+          const a = dataView.getUint32(i, true);
+          const b = dataView.getUint32(i + 4, true);
+          if (a === 0xffffffff && b === 0xffffffff) {
+            values.push(Infinity);
+          } else {
+            values.push(b * NUMBER_2_POW_32 + a);
+          }
+        }
+
+        const nodesRes = packedUintNum64sToLeafNodes(values);
+        expect(onlyHashObject(nodesRes[0].rootHashObject)).to.deep.equal(onlyHashObject(nodes[0].rootHashObject));
+      });
+    }
   }
 });
 

diff --git a/packages/ssz/src/type/listUintNum64.ts b/packages/ssz/src/type/listUintNum64.ts
@@ -0,0 +1,49 @@
+import {LeafNode, Node, packedUintNum64sToLeafNodes, subtreeFillToContents} from "@chainsafe/persistent-merkle-tree";
+
+import {ListBasicTreeViewDU} from "../viewDU/listBasic";
+import {ListBasicOpts, ListBasicType} from "./listBasic";
+import {UintNumberType} from "./uint";
+import {addLengthNode} from "./arrayBasic";
+
+/**
+ * Specific implementation of ListBasicType for UintNumberType with some optimizations.
+ */
+export class ListUintNum64Type extends ListBasicType<UintNumberType> {
+  constructor(limit: number, opts?: ListBasicOpts) {
+    super(new UintNumberType(8), limit, opts);
+  }
+
+  /**
+   * Return a ListBasicTreeViewDU with nodes populated
+   */
+  toViewDU(value: number[]): ListBasicTreeViewDU<UintNumberType> {
+    // no need to serialize and deserialize like in the abstract class
+    const {treeNode, leafNodes} = this.packedUintNum64sToNode(value);
+    // cache leaf nodes in the ViewDU
+    return this.getViewDU(treeNode, {
+      nodes: leafNodes,
+      length: value.length,
+      nodesPopulated: true,
+    });
+  }
+
+  /**
+   * No need to serialize and deserialize like in the abstract class
+   */
+  value_toTree(value: number[]): Node {
+    const {treeNode} = this.packedUintNum64sToNode(value);
+    return treeNode;
+  }
+
+  private packedUintNum64sToNode(value: number[]): {treeNode: Node; leafNodes: LeafNode[]} {
+    if (value.length > this.limit) {
+      throw new Error(`Exceeds limit: ${value.length} > ${this.limit}`);
+    }
+
+    const leafNodes = packedUintNum64sToLeafNodes(value);
+    // subtreeFillToContents mutates the leafNodes array
+    const rootNode = subtreeFillToContents([...leafNodes], this.chunkDepth);
+    const treeNode = addLengthNode(rootNode, value.length);
+    return {treeNode, leafNodes};
+  }
+}
diff --git a/packages/ssz/test/lodestarTypes/phase0/sszTypes.ts b/packages/ssz/test/lodestarTypes/phase0/sszTypes.ts
@@ -8,6 +8,7 @@ import {
   VectorBasicType,
   VectorCompositeType,
 } from "../../../src";
+import {ListUintNum64Type} from "../../../src/type/listUintNum64";
 import {
   preset,
   MAX_REQUEST_BLOCKS,
@@ -250,7 +251,7 @@ export const Validator = ValidatorNodeStruct;
 
 // Export as stand-alone for direct tree optimizations
 export const Validators = new ListCompositeType(ValidatorNodeStruct, VALIDATOR_REGISTRY_LIMIT);
-export const Balances = new ListBasicType(UintNum64, VALIDATOR_REGISTRY_LIMIT);
+export const Balances = new ListUintNum64Type(VALIDATOR_REGISTRY_LIMIT);
 export const RandaoMixes = new VectorCompositeType(Bytes32, EPOCHS_PER_HISTORICAL_VECTOR);
 export const Slashings = new VectorBasicType(Gwei, EPOCHS_PER_SLASHINGS_VECTOR);
 export const JustificationBits = new BitVectorType(JUSTIFICATION_BITS_LENGTH);