@proto-kit/common 0.1.1-develop.211 → 0.1.1-develop.2137

package/src/trees/lmt/InMemoryLinkedLeafStore.ts

@@ -0,0 +1,42 @@
+import { LinkedLeafStore, LinkedLeaf } from "./LinkedLeafStore";
+
+export class InMemoryLinkedLeafStore implements LinkedLeafStore {
+  public leaves: {
+    [key: string]: { leaf: LinkedLeaf; index: bigint };
+  } = {};
+
+  public maximumIndex?: bigint;
+
+  public getLeaf(
+    path: bigint
+  ): { leaf: LinkedLeaf; index: bigint } | undefined {
+    return this.leaves[path.toString()];
+  }
+
+  public setLeaf(index: bigint, value: LinkedLeaf): void {
+    const leaf = this.getLeaf(value.path);
+
+    if (leaf !== undefined && leaf?.index !== index) {
+      throw new Error("Cannot change index of existing leaf");
+    }
+
+    this.leaves[value.path.toString()] = { leaf: value, index: index };
+    if (index > (this.maximumIndex ?? -1)) {
+      this.maximumIndex = index;
+    }
+  }
+
+  public getMaximumIndex(): bigint | undefined {
+    return this.maximumIndex;
+  }
+
+  // This gets the leaf with the closest path.
+  public getPreviousLeaf(
+    path: bigint
+  ): { leaf: LinkedLeaf; index: bigint } | undefined {
+    return Object.values(this.leaves).find(
+      (storedLeaf) =>
+        storedLeaf.leaf.nextPath > path && storedLeaf.leaf.path < path
+    );
+  }
+}

package/src/trees/lmt/LinkedLeafStore.ts

@@ -0,0 +1,13 @@
+export type LinkedLeaf = { value: bigint; path: bigint; nextPath: bigint };
+
+export type StoredLeaf = { leaf: LinkedLeaf; index: bigint };
+
+export interface LinkedLeafStore {
+  setLeaf: (index: bigint, value: LinkedLeaf) => void;
+
+  getLeaf: (path: bigint) => StoredLeaf | undefined;
+
+  getPreviousLeaf: (path: bigint) => StoredLeaf | undefined;
+
+  getMaximumIndex: () => bigint | undefined;
+}

package/src/trees/lmt/LinkedMerkleTree.ts

