@aztec/world-state 0.0.1-commit.2f68f620 → 0.0.1-commit.3100065

package/src/native/merkle_trees_facade.ts

@@ -22,13 +22,14 @@ import {
   PublicDataTreeLeafPreimage,
 } from '@aztec/stdlib/trees';
 import { type BlockHeader, PartialStateReference, StateReference } from '@aztec/stdlib/tx';
-import type { WorldStateRevision } from '@aztec/stdlib/world-state';
+import { type WorldStateRevision, WorldStateRevisionWithHandle } from '@aztec/stdlib/world-state';
 
 import assert from 'assert';
 
 import {
   type SerializedIndexedLeaf,
   type SerializedLeafValue,
+  WorldStateMessageType,
   blockStateReference,
   treeStateReferenceToSnapshot,
 } from './message.js';
@@ -45,12 +46,8 @@ export class MerkleTreesFacade implements MerkleTreeReadOperations {
     return this.initialHeader;
   }
 
-  getRevision(): WorldStateRevision {
-    return this.revision;
-  }
-
-  getIpcPath(): string {
-    return this.instance.getIpcPath();
+  getRevision(): WorldStateRevisionWithHandle {
+    return WorldStateRevisionWithHandle.fromWorldStateRevision(this.revision, this.instance.getHandle());
   }
 
   findLeafIndices(treeId: MerkleTreeId, values: MerkleTreeLeafType<MerkleTreeId>[]): Promise<(bigint | undefined)[]> {
@@ -61,13 +58,13 @@ export class MerkleTreesFacade implements MerkleTreeReadOperations {
     treeId: MerkleTreeId,
     values: MerkleTreeLeafType<MerkleTreeId>[],
   ): Promise<({ path: SiblingPath<N>; index: bigint } | undefined)[]> {
-    const paths = await this.instance.findSiblingPaths(
+    const response = await this.instance.call(WorldStateMessageType.FIND_SIBLING_PATHS, {
+      leaves: values.map(leaf => serializeLeaf(hydrateLeaf(treeId, leaf))),
+      revision: this.revision,
       treeId,
-      this.revision,
-      values.map(leaf => serializeLeaf(hydrateLeaf(treeId, leaf))),
-    );
+    });
 
-    return paths.map(path => {
+    return response.paths.map(path => {
       if (!path) {
         return undefined;
       }
@@ -80,14 +77,14 @@ export class MerkleTreesFacade implements MerkleTreeReadOperations {
     leaves: MerkleTreeLeafType<MerkleTreeId>[],
     startIndex: bigint,
   ): Promise<(bigint | undefined)[]> {
-    const indices = await this.instance.findLeafIndices(
+    const response = await this.instance.call(WorldStateMessageType.FIND_LEAF_INDICES, {
+      leaves: leaves.map(leaf => serializeLeaf(hydrateLeaf(treeId, leaf))),
+      revision: this.revision,
       treeId,
-      this.revision,
-      leaves.map(leaf => serializeLeaf(hydrateLeaf(treeId, leaf))),
       startIndex,
-    );
+    });
 
-    return indices.map(index => {
+    return response.indices.map(index => {
       if (typeof index === 'number' || typeof index === 'bigint') {
         return BigInt(index);
       } else {
@@ -97,7 +94,11 @@ export class MerkleTreesFacade implements MerkleTreeReadOperations {
   }
 
   async getLeafPreimage(treeId: IndexedTreeId, leafIndex: bigint): Promise<IndexedTreeLeafPreimage | undefined> {
-    const resp = await this.instance.getLeafPreimage(treeId, this.revision, leafIndex);
+    const resp = await this.instance.call(WorldStateMessageType.GET_LEAF_PREIMAGE, {
+      leafIndex,
+      revision: this.revision,
+      treeId,
+    });
 
     return resp ? deserializeIndexedLeaf(resp) : undefined;
   }
@@ -106,7 +107,11 @@ export class MerkleTreesFacade implements MerkleTreeReadOperations {
     treeId: ID,
     leafIndex: bigint,
   ): Promise<MerkleTreeLeafType<ID> | undefined> {
-    const resp = await this.instance.getLeafValue(treeId, this.revision, leafIndex);
+    const resp = await this.instance.call(WorldStateMessageType.GET_LEAF_VALUE, {
+      leafIndex,
+      revision: this.revision,
+      treeId,
+    });
 
     if (!resp) {
       return undefined;
@@ -124,7 +129,11 @@ export class MerkleTreesFacade implements MerkleTreeReadOperations {
     treeId: IndexedTreeId,
     value: bigint,
   ): Promise<{ index: bigint; alreadyPresent: boolean } | undefined> {
-    const resp = await this.instance.findLowLeaf(treeId, this.revision, new Fr(value));
+    const resp = await this.instance.call(WorldStateMessageType.FIND_LOW_LEAF, {
+      key: new Fr(value),
+      revision: this.revision,
+      treeId,
+    });
     return {
       alreadyPresent: resp.alreadyPresent,
       index: BigInt(resp.index),
@@ -132,48 +141,68 @@ export class MerkleTreesFacade implements MerkleTreeReadOperations {
   }
 
   async getSiblingPath<N extends number>(treeId: MerkleTreeId, leafIndex: bigint): Promise<SiblingPath<N>> {
-    const siblingPath = await this.instance.getSiblingPath(treeId, this.revision, leafIndex);
+    const siblingPath = await this.instance.call(WorldStateMessageType.GET_SIBLING_PATH, {
+      leafIndex,
+      revision: this.revision,
+      treeId,
+    });
 
     return new SiblingPath(siblingPath.length, siblingPath) as any;
   }
 
   async getStateReference(): Promise<StateReference> {
-    const state = await this.instance.getStateReference(this.revision);
+    const resp = await this.instance.call(WorldStateMessageType.GET_STATE_REFERENCE, {
+      revision: this.revision,
+    });
 
     return new StateReference(
-      treeStateReferenceToSnapshot(state[MerkleTreeId.L1_TO_L2_MESSAGE_TREE]),
+      treeStateReferenceToSnapshot(resp.state[MerkleTreeId.L1_TO_L2_MESSAGE_TREE]),
       new PartialStateReference(
-        treeStateReferenceToSnapshot(state[MerkleTreeId.NOTE_HASH_TREE]),
-        treeStateReferenceToSnapshot(state[MerkleTreeId.NULLIFIER_TREE]),
-        treeStateReferenceToSnapshot(state[MerkleTreeId.PUBLIC_DATA_TREE]),
+        treeStateReferenceToSnapshot(resp.state[MerkleTreeId.NOTE_HASH_TREE]),
+        treeStateReferenceToSnapshot(resp.state[MerkleTreeId.NULLIFIER_TREE]),
+        treeStateReferenceToSnapshot(resp.state[MerkleTreeId.PUBLIC_DATA_TREE]),
       ),
     );
   }
 
   async getInitialStateReference(): Promise<StateReference> {
-    const state = await this.instance.getInitialStateReference();
+    const resp = await this.instance.call(WorldStateMessageType.GET_INITIAL_STATE_REFERENCE, { canonical: true });
 
     return new StateReference(
-      treeStateReferenceToSnapshot(state[MerkleTreeId.L1_TO_L2_MESSAGE_TREE]),
+      treeStateReferenceToSnapshot(resp.state[MerkleTreeId.L1_TO_L2_MESSAGE_TREE]),
       new PartialStateReference(
-        treeStateReferenceToSnapshot(state[MerkleTreeId.NOTE_HASH_TREE]),
-        treeStateReferenceToSnapshot(state[MerkleTreeId.NULLIFIER_TREE]),
-        treeStateReferenceToSnapshot(state[MerkleTreeId.PUBLIC_DATA_TREE]),
+        treeStateReferenceToSnapshot(resp.state[MerkleTreeId.NOTE_HASH_TREE]),
+        treeStateReferenceToSnapshot(resp.state[MerkleTreeId.NULLIFIER_TREE]),
+        treeStateReferenceToSnapshot(resp.state[MerkleTreeId.PUBLIC_DATA_TREE]),
       ),
     );
   }
 
   async getTreeInfo(treeId: MerkleTreeId): Promise<TreeInfo> {
-    return await this.instance.getTreeInfo(treeId, this.revision);
+    const resp = await this.instance.call(WorldStateMessageType.GET_TREE_INFO, {
+      treeId: treeId,
+      revision: this.revision,
+    });
+
+    return {
+      depth: resp.depth,
+      root: resp.root,
+      size: BigInt(resp.size),
+      treeId,
+    };
   }
 
   async getBlockNumbersForLeafIndices<ID extends MerkleTreeId>(
     treeId: ID,
     leafIndices: bigint[],
   ): Promise<(BlockNumber | undefined)[]> {
-    const blockNumbers = await this.instance.getBlockNumbersForLeafIndices(treeId, this.revision, leafIndices);
+    const response = await this.instance.call(WorldStateMessageType.GET_BLOCK_NUMBERS_FOR_LEAF_INDICES, {
+      treeId,
+      revision: this.revision,
+      leafIndices,
+    });
 
-    return blockNumbers.map(x => (x === undefined || x === null ? undefined : BlockNumber(Number(x))));
+    return response.blockNumbers.map(x => (x === undefined || x === null ? undefined : BlockNumber(Number(x))));
   }
 }
 
@@ -192,19 +221,19 @@ export class MerkleTreesForkFacade extends MerkleTreesFacade implements MerkleTr
     super(instance, initialHeader, revision);
   }
   async updateArchive(header: BlockHeader): Promise<void> {
-    await this.instance.updateArchive(
-      this.revision.forkId,
-      blockStateReference(header.state),
-      (await header.hash()).toBuffer(),
-    );
+    await this.instance.call(WorldStateMessageType.UPDATE_ARCHIVE, {
+      forkId: this.revision.forkId,
+      blockHeaderHash: (await header.hash()).toBuffer(),
+      blockStateRef: blockStateReference(header.state),
+    });
   }
 
   async appendLeaves<ID extends MerkleTreeId>(treeId: ID, leaves: MerkleTreeLeafType<ID>[]): Promise<void> {
-    await this.instance.appendLeaves(
+    await this.instance.call(WorldStateMessageType.APPEND_LEAVES, {
+      leaves: leaves.map(leaf => serializeLeaf(hydrateLeaf(treeId, leaf as any))),
+      forkId: this.revision.forkId,
       treeId,
-      this.revision.forkId,
-      leaves.map(leaf => serializeLeaf(hydrateLeaf(treeId, leaf as any))),
-    );
+    });
   }
 
   async batchInsert<TreeHeight extends number, SubtreeSiblingPathHeight extends number, ID extends IndexedTreeId>(
@@ -213,19 +242,24 @@ export class MerkleTreesForkFacade extends MerkleTreesFacade implements MerkleTr
     subtreeHeight: number,
   ): Promise<BatchInsertionResult<TreeHeight, SubtreeSiblingPathHeight>> {
     const leaves = rawLeaves.map((leaf: Buffer) => hydrateLeaf(treeId, leaf)).map(serializeLeaf);
-    const resp = await this.instance.batchInsert(treeId, this.revision.forkId, leaves, subtreeHeight);
+    const resp = await this.instance.call(WorldStateMessageType.BATCH_INSERT, {
+      leaves,
+      treeId,
+      forkId: this.revision.forkId,
+      subtreeDepth: subtreeHeight,
+    });
 
     return {
       newSubtreeSiblingPath: new SiblingPath<SubtreeSiblingPathHeight>(
-        resp.subtreePath.length as any,
-        resp.subtreePath,
+        resp.subtree_path.length as any,
+        resp.subtree_path,
       ),
-      sortedNewLeaves: resp.sortedLeaves
+      sortedNewLeaves: resp.sorted_leaves
         .map(([leaf]) => leaf)
         .map(deserializeLeafValue)
         .map(serializeToBuffer),
-      sortedNewLeavesIndexes: resp.sortedLeaves.map(([, index]) => index),
-      lowLeavesWitnessData: resp.lowLeafWitnessData.map(data => ({
+      sortedNewLeavesIndexes: resp.sorted_leaves.map(([, index]) => index),
+      lowLeavesWitnessData: resp.low_leaf_witness_data.map(data => ({
         index: BigInt(data.index),
         leafPreimage: deserializeIndexedLeaf(data.leaf),
         siblingPath: new SiblingPath<TreeHeight>(data.path.length as any, data.path),
@@ -238,15 +272,19 @@ export class MerkleTreesForkFacade extends MerkleTreesFacade implements MerkleTr
     rawLeaves: Buffer[],
   ): Promise<SequentialInsertionResult<TreeHeight>> {
     const leaves = rawLeaves.map((leaf: Buffer) => hydrateLeaf(treeId, leaf)).map(serializeLeaf);
-    const resp = await this.instance.sequentialInsert(treeId, this.revision.forkId, leaves);
+    const resp = await this.instance.call(WorldStateMessageType.SEQUENTIAL_INSERT, {
+      leaves,
+      treeId,
+      forkId: this.revision.forkId,
+    });
 
     return {
-      lowLeavesWitnessData: resp.lowLeafWitnessData.map(data => ({
+      lowLeavesWitnessData: resp.low_leaf_witness_data.map(data => ({
         index: BigInt(data.index),
         leafPreimage: deserializeIndexedLeaf(data.leaf),
         siblingPath: new SiblingPath<TreeHeight>(data.path.length as any, data.path),
       })),
-      insertionWitnessData: resp.insertionWitnessData.map(data => ({
+      insertionWitnessData: resp.insertion_witness_data.map(data => ({
         index: BigInt(data.index),
         leafPreimage: deserializeIndexedLeaf(data.leaf),
         siblingPath: new SiblingPath<TreeHeight>(data.path.length as any, data.path),
@@ -266,11 +304,11 @@ export class MerkleTreesForkFacade extends MerkleTreesFacade implements MerkleTr
 
   private async doClose(): Promise<void> {
     try {
-      await this.instance.deleteFork(this.revision.forkId);
+      await this.instance.call(WorldStateMessageType.DELETE_FORK, { forkId: this.revision.forkId });
     } catch (err: any) {
       // Ignore errors due to native instance being closed during shutdown.
       // This can happen when validators are still processing block proposals while the node is stopping.
-      if (err?.message === 'Native instance is closed' || err?.message === 'IPC instance is closed') {
+      if (err?.message === 'Native instance is closed') {
         return;
       }
       // Ignore "Fork not found": the native fork was already destroyed by a pending-chain unwind or a
@@ -297,27 +335,34 @@ export class MerkleTreesForkFacade extends MerkleTreesFacade implements MerkleTr
 
   public async createCheckpoint(): Promise<number> {
     assert.notEqual(this.revision.forkId, 0, 'Fork ID must be set');
-    return await this.instance.createCheckpoint(this.revision.forkId);
+    const resp = await this.instance.call(WorldStateMessageType.CREATE_CHECKPOINT, { forkId: this.revision.forkId });
+    return resp.depth;
   }
 
   public async commitCheckpoint(): Promise<void> {
     assert.notEqual(this.revision.forkId, 0, 'Fork ID must be set');
-    await this.instance.commitCheckpoint(this.revision.forkId);
+    await this.instance.call(WorldStateMessageType.COMMIT_CHECKPOINT, { forkId: this.revision.forkId });
   }
 
   public async revertCheckpoint(): Promise<void> {
     assert.notEqual(this.revision.forkId, 0, 'Fork ID must be set');
-    await this.instance.revertCheckpoint(this.revision.forkId);
+    await this.instance.call(WorldStateMessageType.REVERT_CHECKPOINT, { forkId: this.revision.forkId });
   }
 
   public async commitAllCheckpointsTo(depth: number): Promise<void> {
     assert.notEqual(this.revision.forkId, 0, 'Fork ID must be set');
-    await this.instance.commitAllCheckpoints(this.revision.forkId, depth);
+    await this.instance.call(WorldStateMessageType.COMMIT_ALL_CHECKPOINTS, {
+      forkId: this.revision.forkId,
+      depth,
+    });
   }
 
   public async revertAllCheckpointsTo(depth: number): Promise<void> {
     assert.notEqual(this.revision.forkId, 0, 'Fork ID must be set');
-    await this.instance.revertAllCheckpoints(this.revision.forkId, depth);
+    await this.instance.call(WorldStateMessageType.REVERT_ALL_CHECKPOINTS, {
+      forkId: this.revision.forkId,
+      depth,
+    });
   }
 }
 

package/src/native/message.ts

@@ -1,7 +1,59 @@
 import { BlockNumber } from '@aztec/foundation/branded-types';
 import { Fr } from '@aztec/foundation/curves/bn254';
+import type { Tuple } from '@aztec/foundation/serialize';
 import { AppendOnlyTreeSnapshot, MerkleTreeId } from '@aztec/stdlib/trees';
 import type { StateReference } from '@aztec/stdlib/tx';
+import type { UInt32 } from '@aztec/stdlib/types';
+import type { WorldStateRevision } from '@aztec/stdlib/world-state';
+
+export enum WorldStateMessageType {
+  GET_TREE_INFO = 100,
+  GET_STATE_REFERENCE,
+  GET_INITIAL_STATE_REFERENCE,
+
+  GET_LEAF_VALUE,
+  GET_LEAF_PREIMAGE,
+  GET_SIBLING_PATH,
+  GET_BLOCK_NUMBERS_FOR_LEAF_INDICES,
+
+  FIND_LEAF_INDICES,
+  FIND_LOW_LEAF,
+  FIND_SIBLING_PATHS,
+
+  APPEND_LEAVES,
+  BATCH_INSERT,
+  SEQUENTIAL_INSERT,
+
+  UPDATE_ARCHIVE,
+
+  COMMIT,
+  ROLLBACK,
+
+  SYNC_BLOCK,
+
+  CREATE_FORK,
+  DELETE_FORK,
+
+  FINALIZE_BLOCKS,
+  UNWIND_BLOCKS,
+  REMOVE_HISTORICAL_BLOCKS,
+
+  GET_STATUS,
+
+  CREATE_CHECKPOINT,
+  COMMIT_CHECKPOINT,
+  REVERT_CHECKPOINT,
+  COMMIT_ALL_CHECKPOINTS,
+  REVERT_ALL_CHECKPOINTS,
+
+  COPY_STORES,
+
+  CLOSE = 999,
+}
+
+interface WithTreeId {
+  treeId: MerkleTreeId;
+}
 
 export interface WorldStateStatusSummary {
   /** Last block number that can still be unwound. */
@@ -227,6 +279,32 @@ export function sanitizeFullStatus(status: WorldStateStatusFull) {
   return status;
 }
 
+interface WithForkId {
+  forkId: number;
+}
+
+interface CreateCheckpointResponse {
+  depth: number;
+}
+
+/** Request to commit/revert all checkpoints down to a target depth. The resulting depth after the operation equals the given depth. */
+interface CheckpointDepthRequest extends WithForkId {
+  /** The target depth after the operation. All checkpoints above this depth are committed/reverted. */
+  depth: number;
+}
+
+interface WithWorldStateRevision {
+  revision: WorldStateRevision;
+}
+
+interface WithCanonicalForkId {
+  canonical: true;
+}
+
+interface WithLeafIndex {
+  leafIndex: bigint;
+}
+
 export type SerializedLeafValue =
   | Buffer // Fr
   | { nullifier: Buffer } // NullifierLeaf
@@ -238,36 +316,245 @@ export type SerializedIndexedLeaf = {
   nextKey: Buffer; // Fr
 };
 
-export interface SerializedSiblingPathAndIndex {
+interface WithLeafValues {
+  leaves: SerializedLeafValue[];
+}
+
+interface BlockShiftRequest extends WithCanonicalForkId {
+  toBlockNumber: BlockNumber;
+}
+
+interface WithLeaves {
+  leaves: SerializedLeafValue[];
+}
+
+interface GetTreeInfoRequest extends WithTreeId, WithWorldStateRevision {}
+interface GetTreeInfoResponse {
+  treeId: MerkleTreeId;
+  depth: UInt32;
+  size: bigint | number;
+  root: Buffer;
+}
+
+interface GetBlockNumbersForLeafIndicesRequest extends WithTreeId, WithWorldStateRevision {
+  leafIndices: bigint[];
+}
+
+interface GetBlockNumbersForLeafIndicesResponse {
+  blockNumbers: bigint[];
+}
+
+interface GetSiblingPathRequest extends WithTreeId, WithLeafIndex, WithWorldStateRevision {}
+type GetSiblingPathResponse = Buffer[];
+
+interface GetStateReferenceRequest extends WithWorldStateRevision {}
+interface GetStateReferenceResponse {
+  state: Record<MerkleTreeId, TreeStateReference>;
+}
+
+interface GetLeafRequest extends WithTreeId, WithWorldStateRevision, WithLeafIndex {}
+type GetLeafResponse = SerializedLeafValue | undefined;
+
+interface GetLeafPreImageRequest extends WithTreeId, WithLeafIndex, WithWorldStateRevision {}
+type GetLeafPreImageResponse = SerializedIndexedLeaf | undefined;
+
+interface FindLeafIndicesRequest extends WithTreeId, WithLeafValues, WithWorldStateRevision {
+  startIndex: bigint;
+}
+interface FindLeafIndicesResponse {
+  indices: bigint[];
+}
+
+interface FindSiblingPathsRequest extends WithTreeId, WithLeafValues, WithWorldStateRevision {}
+
+interface SiblingPathAndIndex {
   index: bigint;
   path: Buffer[];
 }
+interface FindSiblingPathsResponse {
+  paths: (SiblingPathAndIndex | undefined)[];
+}
 
-export interface SerializedBatchInsertionResult {
-  lowLeafWitnessData: ReadonlyArray<{
+interface FindLowLeafRequest extends WithTreeId, WithWorldStateRevision {
+  key: Fr;
+}
+interface FindLowLeafResponse {
+  index: bigint | number;
+  alreadyPresent: boolean;
+}
+
+interface AppendLeavesRequest extends WithTreeId, WithForkId, WithLeaves {}
+
+interface BatchInsertRequest extends WithTreeId, WithForkId, WithLeaves {
+  subtreeDepth: number;
+}
+
+interface BatchInsertResponse {
+  low_leaf_witness_data: ReadonlyArray<{
     leaf: SerializedIndexedLeaf;
     index: bigint | number;
-    path: Buffer[];
+    path: Tuple<Buffer, number>;
   }>;
-  sortedLeaves: ReadonlyArray<readonly [SerializedLeafValue, number]>;
-  subtreePath: Buffer[];
+  sorted_leaves: ReadonlyArray<[SerializedLeafValue, UInt32]>;
+  subtree_path: Tuple<Buffer, number>;
 }
 
-export interface SerializedSequentialInsertionResult {
-  lowLeafWitnessData: ReadonlyArray<{
+interface SequentialInsertRequest extends WithTreeId, WithForkId, WithLeaves {}
+
+interface SequentialInsertResponse {
+  low_leaf_witness_data: ReadonlyArray<{
     leaf: SerializedIndexedLeaf;
     index: bigint | number;
-    path: Buffer[];
+    path: Tuple<Buffer, number>;
   }>;
-  insertionWitnessData: ReadonlyArray<{
+  insertion_witness_data: ReadonlyArray<{
     leaf: SerializedIndexedLeaf;
     index: bigint | number;
-    path: Buffer[];
+    path: Tuple<Buffer, number>;
   }>;
 }
 
-export type TreeStateReference = readonly [Buffer, number | bigint];
-export type BlockStateReference = Map<Exclude<MerkleTreeId, MerkleTreeId.ARCHIVE>, TreeStateReference>;
+interface UpdateArchiveRequest extends WithForkId {
+  blockStateRef: BlockStateReference;
+  blockHeaderHash: Buffer;
+}
+
+interface SyncBlockRequest extends WithCanonicalForkId {
+  blockNumber: BlockNumber;
+  blockStateRef: BlockStateReference;
+  blockHeaderHash: Buffer;
+  /** Canonical archive root after this block; world state rejects the sync if its computed root differs. */
+  expectedArchiveRoot: Buffer;
+  /** Canonical archive root before this block (the block's lastArchive); verified against the local root. */
+  expectedPreviousArchiveRoot: Buffer;
+  paddedNoteHashes: readonly SerializedLeafValue[];
+  paddedL1ToL2Messages: readonly SerializedLeafValue[];
+  paddedNullifiers: readonly SerializedLeafValue[];
+  publicDataWrites: readonly SerializedLeafValue[];
+}
+
+interface CreateForkRequest extends WithCanonicalForkId {
+  latest: boolean;
+  blockNumber: BlockNumber;
+}
+
+interface CreateForkResponse {
+  forkId: number;
+}
+
+interface DeleteForkRequest extends WithForkId {}
+
+interface CopyStoresRequest extends WithCanonicalForkId {
+  dstPath: string;
+  compact: boolean;
+}
+
+export type WorldStateRequestCategories = WithForkId | WithWorldStateRevision | WithCanonicalForkId;
+
+export function isWithForkId(body: WorldStateRequestCategories): body is WithForkId {
+  return body && 'forkId' in body;
+}
+
+export function isWithRevision(body: WorldStateRequestCategories): body is WithWorldStateRevision {
+  return body && 'revision' in body;
+}
+
+export function isWithCanonical(body: WorldStateRequestCategories): body is WithCanonicalForkId {
+  return body && 'canonical' in body;
+}
+
+export type WorldStateRequest = {
+  [WorldStateMessageType.GET_TREE_INFO]: GetTreeInfoRequest;
+  [WorldStateMessageType.GET_STATE_REFERENCE]: GetStateReferenceRequest;
+  [WorldStateMessageType.GET_INITIAL_STATE_REFERENCE]: WithCanonicalForkId;
+
+  [WorldStateMessageType.GET_LEAF_VALUE]: GetLeafRequest;
+  [WorldStateMessageType.GET_LEAF_PREIMAGE]: GetLeafPreImageRequest;
+  [WorldStateMessageType.GET_SIBLING_PATH]: GetSiblingPathRequest;
+  [WorldStateMessageType.GET_BLOCK_NUMBERS_FOR_LEAF_INDICES]: GetBlockNumbersForLeafIndicesRequest;
+
+  [WorldStateMessageType.FIND_LEAF_INDICES]: FindLeafIndicesRequest;
+  [WorldStateMessageType.FIND_LOW_LEAF]: FindLowLeafRequest;
+  [WorldStateMessageType.FIND_SIBLING_PATHS]: FindSiblingPathsRequest;
+
+  [WorldStateMessageType.APPEND_LEAVES]: AppendLeavesRequest;
+  [WorldStateMessageType.BATCH_INSERT]: BatchInsertRequest;
+  [WorldStateMessageType.SEQUENTIAL_INSERT]: SequentialInsertRequest;
+
+  [WorldStateMessageType.UPDATE_ARCHIVE]: UpdateArchiveRequest;
+
+  [WorldStateMessageType.COMMIT]: WithCanonicalForkId;
+  [WorldStateMessageType.ROLLBACK]: WithCanonicalForkId;
+
+  [WorldStateMessageType.SYNC_BLOCK]: SyncBlockRequest;
+
+  [WorldStateMessageType.CREATE_FORK]: CreateForkRequest;
+  [WorldStateMessageType.DELETE_FORK]: DeleteForkRequest;
+
+  [WorldStateMessageType.REMOVE_HISTORICAL_BLOCKS]: BlockShiftRequest;
+  [WorldStateMessageType.UNWIND_BLOCKS]: BlockShiftRequest;
+  [WorldStateMessageType.FINALIZE_BLOCKS]: BlockShiftRequest;
+
+  [WorldStateMessageType.GET_STATUS]: WithCanonicalForkId;
+
+  [WorldStateMessageType.CREATE_CHECKPOINT]: WithForkId;
+  [WorldStateMessageType.COMMIT_CHECKPOINT]: WithForkId;
+  [WorldStateMessageType.REVERT_CHECKPOINT]: WithForkId;
+  [WorldStateMessageType.COMMIT_ALL_CHECKPOINTS]: CheckpointDepthRequest;
+  [WorldStateMessageType.REVERT_ALL_CHECKPOINTS]: CheckpointDepthRequest;
+
+  [WorldStateMessageType.COPY_STORES]: CopyStoresRequest;
+
+  [WorldStateMessageType.CLOSE]: WithCanonicalForkId;
+};
+
+export type WorldStateResponse = {
+  [WorldStateMessageType.GET_TREE_INFO]: GetTreeInfoResponse;
+  [WorldStateMessageType.GET_STATE_REFERENCE]: GetStateReferenceResponse;
+  [WorldStateMessageType.GET_INITIAL_STATE_REFERENCE]: GetStateReferenceResponse;
+
+  [WorldStateMessageType.GET_LEAF_VALUE]: GetLeafResponse;
+  [WorldStateMessageType.GET_LEAF_PREIMAGE]: GetLeafPreImageResponse;
+  [WorldStateMessageType.GET_SIBLING_PATH]: GetSiblingPathResponse;
+  [WorldStateMessageType.GET_BLOCK_NUMBERS_FOR_LEAF_INDICES]: GetBlockNumbersForLeafIndicesResponse;
+
+  [WorldStateMessageType.FIND_LEAF_INDICES]: FindLeafIndicesResponse;
+  [WorldStateMessageType.FIND_LOW_LEAF]: FindLowLeafResponse;
+  [WorldStateMessageType.FIND_SIBLING_PATHS]: FindSiblingPathsResponse;
+
+  [WorldStateMessageType.APPEND_LEAVES]: void;
+  [WorldStateMessageType.BATCH_INSERT]: BatchInsertResponse;
+  [WorldStateMessageType.SEQUENTIAL_INSERT]: SequentialInsertResponse;
+
+  [WorldStateMessageType.UPDATE_ARCHIVE]: void;
+
+  [WorldStateMessageType.COMMIT]: void;
+  [WorldStateMessageType.ROLLBACK]: void;
+
+  [WorldStateMessageType.SYNC_BLOCK]: WorldStateStatusFull;
+
+  [WorldStateMessageType.CREATE_FORK]: CreateForkResponse;
+  [WorldStateMessageType.DELETE_FORK]: void;
+
+  [WorldStateMessageType.REMOVE_HISTORICAL_BLOCKS]: WorldStateStatusFull;
+  [WorldStateMessageType.UNWIND_BLOCKS]: WorldStateStatusFull;
+  [WorldStateMessageType.FINALIZE_BLOCKS]: WorldStateStatusSummary;
+
+  [WorldStateMessageType.GET_STATUS]: WorldStateStatusSummary;
+
+  [WorldStateMessageType.CREATE_CHECKPOINT]: CreateCheckpointResponse;
+  [WorldStateMessageType.COMMIT_CHECKPOINT]: void;
+  [WorldStateMessageType.REVERT_CHECKPOINT]: void;
+  [WorldStateMessageType.COMMIT_ALL_CHECKPOINTS]: void;
+  [WorldStateMessageType.REVERT_ALL_CHECKPOINTS]: void;
+
+  [WorldStateMessageType.COPY_STORES]: void;
+
+  [WorldStateMessageType.CLOSE]: void;
+};
+
+type TreeStateReference = readonly [Buffer, number | bigint];
+type BlockStateReference = Map<Exclude<MerkleTreeId, MerkleTreeId.ARCHIVE>, TreeStateReference>;
 
 export function treeStateReferenceToSnapshot([root, size]: TreeStateReference): AppendOnlyTreeSnapshot {
   return new AppendOnlyTreeSnapshot(Fr.fromBuffer(root), Number(size));