feat: snapshot apis for EIP-4881 (#400)

* feat: implement indexToFinalizedGindices snapshot util

* feat: implement toSnapshot() and fromSnapshot() apis in persistent-merkle-tree

* feat: implement PartialListCompositeType and ViewDU

* fix: use ssz Snapshot type

* fix: move ZERO_SNAPSHOT to ssz

* fix: use common type for ArrayBasicType and ArrayCompositeType

* fix: throw Error for methods that PartialListCompositeType does not support

* feat: implement type.defaultPartialViewDU() and zeroSnapshot()

* fix: EIP-4881 spec test

* chore: export types

* chore: more comments for Snapshot type

* chore: use addLengthNode() api

packages/persistent-merkle-tree/src/index.ts

@@ -8,3 +8,4 @@ export * from "./subtree";
 export * from "./tree";
 export * from "./zeroNode";
 export * from "./zeroHash";
+export * from "./snapshot";

packages/persistent-merkle-tree/src/snapshot.ts

@@ -0,0 +1,89 @@
+import {Tree, getNode} from "./tree";
+import {zeroNode} from "./zeroNode";
+import {Gindex, toGindex} from "./gindex";
+import {LeafNode, Node} from "./node";
+
+type Snapshot = {
+  finalized: Uint8Array[];
+  count: number;
+};
+
+/**
+ * Given a tree, return a snapshot of the tree with the root, finalized nodes, and count.
+ * Tree could be full tree, or partial tree. See https://github.com/ChainSafe/ssz/issues/293
+ */
+export function toSnapshot(rootNode: Node, depth: number, count: number): Snapshot {
+  if (count < 0) {
+    throw new Error(`Expect count to be non-negative, got ${count}`);
+  }
+
+  const finalizedGindices = count > 0 ? indexToFinalizedGindices(depth, count - 1) : [];
+  const finalized = finalizedGindices.map((gindex) => getNode(rootNode, gindex).root);
+
+  return {
+    finalized,
+    count,
+  };
+}
+
+/**
+ * Given a snapshot, return root node of a tree.
+ * See https://github.com/ChainSafe/ssz/issues/293
+ */
+export function fromSnapshot(snapshot: Snapshot, depth: number): Node {
+  const tree = new Tree(zeroNode(depth));
+  const {count, finalized} = snapshot;
+  if (count < 0) {
+    throw new Error(`Expect count to be non-negative, got ${count}`);
+  }
+
+  const finalizedGindices = count > 0 ? indexToFinalizedGindices(depth, count - 1) : [];
+
+  if (finalizedGindices.length !== finalized.length) {
+    throw new Error(`Expected ${finalizedGindices.length} finalized gindices, got ${finalized.length}`);
+  }
+
+  for (const [i, gindex] of finalizedGindices.entries()) {
+    const node = LeafNode.fromRoot(finalized[i]);
+    tree.setNode(gindex, node);
+  }
+
+  return tree.rootNode;
+}
+
+/**
+ * A finalized gindex means that the gindex is at the root of a subtree of the tree where there is no ZERO_NODE belong to it.
+ * Given a list of depth `depth` and an index `index`, return a list of finalized gindexes.
+ */
+export function indexToFinalizedGindices(depth: number, index: number): Gindex[] {
+  if (index < 0 || depth < 0) {
+    throw new Error(`Expect index and depth to be non-negative, got ${index} and ${depth}`);
+  }
+
+  // given this tree with depth 3 and index 6
+  //        X
+  //    X       X
+  //  X   X   X   0
+  // X X X X X X 0 0
+  // we'll extract the root 4 left most nodes, then root node of the next 2 nodes
+  // need to track the offset at each level to compute gindex of each root node
+  const offsetByDepth = Array.from({length: depth + 1}, () => 0);
+  // count starts with 1
+  let count = index + 1;
+
+  const result: Gindex[] = [];
+  while (count > 0) {
+    const prevLog2 = Math.floor(Math.log2(count));
+    const prevPowerOf2 = 2 ** prevLog2;
+    const depthFromRoot = depth - prevLog2;
+    const finalizedGindex = toGindex(depthFromRoot, BigInt(offsetByDepth[depthFromRoot]));
+    result.push(finalizedGindex);
+    for (let i = 0; i <= prevLog2; i++) {
+      offsetByDepth[depthFromRoot + i] += Math.pow(2, i);
+    }
+
+    count -= prevPowerOf2;
+  }
+
+  return result;
+}

packages/persistent-merkle-tree/test/unit/snapshot.test.ts

@@ -0,0 +1,115 @@
+import { expect } from "chai";
+import {describe, it} from "mocha";
+import {fromSnapshot, indexToFinalizedGindices, toSnapshot} from "../../src/snapshot";
+import {subtreeFillToContents} from "../../src/subtree";
+import { LeafNode } from "../../src/node";
+import { Tree, setNodesAtDepth } from "../../src/tree";
+import { toGindex } from "../../src";
+
+describe("toSnapshot and fromSnapshot", () => {
+  const depth = 4;
+  const maxItems = Math.pow(2, depth);
+
+  for (let count = 0; count <= maxItems; count ++) {
+    it(`toSnapshot and fromSnapshot with count ${count}`, () => {
+      const nodes = Array.from({length: count}, (_, i) => LeafNode.fromRoot(Buffer.alloc(32, i)));
+      const fullListRootNode = subtreeFillToContents(nodes, depth);
+      const snapshot = toSnapshot(fullListRootNode, depth, count);
+      const partialListRootNode = fromSnapshot(snapshot, depth);
+
+      // 1st step - check if the restored root node is the same
+      expect(partialListRootNode.root).to.deep.equal(fullListRootNode.root);
+
+      // 2nd step - make sure we can add more nodes to the restored tree
+      const fullTree = new Tree(fullListRootNode);
+      const partialTree = new Tree(partialListRootNode);
+      for (let i = count; i < maxItems; i++) {
+        const gindex = toGindex(depth, BigInt(i));
+        fullTree.setNode(gindex, LeafNode.fromRoot(Buffer.alloc(32, i)));
+        partialTree.setNode(gindex, LeafNode.fromRoot(Buffer.alloc(32, i)));
+        expect(partialTree.root).to.deep.equal(fullTree.root);
+
+        // and snapshot created from 2 trees are the same
+        const snapshot1 = toSnapshot(fullTree.rootNode, depth, i + 1);
+        const snapshot2 = toSnapshot(partialTree.rootNode, depth, i + 1);
+        expect(snapshot2).to.deep.equal(snapshot1);
+      }
+    });
+
+    // setNodesAtDepth() api is what ssz uses to grow the tree in its commit() phase
+    it(`toSnapshot and fromSnapshot with count ${count} then grow with setNodeAtDepth`, () => {
+      const nodes = Array.from({length: count}, (_, i) => LeafNode.fromRoot(Buffer.alloc(32, i)));
+      const fullListRootNode = subtreeFillToContents(nodes, depth);
+      const snapshot = toSnapshot(fullListRootNode, depth, count);
+      const partialListRootNode = fromSnapshot(snapshot, depth);
+
+      // 1st step - check if the restored root node is the same
+      expect(partialListRootNode.root).to.deep.equal(fullListRootNode.root);
+
+      // 2nd step - grow the tree with setNodesAtDepth
+      for (let i = count; i < maxItems; i++) {
+        const addedNodes = Array.from({length: i - count + 1}, (_, j) => LeafNode.fromRoot(Buffer.alloc(32, j)));
+        const indices = Array.from({length: i - count + 1}, (_, j) => j + count);
+        const root1 = setNodesAtDepth(fullListRootNode, depth, indices, addedNodes);
+        const root2 = setNodesAtDepth(partialListRootNode, depth, indices, addedNodes);
+        expect(root2.root).to.deep.equal(root1.root);
+
+        for (let j = count; j <= i; j++) {
+          const snapshot1 = toSnapshot(root1, depth, j);
+          const snapshot2 = toSnapshot(root2, depth, j);
+          expect(snapshot2).to.deep.equal(snapshot1);
+        }
+      }
+    });
+
+    it(`toSnapshot() multiple times with count ${count}`, () => {
+      const nodes = Array.from({length: count}, (_, i) => LeafNode.fromRoot(Buffer.alloc(32, i)));
+      const fullListRootNode = subtreeFillToContents(nodes, depth);
+      const snapshot = toSnapshot(fullListRootNode, depth, count);
+      const partialListRootNode = fromSnapshot(snapshot, depth);
+
+      // 1st step - check if the restored root node is the same
+      expect(partialListRootNode.root).to.deep.equal(fullListRootNode.root);
+
+      const snapshot2 = toSnapshot(partialListRootNode, depth, count);
+      const restoredRootNode2 = fromSnapshot(snapshot2, depth);
+
+      // 2nd step - check if the restored root node is the same
+      expect(restoredRootNode2.root).to.deep.equal(partialListRootNode.root);
+    });
+  }
+});
+
+describe("indexToFinalizedGindices", () => {
+  // given a tree with depth = 4
+  //                                      1
+  //                     2                                3
+  //         4                   5                 6               7
+  //    8        9         10       11        12       13       14      15
+  // 16  17   18  19    20   21   22  23    24  25   26  27   28  29  30  31
+  const testCases: [number, number, bigint[]][] = [
+    [4, 0, [BigInt(16)]],
+    [4, 1, [BigInt(8)]],
+    [4, 2, [8, 18].map(BigInt)],
+    [4, 3, [4].map(BigInt)],
+    [4, 4, [4, 20].map(BigInt)],
+    [4, 5, [4, 10].map(BigInt)],
+    [4, 6, [4, 10, 22].map(BigInt)],
+    [4, 7, [2].map(BigInt)],
+    [4, 8, [2, 24].map(BigInt)],
+    [4, 9, [2, 12].map(BigInt)],
+    [4, 10, [2, 12, 26].map(BigInt)],
+    [4, 11, [2, 6].map(BigInt)],
+    [4, 12, [2, 6, 28].map(BigInt)],
+    [4, 13, [2, 6, 14].map(BigInt)],
+    [4, 14, [2, 6, 14, 30].map(BigInt)],
+    [4, 15, [1].map(BigInt)],
+  ];
+
+  for (const [depth, index, finalizeGindices] of testCases) {
+    it(`should correctly get finalized gindexes for index ${index} and depth ${depth}`, () => {
+      const actual = indexToFinalizedGindices(depth, index);
+      expect(actual).to.deep.equal(finalizeGindices);
+    });
+  }
+});

packages/ssz/package.json

@@ -26,11 +26,12 @@
     "benchmark:local": "yarn benchmark --local",
     "test:perf": "mocha  \"test/perf/**/*.test.ts\"",
     "test:unit": "nyc mocha  \"test/unit/**/*.test.ts\"",
-    "test:spec": "yarn test:spec-generic && yarn test:spec-static",
+    "test:spec": "yarn test:spec-generic && yarn test:spec-static test:spec-eip-4881",
     "test:spec-generic": "mocha  \"test/spec/generic/**/*.test.ts\"",
     "test:spec-static": "yarn test:spec-static-minimal && yarn test:spec-static-mainnet",
     "test:spec-static-minimal": "LODESTAR_PRESET=minimal  mocha test/spec/ssz_static.test.ts",
     "test:spec-static-mainnet": "LODESTAR_PRESET=mainnet  mocha test/spec/ssz_static.test.ts",
+    "test:spec-eip-4881": "mocha  \"test/spec/eip-4881/**/*.test.ts\"",
     "download-spec-tests": "node -r ts-node/register test/spec/downloadTests.ts"
   },
   "types": "lib/index.d.ts",