@@ -0,0 +1,335 @@
+import { Field, Struct } from "o1js";
+
+import { createMerkleTree, RollupMerkleTree } from "../sparse/RollupMerkleTree";
+import { MerkleTreeStore } from "../sparse/MerkleTreeStore";
+import { InMemoryMerkleTreeStorage } from "../sparse/InMemoryMerkleTreeStorage";
+
+import { InMemoryLinkedLeafStore } from "./InMemoryLinkedLeafStore";
+import { LinkedLeaf, LinkedLeafStore } from "./LinkedLeafStore";
+import { LinkedLeafStruct } from "./LinkedMerkleTreeTypes";
+import {
+  AbstractLinkedMerkleTree,
+  AbstractLinkedMerkleTreeClass,
+} from "./AbstractLinkedMerkleTree";
+
+type LeafOperationInstruction = {
+  witness: bigint;
+  witnessLeaf: LinkedLeafStruct;
+  write: { index: bigint; leaf: Field };
+};
+
+type SetLeafMetadata = {
+  leafPrevious: LeafOperationInstruction | "dummy";
+  leafCurrent: LeafOperationInstruction;
+};
+
+export function createLinkedMerkleTree(
+  height: number
+): AbstractLinkedMerkleTreeClass {
+  const SparseTreeClass = createMerkleTree(height);
+
+  class LinkedLeafAndMerkleWitness extends Struct({
+    leaf: LinkedLeafStruct,
+    merkleWitness: SparseTreeClass.WITNESS,
+  }) {
+    public checkMembership(root: Field, path: Field, value: Field) {
+      const pathEquals = path.equals(this.leaf.path);
+
+      return this.merkleWitness
+        .calculateRoot(
+          new LinkedLeafStruct({
+            ...this.leaf,
+            value,
+          }).hash()
+        )
+        .equals(root)
+        .and(pathEquals);
+    }
+  }
+
+  class LinkedOperationWitness extends Struct({
+    leafPrevious: LinkedLeafAndMerkleWitness,
+    leafCurrent: LinkedLeafAndMerkleWitness,
+  }) {
+    // implements LinkedStructTemplate
+    public static fromReadWitness(readWitness: LinkedLeafAndMerkleWitness) {
+      return new LinkedOperationWitness({
+        leafPrevious: new LinkedLeafAndMerkleWitness({
+          merkleWitness: SparseTreeClass.WITNESS.dummy(),
+          leaf: LinkedLeafStruct.dummy(),
+        }),
+        leafCurrent: readWitness,
+      });
+    }
+  }
+
+  return class AbstractLinkedRollupMerkleTree
+    implements AbstractLinkedMerkleTree
+  {
+    public static HEIGHT = height;
+
+    public static EMPTY_ROOT = new AbstractLinkedRollupMerkleTree(
+      new InMemoryMerkleTreeStorage(),
+      new InMemoryLinkedLeafStore()
+    ).getRoot();
+
+    public static READ_WITNESS = LinkedLeafAndMerkleWitness;
+
+    public static WITNESS = LinkedOperationWitness;
+
+    readonly tree: RollupMerkleTree;
+
+    readonly leafStore: LinkedLeafStore;
+
+    public constructor(store: MerkleTreeStore, leafStore: LinkedLeafStore) {
+      this.leafStore = leafStore;
+
+      this.tree = new SparseTreeClass(store);
+
+      // We only do the leaf initialisation when the store
+      // has no values. Otherwise, we leave the store
+      // as is to not overwrite any data.
+      if (this.leafStore.getLeaf(0n) === undefined) {
+        this.setLeafInitialisation();
+      }
+    }
+
+    /**
+     * Returns leaf which lives at a given path.
+     * Errors if the path is not defined.
+     * @param path path of the node.
+     * @returns The data of the node.
+     */
+    public getLeaf(path: bigint): LinkedLeafStruct | undefined {
+      const storedLeaf = this.leafStore.getLeaf(path);
+      if (storedLeaf === undefined) {
+        return undefined;
+      }
+      return new LinkedLeafStruct({
+        value: Field(storedLeaf.leaf.value),
+        path: Field(storedLeaf.leaf.path),
+        nextPath: Field(storedLeaf.leaf.nextPath),
+      });
+    }
+
+    /**
+     * Returns the root of the [Merkle Tree](https://en.wikipedia.org/wiki/Merkle_tree).
+     * @returns The root of the Merkle Tree.
+     */
+    public getRoot(): Field {
+      return this.tree.getRoot().toConstant();
+    }
+
+    private writeLeaf(index: bigint, leaf: LinkedLeaf) {
+      this.leafStore.setLeaf(index, leaf);
+
+      const leafHash = new LinkedLeafStruct(
+        LinkedLeafStruct.fromValue(leaf)
+      ).hash();
+
+      return {
+        index,
+        leaf: leafHash,
+      };
+    }
+
+    private setLeafInternal(path: bigint, value: bigint): SetLeafMetadata {
+      const storedLeaf = this.leafStore.getLeaf(path);
+
+      if (storedLeaf === undefined) {
+        // Insert case
+        // The above means the path doesn't already exist, and we are inserting, not updating.
+        // This requires us to update the node with the previous path, as well.
+        const tempIndex = this.leafStore.getMaximumIndex();
+        if (tempIndex === undefined) {
+          throw Error("Store Max Index not defined");
+        }
+        if (tempIndex + 1n >= 2 ** height) {
+          throw new Error("Index greater than maximum leaf number");
+        }
+        const nextFreeIndex = tempIndex + 1n;
+
+        const previousLeaf = this.leafStore.getPreviousLeaf(path);
+
+        if (previousLeaf === undefined) {
+          throw Error(`Prev leaf shouldn't be undefined (path ${path})`);
+        }
+
+        const newPrevLeaf = {
+          ...previousLeaf.leaf,
+          nextPath: path,
+        };
+        const treeWrite1 = this.writeLeaf(previousLeaf.index, newPrevLeaf);
+
+        const newLeaf = {
+          path,
+          value,
+          nextPath: previousLeaf.leaf.nextPath,
+        };
+        const treeWrite2 = this.writeLeaf(nextFreeIndex, newLeaf);
+
+        return {
+          leafPrevious: {
+            witness: previousLeaf.index,
+            witnessLeaf: new LinkedLeafStruct(
+              LinkedLeafStruct.fromValue(previousLeaf.leaf)
+            ),
+            write: treeWrite1,
+          },
+          leafCurrent: {
+            witness: nextFreeIndex,
+            witnessLeaf: LinkedLeafStruct.dummy(),
+            write: treeWrite2,
+          },
+        };
+      } else {
+        // Update case
+        const updatedLeaf = {
+          ...storedLeaf.leaf,
+          value: value,
+        };
+
+        const treeWrite = this.writeLeaf(storedLeaf.index, updatedLeaf);
+
+        return {
+          leafPrevious: "dummy",
+          leafCurrent: {
+            witness: storedLeaf.index,
+            witnessLeaf: new LinkedLeafStruct(
+              LinkedLeafStruct.fromValue(storedLeaf.leaf)
+            ),
+            write: treeWrite,
+          },
+        };
+      }
+    }
+
+    private applyOperationInstruction(
+      instruction: LeafOperationInstruction | "dummy"
+    ): LinkedLeafAndMerkleWitness {
+      if (instruction === "dummy") {
+        return AbstractLinkedRollupMerkleTree.dummyReadWitness();
+      }
+
+      const merkleWitness = this.tree.getWitness(instruction.witness);
+
+      this.tree.setLeaf(instruction.write.index, instruction.write.leaf);
+
+      return new LinkedLeafAndMerkleWitness({
+        merkleWitness,
+        leaf: instruction.witnessLeaf,
+      });
+    }
+
+    /**
+     * Sets the value of a node at a given index to a given value.
+     * @param path Position of the leaf node.
+     * @param value New value.
+     */
+    public setLeaf(path: bigint, value: bigint): LinkedOperationWitness {
+      const {
+        leafPrevious: previousLeafInstruction,
+        leafCurrent: currentLeafInstruction,
+      } = this.setLeafInternal(path, value);
+
+      const leafPrevious = this.applyOperationInstruction(
+        previousLeafInstruction
+      );
+      const leafCurrent = this.applyOperationInstruction(
+        currentLeafInstruction
+      );
+
+      return { leafPrevious, leafCurrent };
+    }
+
+    public setLeaves(batch: { path: bigint; value: bigint }[]) {
+      if (batch.length > 0) {
+        const witnesses = batch.map(({ path, value }) =>
+          this.setLeafInternal(path, value)
+        );
+
+        // tree.setLeafBatch internally takes care of making the writes unique to optimize
+        this.tree.setLeaves(
+          witnesses.flatMap(({ leafPrevious, leafCurrent }) =>
+            (leafPrevious === "dummy" ? [] : [leafPrevious.write]).concat(
+              leafCurrent.write
+            )
+          )
+        );
+      }
+    }
+
+    /**
+     * Sets the value of a leaf node at initialisation,
+     * i.e.  {vale: 0, path: 0, nextPath: Field.Max}
+     */
+    private setLeafInitialisation() {
+      // This is the maximum value of the hash
+      const MAX_FIELD_VALUE: bigint = Field.ORDER - 1n;
+      const zeroLeaf = {
+        value: 0n,
+        path: 0n,
+        nextPath: MAX_FIELD_VALUE,
+      };
+      this.leafStore.setLeaf(0n, zeroLeaf);
+      // We now set the leafs in the merkle tree to cascade the values up
+      // the tree.
+      this.tree.setLeaf(
+        0n,
+        new LinkedLeafStruct(LinkedLeafStruct.fromValue(zeroLeaf)).hash()
+      );
+    }
+
+    /**
+     * Returns the witness (also known as
+     * [Merkle Proof or Merkle Witness](https://computersciencewiki.org/index.php/Merkle_proof))
+     * for the leaf at the given path, otherwise returns a witness for the first unused index.
+     * @param path of the leaf node.
+     * @returns The witness that belongs to the leaf.
+     */
+    public getReadWitness(path: bigint): LinkedLeafAndMerkleWitness {
+      const storedLeaf = this.leafStore.getLeaf(path);
+      let leaf;
+      let currentIndex: bigint;
+
+      if (storedLeaf === undefined) {
+        const storeIndex = this.leafStore.getMaximumIndex();
+        if (storeIndex === undefined) {
+          throw new Error("Store undefined");
+        }
+        currentIndex = storeIndex + 1n;
+        leaf = LinkedLeafStruct.dummy();
+      } else {
+        leaf = new LinkedLeafStruct(
+          LinkedLeafStruct.fromValue(storedLeaf.leaf)
+        );
+        currentIndex = storedLeaf.index;
+      }
+
+      const merkleWitness = this.tree.getWitness(currentIndex);
+
+      return new LinkedLeafAndMerkleWitness({
+        merkleWitness,
+        leaf,
+      });
+    }
+
+    public static dummyReadWitness(): LinkedLeafAndMerkleWitness {
+      return new LinkedLeafAndMerkleWitness({
+        merkleWitness: SparseTreeClass.WITNESS.dummy(),
+        leaf: LinkedLeafStruct.dummy(),
+      });
+    }
+
+    public static dummyWitness() {
+      return new LinkedOperationWitness({
+        leafPrevious: AbstractLinkedRollupMerkleTree.dummyReadWitness(),
+        leafCurrent: AbstractLinkedRollupMerkleTree.dummyReadWitness(),
+      });
+    }
+  };
+}
+
+export class LinkedMerkleTree extends createLinkedMerkleTree(40) {}
+export class LinkedMerkleTreeWitness extends LinkedMerkleTree.WITNESS {}
+export class LinkedMerkleTreeReadWitness extends LinkedMerkleTree.READ_WITNESS {}

