Add composite type leaves to proof

src/types/composite/abstract.ts

@@ -173,7 +173,16 @@ export abstract class CompositeType<T extends CompositeValue> extends Type<T> {
   abstract tree_getPropertyNames(tree: Tree): (string | number)[];
 
   abstract getPropertyGindex(property: PropertyKey): Gindex;
-  getPathGindex(path: Path): Gindex {
+  abstract getPropertyType(property: PropertyKey): Type<unknown> | undefined;
+  abstract tree_getProperty(tree: Tree, property: PropertyKey): Tree | unknown;
+  abstract tree_setProperty(tree: Tree, property: PropertyKey, value: Tree | unknown): boolean;
+  abstract tree_deleteProperty(tree: Tree, property: PropertyKey): boolean;
+  abstract tree_iterateValues(tree: Tree): IterableIterator<Tree | unknown>;
+  abstract tree_readonlyIterateValues(tree: Tree): IterableIterator<Tree | unknown>;
+  /**
+   * Navigate to a subtype & gindex using a path
+   */
+  getPathInfo(path: Path): {gindex: Gindex; type: Type<unknown>} {
     const gindices = [];
     let type = this as CompositeType<CompositeValue>;
     for (const prop of path) {
@@ -183,17 +192,42 @@ export abstract class CompositeType<T extends CompositeValue> extends Type<T> {
       gindices.push(type.getPropertyGindex(prop));
       type = type.getPropertyType(prop) as CompositeType<CompositeValue>;
     }
-    return concatGindices(gindices);
+    return {
+      type,
+      gindex: concatGindices(gindices),
+    };
+  }
+  getPathGindex(path: Path): Gindex {
+    return this.getPathInfo(path).gindex;
   }
+  /**
+   * Get leaf gindices
+   *
+   * Note: This is a recursively called method.
+   * Subtypes recursively call this method until basic types / leaf data is hit.
+   *
+   * @param target Used for variable-length types.
+   * @param root Used to anchor the returned gindices to a non-root gindex.
+   * This is used to augment leaf gindices in recursively-called subtypes relative to the type.
+   * @returns The gindices corresponding to leaf data.
+   */
+  abstract tree_getLeafGindices(target?: Tree, root?: Gindex): Gindex[];
 
-  abstract getPropertyType(property: PropertyKey): Type<unknown> | undefined;
-  abstract tree_getProperty(tree: Tree, property: PropertyKey): Tree | unknown;
-  abstract tree_setProperty(tree: Tree, property: PropertyKey, value: Tree | unknown): boolean;
-  abstract tree_deleteProperty(tree: Tree, property: PropertyKey): boolean;
-  abstract tree_iterateValues(tree: Tree): IterableIterator<Tree | unknown>;
-  abstract tree_readonlyIterateValues(tree: Tree): IterableIterator<Tree | unknown>;
   tree_createProof(target: Tree, paths: Path[]): Proof {
-    const gindices = paths.map((path) => this.getPathGindex(path));
+    const gindices = paths
+      .map((path) => {
+        const {type, gindex} = this.getPathInfo(path);
+        if (!isCompositeType(type)) {
+          return gindex;
+        } else {
+          // if the path subtype is composite, include the gindices of all the leaves
+          return type.tree_getLeafGindices(
+            type.hasVariableSerializedLength() ? target.getSubtree(gindex) : undefined,
+            gindex
+          );
+        }
+      })
+      .flat(1);
     return target.getProof({
       type: ProofType.treeOffset,
       gindices,

src/types/composite/array.ts

@@ -5,7 +5,16 @@ import {IJsonOptions, Type} from "../type";
 import {BasicType} from "../basic";
 import {CompositeType} from "./abstract";
 import {SszErrorPath} from "../../util/errorPath";
-import {Gindex, iterateAtDepth, LeafNode, Node, subtreeFillToContents, Tree} from "@chainsafe/persistent-merkle-tree";
+import {
+  concatGindices,
+  Gindex,
+  iterateAtDepth,
+  LeafNode,
+  Node,
+  subtreeFillToContents,
+  toGindex,
+  Tree,
+} from "@chainsafe/persistent-merkle-tree";
 import {isTreeBacked} from "../../backings/tree/treeValue";
 
 export interface IArrayOptions {
@@ -265,6 +274,16 @@ export abstract class BasicArrayType<T extends ArrayLike<unknown>> extends Compo
   bytes_getVariableOffsets(target: Uint8Array): [number, number][] {
     return [];
   }
+
+  tree_getLeafGindices(target?: Tree, root: Gindex = BigInt(1)): Gindex[] {
+    const chunkCount = this.tree_getChunkCount(target);
+    const startIndex = concatGindices([root, toGindex(this.getChunkDepth(), BigInt(0))]);
+    const gindices: Gindex[] = [];
+    for (let i = 0, gindex = startIndex; i < chunkCount; i++, gindex++) {
+      gindices.push(gindex);
+    }
+    return gindices;
+  }
 }
 
 export abstract class CompositeArrayType<T extends ArrayLike<unknown>> extends CompositeType<T> {
@@ -564,4 +583,31 @@ export abstract class CompositeArrayType<T extends ArrayLike<unknown>> extends C
       return [];
     }
   }
+
+  tree_getLeafGindices(target?: Tree, root: Gindex = BigInt(1)): Gindex[] {
+    // Underlying elements exist one per chunk
+    // Iterate through chunk gindices, recursively fetching leaf gindices from each chunk
+    const chunkCount = this.tree_getChunkCount(target);
+    const gindices: Gindex[] = [];
+    const startIndex = toGindex(this.getChunkDepth(), BigInt(0));
+    const extendedStartIndex = concatGindices([root, startIndex]);
+    if (this.elementType.hasVariableSerializedLength()) {
+      if (!target) {
+        throw new Error("variable type requires tree argument to get leaves");
+      }
+      // variable-length elements must pass the underlying subtrees to determine the length
+      for (
+        let i = 0, gindex = startIndex, extendedGindex = extendedStartIndex;
+        i < chunkCount;
+        i++, gindex++, extendedGindex++
+      ) {
+        gindices.push(...this.elementType.tree_getLeafGindices(target.getSubtree(gindex), extendedGindex));
+      }
+    } else {
+      for (let i = 0, gindex = extendedStartIndex; i < chunkCount; i++, gindex++) {
+        gindices.push(...this.elementType.tree_getLeafGindices(undefined, gindex));
+      }
+    }
+    return gindices;
+  }
 }

src/types/composite/container.ts

@@ -3,6 +3,7 @@ import {Json, ObjectLike} from "../../interface";
 import {CompositeType, isCompositeType} from "./abstract";
 import {IJsonOptions, isTypeOf, Type} from "../type";
 import {
+  concatGindices,
   Gindex,
   iterateAtDepth,
   LeafNode,
@@ -510,4 +511,25 @@ export class ContainerType<T extends ObjectLike = ObjectLike> extends CompositeT
   getMaxChunkCount(): number {
     return Object.keys(this.fields).length;
   }
+
+  tree_getLeafGindices(target?: Tree, root: Gindex = BigInt(1)): Gindex[] {
+    const gindices: Gindex[] = [];
+    for (const [fieldName, fieldType] of Object.entries(this.fields)) {
+      const fieldGindex = this.getPropertyGindex(fieldName);
+      const extendedFieldGindex = concatGindices([root, fieldGindex]);
+      if (!isCompositeType(fieldType)) {
+        gindices.push(extendedFieldGindex);
+      } else {
+        if (fieldType.hasVariableSerializedLength()) {
+          if (!target) {
+            throw new Error("variable type requires tree argument to get leaves");
+          }
+          gindices.push(...fieldType.tree_getLeafGindices(target.getSubtree(fieldGindex), extendedFieldGindex));
+        } else {
+          gindices.push(...fieldType.tree_getLeafGindices(undefined, extendedFieldGindex));
+        }
+      }
+    }
+    return gindices;
+  }
 }

src/types/composite/list.ts

@@ -4,9 +4,20 @@ import {IArrayOptions, BasicArrayType, CompositeArrayType} from "./array";
 import {isBasicType, number32Type} from "../basic";
 import {IJsonOptions, isTypeOf, Type} from "../type";
 import {mixInLength} from "../../util/compat";
-import {BranchNode, Node, Tree, zeroNode} from "@chainsafe/persistent-merkle-tree";
+import {BranchNode, concatGindices, Gindex, Node, Tree, zeroNode} from "@chainsafe/persistent-merkle-tree";
 import {isTreeBacked} from "../../backings/tree/treeValue";
 
+/**
+ * SSZ Lists (variable-length arrays) include the length of the list in the tree
+ * This length is always in the same index in the tree
+ * ```
+ *   1
+ *  / \
+ * 2   3 // <-here
+ * ```
+ */
+export const LENGTH_GINDEX = BigInt(3);
+
 export interface IListOptions extends IArrayOptions {
   limit: number;
 }
@@ -112,13 +123,13 @@ export class BasicListType<T extends List<unknown> = List<unknown>> extends Basi
   }
 
   tree_getLength(target: Tree): number {
-    return number32Type.struct_deserializeFromBytes(target.getRoot(BigInt(3)), 0);
+    return number32Type.struct_deserializeFromBytes(target.getRoot(LENGTH_GINDEX), 0);
   }
 
   tree_setLength(target: Tree, length: number): void {
     const chunk = new Uint8Array(32);
     number32Type.toBytes(length, chunk, 0);
-    target.setRoot(BigInt(3), chunk);
+    target.setRoot(LENGTH_GINDEX, chunk);
   }
 
   tree_deserializeFromBytes(data: Uint8Array, start: number, end: number): Tree {
@@ -195,6 +206,15 @@ export class BasicListType<T extends List<unknown> = List<unknown>> extends Basi
   getMaxChunkCount(): number {
     return Math.ceil((this.limit * this.elementType.size()) / 32);
   }
+  tree_getLeafGindices(target?: Tree, root: Gindex = BigInt(1)): Gindex[] {
+    if (!target) {
+      throw new Error("variable type requires tree argument to get leaves");
+    }
+    const gindices = super.tree_getLeafGindices(target, root);
+    // include the length chunk
+    gindices.push(concatGindices([root, LENGTH_GINDEX]));
+    return gindices;
+  }
 }
 
 export class CompositeListType<T extends List<object> = List<object>> extends CompositeArrayType<T> {
@@ -277,13 +297,13 @@ export class CompositeListType<T extends List<object> = List<object>> extends Co
   }
 
   tree_getLength(target: Tree): number {
-    return number32Type.struct_deserializeFromBytes(target.getRoot(BigInt(3)), 0);
+    return number32Type.struct_deserializeFromBytes(target.getRoot(LENGTH_GINDEX), 0);
   }
 
   tree_setLength(target: Tree, length: number): void {
     const chunk = new Uint8Array(32);
     number32Type.struct_serializeToBytes(length, chunk, 0);
-    target.setRoot(BigInt(3), chunk);
+    target.setRoot(LENGTH_GINDEX, chunk);
   }
 
   tree_deserializeFromBytes(data: Uint8Array, start: number, end: number): Tree {
@@ -376,4 +396,13 @@ export class CompositeListType<T extends List<object> = List<object>> extends Co
     this.tree_setLength(target, length - 1);
     return value as T[number];
   }
+  tree_getLeafGindices(target?: Tree, root: Gindex = BigInt(1)): Gindex[] {
+    if (!target) {
+      throw new Error("variable type requires tree argument to get leaves");
+    }
+    const gindices = super.tree_getLeafGindices(target, root);
+    // include the length chunk
+    gindices.push(concatGindices([root, LENGTH_GINDEX]));
+    return gindices;
+  }
 }

test/unit/proof.test.ts

@@ -0,0 +1,14 @@
+import { expect } from "chai";
+import { ArrayObject, SimpleObject } from "./objects";
+
+describe("create proof", () => {
+  it("should include all leaves of path to composite type", () => {
+    const arrayObj = ArrayObject.defaultTreeBacked();
+    const simpleObj = SimpleObject.defaultTreeBacked();
+    arrayObj.v.push(simpleObj);
+    const proof = arrayObj.createProof([["v"]]);
+    const arrayObj2 = ArrayObject.createTreeBackedFromProofUnsafe(proof);
+    expect(arrayObj2.v[0].valueOf()).to.deep.equal(arrayObj.v[0].valueOf());
+    expect(arrayObj2.v.valueOf()).to.deep.equal(arrayObj.v.valueOf());
+  });
+});