feat: compact multiproof (#292)

* feat: add dynamic multiproofs

* Fix logic and rename to compact proof

* Add perf tests

* Fix benchmark

* Update packages/persistent-merkle-tree/src/proof/compactMulti.ts

Co-authored-by: Lion - dapplion <35266934+dapplion@users.noreply.github.com>

* Fix linter errors

* PR review

Co-authored-by: Lion - dapplion <35266934+dapplion@users.noreply.github.com>

packages/persistent-merkle-tree/src/proof/compactMulti.ts

@@ -0,0 +1,151 @@
+import {convertGindexToBitstring, Gindex, GindexBitstring} from "../gindex";
+import {BranchNode, LeafNode, Node} from "../node";
+import {computeProofBitstrings} from "./util";
+
+export function computeDescriptor(indices: Gindex[]): Uint8Array {
+  // include all helper indices
+  const proofBitstrings = new Set<GindexBitstring>();
+  const pathBitstrings = new Set<GindexBitstring>();
+  for (const leafIndex of indices) {
+    const leafBitstring = convertGindexToBitstring(leafIndex);
+    proofBitstrings.add(leafBitstring);
+    const {branch, path} = computeProofBitstrings(leafBitstring);
+    path.delete(leafBitstring);
+    for (const pathIndex of path) {
+      pathBitstrings.add(pathIndex);
+    }
+    for (const branchIndex of branch) {
+      proofBitstrings.add(branchIndex);
+    }
+  }
+  for (const pathIndex of pathBitstrings) {
+    proofBitstrings.delete(pathIndex);
+  }
+
+  // sort gindex bitstrings in-order
+  const allBitstringsSorted = Array.from(proofBitstrings).sort((a, b) => a.localeCompare(b));
+
+  // convert gindex bitstrings into descriptor bitstring
+  let descriptorBitstring = "";
+  for (const gindexBitstring of allBitstringsSorted) {
+    for (let i = 0; i < gindexBitstring.length; i++) {
+      if (gindexBitstring[gindexBitstring.length - 1 - i] === "1") {
+        descriptorBitstring += "1".padStart(i + 1, "0");
+        break;
+      }
+    }
+  }
+
+  // append zero bits to byte-alignt
+  if (descriptorBitstring.length % 8 != 0) {
+    descriptorBitstring = descriptorBitstring.padEnd(
+      8 - (descriptorBitstring.length % 8) + descriptorBitstring.length,
+      "0"
+    );
+  }
+
+  // convert descriptor bitstring to bytes
+  const descriptor = new Uint8Array(descriptorBitstring.length / 8);
+  for (let i = 0; i < descriptor.length; i++) {
+    descriptor[i] = Number("0b" + descriptorBitstring.substring(i * 8, (i + 1) * 8));
+  }
+  return descriptor;
+}
+
+function getBit(bitlist: Uint8Array, bitIndex: number): boolean {
+  const bit = bitIndex % 8;
+  const byteIdx = Math.floor(bitIndex / 8);
+  const byte = bitlist[byteIdx];
+  switch (bit) {
+    case 0:
+      return (byte & 0b1000_0000) !== 0;
+    case 1:
+      return (byte & 0b0100_0000) !== 0;
+    case 2:
+      return (byte & 0b0010_0000) !== 0;
+    case 3:
+      return (byte & 0b0001_0000) !== 0;
+    case 4:
+      return (byte & 0b0000_1000) !== 0;
+    case 5:
+      return (byte & 0b0000_0100) !== 0;
+    case 6:
+      return (byte & 0b0000_0010) !== 0;
+    case 7:
+      return (byte & 0b0000_0001) !== 0;
+    default:
+      throw new Error("unreachable");
+  }
+}
+
+export function descriptorToBitlist(descriptor: Uint8Array): boolean[] {
+  const bools: boolean[] = [];
+  const maxBitLength = descriptor.length * 8;
+  let count0 = 0;
+  let count1 = 0;
+  for (let i = 0; i < maxBitLength; i++) {
+    const bit = getBit(descriptor, i);
+    bools.push(bit);
+    if (bit) {
+      count1++;
+    } else {
+      count0++;
+    }
+    if (count1 > count0) {
+      i++;
+      if (i + 7 < maxBitLength) {
+        throw new Error("Invalid descriptor: too many bytes");
+      }
+      for (; i < maxBitLength; i++) {
+        const bit = getBit(descriptor, i);
+        if (bit) {
+          throw new Error("Invalid descriptor: too many 1 bits");
+        }
+      }
+      return bools;
+    }
+  }
+  throw new Error("Invalid descriptor: not enough 1 bits");
+}
+
+export function nodeToCompactMultiProof(node: Node, bitlist: boolean[], bitIndex: number): Uint8Array[] {
+  if (bitlist[bitIndex]) {
+    return [node.root];
+  } else {
+    const left = nodeToCompactMultiProof(node.left, bitlist, bitIndex + 1);
+    const right = nodeToCompactMultiProof(node.right, bitlist, bitIndex + left.length * 2);
+    return [...left, ...right];
+  }
+}
+
+/**
+ * Create a Node given a validated bitlist, leaves, and a pointer into the bitlist and leaves
+ *
+ * Recursive definition
+ */
+export function compactMultiProofToNode(
+  bitlist: boolean[],
+  leaves: Uint8Array[],
+  pointer: {bitIndex: number; leafIndex: number}
+): Node {
+  if (bitlist[pointer.bitIndex++]) {
+    return LeafNode.fromRoot(leaves[pointer.leafIndex++]);
+  } else {
+    return new BranchNode(
+      compactMultiProofToNode(bitlist, leaves, pointer),
+      compactMultiProofToNode(bitlist, leaves, pointer)
+    );
+  }
+}
+
+export function createCompactMultiProof(rootNode: Node, descriptor: Uint8Array): Uint8Array[] {
+  return nodeToCompactMultiProof(rootNode, descriptorToBitlist(descriptor), 0);
+}
+
+export function createNodeFromCompactMultiProof(leaves: Uint8Array[], descriptor: Uint8Array): Node {
+  const bools = descriptorToBitlist(descriptor);
+  if (bools.length !== leaves.length * 2 - 1) {
+    throw new Error("Invalid multiproof: invalid number of leaves");
+  }
+  return compactMultiProofToNode(bools, leaves, {bitIndex: 0, leafIndex: 0});
+}

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

@@ -1,6 +1,7 @@
 import {Gindex} from "../gindex";
 import {Node} from "../node";
 import {createMultiProof, createNodeFromMultiProof} from "./multi";
+import {createNodeFromCompactMultiProof, createCompactMultiProof} from "./compactMulti";
 import {createNodeFromSingleProof, createSingleProof} from "./single";
 import {
   computeTreeOffsetProofSerializedLength,
@@ -10,10 +11,13 @@ import {
   serializeTreeOffsetProof,
 } from "./treeOffset";
 
+export {computeDescriptor, descriptorToBitlist} from "./compactMulti";
+
 export enum ProofType {
   single = "single",
   treeOffset = "treeOffset",
   multi = "multi",
+  compactMulti = "compactMulti",
 }
 
 /**
@@ -23,6 +27,7 @@ export const ProofTypeSerialized = [
   ProofType.single, // 0
   ProofType.treeOffset, // 1
   ProofType.multi, // 2
+  ProofType.compactMulti, // 3
 ];
 
 /**
@@ -58,7 +63,13 @@ export interface MultiProof {
   gindices: Gindex[];
 }
 
-export type Proof = SingleProof | TreeOffsetProof | MultiProof;
+export interface CompactMultiProof {
+  type: ProofType.compactMulti;
+  leaves: Uint8Array[];
+  descriptor: Uint8Array;
+}
+
+export type Proof = SingleProof | TreeOffsetProof | MultiProof | CompactMultiProof;
 
 export interface SingleProofInput {
   type: ProofType.single;
@@ -74,7 +85,12 @@ export interface MultiProofInput {
   gindices: Gindex[];
 }
 
-export type ProofInput = SingleProofInput | TreeOffsetProofInput | MultiProofInput;
+export interface CompactMultiProofInput {
+  type: ProofType.compactMulti;
+  descriptor: Uint8Array;
+}
+
+export type ProofInput = SingleProofInput | TreeOffsetProofInput | MultiProofInput | CompactMultiProofInput;
 
 export function createProof(rootNode: Node, input: ProofInput): Proof {
   switch (input.type) {
@@ -104,6 +120,14 @@ export function createProof(rootNode: Node, input: ProofInput): Proof {
         gindices,
       };
     }
+    case ProofType.compactMulti: {
+      const leaves = createCompactMultiProof(rootNode, input.descriptor);
+      return {
+        type: ProofType.compactMulti,
+        leaves,
+        descriptor: input.descriptor,
+      };
+    }
     default:
       throw new Error("Invalid proof type");
   }
@@ -117,6 +141,8 @@ export function createNodeFromProof(proof: Proof): Node {
       return createNodeFromTreeOffsetProof(proof.offsets, proof.leaves);
     case ProofType.multi:
       return createNodeFromMultiProof(proof.leaves, proof.witnesses, proof.gindices);
+    case ProofType.compactMulti:
+      return createNodeFromCompactMultiProof(proof.leaves, proof.descriptor);
     default:
       throw new Error("Invalid proof type");
   }

packages/persistent-merkle-tree/test/perf/proof.test.ts

@@ -0,0 +1,40 @@
+import {itBench} from "@dapplion/benchmark";
+import {computeDescriptor, createProof, ProofType} from "../../src/proof";
+import {createTree} from "../utils/tree";
+
+describe("Proofs", () => {
+  const depth = 15;
+  const tree = createTree(depth);
+  const maxNumLeaves = 10;
+  const allLeafIndices = Array.from(
+    {length: maxNumLeaves},
+    (_, i) => BigInt(2) ** BigInt(depth) + BigInt(i) ** BigInt(2)
+  );
+  for (let numLeaves = 1; numLeaves < 5; numLeaves++) {
+    const leafIndices = allLeafIndices.slice(0, numLeaves);
+
+    itBench({
+      id: `multiproof - depth ${depth}, ${numLeaves} requested leaves`,
+      fn: () => {
+        createProof(tree, {type: ProofType.multi, gindices: leafIndices});
+      },
+    });
+
+    itBench({
+      id: `tree offset multiproof - depth ${depth}, ${numLeaves} requested leaves`,
+      fn: () => {
+        createProof(tree, {type: ProofType.treeOffset, gindices: leafIndices});
+      },
+    });
+
+    itBench({
+      id: `compact multiproof - depth ${depth}, ${numLeaves} requested leaves`,
+      beforeEach: () => {
+        return computeDescriptor(leafIndices);
+      },
+      fn: (descriptor) => {
+        createProof(tree, {type: ProofType.compactMulti, descriptor});
+      },
+    });
+  }
+});

packages/persistent-merkle-tree/test/unit/proof/compactMulti.test.ts

@@ -0,0 +1,63 @@
+import {expect} from "chai";
+import {
+  createNodeFromCompactMultiProof,
+  createCompactMultiProof,
+  descriptorToBitlist,
+  computeDescriptor,
+} from "../../../src/proof/compactMulti";
+import {createTree} from "../../utils/tree";
+
+describe("CompactMultiProof", () => {
+  const descriptorTestCases = [
+    {
+      input: Uint8Array.from([0b1000_0000]),
+      output: [1].map(Boolean),
+    },
+    {
+      input: Uint8Array.from([0b0010_0101, 0b1110_0000]),
+      output: [0, 0, 1, 0, 0, 1, 0, 1, 1, 1, 1].map(Boolean),
+    },
+    {
+      input: Uint8Array.from([0b0101_0101, 0b1000_0000]),
+      output: [0, 1, 0, 1, 0, 1, 0, 1, 1].map(Boolean),
+    },
+    {
+      input: Uint8Array.from([0b0101_0110]),
+      output: [0, 1, 0, 1, 0, 1, 1].map(Boolean),
+    },
+  ];
+  describe("descriptorToBitlist", () => {
+    it("should convert valid descriptor to a bitlist", () => {
+      for (const {input, output} of descriptorTestCases) {
+        expect(descriptorToBitlist(input)).to.deep.equal(output);
+      }
+    });
+    it("should throw on invalid descriptors", () => {
+      const errorCases = [
+        Uint8Array.from([0b1000_0000, 0]),
+        Uint8Array.from([0b0000_0001, 0]),
+        Uint8Array.from([0b0101_0111]),
+        Uint8Array.from([0b0101_0110, 0]),
+      ];
+      for (const input of errorCases) {
+        expect(() => descriptorToBitlist(input)).to.throw();
+      }
+    });
+  });
+  describe("computeDescriptor", () => {
+    it("should convert gindices to a descriptor", () => {
+      const index = 42n;
+      const expected = Uint8Array.from([0x25, 0xe0]);
+      expect(computeDescriptor([index])).to.deep.equal(expected);
+    });
+  });
+
+  const tree = createTree(5);
+  it("should roundtrip node -> proof -> node", () => {
+    for (const {input} of descriptorTestCases) {
+      const proof = createCompactMultiProof(tree, input);
+      const newNode = createNodeFromCompactMultiProof(proof, input);
+      expect(newNode.root).to.deep.equal(tree.root);
+    }
+  });
+});

packages/persistent-merkle-tree/test/unit/proof/index.test.ts

@@ -1,15 +1,14 @@
 import {expect} from "chai";
 import {describe, it} from "mocha";
-import {createNodeFromProof, createProof, deserializeProof, ProofType, serializeProof} from "../../../src/proof";
-import {Node, LeafNode, BranchNode} from "../../../src/node";
-
-// Create a tree with leaves of different values
-function createTree(depth: number, index = 0): Node {
-  if (!depth) {
-    return LeafNode.fromRoot(Buffer.alloc(32, index));
-  }
-  return new BranchNode(createTree(depth - 1, 2 ** depth + index), createTree(depth - 1, 2 ** depth + index + 1));
-}
+import {
+  computeDescriptor,
+  createNodeFromProof,
+  createProof,
+  deserializeProof,
+  ProofType,
+  serializeProof,
+} from "../../../src/proof";
+import {createTree} from "../../utils/tree";
 
 describe("proof equivalence", () => {
   it("should compute the same root from different proof types - single leaf", () => {
@@ -19,9 +18,14 @@ describe("proof equivalence", () => {
       const singleProof = createProof(node, {type: ProofType.single, gindex});
       const treeOffsetProof = createProof(node, {type: ProofType.treeOffset, gindices: [gindex]});
       const multiProof = createProof(node, {type: ProofType.multi, gindices: [gindex]});
+      const compactMultiProof = createProof(node, {
+        type: ProofType.compactMulti,
+        descriptor: computeDescriptor([gindex]),
+      });
       expect(node.root).to.deep.equal(createNodeFromProof(singleProof).root);
       expect(node.root).to.deep.equal(createNodeFromProof(treeOffsetProof).root);
       expect(node.root).to.deep.equal(createNodeFromProof(multiProof).root);
+      expect(node.root).to.deep.equal(createNodeFromProof(compactMultiProof).root);
     }
   });
   it("should compute the same root from different proof types - multiple leaves", function () {
@@ -43,9 +47,14 @@ describe("proof equivalence", () => {
 
               const treeOffsetProof = createProof(node, {type: ProofType.treeOffset, gindices});
               const multiProof = createProof(node, {type: ProofType.multi, gindices});
+              const compactMultiProof = createProof(node, {
+                type: ProofType.compactMulti,
+                descriptor: computeDescriptor(gindices),
+              });
 
               expect(node.root).to.deep.equal(createNodeFromProof(treeOffsetProof).root);
               expect(node.root).to.deep.equal(createNodeFromProof(multiProof).root);
+              expect(node.root).to.deep.equal(createNodeFromProof(compactMultiProof).root);
             }
           }
         }