@aztec/merkle-tree 0.23.0 → 0.24.0

package/src/snapshots/base_full_snapshot.ts

@@ -0,0 +1,215 @@
+import { SiblingPath } from '@aztec/circuit-types';
+import { AztecKVStore, AztecMap } from '@aztec/kv-store';
+
+import { TreeBase } from '../tree_base.js';
+import { TreeSnapshot, TreeSnapshotBuilder } from './snapshot_builder.js';
+
+/**
+ * Metadata for a snapshot, per block
+ */
+type SnapshotMetadata = {
+  /** The tree root at the time */
+  root: Buffer;
+  /** The number of filled leaves */
+  numLeaves: bigint;
+};
+
+/**
+ * Builds a full snapshot of a tree. This implementation works for any Merkle tree and stores
+ * it in a database in a similar way to how a tree is stored in memory, using pointers.
+ *
+ * Sharing the same database between versions and trees is recommended as the trees would share
+ * structure.
+ *
+ * Implement the protected method `handleLeaf` to store any additional data you need for each leaf.
+ *
+ * Complexity:
+ * N - count of non-zero nodes in tree
+ * M - count of snapshots
+ * H - tree height
+ * Worst case space complexity: O(N * M)
+ * Sibling path access: O(H) database reads
+ */
+export abstract class BaseFullTreeSnapshotBuilder<T extends TreeBase, S extends TreeSnapshot>
+  implements TreeSnapshotBuilder<S>
+{
+  protected nodes: AztecMap<string, [Buffer, Buffer]>;
+  protected snapshotMetadata: AztecMap<number, SnapshotMetadata>;
+
+  constructor(protected db: AztecKVStore, protected tree: T) {
+    this.nodes = db.openMap(`full_snapshot:${tree.getName()}:node`);
+    this.snapshotMetadata = db.openMap(`full_snapshot:${tree.getName()}:metadata`);
+  }
+
+  snapshot(block: number): Promise<S> {
+    return this.db.transaction(() => {
+      const snapshotMetadata = this.#getSnapshotMeta(block);
+
+      if (snapshotMetadata) {
+        return this.openSnapshot(snapshotMetadata.root, snapshotMetadata.numLeaves);
+      }
+
+      const root = this.tree.getRoot(false);
+      const numLeaves = this.tree.getNumLeaves(false);
+      const depth = this.tree.getDepth();
+      const queue: [Buffer, number, bigint][] = [[root, 0, 0n]];
+
+      // walk the tree breadth-first and store each of its nodes in the database
+      // for each node we save two keys
+      //   <node hash>:0 -> <left child's hash>
+      //   <node hash>:1 -> <right child's hash>
+      while (queue.length > 0) {
+        const [node, level, i] = queue.shift()!;
+        const nodeKey = node.toString('hex');
+        // check if the database already has a child for this tree
+        // if it does, then we know we've seen the whole subtree below it before
+        // and we don't have to traverse it anymore
+        // we use the left child here, but it could be anything that shows we've stored the node before
+        if (this.nodes.has(nodeKey)) {
+          continue;
+        }
+
+        if (level + 1 > depth) {
+          // short circuit if we've reached the leaf level
+          // otherwise getNode might throw if we ask for the children of a leaf
+          this.handleLeaf(i, node);
+          continue;
+        }
+
+        const [lhs, rhs] = [this.tree.getNode(level + 1, 2n * i), this.tree.getNode(level + 1, 2n * i + 1n)];
+
+        // we want the zero hash at the children's level, not the node's level
+        const zeroHash = this.tree.getZeroHash(level + 1);
+
+        void this.nodes.set(nodeKey, [lhs ?? zeroHash, rhs ?? zeroHash]);
+        // enqueue the children only if they're not zero hashes
+        if (lhs) {
+          queue.push([lhs, level + 1, 2n * i]);
+        }
+
+        if (rhs) {
+          queue.push([rhs, level + 1, 2n * i + 1n]);
+        }
+      }
+
+      void this.snapshotMetadata.set(block, { root, numLeaves });
+      return this.openSnapshot(root, numLeaves);
+    });
+  }
+
+  protected handleLeaf(_index: bigint, _node: Buffer): void {}
+
+  getSnapshot(version: number): Promise<S> {
+    const snapshotMetadata = this.#getSnapshotMeta(version);
+
+    if (!snapshotMetadata) {
+      return Promise.reject(new Error(`Version ${version} does not exist for tree ${this.tree.getName()}`));
+    }
+
+    return Promise.resolve(this.openSnapshot(snapshotMetadata.root, snapshotMetadata.numLeaves));
+  }
+
+  protected abstract openSnapshot(root: Buffer, numLeaves: bigint): S;
+
+  #getSnapshotMeta(block: number): SnapshotMetadata | undefined {
+    return this.snapshotMetadata.get(block);
+  }
+}
+
+/**
+ * A source of sibling paths from a snapshot tree
+ */
+export class BaseFullTreeSnapshot implements TreeSnapshot {
+  constructor(
+    protected db: AztecMap<string, [Buffer, Buffer]>,
+    protected historicRoot: Buffer,
+    protected numLeaves: bigint,
+    protected tree: TreeBase,
+  ) {}
+
+  getSiblingPath<N extends number>(index: bigint): SiblingPath<N> {
+    const siblings: Buffer[] = [];
+
+    for (const [_node, sibling] of this.pathFromRootToLeaf(index)) {
+      siblings.push(sibling);
+    }
+
+    // we got the siblings we were looking for, but they are in root-leaf order
+    // reverse them here so we have leaf-root (what SiblingPath expects)
+    siblings.reverse();
+
+    return new SiblingPath<N>(this.tree.getDepth() as N, siblings);
+  }
+
+  getLeafValue(index: bigint): Buffer | undefined {
+    let leafNode: Buffer | undefined = undefined;
+    for (const [node, _sibling] of this.pathFromRootToLeaf(index)) {
+      leafNode = node;
+    }
+
+    return leafNode;
+  }
+
+  getDepth(): number {
+    return this.tree.getDepth();
+  }
+
+  getRoot(): Buffer {
+    return this.historicRoot;
+  }
+
+  getNumLeaves(): bigint {
+    return this.numLeaves;
+  }
+
+  protected *pathFromRootToLeaf(leafIndex: bigint) {
+    const root = this.historicRoot;
+    const pathFromRoot = this.#getPathFromRoot(leafIndex);
+
+    let node: Buffer = root;
+    for (let i = 0; i < pathFromRoot.length; i++) {
+      // get both children. We'll need both anyway (one to keep track of, the other to walk down to)
+      const children: [Buffer, Buffer] = this.db.get(node.toString('hex')) ?? [
+        this.tree.getZeroHash(i + 1),
+        this.tree.getZeroHash(i + 1),
+      ];
+      const next = children[pathFromRoot[i]];
+      const sibling = children[(pathFromRoot[i] + 1) % 2];
+
+      yield [next, sibling];
+
+      node = next;
+    }
+  }
+
+  /**
+   * Calculates the path from the root to the target leaf. Returns an array of 0s and 1s,
+   * each 0 represents walking down a left child and each 1 walking down to the child on the right.
+   *
+   * @param leafIndex - The target leaf
+   * @returns An array of 0s and 1s
+   */
+  #getPathFromRoot(leafIndex: bigint): ReadonlyArray<0 | 1> {
+    const path: Array<0 | 1> = [];
+    let level = this.tree.getDepth();
+    while (level > 0) {
+      path.push(leafIndex & 0x01n ? 1 : 0);
+      leafIndex >>= 1n;
+      level--;
+    }
+
+    path.reverse();
+    return path;
+  }
+
+  findLeafIndex(value: Buffer): bigint | undefined {
+    const numLeaves = this.getNumLeaves();
+    for (let i = 0n; i < numLeaves; i++) {
+      const currentValue = this.getLeafValue(i);
+      if (currentValue && currentValue.equals(value)) {
+        return i;
+      }
+    }
+    return undefined;
+  }
+}