packages/ssz/src/index.ts

@@ -8,6 +8,7 @@ export {ContainerType} from "./type/container";
 export {ContainerNodeStructType} from "./type/containerNodeStruct";
 export {ListBasicType} from "./type/listBasic";
 export {ListCompositeType} from "./type/listComposite";
+export {PartialListCompositeType} from "./type/partialListComposite";
 export {NoneType} from "./type/none";
 export {UintBigintType, UintNumberType} from "./type/uint";
 export {UnionType} from "./type/union";
@@ -34,5 +35,5 @@ export {BitArray, getUint8ByteToBitBooleanArray} from "./value/bitArray";
 
 // Utils
 export {fromHexString, toHexString, byteArrayEquals} from "./util/byteArray";
-
+export {Snapshot} from "./util/types";
 export {hash64, symbolCachedPermanentRoot} from "./util/merkleize";

packages/ssz/src/type/partialListComposite.ts

@@ -0,0 +1,68 @@
+import {fromSnapshot, zeroNode} from "@chainsafe/persistent-merkle-tree";
+import {CompositeType, CompositeView, CompositeViewDU} from "./composite";
+import {ListCompositeOpts, ListCompositeType} from "./listComposite";
+import {PartialListCompositeTreeViewDU} from "../viewDU/partialListComposite";
+import {Snapshot} from "../util/types";
+import {byteArrayEquals} from "../util/byteArray";
+import {zeroSnapshot} from "../util/snapshot";
+import {addLengthNode} from "./arrayBasic";
+
+/**
+ * Similar to ListCompositeType, this is mainly used to create a PartialListCompositeTreeViewDU from a snapshot.
+ * The ViewDU created is a partial tree created from a snapshot, not a full tree.
+ * Note that this class only inherits minimal methods as defined in ArrayType of ../view/arrayBasic.ts
+ * It'll throw errors for all other methods, most of the usage is in the ViewDU class.
+ */
+export class PartialListCompositeType<
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  ElementType extends CompositeType<any, CompositeView<ElementType>, CompositeViewDU<ElementType>>
+> extends ListCompositeType<ElementType> {
+  constructor(readonly elementType: ElementType, readonly limit: number, opts?: ListCompositeOpts) {
+    super(elementType, limit, opts);
+
+    // only inherit methods in ArrayType of ../view/arrayBasic.ts
+    const inheritedMethods = [
+      "tree_getLength",
+      "tree_setLength",
+      "tree_getChunksNode",
+      "tree_chunksNodeOffset",
+      "tree_setChunksNode",
+    ];
+    const methodNames = Object.getOwnPropertyNames(ListCompositeType.prototype).filter(
+      (prop) =>
+        prop !== "constructor" &&
+        typeof (this as unknown as Record<string, unknown>)[prop] === "function" &&
+        !inheritedMethods.includes(prop)
+    );
+
+    // throw errors for all remaining methods
+    for (const methodName of methodNames) {
+      (this as unknown as Record<string, unknown>)[methodName] = () => {
+        throw new Error(`Method ${methodName} is not implemented for PartialListCompositeType`);
+      };
+    }
+  }
+
+  /**
+   * Create a PartialListCompositeTreeViewDU from a snapshot.
+   */
+  toPartialViewDU(snapshot: Snapshot): PartialListCompositeTreeViewDU<ElementType> {
+    const chunksNode = fromSnapshot(snapshot, this.chunkDepth);
+    const rootNode = addLengthNode(chunksNode, snapshot.count);
+
+    if (!byteArrayEquals(rootNode.root, snapshot.root)) {
+      throw new Error(`Snapshot root is incorrect, expected ${snapshot.root}, got ${rootNode.root}`);
+    }
+
+    return new PartialListCompositeTreeViewDU(this, rootNode, snapshot);
+  }
+
+  /**
+   * Creates a PartialListCompositeTreeViewDU from a zero snapshot.
+   */
+  defaultPartialViewDU(): PartialListCompositeTreeViewDU<ElementType> {
+    const rootNode = addLengthNode(zeroNode(this.chunkDepth), 0);
+
+    return new PartialListCompositeTreeViewDU(this, rootNode, zeroSnapshot(this.chunkDepth));
+  }
+}