package/src/trees/lmt/LinkedMerkleTreeCircuitOps.ts

@@ -0,0 +1,188 @@
+import { Bool, Field, Provable, Struct } from "o1js";
+
+import { LinkedMerkleTreeWitness } from "./LinkedMerkleTree";
+import { LinkedLeafStruct } from "./LinkedMerkleTreeTypes";
+
+/* eslint-disable no-inner-declarations */
+// TODO Add a struct that captures the errors monad-style
+
+export type TreeWrite = {
+  path: Field;
+  from: Field;
+  to: Field;
+};
+
+export namespace LinkedMerkleTreeCircuitOps {
+  function boolAllTrue(...args: Bool[]): Bool {
+    return args.reduce((a, b, i) => {
+      // if (!b.toBoolean()) {
+      //   console.log();
+      // }
+      return a.and(b);
+    });
+  }
+
+  export class ComputeRootInstruction extends Struct({
+    newPreviousLeaf: LinkedLeafStruct,
+    newCurrentLeaf: LinkedLeafStruct,
+    allChecksMet: Bool,
+  }) {}
+
+  function chooseInstruction(
+    isUpdate: Bool,
+    updateInstruction: ComputeRootInstruction,
+    insertInstruction: ComputeRootInstruction
+  ): ComputeRootInstruction {
+    return Provable.if(
+      isUpdate,
+      ComputeRootInstruction,
+      updateInstruction,
+      insertInstruction
+    );
+  }
+
+  /**
+   * verifyWitness(current.merkleWitness, current.leaf)
+   *
+   * st.path == current.leaf.path
+   * st.from.value == current.leaf.value
+   * ```
+   *
+   * updates:
+   * ```
+   * current := { current.path, current.nextPath, value: st.to.value }
+   */
+  function update(
+    { leafCurrent, leafPrevious }: LinkedMerkleTreeWitness,
+    { to, from, path }: TreeWrite
+  ): ComputeRootInstruction {
+    const allChecksMet = boolAllTrue(
+      path.equals(leafCurrent.leaf.path),
+      leafCurrent.leaf.value.equals(from)
+    );
+
+    return {
+      newPreviousLeaf: leafPrevious.leaf,
+      newCurrentLeaf: new LinkedLeafStruct({
+        ...leafCurrent.leaf,
+        value: to,
+      }),
+      allChecksMet,
+    };
+  }
+
+  /**
+   * st.path == current.leaf.path
+   *
+   * previous.leaf.nextPath > current.leaf.path
+   * previous.leaf.path < current.leaf.path
+   * previous.leaf.nextPath == current.leaf.nextPath
+   *
+   * index(current.merkleWitness) == nextFreeIndex
+   *
+   * updates:
+   * previous := { previous.path, previous.value, nextPath: current.path }
+   * current := current.leaf
+   */
+  function insert(
+    witness: LinkedMerkleTreeWitness,
+    { path, to }: TreeWrite
+  ): ComputeRootInstruction {
+    const { leafPrevious: previous, leafCurrent: current } = witness;
+
+    const allChecksMet = boolAllTrue(
+      // Already covered in general checks
+      // path.equals(current.leaf.path),
+      current.leaf.isDummy(),
+      previous.leaf.nextPath.greaterThan(path),
+      previous.leaf.path.lessThan(path)
+    );
+
+    return {
+      newPreviousLeaf: new LinkedLeafStruct({
+        ...previous.leaf,
+        nextPath: path,
+      }),
+      newCurrentLeaf: new LinkedLeafStruct({
+        path,
+        value: to,
+        nextPath: previous.leaf.nextPath,
+      }),
+      allChecksMet,
+    };
+  }
+
+  function computeRoot(
+    witness: LinkedMerkleTreeWitness,
+    newPreviousLeaf: LinkedLeafStruct,
+    newCurrentLeaf: LinkedLeafStruct,
+    isUpdate: Bool,
+    isDummy: Bool,
+    root: Field
+  ) {
+    const { leafPrevious, leafCurrent } = witness;
+
+    leafPrevious.merkleWitness
+      .calculateRoot(leafPrevious.leaf.hash())
+      .equals(root)
+      .or(isUpdate)
+      .assertTrue("Previous leaf calculation not matching");
+
+    const root1 = leafPrevious.merkleWitness.calculateRoot(
+      newPreviousLeaf.hash()
+    );
+
+    const intermediateRoot = Provable.if(isUpdate, root, root1);
+
+    // TODO Make this Provable.if more efficient
+    const leafCurrentLeaf = Provable.if(
+      isUpdate,
+      leafCurrent.leaf.hash(),
+      Field(0)
+    );
+    leafCurrent.merkleWitness
+      .calculateRoot(leafCurrentLeaf)
+      .equals(intermediateRoot)
+      .or(isDummy)
+      .assertTrue("Current leaf witness invalid");
+
+    return leafCurrent.merkleWitness.calculateRoot(newCurrentLeaf.hash());
+  }
+
+  export function applyTreeWrite(
+    root: Field,
+    witness: LinkedMerkleTreeWitness,
+    treeWrite: TreeWrite,
+    index: number
+  ): Field {
+    const { leafPrevious, leafCurrent } = witness;
+
+    const isUpdate = leafPrevious.leaf.isDummy();
+    const isDummy = leafCurrent.leaf.isDummy().and(isUpdate);
+
+    // For read-only and update
+    const updateState = update(witness, treeWrite);
+
+    // For insert
+    const insertState = insert(witness, treeWrite);
+
+    const instruction = chooseInstruction(isUpdate, updateState, insertState);
+
+    instruction.allChecksMet
+      .or(isDummy)
+      .assertTrue(`Not all witness checks have been met: ${index}`);
+
+    const newRoot = computeRoot(
+      witness,
+      instruction.newPreviousLeaf,
+      instruction.newCurrentLeaf,
+      isUpdate,
+      isDummy,
+      root
+    );
+
+    return Provable.if(isDummy, root, newRoot);
+  }
+}
+
+/* eslint-enable no-inner-declarations */