package/src/snapshots/full_snapshot.ts

@@ -0,0 +1,26 @@
+import { TreeBase } from '../tree_base.js';
+import { BaseFullTreeSnapshot, BaseFullTreeSnapshotBuilder } from './base_full_snapshot.js';
+import { TreeSnapshot, TreeSnapshotBuilder } from './snapshot_builder.js';
+
+/**
+ * Builds a full snapshot of a tree. This implementation works for any Merkle tree and stores
+ * it in a database in a similar way to how a tree is stored in memory, using pointers.
+ *
+ * Sharing the same database between versions and trees is recommended as the trees would share
+ * structure.
+ *
+ * Complexity:
+ * N - count of non-zero nodes in tree
+ * M - count of snapshots
+ * H - tree height
+ * Worst case space complexity: O(N * M)
+ * Sibling path access: O(H) database reads
+ */
+export class FullTreeSnapshotBuilder
+  extends BaseFullTreeSnapshotBuilder<TreeBase, TreeSnapshot>
+  implements TreeSnapshotBuilder<TreeSnapshot>
+{
+  protected openSnapshot(root: Buffer, numLeaves: bigint): TreeSnapshot {
+    return new BaseFullTreeSnapshot(this.nodes, root, numLeaves, this.tree);
+  }
+}