packages/ssz/src/util/snapshot.ts

@@ -0,0 +1,14 @@
+import {zeroHash} from "@chainsafe/persistent-merkle-tree";
+import {hash64} from "./merkleize";
+import {Snapshot} from "./types";
+
+/**
+ * Create a zero snapshot with the given chunksDepth.
+ */
+export function zeroSnapshot(chunkDepth: number): Snapshot {
+  return {
+    finalized: [],
+    count: 0,
+    root: hash64(zeroHash(chunkDepth), zeroHash(0)),
+  };
+}

packages/ssz/src/util/types.ts

@@ -1 +1,12 @@
 export type Require<T, K extends keyof T> = T & Required<Pick<T, K>>;
+
+/**
+ * A snapshot contains the minimum amount of information needed to reconstruct a merkleized list, for the purposes of appending more items.
+ * Note: This does not contain list elements, rather only contains intermediate merkle nodes.
+ * This is used primarily for PartialListCompositeType.
+ */
+export type Snapshot = {
+  finalized: Uint8Array[];
+  root: Uint8Array;
+  count: number;
+};

packages/ssz/src/view/arrayBasic.ts

@@ -10,11 +10,15 @@ export type ArrayBasicType<ElementType extends BasicType<unknown>> = CompositeTy
   ValueOf<ElementType>[],
   TreeView<ArrayBasicType<ElementType>>,
   TreeViewDU<ArrayBasicType<ElementType>>