package/src/trees/lmt/LinkedMerkleTreeTypes.ts

@@ -0,0 +1,53 @@
+import { Field, Poseidon, Provable, Struct } from "o1js";
+
+export class LinkedLeafStruct extends Struct({
+  value: Field,
+  path: Field,
+  nextPath: Field,
+}) {
+  public isDummy() {
+    return this.path.equals(0).and(this.nextPath.equals(0));
+  }
+
+  public hash(): Field {
+    const hash = Poseidon.hash(LinkedLeafStruct.toFields(this));
+    return Provable.if(this.isDummy(), Field(0), hash);
+  }
+
+  public static dummy(): LinkedLeafStruct {
+    return new LinkedLeafStruct({
+      value: Field(0),
+      path: Field(0),
+      nextPath: Field(0),
+    });
+  }
+}
+
+// export class LinkedMerkleTreeGlobalState extends Struct({
+//   root: Field,
+//   lastOccupiedIndex: Field,
+// }) {
+//   public static equals(
+//     state1: LinkedMerkleTreeGlobalState,
+//     state2: LinkedMerkleTreeGlobalState
+//   ): Bool {
+//     return state1.root
+//       .equals(state2.root)
+//       .and(state1.lastOccupiedIndex.equals(state2.lastOccupiedIndex));
+//   }
+//
+//   public static assertEquals(
+//     state1: LinkedMerkleTreeGlobalState,
+//     state2: LinkedMerkleTreeGlobalState,
+//     msg?: string
+//   ) {
+//     state1.root.assertEquals(
+//       state2.root,
+//       msg !== undefined ? `${msg}: root` : msg
+//     );
+//     state1.lastOccupiedIndex.assertEquals(
+//       state2.lastOccupiedIndex,
+//       msg !== undefined ? `${msg}: lastOccupiedIndex` : msg
+//     );
+//   }
+// }