package/src/snapshots/indexed_tree_snapshot.ts

@@ -0,0 +1,108 @@
+import { IndexedTreeLeafPreimage } from '@aztec/foundation/trees';
+import { AztecKVStore, AztecMap } from '@aztec/kv-store';
+
+import { IndexedTree, PreimageFactory } from '../interfaces/indexed_tree.js';
+import { TreeBase } from '../tree_base.js';
+import { BaseFullTreeSnapshot, BaseFullTreeSnapshotBuilder } from './base_full_snapshot.js';
+import { IndexedTreeSnapshot, TreeSnapshotBuilder } from './snapshot_builder.js';
+
+const snapshotLeafValue = (node: Buffer, index: bigint) => 'snapshot:leaf:' + node.toString('hex') + ':' + index;
+
+/** a */
+export class IndexedTreeSnapshotBuilder
+  extends BaseFullTreeSnapshotBuilder<IndexedTree & TreeBase, IndexedTreeSnapshot>
+  implements TreeSnapshotBuilder<IndexedTreeSnapshot>
+{
+  leaves: AztecMap<string, Buffer>;
+  constructor(store: AztecKVStore, tree: IndexedTree & TreeBase, private leafPreimageBuilder: PreimageFactory) {
+    super(store, tree);
+    this.leaves = store.openMap('indexed_tree_snapshot:' + tree.getName());
+  }
+
+  protected openSnapshot(root: Buffer, numLeaves: bigint): IndexedTreeSnapshot {
+    return new IndexedTreeSnapshotImpl(this.nodes, this.leaves, root, numLeaves, this.tree, this.leafPreimageBuilder);
+  }
+
+  protected handleLeaf(index: bigint, node: Buffer) {
+    const leafPreimage = this.tree.getLatestLeafPreimageCopy(index, false);
+    if (leafPreimage) {
+      void this.leaves.set(snapshotLeafValue(node, index), leafPreimage.toBuffer());
+    }
+  }
+}
+
+/** A snapshot of an indexed tree at a particular point in time */
+class IndexedTreeSnapshotImpl extends BaseFullTreeSnapshot implements IndexedTreeSnapshot {
+  constructor(
+    db: AztecMap<string, [Buffer, Buffer]>,
+    private leaves: AztecMap<string, Buffer>,
+    historicRoot: Buffer,
+    numLeaves: bigint,
+    tree: IndexedTree & TreeBase,
+    private leafPreimageBuilder: PreimageFactory,
+  ) {
+    super(db, historicRoot, numLeaves, tree);
+  }
+
+  getLeafValue(index: bigint): Buffer | undefined {
+    const leafPreimage = this.getLatestLeafPreimageCopy(index);
+    return leafPreimage?.toBuffer();
+  }
+
+  getLatestLeafPreimageCopy(index: bigint): IndexedTreeLeafPreimage | undefined {
+    const leafNode = super.getLeafValue(index);
+    const leafValue = this.leaves.get(snapshotLeafValue(leafNode!, index));
+    if (leafValue) {
+      return this.leafPreimageBuilder.fromBuffer(leafValue);
+    } else {
+      return undefined;
+    }
+  }
+
+  findIndexOfPreviousKey(newValue: bigint): {
+    /**
+     * The index of the found leaf.
+     */
+    index: bigint;
+    /**
+     * A flag indicating if the corresponding leaf's value is equal to `newValue`.
+     */
+    alreadyPresent: boolean;
+  } {
+    const numLeaves = this.getNumLeaves();
+    const diff: bigint[] = [];
+
+    for (let i = 0; i < numLeaves; i++) {
+      // this is very inefficient
+      const storedLeaf = this.getLatestLeafPreimageCopy(BigInt(i))!;
+
+      // The stored leaf can be undefined if it addresses an empty leaf
+      // If the leaf is empty we do the same as if the leaf was larger
+      if (storedLeaf === undefined) {
+        diff.push(newValue);
+      } else if (storedLeaf.getKey() > newValue) {
+        diff.push(newValue);
+      } else if (storedLeaf.getKey() === newValue) {
+        return { index: BigInt(i), alreadyPresent: true };
+      } else {
+        diff.push(newValue - storedLeaf.getKey());
+      }
+    }
+
+    let minIndex = 0;
+    for (let i = 1; i < diff.length; i++) {
+      if (diff[i] < diff[minIndex]) {
+        minIndex = i;
+      }
+    }
+
+    return { index: BigInt(minIndex), alreadyPresent: false };
+  }
+
+  findLeafIndex(value: Buffer): bigint | undefined {
+    const index = this.tree.findLeafIndex(value, false);
+    if (index !== undefined && index < this.getNumLeaves()) {
+      return index;
+    }
+  }
+}