-> & {
-  readonly elementType: ElementType;
-  readonly itemsPerChunk: number;
-  readonly chunkDepth: number;
-
+> &
+  ArrayType & {
+    readonly elementType: ElementType;
+    readonly itemsPerChunk: number;
+    readonly chunkDepth: number;
+  };
+
+/** Common type for both ArrayBasicType and ArrayCompositeTypesrc/view/arrayBasic.ts */
+export type ArrayType = {
   /** INTERNAL METHOD: Return the length of this type from an Array's root node */
   tree_getLength(node: Node): number;
   /** INTERNAL METHOD: Mutate a tree's rootNode with a new length value */

packages/ssz/src/view/arrayComposite.ts

@@ -1,32 +1,17 @@
-import {getNodesAtDepth, Node, toGindexBitstring, Tree, HashComputationLevel} from "@chainsafe/persistent-merkle-tree";
+import {getNodesAtDepth, Node, toGindexBitstring, Tree} from "@chainsafe/persistent-merkle-tree";
 import {ValueOf} from "../type/abstract";
 import {CompositeType, CompositeView, CompositeViewDU} from "../type/composite";
 import {TreeView} from "./abstract";
+import {ArrayType} from "./arrayBasic";
 
 /** Expected API of this View's type. This interface allows to break a recursive dependency between types and views */
 export type ArrayCompositeType<
   ElementType extends CompositeType<unknown, CompositeView<ElementType>, CompositeViewDU<ElementType>>
-> = CompositeType<ValueOf<ElementType>[], unknown, unknown> & {
-  readonly elementType: ElementType;
-  readonly chunkDepth: number;
-
-  /** INTERNAL METHOD: Return the length of this type from an Array's root node */
-  tree_getLength(node: Node): number;
-  /** INTERNAL METHOD: Mutate a tree's rootNode with a new length value */
-  tree_setLength(tree: Tree, length: number): void;
-  /** INTERNAL METHOD: Return the chunks node from a root node */
-  tree_getChunksNode(rootNode: Node): Node;
-  /** INTERNAL METHOD: Return the offset from root for HashComputation */
-  tree_chunksNodeOffset(): number;
-  /** INTERNAL METHOD: Return a new root node with changed chunks node and length */
-  tree_setChunksNode(
-    rootNode: Node,
-    chunksNode: Node,
-    newLength: number | null,
-    hcOffset?: number,
-    hcByLevel?: HashComputationLevel[] | null
-  ): Node;
-};
+> = CompositeType<ValueOf<ElementType>[], unknown, unknown> &
+  ArrayType & {
+    readonly elementType: ElementType;
+    readonly chunkDepth: number;
+  };
 
 export class ArrayCompositeTreeView<
   ElementType extends CompositeType<ValueOf<ElementType>, CompositeView<ElementType>, CompositeViewDU<ElementType>>