package/src/trees/sparse/InMemoryMerkleTreeStorage.ts

@@ -0,0 +1,17 @@
+import { MerkleTreeStore } from "./MerkleTreeStore";
+
+export class InMemoryMerkleTreeStorage implements MerkleTreeStore {
+  public nodes: {
+    [key: number]: {
+      [key: string]: bigint;
+    };
+  } = {};
+
+  public getNode(key: bigint, level: number): bigint | undefined {
+    return this.nodes[level]?.[key.toString()];
+  }
+
+  public setNode(key: bigint, level: number, value: bigint): void {
+    (this.nodes[level] ??= {})[key.toString()] = value;
+  }
+}

package/src/trees/sparse/MerkleTreeStore.ts

@@ -0,0 +1,5 @@
+export interface MerkleTreeStore {
+  setNode: (key: bigint, level: number, value: bigint) => void;
+
+  getNode: (key: bigint, level: number) => bigint | undefined;
+}

package/src/trees/sparse/MockAsyncMerkleStore.ts

@@ -0,0 +1,30 @@
+import { noop } from "../../utils";
+
+import { InMemoryMerkleTreeStorage } from "./InMemoryMerkleTreeStorage";
+
+export class MockAsyncMerkleTreeStore {
+  public readonly store = new InMemoryMerkleTreeStorage();
+
+  public commit(): void {
+    noop();
+  }
+
+  public openTransaction(): void {
+    noop();
+  }
+
+  public async getNodeAsync(
+    key: bigint,
+    level: number
+  ): Promise<bigint | undefined> {
+    return this.store.getNode(key, level);
+  }
+
+  public async setNodeAsync(
+    key: bigint,
+    level: number,
+    value: bigint
+  ): Promise<void> {
+    this.store.setNode(key, level, value);
+  }
+}