package/src/snapshots/snapshot_builder.ts

@@ -0,0 +1,84 @@
+import { SiblingPath } from '@aztec/circuit-types';
+import { IndexedTreeLeafPreimage } from '@aztec/foundation/trees';
+
+/**
+ * An interface for a tree that can record snapshots of its contents.
+ */
+export interface TreeSnapshotBuilder<S extends TreeSnapshot = TreeSnapshot> {
+  /**
+   * Creates a snapshot of the tree at the given version.
+   * @param block - The version to snapshot the tree at.
+   */
+  snapshot(block: number): Promise<S>;
+
+  /**
+   * Returns a snapshot of the tree at the given version.
+   * @param block - The version of the snapshot to return.
+   */
+  getSnapshot(block: number): Promise<S>;
+}
+
+/**
+ * A tree snapshot
+ */
+export interface TreeSnapshot {
+  /**
+   * Returns the current root of the tree.
+   */
+  getRoot(): Buffer;
+
+  /**
+   * Returns the number of leaves in the tree.
+   */
+  getDepth(): number;
+
+  /**
+   * Returns the number of leaves in the tree.
+   */
+  getNumLeaves(): bigint;
+
+  /**
+   * Returns the value of a leaf at the specified index.
+   * @param index - The index of the leaf value to be returned.
+   */
+  getLeafValue(index: bigint): Buffer | undefined;
+
+  /**
+   * Returns the sibling path for a requested leaf index.
+   * @param index - The index of the leaf for which a sibling path is required.
+   */
+  getSiblingPath<N extends number>(index: bigint): SiblingPath<N>;
+
+  /**
+   * Returns the index of a leaf given its value, or undefined if no leaf with that value is found.
+   * @param treeId - The ID of the tree.
+   * @param value - The leaf value to look for.
+   * @returns The index of the first leaf found with a given value (undefined if not found).
+   */
+  findLeafIndex(value: Buffer): bigint | undefined;
+}
+
+/** A snapshot of an indexed tree */
+export interface IndexedTreeSnapshot extends TreeSnapshot {
+  /**
+   * Gets the historical data for a leaf
+   * @param index - The index of the leaf to get the data for
+   */
+  getLatestLeafPreimageCopy(index: bigint): IndexedTreeLeafPreimage | undefined;
+
+  /**
+   * Finds the index of the largest leaf whose value is less than or equal to the provided value.
+   * @param newValue - The new value to be inserted into the tree.
+   * @returns The found leaf index and a flag indicating if the corresponding leaf's value is equal to `newValue`.
+   */
+  findIndexOfPreviousKey(newValue: bigint): {
+    /**
+     * The index of the found leaf.
+     */
+    index: bigint;
+    /**
+     * A flag indicating if the corresponding leaf's value is equal to `newValue`.
+     */
+    alreadyPresent: boolean;
+  };
+}

