@aztec/world-state 0.0.0-test.0 → 0.0.1-commit.24de95ac

package/src/native/native_world_state.ts

@@ -1,8 +1,9 @@
 import { MAX_NOTE_HASHES_PER_TX, MAX_NULLIFIERS_PER_TX, NUMBER_OF_L1_L2_MESSAGES_PER_ROLLUP } from '@aztec/constants';
-import { padArrayEnd } from '@aztec/foundation/collection';
+import { fromEntries, padArrayEnd } from '@aztec/foundation/collection';
 import { EthAddress } from '@aztec/foundation/eth-address';
 import { Fr } from '@aztec/foundation/fields';
-import { createLogger } from '@aztec/foundation/log';
+import { tryRmDir } from '@aztec/foundation/fs';
+import { type Logger, createLogger } from '@aztec/foundation/log';
 import type { L2Block } from '@aztec/stdlib/block';
 import { DatabaseVersionManager } from '@aztec/stdlib/database-version';
 import type {
@@ -10,8 +11,10 @@ import type {
   MerkleTreeReadOperations,
   MerkleTreeWriteOperations,
 } from '@aztec/stdlib/interfaces/server';
+import type { SnapshotDataKeys } from '@aztec/stdlib/snapshots';
 import { MerkleTreeId, NullifierLeaf, type NullifierLeafPreimage, PublicDataTreeLeaf } from '@aztec/stdlib/trees';
 import { BlockHeader, PartialStateReference, StateReference } from '@aztec/stdlib/tx';
+import { WorldStateRevision } from '@aztec/stdlib/world-state';
 import { getTelemetryClient } from '@aztec/telemetry-client';
 
 import assert from 'assert/strict';
@@ -20,6 +23,7 @@ import { tmpdir } from 'os';
 import { join } from 'path';
 
 import { WorldStateInstrumentation } from '../instrumentation/instrumentation.js';
+import type { WorldStateTreeMapSizes } from '../synchronizer/factory.js';
 import type { MerkleTreeAdminDatabase as MerkleTreeDatabase } from '../world-state-db/merkle_tree_db.js';
 import { MerkleTreesFacade, MerkleTreesForkFacade, serializeLeaf } from './merkle_trees_facade.js';
 import {
@@ -27,16 +31,17 @@ import {
   type WorldStateStatusFull,
   type WorldStateStatusSummary,
   blockStateReference,
-  sanitiseFullStatus,
-  sanitiseSummary,
+  sanitizeFullStatus,
+  sanitizeSummary,
   treeStateReferenceToSnapshot,
-  worldStateRevision,
 } from './message.js';
 import { NativeWorldState } from './native_world_state_instance.js';
 
 // The current version of the world state database schema
 // Increment this when making incompatible changes to the database schema
-export const WORLD_STATE_DB_VERSION = 1; // The initial version
+export const WORLD_STATE_DB_VERSION = 2; // The initial version
+
+export const WORLD_STATE_DIR = 'world_state';
 
 export class NativeWorldStateService implements MerkleTreeDatabase {
   protected initialHeader: BlockHeader | undefined;
@@ -44,38 +49,38 @@ export class NativeWorldStateService implements MerkleTreeDatabase {
   private cachedStatusSummary: WorldStateStatusSummary | undefined;
 
   protected constructor(
-    protected readonly instance: NativeWorldState,
+    protected instance: NativeWorldState,
     protected readonly worldStateInstrumentation: WorldStateInstrumentation,
-    protected readonly log = createLogger('world-state:database'),
+    protected readonly log: Logger = createLogger('world-state:database'),
     private readonly cleanup = () => Promise.resolve(),
   ) {}
 
   static async new(
     rollupAddress: EthAddress,
     dataDir: string,
-    dbMapSizeKb: number,
+    wsTreeMapSizes: WorldStateTreeMapSizes,
     prefilledPublicData: PublicDataTreeLeaf[] = [],
     instrumentation = new WorldStateInstrumentation(getTelemetryClient()),
     log = createLogger('world-state:database'),
     cleanup = () => Promise.resolve(),
   ): Promise<NativeWorldStateService> {
-    const worldStateDirectory = join(dataDir, 'world_state');
+    const worldStateDirectory = join(dataDir, WORLD_STATE_DIR);
     // Create a version manager to handle versioning
-    const versionManager = new DatabaseVersionManager(
-      WORLD_STATE_DB_VERSION,
+    const versionManager = new DatabaseVersionManager({
+      schemaVersion: WORLD_STATE_DB_VERSION,
       rollupAddress,
-      worldStateDirectory,
-      (dir: string) => {
-        return Promise.resolve(new NativeWorldState(dir, dbMapSizeKb, prefilledPublicData, instrumentation));
+      dataDirectory: worldStateDirectory,
+      onOpen: (dir: string) => {
+        return Promise.resolve(new NativeWorldState(dir, wsTreeMapSizes, prefilledPublicData, instrumentation));
       },
-    );
+    });
 
     const [instance] = await versionManager.open();
     const worldState = new this(instance, instrumentation, log, cleanup);
     try {
       await worldState.init();
     } catch (e) {
-      log.error(`Error initialising world state: ${e}`);
+      log.error(`Error initializing world state: ${e}`);
       throw e;
     }
 
@@ -91,7 +96,14 @@ export class NativeWorldStateService implements MerkleTreeDatabase {
     const log = createLogger('world-state:database');
     const dataDir = await mkdtemp(join(tmpdir(), 'aztec-world-state-'));
     const dbMapSizeKb = 10 * 1024 * 1024;
-    log.debug(`Created temporary world state database at: ${dataDir} with size: ${dbMapSizeKb}`);
+    const worldStateTreeMapSizes: WorldStateTreeMapSizes = {
+      archiveTreeMapSizeKb: dbMapSizeKb,
+      nullifierTreeMapSizeKb: dbMapSizeKb,
+      noteHashTreeMapSizeKb: dbMapSizeKb,
+      messageTreeMapSizeKb: dbMapSizeKb,
+      publicDataTreeMapSizeKb: dbMapSizeKb,
+    };
+    log.debug(`Created temporary world state database at: ${dataDir} with tree map size: ${dbMapSizeKb}`);
 
     // pass a cleanup callback because process.on('beforeExit', cleanup) does not work under Jest
     const cleanup = async () => {
@@ -103,7 +115,7 @@ export class NativeWorldStateService implements MerkleTreeDatabase {
       }
     };
 
-    return this.new(rollupAddress, dataDir, dbMapSizeKb, prefilledPublicData, instrumentation, log, cleanup);
+    return this.new(rollupAddress, dataDir, worldStateTreeMapSizes, prefilledPublicData, instrumentation, log, cleanup);
   }
 
   protected async init() {
@@ -127,12 +139,23 @@ export class NativeWorldStateService implements MerkleTreeDatabase {
     assert.strictEqual(initialHeaderIndex, 0n, 'Invalid initial archive state');
   }
 
+  public async clear() {
+    await this.instance.close();
+    this.cachedStatusSummary = undefined;
+    await tryRmDir(this.instance.getDataDir(), this.log);
+    this.instance = this.instance.clone();
+  }
+
   public getCommitted(): MerkleTreeReadOperations {
-    return new MerkleTreesFacade(this.instance, this.initialHeader!, worldStateRevision(false, 0, 0));
+    return new MerkleTreesFacade(this.instance, this.initialHeader!, WorldStateRevision.empty());
   }
 
   public getSnapshot(blockNumber: number): MerkleTreeReadOperations {
-    return new MerkleTreesFacade(this.instance, this.initialHeader!, worldStateRevision(false, 0, blockNumber));
+    return new MerkleTreesFacade(
+      this.instance,
+      this.initialHeader!,
+      new WorldStateRevision(/*forkId=*/ 0, /* blockNumber=*/ blockNumber, /* includeUncommitted=*/ false),
+    );
   }
 
   public async fork(blockNumber?: number): Promise<MerkleTreeWriteOperations> {
@@ -141,19 +164,31 @@ export class NativeWorldStateService implements MerkleTreeDatabase {
       blockNumber: blockNumber ?? 0,
       canonical: true,
     });
-    return new MerkleTreesForkFacade(this.instance, this.initialHeader!, worldStateRevision(true, resp.forkId, 0));
+    return new MerkleTreesForkFacade(
+      this.instance,
+      this.initialHeader!,
+      new WorldStateRevision(/*forkId=*/ resp.forkId, /* blockNumber=*/ 0, /* includeUncommitted=*/ true),
+    );
   }
 
   public getInitialHeader(): BlockHeader {
     return this.initialHeader!;
   }
 
-  public async handleL2BlockAndMessages(l2Block: L2Block, l1ToL2Messages: Fr[]): Promise<WorldStateStatusFull> {
+  public async handleL2BlockAndMessages(
+    l2Block: L2Block,
+    l1ToL2Messages: Fr[],
+    // TODO(#17027)
+    // Temporary hack to only insert l1 to l2 messages for the first block in a checkpoint.
+    isFirstBlock = true,
+  ): Promise<WorldStateStatusFull> {
     // We have to pad both the values within tx effects because that's how the trees are built by circuits.
     const paddedNoteHashes = l2Block.body.txEffects.flatMap(txEffect =>
       padArrayEnd(txEffect.noteHashes, Fr.ZERO, MAX_NOTE_HASHES_PER_TX),
     );
-    const paddedL1ToL2Messages = padArrayEnd(l1ToL2Messages, Fr.ZERO, NUMBER_OF_L1_L2_MESSAGES_PER_ROLLUP);
+    const paddedL1ToL2Messages = isFirstBlock
+      ? padArrayEnd<Fr, number>(l1ToL2Messages, Fr.ZERO, NUMBER_OF_L1_L2_MESSAGES_PER_ROLLUP)
+      : [];
 
     const paddedNullifiers = l2Block.body.txEffects
       .flatMap(txEffect => padArrayEnd(txEffect.nullifiers, Fr.ZERO, MAX_NULLIFIERS_PER_TX))
@@ -173,7 +208,7 @@ export class NativeWorldStateService implements MerkleTreeDatabase {
         WorldStateMessageType.SYNC_BLOCK,
         {
           blockNumber: l2Block.number,
-          blockHeaderHash: await l2Block.header.hash(),
+          blockHeaderHash: await l2Block.getBlockHeader().hash(),
           paddedL1ToL2Messages: paddedL1ToL2Messages.map(serializeLeaf),
           paddedNoteHashes: paddedNoteHashes.map(serializeLeaf),
           paddedNullifiers: paddedNullifiers.map(serializeLeaf),
@@ -181,7 +216,7 @@ export class NativeWorldStateService implements MerkleTreeDatabase {
           blockStateRef: blockStateReference(l2Block.header.state),
           canonical: true,
         },
-        this.sanitiseAndCacheSummaryFromFull.bind(this),
+        this.sanitizeAndCacheSummaryFromFull.bind(this),
         this.deleteCachedSummary.bind(this),
       );
     } catch (err) {
@@ -200,16 +235,16 @@ export class NativeWorldStateService implements MerkleTreeDatabase {
     return BlockHeader.empty({ state });
   }
 
-  private sanitiseAndCacheSummaryFromFull(response: WorldStateStatusFull) {
-    const sanitised = sanitiseFullStatus(response);
-    this.cachedStatusSummary = { ...sanitised.summary };
-    return sanitised;
+  private sanitizeAndCacheSummaryFromFull(response: WorldStateStatusFull) {
+    const sanitized = sanitizeFullStatus(response);
+    this.cachedStatusSummary = { ...sanitized.summary };
+    return sanitized;
   }
 
-  private sanitiseAndCacheSummary(response: WorldStateStatusSummary) {
-    const sanitised = sanitiseSummary(response);
-    this.cachedStatusSummary = { ...sanitised };
-    return sanitised;
+  private sanitizeAndCacheSummary(response: WorldStateStatusSummary) {
+    const sanitized = sanitizeSummary(response);
+    this.cachedStatusSummary = { ...sanitized };
+    return sanitized;
   }
 
   private deleteCachedSummary(_: string) {
@@ -217,19 +252,19 @@ export class NativeWorldStateService implements MerkleTreeDatabase {
   }
 
   /**
-   * Advances the finalised block number to be the number provided
-   * @param toBlockNumber The block number that is now the tip of the finalised chain
+   * Advances the finalized block number to be the number provided
+   * @param toBlockNumber The block number that is now the tip of the finalized chain
    * @returns The new WorldStateStatus
    */
-  public async setFinalised(toBlockNumber: bigint) {
+  public async setFinalized(toBlockNumber: bigint) {
     try {
       await this.instance.call(
-        WorldStateMessageType.FINALISE_BLOCKS,
+        WorldStateMessageType.FINALIZE_BLOCKS,
         {
           toBlockNumber,
           canonical: true,
         },
-        this.sanitiseAndCacheSummary.bind(this),
+        this.sanitizeAndCacheSummary.bind(this),
         this.deleteCachedSummary.bind(this),
       );
     } catch (err) {
@@ -252,7 +287,7 @@ export class NativeWorldStateService implements MerkleTreeDatabase {
           toBlockNumber,
           canonical: true,
         },
-        this.sanitiseAndCacheSummaryFromFull.bind(this),
+        this.sanitizeAndCacheSummaryFromFull.bind(this),
         this.deleteCachedSummary.bind(this),
       );
     } catch (err) {
@@ -274,7 +309,7 @@ export class NativeWorldStateService implements MerkleTreeDatabase {
           toBlockNumber,
           canonical: true,
         },
-        this.sanitiseAndCacheSummaryFromFull.bind(this),
+        this.sanitizeAndCacheSummaryFromFull.bind(this),
         this.deleteCachedSummary.bind(this),
       );
     } catch (err) {
@@ -290,7 +325,7 @@ export class NativeWorldStateService implements MerkleTreeDatabase {
     return await this.instance.call(
       WorldStateMessageType.GET_STATUS,
       { canonical: true },
-      this.sanitiseAndCacheSummary.bind(this),
+      this.sanitizeAndCacheSummary.bind(this),
     );
   }
 
@@ -314,4 +349,25 @@ export class NativeWorldStateService implements MerkleTreeDatabase {
       ),
     );
   }
+
+  public async backupTo(
+    dstPath: string,
+    compact: boolean = true,
+  ): Promise<Record<Exclude<SnapshotDataKeys, 'archiver'>, string>> {
+    await this.instance.call(WorldStateMessageType.COPY_STORES, {
+      dstPath,
+      compact,
+      canonical: true,
+    });
+    return fromEntries(NATIVE_WORLD_STATE_DBS.map(([name, dir]) => [name, join(dstPath, dir, 'data.mdb')] as const));
+  }
 }
+
+// The following paths are defined in cpp-land
+export const NATIVE_WORLD_STATE_DBS = [
+  ['l1-to-l2-message-tree', 'L1ToL2MessageTree'],
+  ['archive-tree', 'ArchiveTree'],
+  ['public-data-tree', 'PublicDataTree'],
+  ['note-hash-tree', 'NoteHashTree'],
+  ['nullifier-tree', 'NullifierTree'],
+] as const;

package/src/native/native_world_state_instance.ts

@@ -8,7 +8,7 @@ import {
   NULLIFIER_TREE_HEIGHT,
   PUBLIC_DATA_TREE_HEIGHT,
 } from '@aztec/constants';
-import { createLogger } from '@aztec/foundation/log';
+import { type Logger, createLogger } from '@aztec/foundation/log';
 import { NativeWorldState as BaseNativeWorldState, MsgpackChannel } from '@aztec/native';
 import { MerkleTreeId } from '@aztec/stdlib/trees';
 import type { PublicDataTreeLeaf } from '@aztec/stdlib/trees';
@@ -17,6 +17,7 @@ import assert from 'assert';
 import { cpus } from 'os';
 
 import type { WorldStateInstrumentation } from '../instrumentation/instrumentation.js';
+import type { WorldStateTreeMapSizes } from '../synchronizer/factory.js';
 import {
   WorldStateMessageType,
   type WorldStateRequest,
@@ -50,15 +51,17 @@ export class NativeWorldState implements NativeWorldStateInstance {
 
   /** Creates a new native WorldState instance */
   constructor(
-    dataDir: string,
-    dbMapSizeKb: number,
-    prefilledPublicData: PublicDataTreeLeaf[] = [],
-    private instrumentation: WorldStateInstrumentation,
-    private log = createLogger('world-state:database'),
+    private readonly dataDir: string,
+    private readonly wsTreeMapSizes: WorldStateTreeMapSizes,
+    private readonly prefilledPublicData: PublicDataTreeLeaf[] = [],
+    private readonly instrumentation: WorldStateInstrumentation,
+    private readonly log: Logger = createLogger('world-state:database'),
   ) {
     const threads = Math.min(cpus().length, MAX_WORLD_STATE_THREADS);
     log.info(
-      `Creating world state data store at directory ${dataDir} with map size ${dbMapSizeKb} KB and ${threads} threads.`,
+      `Creating world state data store at directory ${dataDir} with map sizes ${JSON.stringify(
+        wsTreeMapSizes,
+      )} and ${threads} threads.`,
     );
     const prefilledPublicDataBufferArray = prefilledPublicData.map(d => [d.slot.toBuffer(), d.value.toBuffer()]);
     const ws = new BaseNativeWorldState(
@@ -76,7 +79,13 @@ export class NativeWorldState implements NativeWorldStateInstance {
       },
       prefilledPublicDataBufferArray,
       GeneratorIndex.BLOCK_HASH,
-      dbMapSizeKb,
+      {
+        [MerkleTreeId.NULLIFIER_TREE]: wsTreeMapSizes.nullifierTreeMapSizeKb,
+        [MerkleTreeId.NOTE_HASH_TREE]: wsTreeMapSizes.noteHashTreeMapSizeKb,
+        [MerkleTreeId.PUBLIC_DATA_TREE]: wsTreeMapSizes.publicDataTreeMapSizeKb,
+        [MerkleTreeId.L1_TO_L2_MESSAGE_TREE]: wsTreeMapSizes.messageTreeMapSizeKb,
+        [MerkleTreeId.ARCHIVE]: wsTreeMapSizes.archiveTreeMapSizeKb,
+      },
       threads,
     );
     this.instance = new MsgpackChannel(ws);
@@ -84,6 +93,20 @@ export class NativeWorldState implements NativeWorldStateInstance {
     this.queues.set(0, new WorldStateOpsQueue());
   }
 
+  public getDataDir() {
+    return this.dataDir;
+  }
+
+  public clone() {
+    return new NativeWorldState(
+      this.dataDir,
+      this.wsTreeMapSizes,
+      this.prefilledPublicData,
+      this.instrumentation,
+      this.log,
+    );
+  }
+
   /**
    * Sends a message to the native instance and returns the response.
    * @param messageType - The type of message to send

package/src/native/world_state_ops_queue.ts

@@ -35,7 +35,7 @@ export const MUTATING_MSG_TYPES = new Set([
   WorldStateMessageType.SYNC_BLOCK,
   WorldStateMessageType.CREATE_FORK,
   WorldStateMessageType.DELETE_FORK,
-  WorldStateMessageType.FINALISE_BLOCKS,
+  WorldStateMessageType.FINALIZE_BLOCKS,
   WorldStateMessageType.UNWIND_BLOCKS,
   WorldStateMessageType.REMOVE_HISTORICAL_BLOCKS,
   WorldStateMessageType.CREATE_CHECKPOINT,

package/src/synchronizer/config.ts

@@ -16,9 +16,24 @@ export interface WorldStateConfig {
   /** Size of the batch for each get-blocks request from the synchronizer to the archiver. */
   worldStateBlockRequestBatchSize?: number;
 
-  /** The map size to be provided to LMDB for each world state tree DB, optional, will inherit from the general dataStoreMapSizeKB if not specified*/
+  /** The map size to be provided to LMDB for each world state tree DB, optional, will inherit from the general dataStoreMapSizeKb if not specified*/
   worldStateDbMapSizeKb?: number;
 
+  /** The map size to be provided to LMDB for each world state archive tree, optional, will inherit from the general worldStateDbMapSizeKb if not specified*/
+  archiveTreeMapSizeKb?: number;
+
+  /** The map size to be provided to LMDB for each world state nullifier tree, optional, will inherit from the general worldStateDbMapSizeKb if not specified*/
+  nullifierTreeMapSizeKb?: number;
+
+  /** The map size to be provided to LMDB for each world state note hash tree, optional, will inherit from the general worldStateDbMapSizeKb if not specified*/
+  noteHashTreeMapSizeKb?: number;
+
+  /** The map size to be provided to LMDB for each world state message tree, optional, will inherit from the general worldStateDbMapSizeKb if not specified*/
+  messageTreeMapSizeKb?: number;
+
+  /** The map size to be provided to LMDB for each world state public data tree, optional, will inherit from the general worldStateDbMapSizeKb if not specified*/
+  publicDataTreeMapSizeKb?: number;
+
   /** Optional directory for the world state DB, if unspecified will default to the general data directory */
   worldStateDataDirectory?: string;
 
@@ -46,7 +61,37 @@ export const worldStateConfigMappings: ConfigMappingsType<WorldStateConfig> = {
   worldStateDbMapSizeKb: {
     env: 'WS_DB_MAP_SIZE_KB',
     parseEnv: (val: string | undefined) => (val ? +val : undefined),
-    description: 'The maximum possible size of the world state DB',
+    description: 'The maximum possible size of the world state DB in KB. Overwrites the general dataStoreMapSizeKb.',
+  },
+  archiveTreeMapSizeKb: {
+    env: 'ARCHIVE_TREE_MAP_SIZE_KB',
+    parseEnv: (val: string | undefined) => (val ? +val : undefined),
+    description:
+      'The maximum possible size of the world state archive tree in KB. Overwrites the general worldStateDbMapSizeKb.',
+  },
+  nullifierTreeMapSizeKb: {
+    env: 'NULLIFIER_TREE_MAP_SIZE_KB',
+    parseEnv: (val: string | undefined) => (val ? +val : undefined),
+    description:
+      'The maximum possible size of the world state nullifier tree in KB. Overwrites the general worldStateDbMapSizeKb.',
+  },
+  noteHashTreeMapSizeKb: {
+    env: 'NOTE_HASH_TREE_MAP_SIZE_KB',
+    parseEnv: (val: string | undefined) => (val ? +val : undefined),
+    description:
+      'The maximum possible size of the world state note hash tree in KB. Overwrites the general worldStateDbMapSizeKb.',
+  },
+  messageTreeMapSizeKb: {
+    env: 'MESSAGE_TREE_MAP_SIZE_KB',
+    parseEnv: (val: string | undefined) => (val ? +val : undefined),
+    description:
+      'The maximum possible size of the world state message tree in KB. Overwrites the general worldStateDbMapSizeKb.',
+  },
+  publicDataTreeMapSizeKb: {
+    env: 'PUBLIC_DATA_TREE_MAP_SIZE_KB',
+    parseEnv: (val: string | undefined) => (val ? +val : undefined),
+    description:
+      'The maximum possible size of the world state public data tree in KB. Overwrites the general worldStateDbMapSizeKb.',
   },
   worldStateDataDirectory: {
     env: 'WS_DATA_DIRECTORY',

package/src/synchronizer/errors.ts

@@ -0,0 +1,5 @@
+export class WorldStateSynchronizerError extends Error {
+  constructor(message: string, options?: ErrorOptions) {
+    super(message, options);
+  }
+}

package/src/synchronizer/factory.ts

@@ -9,6 +9,14 @@ import { NativeWorldStateService } from '../native/native_world_state.js';
 import type { WorldStateConfig } from './config.js';
 import { ServerWorldStateSynchronizer } from './server_world_state_synchronizer.js';
 
+export interface WorldStateTreeMapSizes {
+  archiveTreeMapSizeKb: number;
+  nullifierTreeMapSizeKb: number;
+  noteHashTreeMapSizeKb: number;
+  messageTreeMapSizeKb: number;
+  publicDataTreeMapSizeKb: number;
+}
+
 export async function createWorldStateSynchronizer(
   config: WorldStateConfig & DataStoreConfig,
   l2BlockSource: L2BlockSource & L1ToL2MessageSource,
@@ -21,25 +29,40 @@ export async function createWorldStateSynchronizer(
 }
 
 export async function createWorldState(
-  config: WorldStateConfig & DataStoreConfig,
+  config: Pick<
+    WorldStateConfig,
+    | 'worldStateDataDirectory'
+    | 'worldStateDbMapSizeKb'
+    | 'archiveTreeMapSizeKb'
+    | 'nullifierTreeMapSizeKb'
+    | 'noteHashTreeMapSizeKb'
+    | 'messageTreeMapSizeKb'
+    | 'publicDataTreeMapSizeKb'
+  > &
+    Pick<DataStoreConfig, 'dataDirectory' | 'dataStoreMapSizeKb' | 'l1Contracts'>,
   prefilledPublicData: PublicDataTreeLeaf[] = [],
   instrumentation: WorldStateInstrumentation = new WorldStateInstrumentation(getTelemetryClient()),
 ) {
-  const newConfig = {
-    dataDirectory: config.worldStateDataDirectory ?? config.dataDirectory,
-    dataStoreMapSizeKB: config.worldStateDbMapSizeKb ?? config.dataStoreMapSizeKB,
-  } as DataStoreConfig;
+  const dataDirectory = config.worldStateDataDirectory ?? config.dataDirectory;
+  const dataStoreMapSizeKb = config.worldStateDbMapSizeKb ?? config.dataStoreMapSizeKb;
+  const wsTreeMapSizes: WorldStateTreeMapSizes = {
+    archiveTreeMapSizeKb: config.archiveTreeMapSizeKb ?? dataStoreMapSizeKb,
+    nullifierTreeMapSizeKb: config.nullifierTreeMapSizeKb ?? dataStoreMapSizeKb,
+    noteHashTreeMapSizeKb: config.noteHashTreeMapSizeKb ?? dataStoreMapSizeKb,
+    messageTreeMapSizeKb: config.messageTreeMapSizeKb ?? dataStoreMapSizeKb,
+    publicDataTreeMapSizeKb: config.publicDataTreeMapSizeKb ?? dataStoreMapSizeKb,
+  };
 
   if (!config.l1Contracts?.rollupAddress) {
     throw new Error('Rollup address is required to create a world state synchronizer.');
   }
 
   // If a data directory is provided in config, then create a persistent store.
-  const merkleTrees = newConfig.dataDirectory
+  const merkleTrees = dataDirectory
     ? await NativeWorldStateService.new(
         config.l1Contracts.rollupAddress,
-        newConfig.dataDirectory,
-        newConfig.dataStoreMapSizeKB,
+        dataDirectory,
+        wsTreeMapSizes,
         prefilledPublicData,
         instrumentation,
       )