refactor: replace leveldb with lmdb for merkle trees

yarn-project/acir-simulator/package.json

@@ -39,6 +39,7 @@
     "tslib": "^2.4.0"
   },
   "devDependencies": {
+    "@aztec/kv-store": "workspace:^",
     "@aztec/merkle-tree": "workspace:^",
     "@aztec/noir-contracts": "workspace:^",
     "@jest/globals": "^29.5.0",

yarn-project/acir-simulator/src/client/private_execution.test.ts

@@ -35,6 +35,7 @@ import { EthAddress } from '@aztec/foundation/eth-address';
 import { Fr, GrumpkinScalar } from '@aztec/foundation/fields';
 import { DebugLogger, createDebugLogger } from '@aztec/foundation/log';
 import { FieldsOf } from '@aztec/foundation/types';
+import { AztecLmdbStore } from '@aztec/kv-store';
 import { AppendOnlyTree, Pedersen, StandardTree, newTree } from '@aztec/merkle-tree';
 import {
   ChildContractArtifact,
@@ -48,8 +49,6 @@ import {
 
 import { jest } from '@jest/globals';
 import { MockProxy, mock } from 'jest-mock-extended';
-import { default as levelup } from 'levelup';
-import { type MemDown, default as memdown } from 'memdown';
 import { getFunctionSelector } from 'viem';
 
 import { buildL1ToL2Message } from '../test/utils.js';
@@ -59,8 +58,6 @@ import { AcirSimulator } from './simulator.js';
 
 jest.setTimeout(60_000);
 
-const createMemDown = () => (memdown as any)() as MemDown<any, any>;
-
 describe('Private Execution test suite', () => {
   let oracle: MockProxy<DBOracle>;
   let acirSimulator: AcirSimulator;
@@ -131,7 +128,7 @@ describe('Private Execution test suite', () => {
       throw new Error(`Unknown tree ${name}`);
     }
     if (!trees[name]) {
-      const db = levelup(createMemDown());
+      const db = await AztecLmdbStore.openTmp();
       const pedersen = new Pedersen();
       trees[name] = await newTree(StandardTree, db, pedersen, name, treeHeights[name]);
     }

yarn-project/acir-simulator/tsconfig.json

@@ -15,6 +15,9 @@
     {
       "path": "../foundation"
     },
+    {
+      "path": "../kv-store"
+    },
     {
       "path": "../merkle-tree"
     },

yarn-project/archiver/src/archiver/kv_archiver_store/kv_archiver_store.test.ts

@@ -8,7 +8,7 @@ describe('KVArchiverDataStore', () => {
   let archiverStore: KVArchiverDataStore;
 
   beforeEach(async () => {
-    archiverStore = new KVArchiverDataStore(await AztecLmdbStore.create(EthAddress.random()));
+    archiverStore = new KVArchiverDataStore(await AztecLmdbStore.openTmp());
   });
 
   describeArchiverDataStore('ArchiverStore', () => archiverStore);

yarn-project/aztec-node/src/aztec-node/server.ts

@@ -39,7 +39,7 @@ import { computeGlobalsHash, computePublicDataTreeLeafSlot } from '@aztec/circui
 import { L1ContractAddresses, createEthereumChain } from '@aztec/ethereum';
 import { AztecAddress } from '@aztec/foundation/aztec-address';
 import { createDebugLogger } from '@aztec/foundation/log';
-import { AztecLmdbStore } from '@aztec/kv-store';
+import { AztecKVStore, AztecLmdbStore } from '@aztec/kv-store';
 import { AztecKVTxPool, P2P, createP2PClient } from '@aztec/p2p';
 import {
   GlobalVariableBuilder,
@@ -56,10 +56,7 @@ import {
   getConfigEnvVars as getWorldStateConfig,
 } from '@aztec/world-state';
 
-import { LevelUp } from 'levelup';
-
 import { AztecNodeConfig } from './config.js';
-import { openDb } from './db.js';
 
 /**
  * The aztec node.
@@ -78,7 +75,7 @@ export class AztecNodeService implements AztecNode {
     protected readonly chainId: number,
     protected readonly version: number,
     protected readonly globalVariableBuilder: GlobalVariableBuilder,
-    protected readonly merkleTreesDb: LevelUp,
+    protected readonly merkleTreesDb: AztecKVStore,
     private log = createDebugLogger('aztec:node'),
   ) {
     const message =
@@ -107,7 +104,6 @@ export class AztecNodeService implements AztecNode {
 
     const log = createDebugLogger('aztec:node');
     const store = await AztecLmdbStore.create(config.l1Contracts.rollupAddress, config.dataDirectory);
-    const [_, worldStateDb] = await openDb(config, log);
 
     // first create and sync the archiver
     const archiverStore = new KVArchiverDataStore(store, config.maxLogs);
@@ -121,14 +117,9 @@ export class AztecNodeService implements AztecNode {
     const p2pClient = await createP2PClient(store, config, new AztecKVTxPool(store), archiver);
 
     // now create the merkle trees and the world state synchronizer
-    const merkleTrees = await MerkleTrees.new(worldStateDb);
+    const merkleTrees = await MerkleTrees.new(store);
     const worldStateConfig: WorldStateConfig = getWorldStateConfig();
-    const worldStateSynchronizer = await ServerWorldStateSynchronizer.new(
-      worldStateDb,
-      merkleTrees,
-      archiver,
-      worldStateConfig,
-    );
+    const worldStateSynchronizer = new ServerWorldStateSynchronizer(store, merkleTrees, archiver, worldStateConfig);
 
     // start both and wait for them to sync from the block source
     await Promise.all([p2pClient.start(), worldStateSynchronizer.start()]);
@@ -151,7 +142,7 @@ export class AztecNodeService implements AztecNode {
       ethereumChain.chainInfo.id,
       config.version,
       getGlobalVariableBuilder(config),
-      worldStateDb,
+      store,
       log,
     );
   }
@@ -285,8 +276,6 @@ export class AztecNodeService implements AztecNode {
     await this.p2pClient.stop();
     await this.worldStateSynchronizer.stop();
     await this.blockSource.stop();
-    this.log('Closing Merkle Trees');
-    await this.merkleTreesDb.close();
     this.log.info(`Stopped`);
   }
 

yarn-project/end-to-end/src/e2e_blacklist_token_contract.test.ts

@@ -14,18 +14,15 @@ import {
   computeAuthWitMessageHash,
   computeMessageSecretHash,
 } from '@aztec/aztec.js';
+import { AztecLmdbStore } from '@aztec/kv-store';
 import { Pedersen, SparseTree, newTree } from '@aztec/merkle-tree';
 import { SlowTreeContract, TokenBlacklistContract, TokenContract } from '@aztec/noir-contracts';
 
 import { jest } from '@jest/globals';
-import levelup from 'levelup';
-import { type MemDown, default as memdown } from 'memdown';
 
 import { setup } from './fixtures/utils.js';
 import { TokenSimulator } from './simulators/token_simulator.js';
 
-export const createMemDown = () => (memdown as any)() as MemDown<any, any>;
-
 const TIMEOUT = 90_000;
 
 describe('e2e_blacklist_token_contract', () => {
@@ -48,7 +45,7 @@ describe('e2e_blacklist_token_contract', () => {
   const getMembershipProof = async (index: bigint, includeUncommitted: boolean) => {
     return {
       index,
-      value: Fr.fromBuffer((await slowUpdateTreeSimulator.getLeafValue(index, includeUncommitted))!),
+      value: Fr.fromBuffer(slowUpdateTreeSimulator.getLeafValue(index, includeUncommitted)!),
       // eslint-disable-next-line camelcase
       sibling_path: (await slowUpdateTreeSimulator.getSiblingPath(index, includeUncommitted)).toFieldArray(),
     };
@@ -107,7 +104,7 @@ describe('e2e_blacklist_token_contract', () => {
     slowTree = await SlowTreeContract.deploy(wallets[0]).send().deployed();
 
     const depth = 254;
-    slowUpdateTreeSimulator = await newTree(SparseTree, levelup(createMemDown()), new Pedersen(), 'test', depth);
+    slowUpdateTreeSimulator = await newTree(SparseTree, await AztecLmdbStore.openTmp(), new Pedersen(), 'test', depth);
 
     const deployTx = TokenBlacklistContract.deploy(wallets[0], accounts[0], slowTree.address).send({});
     const receipt = await deployTx.wait();

yarn-project/end-to-end/src/e2e_slow_tree.test.ts

@@ -1,15 +1,11 @@
 /* eslint-disable camelcase */
 import { CheatCodes, DebugLogger, Fr, Wallet } from '@aztec/aztec.js';
+import { AztecLmdbStore } from '@aztec/kv-store';
 import { Pedersen, SparseTree, newTree } from '@aztec/merkle-tree';
 import { SlowTreeContract } from '@aztec/noir-contracts/SlowTree';
 
-import { default as levelup } from 'levelup';
-import { type MemDown, default as memdown } from 'memdown';
-
 import { setup } from './fixtures/utils.js';
 
-export const createMemDown = () => (memdown as any)() as MemDown<any, any>;
-
 describe('e2e_slow_tree', () => {
   let logger: DebugLogger;
   let wallet: Wallet;
@@ -27,11 +23,17 @@ describe('e2e_slow_tree', () => {
 
   it('Messing around with noir slow tree', async () => {
     const depth = 254;
-    const slowUpdateTreeSimulator = await newTree(SparseTree, levelup(createMemDown()), new Pedersen(), 'test', depth);
+    const slowUpdateTreeSimulator = await newTree(
+      SparseTree,
+      await AztecLmdbStore.openTmp(),
+      new Pedersen(),
+      'test',
+      depth,
+    );
     const getMembershipProof = async (index: bigint, includeUncommitted: boolean) => {
       return {
         index,
-        value: Fr.fromBuffer((await slowUpdateTreeSimulator.getLeafValue(index, includeUncommitted))!),
+        value: Fr.fromBuffer(slowUpdateTreeSimulator.getLeafValue(index, includeUncommitted)!),
         // eslint-disable-next-line camelcase
         sibling_path: (await slowUpdateTreeSimulator.getSiblingPath(index, includeUncommitted)).toFieldArray(),
       };

yarn-project/end-to-end/src/integration_l1_publisher.test.ts

@@ -28,6 +28,7 @@ import {
 } from '@aztec/circuits.js/factories';
 import { createEthereumChain } from '@aztec/ethereum';
 import { makeTuple, range } from '@aztec/foundation/array';
+import { AztecLmdbStore } from '@aztec/kv-store';
 import { InboxAbi, OutboxAbi, RollupAbi } from '@aztec/l1-artifacts';
 import {
   EmptyRollupProver,
@@ -44,8 +45,6 @@ import { MerkleTreeOperations, MerkleTrees } from '@aztec/world-state';
 
 import { beforeEach, describe, expect, it } from '@jest/globals';
 import * as fs from 'fs';
-import { default as levelup } from 'levelup';
-import memdown from 'memdown';
 import {
   Address,
   Chain,
@@ -133,7 +132,7 @@ describe('L1Publisher integration', () => {
       publicClient,
     });
 
-    builderDb = await MerkleTrees.new(levelup((memdown as any)())).then(t => t.asLatest());
+    builderDb = await MerkleTrees.new(await AztecLmdbStore.openTmp()).then(t => t.asLatest());
     const vks = getVerificationKeys();
     const simulator = new RealRollupCircuitSimulator();
     const prover = new EmptyRollupProver();

yarn-project/kv-store/src/lmdb/store.ts

@@ -72,6 +72,10 @@ export class AztecLmdbStore implements AztecKVStore {
     return db;
   }
 
+  static openTmp(): Promise<AztecLmdbStore> {
+    return AztecLmdbStore.create(EthAddress.random());
+  }
+
   /**
    * Creates a new AztecMap in the store.
    * @param name - Name of the map

yarn-project/merkle-tree/package.json

@@ -34,18 +34,15 @@
   "dependencies": {
     "@aztec/circuit-types": "workspace:^",
     "@aztec/foundation": "workspace:^",
+    "@aztec/kv-store": "workspace:^",
     "@aztec/types": "workspace:^",
-    "levelup": "^5.1.1",
-    "memdown": "^6.1.1",
     "sha256": "^0.2.0",
     "tslib": "^2.4.0"
   },
   "devDependencies": {
     "@aztec/circuits.js": "workspace:^",
     "@jest/globals": "^29.5.0",
     "@types/jest": "^29.5.0",
-    "@types/levelup": "^5.1.2",
-    "@types/memdown": "^3.0.1",
     "@types/node": "^18.15.3",
     "@types/sha256": "^0.2.0",
     "jest": "^29.5.0",

yarn-project/merkle-tree/src/interfaces/indexed_tree.ts

@@ -1,8 +1,35 @@
-import { IndexedTreeLeafPreimage } from '@aztec/foundation/trees';
+import { IndexedTreeLeaf, IndexedTreeLeafPreimage } from '@aztec/foundation/trees';
 import { SiblingPath } from '@aztec/types/membership';
 
 import { AppendOnlyTree } from './append_only_tree.js';
 
+/**
+ * Factory for creating leaf preimages.
+ */
+export interface PreimageFactory {
+  /**
+   * Creates a new preimage from a leaf.
+   * @param leaf - Leaf to create a preimage from.
+   * @param nextKey - Next key of the leaf.
+   * @param nextIndex - Next index of the leaf.
+   */
+  fromLeaf(leaf: IndexedTreeLeaf, nextKey: bigint, nextIndex: bigint): IndexedTreeLeafPreimage;
+  /**
+   * Creates a new preimage from a buffer.
+   * @param buffer - Buffer to create a preimage from.
+   */
+  fromBuffer(buffer: Buffer): IndexedTreeLeafPreimage;
+  /**
+   * Creates an empty preimage.
+   */
+  empty(): IndexedTreeLeafPreimage;
+  /**
+   * Creates a copy of a preimage.
+   * @param preimage - Preimage to be cloned.
+   */
+  clone(preimage: IndexedTreeLeafPreimage): IndexedTreeLeafPreimage;
+}
+
 /**
  * All of the data to be return during batch insertion.
  */
@@ -56,7 +83,7 @@ export interface IndexedTree extends AppendOnlyTree {
   findIndexOfPreviousKey(
     newValue: bigint,
     includeUncommitted: boolean,
-  ): Promise<
+  ):
     | {
         /**
          * The index of the found leaf.
@@ -67,16 +94,15 @@ export interface IndexedTree extends AppendOnlyTree {
          */
         alreadyPresent: boolean;
       }
-    | undefined
-  >;
+    | undefined;
 
   /**
    * Gets the latest LeafPreimage copy.
    * @param index - Index of the leaf of which to obtain the LeafPreimage copy.
    * @param includeUncommitted - If true, the uncommitted changes are included in the search.
    * @returns A copy of the leaf preimage at the given index or undefined if the leaf was not found.
    */
-  getLatestLeafPreimageCopy(index: bigint, includeUncommitted: boolean): Promise<IndexedTreeLeafPreimage | undefined>;
+  getLatestLeafPreimageCopy(index: bigint, includeUncommitted: boolean): IndexedTreeLeafPreimage | undefined;
 
   /**
    * Batch insert multiple leaves into the tree.

yarn-project/merkle-tree/src/interfaces/merkle_tree.ts

@@ -48,13 +48,13 @@ export interface MerkleTree extends SiblingPathSource {
    * @param index - The index of the leaf value to be returned.
    * @param includeUncommitted - Set to true to include uncommitted updates in the data set.
    */
-  getLeafValue(index: bigint, includeUncommitted: boolean): Promise<Buffer | undefined>;
+  getLeafValue(index: bigint, includeUncommitted: boolean): Buffer | undefined;
 
   /**
    * Returns the index of a leaf given its value, or undefined if no leaf with that value is found.
    * @param leaf - The leaf value to look for.
    * @param includeUncommitted - Indicates whether to include uncommitted data.
    * @returns The index of the first leaf found with a given value (undefined if not found).
    */
-  findLeafIndex(leaf: Buffer, includeUncommitted: boolean): Promise<bigint | undefined>;
+  findLeafIndex(leaf: Buffer, includeUncommitted: boolean): bigint | undefined;
 }

yarn-project/merkle-tree/src/load_tree.ts

@@ -1,8 +1,7 @@
+import { AztecKVStore } from '@aztec/kv-store';
 import { Hasher } from '@aztec/types/interfaces';
 
-import { LevelUp } from 'levelup';
-
-import { TreeBase, decodeMeta } from './tree_base.js';
+import { TreeBase, getTreeMeta } from './tree_base.js';
 
 /**
  * Creates a new tree and sets its root, depth and size based on the meta data which are associated with the name.
@@ -12,15 +11,13 @@ import { TreeBase, decodeMeta } from './tree_base.js';
  * @param name - Name of the tree.
  * @returns The newly created tree.
  */
-export async function loadTree<T extends TreeBase>(
-  c: new (db: LevelUp, hasher: Hasher, name: string, depth: number, size: bigint, root: Buffer) => T,
-  db: LevelUp,
+export function loadTree<T extends TreeBase>(
+  c: new (store: AztecKVStore, hasher: Hasher, name: string, depth: number, size: bigint, root: Buffer) => T,
+  store: AztecKVStore,
   hasher: Hasher,
   name: string,
 ): Promise<T> {
-  const meta: Buffer = await db.get(name);
-  const { root, depth, size } = decodeMeta(meta);
-
-  const tree = new c(db, hasher, name, depth, size, root);
-  return tree;
+  const { root, depth, size } = getTreeMeta(store, name);
+  const tree = new c(store, hasher, name, depth, size, root);
+  return Promise.resolve(tree);
 }