package/src/snapshots/snapshot_builder_test_suite.ts

@@ -0,0 +1,218 @@
+import { TreeBase } from '../tree_base.js';
+import { TreeSnapshotBuilder } from './snapshot_builder.js';
+
+/** Creates a test suit for snapshots */
+export function describeSnapshotBuilderTestSuite<T extends TreeBase, S extends TreeSnapshotBuilder>(
+  getTree: () => T,
+  getSnapshotBuilder: () => S,
+  modifyTree: (tree: T) => Promise<void>,
+) {
+  describe('SnapshotBuilder', () => {
+    let tree: T;
+    let snapshotBuilder: S;
+    let leaves: bigint[];
+
+    beforeEach(() => {
+      tree = getTree();
+      snapshotBuilder = getSnapshotBuilder();
+
+      leaves = Array.from({ length: 4 }).map(() => BigInt(Math.floor(Math.random() * 2 ** tree.getDepth())));
+    });
+
+    describe('snapshot', () => {
+      it('takes snapshots', async () => {
+        await modifyTree(tree);
+        await tree.commit();
+        await expect(snapshotBuilder.snapshot(1)).resolves.toBeDefined();
+      });
+
+      it('is idempotent', async () => {
+        await modifyTree(tree);
+        await tree.commit();
+
+        const block = 1;
+        const snapshot = await snapshotBuilder.snapshot(block);
+        await expect(snapshotBuilder.snapshot(block)).resolves.toEqual(snapshot);
+      });
+
+      it('returns the same path if tree has not diverged', async () => {
+        await modifyTree(tree);
+        await tree.commit();
+        const snapshot = await snapshotBuilder.snapshot(1);
+
+        const historicPaths = await Promise.all(leaves.map(leaf => snapshot.getSiblingPath(leaf)));
+        const expectedPaths = await Promise.all(leaves.map(leaf => tree.getSiblingPath(leaf, false)));
+
+        for (const [index, path] of historicPaths.entries()) {
+          expect(path).toEqual(expectedPaths[index]);
+        }
+      });
+
+      it('returns historic paths if tree has diverged and no new snapshots have been taken', async () => {
+        await modifyTree(tree);
+        await tree.commit();
+        const snapshot = await snapshotBuilder.snapshot(1);
+
+        const expectedPaths = await Promise.all(leaves.map(leaf => tree.getSiblingPath(leaf, false)));
+
+        await modifyTree(tree);
+        await tree.commit();
+
+        const historicPaths = await Promise.all(leaves.map(leaf => snapshot.getSiblingPath(leaf)));
+
+        for (const [index, path] of historicPaths.entries()) {
+          expect(path).toEqual(expectedPaths[index]);
+        }
+      });
+
+      it('retains old snapshots even if new one are created', async () => {
+        await modifyTree(tree);
+        await tree.commit();
+
+        const expectedPaths = await Promise.all(leaves.map(leaf => tree.getSiblingPath(leaf, false)));
+
+        const snapshot = await snapshotBuilder.snapshot(1);
+
+        await modifyTree(tree);
+        await tree.commit();
+
+        await snapshotBuilder.snapshot(2);
+
+        // check that snapshot 2 has not influenced snapshot(1) at all
+        const historicPaths = await Promise.all(leaves.map(leaf => snapshot.getSiblingPath(leaf)));
+
+        for (const [index, path] of historicPaths.entries()) {
+          expect(path).toEqual(expectedPaths[index]);
+        }
+      });
+
+      it('retains old snapshots even if new one are created and the tree diverges', async () => {
+        await modifyTree(tree);
+        await tree.commit();
+
+        const expectedPaths = await Promise.all(leaves.map(leaf => tree.getSiblingPath(leaf, false)));
+
+        const snapshot = await snapshotBuilder.snapshot(1);
+
+        await modifyTree(tree);
+        await tree.commit();
+
+        await snapshotBuilder.snapshot(2);
+
+        await modifyTree(tree);
+        await tree.commit();
+
+        // check that snapshot 2 has not influenced snapshot(1) at all
+        // and that the diverging tree does not influence the old snapshot
+        const historicPaths = await Promise.all(leaves.map(leaf => snapshot.getSiblingPath(leaf)));
+
+        for (const [index, path] of historicPaths.entries()) {
+          expect(path).toEqual(expectedPaths[index]);
+        }
+      });
+    });
+
+    describe('getSnapshot', () => {
+      it('returns old snapshots', async () => {
+        await modifyTree(tree);
+        await tree.commit();
+        const expectedPaths = await Promise.all(leaves.map(leaf => tree.getSiblingPath(leaf, false)));
+        await snapshotBuilder.snapshot(1);
+
+        for (let i = 2; i < 5; i++) {
+          await modifyTree(tree);
+          await tree.commit();
+          await snapshotBuilder.snapshot(i);
+        }
+
+        const firstSnapshot = await snapshotBuilder.getSnapshot(1);
+        const historicPaths = await Promise.all(leaves.map(leaf => firstSnapshot.getSiblingPath(leaf)));
+
+        for (const [index, path] of historicPaths.entries()) {
+          expect(path).toEqual(expectedPaths[index]);
+        }
+      });
+
+      it('throws if an unknown snapshot is requested', async () => {
+        await modifyTree(tree);
+        await tree.commit();
+        await snapshotBuilder.snapshot(1);
+
+        await expect(snapshotBuilder.getSnapshot(2)).rejects.toThrow();
+      });
+    });
+
+    describe('getRoot', () => {
+      it('returns the historical root of the tree when the snapshot was taken', async () => {
+        await modifyTree(tree);
+        await tree.commit();
+        const snapshot = await snapshotBuilder.snapshot(1);
+        const historicalRoot = tree.getRoot(false);
+
+        await modifyTree(tree);
+        await tree.commit();
+
+        expect(snapshot.getRoot()).toEqual(historicalRoot);
+        expect(snapshot.getRoot()).not.toEqual(tree.getRoot(false));
+      });
+    });
+
+    describe('getDepth', () => {
+      it('returns the same depth as the tree', async () => {
+        await modifyTree(tree);
+        await tree.commit();
+        const snapshot = await snapshotBuilder.snapshot(1);
+        expect(snapshot.getDepth()).toEqual(tree.getDepth());
+      });
+    });
+
+    describe('getNumLeaves', () => {
+      it('returns the historical leaves count when the snapshot was taken', async () => {
+        await modifyTree(tree);
+        await tree.commit();
+        const snapshot = await snapshotBuilder.snapshot(1);
+        const historicalNumLeaves = tree.getNumLeaves(false);
+
+        await modifyTree(tree);
+        await tree.commit();
+
+        expect(snapshot.getNumLeaves()).toEqual(historicalNumLeaves);
+      });
+    });
+
+    describe('getLeafValue', () => {
+      it('returns the historical leaf value when the snapshot was taken', async () => {
+        await modifyTree(tree);
+        await tree.commit();
+        const snapshot = await snapshotBuilder.snapshot(1);
+        const historicalLeafValue = tree.getLeafValue(0n, false);
+        expect(snapshot.getLeafValue(0n)).toEqual(historicalLeafValue);
+
+        await modifyTree(tree);
+        await tree.commit();
+
+        expect(snapshot.getLeafValue(0n)).toEqual(historicalLeafValue);
+      });
+    });
+
+    describe('findLeafIndex', () => {
+      it('returns the historical leaf index when the snapshot was taken', async () => {
+        await modifyTree(tree);
+        await tree.commit();
+        const snapshot = await snapshotBuilder.snapshot(1);
+
+        const initialLastLeafIndex = tree.getNumLeaves(false) - 1n;
+        let lastLeaf = tree.getLeafValue(initialLastLeafIndex, false);
+        expect(snapshot.findLeafIndex(lastLeaf!)).toBe(initialLastLeafIndex);
+
+        await modifyTree(tree);
+        await tree.commit();
+
+        const newLastLeafIndex = tree.getNumLeaves(false) - 1n;
+        lastLeaf = tree.getLeafValue(newLastLeafIndex, false);
+
+        expect(snapshot.findLeafIndex(lastLeaf!)).toBe(undefined);
+      });
+    });
+  });
+}

