From 5f1ea9914d0796cf0c9a8c2f9622fc5e459a12f2 Mon Sep 17 00:00:00 2001 From: Cayman Date: Fri, 6 Jan 2023 15:50:17 -0600 Subject: [PATCH] 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> --- .../src/proof/compactMulti.ts | 151 ++++++++++++++++++ .../persistent-merkle-tree/src/proof/index.ts | 30 +++- .../test/perf/proof.test.ts | 40 +++++ .../test/unit/proof/compactMulti.test.ts | 63 ++++++++ .../test/unit/proof/index.test.ts | 29 ++-- .../persistent-merkle-tree/test/utils/tree.ts | 8 + 6 files changed, 309 insertions(+), 12 deletions(-) create mode 100644 packages/persistent-merkle-tree/src/proof/compactMulti.ts create mode 100644 packages/persistent-merkle-tree/test/perf/proof.test.ts create mode 100644 packages/persistent-merkle-tree/test/unit/proof/compactMulti.test.ts create mode 100644 packages/persistent-merkle-tree/test/utils/tree.ts diff --git a/packages/persistent-merkle-tree/src/proof/compactMulti.ts b/packages/persistent-merkle-tree/src/proof/compactMulti.ts new file mode 100644 index 00000000..4be7ad7d --- /dev/null +++ b/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(); + const pathBitstrings = new Set(); + 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}); +} diff --git a/packages/persistent-merkle-tree/src/proof/index.ts b/packages/persistent-merkle-tree/src/proof/index.ts index 4523c272..90118ba2 100644 --- a/packages/persistent-merkle-tree/src/proof/index.ts +++ b/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"); } diff --git a/packages/persistent-merkle-tree/test/perf/proof.test.ts b/packages/persistent-merkle-tree/test/perf/proof.test.ts new file mode 100644 index 00000000..bc52f84f --- /dev/null +++ b/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}); + }, + }); + } +}); diff --git a/packages/persistent-merkle-tree/test/unit/proof/compactMulti.test.ts b/packages/persistent-merkle-tree/test/unit/proof/compactMulti.test.ts new file mode 100644 index 00000000..be317414 --- /dev/null +++ b/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); + } + }); +}); diff --git a/packages/persistent-merkle-tree/test/unit/proof/index.test.ts b/packages/persistent-merkle-tree/test/unit/proof/index.test.ts index d6b15fce..2a7da945 100644 --- a/packages/persistent-merkle-tree/test/unit/proof/index.test.ts +++ b/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); } } } diff --git a/packages/persistent-merkle-tree/test/utils/tree.ts b/packages/persistent-merkle-tree/test/utils/tree.ts new file mode 100644 index 00000000..2ecec3e8 --- /dev/null +++ b/packages/persistent-merkle-tree/test/utils/tree.ts @@ -0,0 +1,8 @@ +import {BranchNode, LeafNode, Node} from "../../src/node"; + +export 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)); +}