package/src/sparse_tree/sparse_tree.ts

@@ -0,0 +1,49 @@
+import { UpdateOnlyTree } from '../interfaces/update_only_tree.js';
+import { FullTreeSnapshotBuilder } from '../snapshots/full_snapshot.js';
+import { TreeSnapshot } from '../snapshots/snapshot_builder.js';
+import { INITIAL_LEAF, TreeBase } from '../tree_base.js';
+
+/**
+ * A Merkle tree implementation that uses a LevelDB database to store the tree.
+ */
+export class SparseTree extends TreeBase implements UpdateOnlyTree {
+  #snapshotBuilder = new FullTreeSnapshotBuilder(this.store, this);
+  /**
+   * Updates a leaf in the tree.
+   * @param leaf - New contents of the leaf.
+   * @param index - Index of the leaf to be updated.
+   */
+  public updateLeaf(leaf: Buffer, index: bigint): Promise<void> {
+    if (index > this.maxIndex) {
+      throw Error(`Index out of bounds. Index ${index}, max index: ${this.maxIndex}.`);
+    }
+
+    const insertingZeroElement = leaf.equals(INITIAL_LEAF);
+    const originallyZeroElement = this.getLeafValue(index, true)?.equals(INITIAL_LEAF);
+    if (insertingZeroElement && originallyZeroElement) {
+      return Promise.resolve();
+    }
+    this.addLeafToCacheAndHashToRoot(leaf, index);
+    if (insertingZeroElement) {
+      // Deleting element (originally non-zero and new value is zero)
+      this.cachedSize = (this.cachedSize ?? this.size) - 1n;
+    } else if (originallyZeroElement) {
+      // Inserting new element (originally zero and new value is non-zero)
+      this.cachedSize = (this.cachedSize ?? this.size) + 1n;
+    }
+
+    return Promise.resolve();
+  }
+
+  public snapshot(block: number): Promise<TreeSnapshot> {
+    return this.#snapshotBuilder.snapshot(block);
+  }
+
+  public getSnapshot(block: number): Promise<TreeSnapshot> {
+    return this.#snapshotBuilder.getSnapshot(block);
+  }
+
+  public findLeafIndex(_value: Buffer, _includeUncommitted: boolean): bigint | undefined {
+    throw new Error('Finding leaf index is not supported for sparse trees');
+  }
+}