@aztec/world-state 0.0.1-commit.d3ec352c → 0.0.1-commit.d58ff9d0

package/src/native/merkle_trees_facade.ts

@@ -1,7 +1,10 @@
 import { BlockNumber } from '@aztec/foundation/branded-types';
-import { Fr } from '@aztec/foundation/fields';
+import { Fr } from '@aztec/foundation/curves/bn254';
+import { createLogger } from '@aztec/foundation/log';
 import { serializeToBuffer } from '@aztec/foundation/serialize';
+import { sleep } from '@aztec/foundation/sleep';
 import { type IndexedTreeLeafPreimage, SiblingPath } from '@aztec/foundation/trees';
+import { BlockHash } from '@aztec/stdlib/block';
 import type {
   BatchInsertionResult,
   IndexedTreeId,
@@ -19,14 +22,13 @@ import {
   PublicDataTreeLeafPreimage,
 } from '@aztec/stdlib/trees';
 import { type BlockHeader, PartialStateReference, StateReference } from '@aztec/stdlib/tx';
-import { type WorldStateRevision, WorldStateRevisionWithHandle } from '@aztec/stdlib/world-state';
+import type { WorldStateRevision } from '@aztec/stdlib/world-state';
 
 import assert from 'assert';
 
 import {
   type SerializedIndexedLeaf,
   type SerializedLeafValue,
-  WorldStateMessageType,
   blockStateReference,
   treeStateReferenceToSnapshot,
 } from './message.js';
@@ -43,8 +45,12 @@ export class MerkleTreesFacade implements MerkleTreeReadOperations {
     return this.initialHeader;
   }
 
-  getRevision(): WorldStateRevisionWithHandle {
-    return WorldStateRevisionWithHandle.fromWorldStateRevision(this.revision, this.instance.getHandle());
+  getRevision(): WorldStateRevision {
+    return this.revision;
+  }
+
+  getIpcPath(): string {
+    return this.instance.getIpcPath();
   }
 
   findLeafIndices(treeId: MerkleTreeId, values: MerkleTreeLeafType<MerkleTreeId>[]): Promise<(bigint | undefined)[]> {
@@ -55,13 +61,13 @@ export class MerkleTreesFacade implements MerkleTreeReadOperations {
     treeId: MerkleTreeId,
     values: MerkleTreeLeafType<MerkleTreeId>[],
   ): Promise<({ path: SiblingPath<N>; index: bigint } | undefined)[]> {
-    const response = await this.instance.call(WorldStateMessageType.FIND_SIBLING_PATHS, {
-      leaves: values.map(leaf => serializeLeaf(hydrateLeaf(treeId, leaf))),
-      revision: this.revision,
+    const paths = await this.instance.findSiblingPaths(
       treeId,
-    });
+      this.revision,
+      values.map(leaf => serializeLeaf(hydrateLeaf(treeId, leaf))),
+    );
 
-    return response.paths.map(path => {
+    return paths.map(path => {
       if (!path) {
         return undefined;
       }
@@ -74,14 +80,14 @@ export class MerkleTreesFacade implements MerkleTreeReadOperations {
     leaves: MerkleTreeLeafType<MerkleTreeId>[],
     startIndex: bigint,
   ): Promise<(bigint | undefined)[]> {
-    const response = await this.instance.call(WorldStateMessageType.FIND_LEAF_INDICES, {
-      leaves: leaves.map(leaf => serializeLeaf(hydrateLeaf(treeId, leaf))),
-      revision: this.revision,
+    const indices = await this.instance.findLeafIndices(
       treeId,
+      this.revision,
+      leaves.map(leaf => serializeLeaf(hydrateLeaf(treeId, leaf))),
       startIndex,
-    });
+    );
 
-    return response.indices.map(index => {
+    return indices.map(index => {
       if (typeof index === 'number' || typeof index === 'bigint') {
         return BigInt(index);
       } else {
@@ -91,11 +97,7 @@ export class MerkleTreesFacade implements MerkleTreeReadOperations {
   }
 
   async getLeafPreimage(treeId: IndexedTreeId, leafIndex: bigint): Promise<IndexedTreeLeafPreimage | undefined> {
-    const resp = await this.instance.call(WorldStateMessageType.GET_LEAF_PREIMAGE, {
-      leafIndex,
-      revision: this.revision,
-      treeId,
-    });
+    const resp = await this.instance.getLeafPreimage(treeId, this.revision, leafIndex);
 
     return resp ? deserializeIndexedLeaf(resp) : undefined;
   }
@@ -104,11 +106,7 @@ export class MerkleTreesFacade implements MerkleTreeReadOperations {
     treeId: ID,
     leafIndex: bigint,
   ): Promise<MerkleTreeLeafType<ID> | undefined> {
-    const resp = await this.instance.call(WorldStateMessageType.GET_LEAF_VALUE, {
-      leafIndex,
-      revision: this.revision,
-      treeId,
-    });
+    const resp = await this.instance.getLeafValue(treeId, this.revision, leafIndex);
 
     if (!resp) {
       return undefined;
@@ -116,7 +114,7 @@ export class MerkleTreesFacade implements MerkleTreeReadOperations {
 
     const leaf = deserializeLeafValue(resp);
     if (leaf instanceof Fr) {
-      return leaf as any;
+      return treeId === MerkleTreeId.ARCHIVE ? (new BlockHash(leaf) as any) : (leaf as any);
     } else {
       return leaf.toBuffer() as any;
     }
@@ -126,11 +124,7 @@ export class MerkleTreesFacade implements MerkleTreeReadOperations {
     treeId: IndexedTreeId,
     value: bigint,
   ): Promise<{ index: bigint; alreadyPresent: boolean } | undefined> {
-    const resp = await this.instance.call(WorldStateMessageType.FIND_LOW_LEAF, {
-      key: new Fr(value),
-      revision: this.revision,
-      treeId,
-    });
+    const resp = await this.instance.findLowLeaf(treeId, this.revision, new Fr(value));
     return {
       alreadyPresent: resp.alreadyPresent,
       index: BigInt(resp.index),
@@ -138,91 +132,79 @@ export class MerkleTreesFacade implements MerkleTreeReadOperations {
   }
 
   async getSiblingPath<N extends number>(treeId: MerkleTreeId, leafIndex: bigint): Promise<SiblingPath<N>> {
-    const siblingPath = await this.instance.call(WorldStateMessageType.GET_SIBLING_PATH, {
-      leafIndex,
-      revision: this.revision,
-      treeId,
-    });
+    const siblingPath = await this.instance.getSiblingPath(treeId, this.revision, leafIndex);
 
     return new SiblingPath(siblingPath.length, siblingPath) as any;
   }
 
   async getStateReference(): Promise<StateReference> {
-    const resp = await this.instance.call(WorldStateMessageType.GET_STATE_REFERENCE, {
-      revision: this.revision,
-    });
+    const state = await this.instance.getStateReference(this.revision);
 
     return new StateReference(
-      treeStateReferenceToSnapshot(resp.state[MerkleTreeId.L1_TO_L2_MESSAGE_TREE]),
+      treeStateReferenceToSnapshot(state[MerkleTreeId.L1_TO_L2_MESSAGE_TREE]),
       new PartialStateReference(
-        treeStateReferenceToSnapshot(resp.state[MerkleTreeId.NOTE_HASH_TREE]),
-        treeStateReferenceToSnapshot(resp.state[MerkleTreeId.NULLIFIER_TREE]),
-        treeStateReferenceToSnapshot(resp.state[MerkleTreeId.PUBLIC_DATA_TREE]),
+        treeStateReferenceToSnapshot(state[MerkleTreeId.NOTE_HASH_TREE]),
+        treeStateReferenceToSnapshot(state[MerkleTreeId.NULLIFIER_TREE]),
+        treeStateReferenceToSnapshot(state[MerkleTreeId.PUBLIC_DATA_TREE]),
       ),
     );
   }
 
   async getInitialStateReference(): Promise<StateReference> {
-    const resp = await this.instance.call(WorldStateMessageType.GET_INITIAL_STATE_REFERENCE, { canonical: true });
+    const state = await this.instance.getInitialStateReference();
 
     return new StateReference(
-      treeStateReferenceToSnapshot(resp.state[MerkleTreeId.L1_TO_L2_MESSAGE_TREE]),
+      treeStateReferenceToSnapshot(state[MerkleTreeId.L1_TO_L2_MESSAGE_TREE]),
       new PartialStateReference(
-        treeStateReferenceToSnapshot(resp.state[MerkleTreeId.NOTE_HASH_TREE]),
-        treeStateReferenceToSnapshot(resp.state[MerkleTreeId.NULLIFIER_TREE]),
-        treeStateReferenceToSnapshot(resp.state[MerkleTreeId.PUBLIC_DATA_TREE]),
+        treeStateReferenceToSnapshot(state[MerkleTreeId.NOTE_HASH_TREE]),
+        treeStateReferenceToSnapshot(state[MerkleTreeId.NULLIFIER_TREE]),
+        treeStateReferenceToSnapshot(state[MerkleTreeId.PUBLIC_DATA_TREE]),
       ),
     );
   }
 
   async getTreeInfo(treeId: MerkleTreeId): Promise<TreeInfo> {
-    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,
-    };
+    return await this.instance.getTreeInfo(treeId, this.revision);
   }
 
   async getBlockNumbersForLeafIndices<ID extends MerkleTreeId>(
     treeId: ID,
     leafIndices: bigint[],
   ): Promise<(BlockNumber | undefined)[]> {
-    const response = await this.instance.call(WorldStateMessageType.GET_BLOCK_NUMBERS_FOR_LEAF_INDICES, {
-      treeId,
-      revision: this.revision,
-      leafIndices,
-    });
+    const blockNumbers = await this.instance.getBlockNumbersForLeafIndices(treeId, this.revision, leafIndices);
 
-    return response.blockNumbers.map(x => (x === undefined || x === null ? undefined : BlockNumber(Number(x))));
+    return blockNumbers.map(x => (x === undefined || x === null ? undefined : BlockNumber(Number(x))));
   }
 }
 
 export class MerkleTreesForkFacade extends MerkleTreesFacade implements MerkleTreeWriteOperations {
-  constructor(instance: NativeWorldStateInstance, initialHeader: BlockHeader, revision: WorldStateRevision) {
+  private log = createLogger('world-state:merkle-trees-fork-facade');
+  private closePromise: Promise<void> | undefined;
+
+  constructor(
+    instance: NativeWorldStateInstance,
+    initialHeader: BlockHeader,
+    revision: WorldStateRevision,
+    private opts: { closeDelayMs?: number },
+  ) {
     assert.notEqual(revision.forkId, 0, 'Fork ID must be set');
     assert.equal(revision.includeUncommitted, true, 'Fork must include uncommitted data');
     super(instance, initialHeader, revision);
   }
   async updateArchive(header: BlockHeader): Promise<void> {
-    await this.instance.call(WorldStateMessageType.UPDATE_ARCHIVE, {
-      forkId: this.revision.forkId,
-      blockHeaderHash: (await header.hash()).toBuffer(),
-      blockStateRef: blockStateReference(header.state),
-    });
+    await this.instance.updateArchive(
+      this.revision.forkId,
+      blockStateReference(header.state),
+      (await header.hash()).toBuffer(),
+    );
   }
 
   async appendLeaves<ID extends MerkleTreeId>(treeId: ID, leaves: MerkleTreeLeafType<ID>[]): Promise<void> {
-    await this.instance.call(WorldStateMessageType.APPEND_LEAVES, {
-      leaves: leaves.map(leaf => leaf as any),
-      forkId: this.revision.forkId,
+    await this.instance.appendLeaves(
       treeId,
-    });
+      this.revision.forkId,
+      leaves.map(leaf => serializeLeaf(hydrateLeaf(treeId, leaf as any))),
+    );
   }
 
   async batchInsert<TreeHeight extends number, SubtreeSiblingPathHeight extends number, ID extends IndexedTreeId>(
@@ -231,24 +213,19 @@ 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.call(WorldStateMessageType.BATCH_INSERT, {
-      leaves,
-      treeId,
-      forkId: this.revision.forkId,
-      subtreeDepth: subtreeHeight,
-    });
+    const resp = await this.instance.batchInsert(treeId, this.revision.forkId, leaves, subtreeHeight);
 
     return {
       newSubtreeSiblingPath: new SiblingPath<SubtreeSiblingPathHeight>(
-        resp.subtree_path.length as any,
-        resp.subtree_path,
+        resp.subtreePath.length as any,
+        resp.subtreePath,
       ),
-      sortedNewLeaves: resp.sorted_leaves
+      sortedNewLeaves: resp.sortedLeaves
         .map(([leaf]) => leaf)
         .map(deserializeLeafValue)
         .map(serializeToBuffer),
-      sortedNewLeavesIndexes: resp.sorted_leaves.map(([, index]) => index),
-      lowLeavesWitnessData: resp.low_leaf_witness_data.map(data => ({
+      sortedNewLeavesIndexes: resp.sortedLeaves.map(([, index]) => index),
+      lowLeavesWitnessData: resp.lowLeafWitnessData.map(data => ({
         index: BigInt(data.index),
         leafPreimage: deserializeIndexedLeaf(data.leaf),
         siblingPath: new SiblingPath<TreeHeight>(data.path.length as any, data.path),
@@ -261,19 +238,15 @@ 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.call(WorldStateMessageType.SEQUENTIAL_INSERT, {
-      leaves,
-      treeId,
-      forkId: this.revision.forkId,
-    });
+    const resp = await this.instance.sequentialInsert(treeId, this.revision.forkId, leaves);
 
     return {
-      lowLeavesWitnessData: resp.low_leaf_witness_data.map(data => ({
+      lowLeavesWitnessData: resp.lowLeafWitnessData.map(data => ({
         index: BigInt(data.index),
         leafPreimage: deserializeIndexedLeaf(data.leaf),
         siblingPath: new SiblingPath<TreeHeight>(data.path.length as any, data.path),
       })),
-      insertionWitnessData: resp.insertion_witness_data.map(data => ({
+      insertionWitnessData: resp.insertionWitnessData.map(data => ({
         index: BigInt(data.index),
         leafPreimage: deserializeIndexedLeaf(data.leaf),
         siblingPath: new SiblingPath<TreeHeight>(data.path.length as any, data.path),
@@ -281,40 +254,78 @@ export class MerkleTreesForkFacade extends MerkleTreesFacade implements MerkleTr
     };
   }
 
-  public async close(): Promise<void> {
+  public close(): Promise<void> {
     assert.notEqual(this.revision.forkId, 0, 'Fork ID must be set');
-    await this.instance.call(WorldStateMessageType.DELETE_FORK, { forkId: this.revision.forkId });
+    // Share the in-flight close promise across duplicate dispose calls so DELETE_FORK is sent at most once.
+    if (this.closePromise) {
+      return this.closePromise;
+    }
+    this.closePromise = this.doClose();
+    return this.closePromise;
   }
 
-  public async createCheckpoint(): Promise<void> {
+  private async doClose(): Promise<void> {
+    try {
+      await this.instance.deleteFork(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') {
+        return;
+      }
+      // Ignore "Fork not found": the native fork was already destroyed by a pending-chain unwind or a
+      // historical prune (both call C++ remove_forks_for_block). Fork IDs are monotonic and never reused,
+      // so swallowing this on close cannot mask a deletion of a different fork.
+      if (err?.message === 'Fork not found') {
+        return;
+      }
+      throw err;
+    }
+  }
+
+  async [Symbol.asyncDispose](): Promise<void> {
+    if (this.opts.closeDelayMs) {
+      void sleep(this.opts.closeDelayMs)
+        .then(() => this.close())
+        .catch(err => {
+          this.log.warn('Error closing MerkleTreesForkFacade after delay', { err });
+        });
+    } else {
+      await this.close();
+    }
+  }
+
+  public async createCheckpoint(): Promise<number> {
     assert.notEqual(this.revision.forkId, 0, 'Fork ID must be set');
-    await this.instance.call(WorldStateMessageType.CREATE_CHECKPOINT, { forkId: this.revision.forkId });
+    return await this.instance.createCheckpoint(this.revision.forkId);
   }
 
   public async commitCheckpoint(): Promise<void> {
     assert.notEqual(this.revision.forkId, 0, 'Fork ID must be set');
-    await this.instance.call(WorldStateMessageType.COMMIT_CHECKPOINT, { forkId: this.revision.forkId });
+    await this.instance.commitCheckpoint(this.revision.forkId);
   }
 
   public async revertCheckpoint(): Promise<void> {
     assert.notEqual(this.revision.forkId, 0, 'Fork ID must be set');
-    await this.instance.call(WorldStateMessageType.REVERT_CHECKPOINT, { forkId: this.revision.forkId });
+    await this.instance.revertCheckpoint(this.revision.forkId);
   }
 
-  public async commitAllCheckpoints(): Promise<void> {
+  public async commitAllCheckpointsTo(depth: number): Promise<void> {
     assert.notEqual(this.revision.forkId, 0, 'Fork ID must be set');
-    await this.instance.call(WorldStateMessageType.COMMIT_ALL_CHECKPOINTS, { forkId: this.revision.forkId });
+    await this.instance.commitAllCheckpoints(this.revision.forkId, depth);
   }
 
-  public async revertAllCheckpoints(): Promise<void> {
+  public async revertAllCheckpointsTo(depth: number): Promise<void> {
     assert.notEqual(this.revision.forkId, 0, 'Fork ID must be set');
-    await this.instance.call(WorldStateMessageType.REVERT_ALL_CHECKPOINTS, { forkId: this.revision.forkId });
+    await this.instance.revertAllCheckpoints(this.revision.forkId, depth);
   }
 }
 
-function hydrateLeaf<ID extends MerkleTreeId>(treeId: ID, leaf: Fr | Buffer) {
+function hydrateLeaf(treeId: MerkleTreeId, leaf: Fr | BlockHash | Buffer) {
   if (leaf instanceof Fr) {
     return leaf;
+  } else if (leaf instanceof BlockHash) {
+    return leaf.toFr();
   } else if (treeId === MerkleTreeId.NULLIFIER_TREE) {
     return NullifierLeaf.fromBuffer(leaf);
   } else if (treeId === MerkleTreeId.PUBLIC_DATA_TREE) {
@@ -324,8 +335,10 @@ function hydrateLeaf<ID extends MerkleTreeId>(treeId: ID, leaf: Fr | Buffer) {
   }
 }
 
-export function serializeLeaf(leaf: Fr | NullifierLeaf | PublicDataTreeLeaf): SerializedLeafValue {
-  if (leaf instanceof Fr) {
+export function serializeLeaf(leaf: Fr | BlockHash | NullifierLeaf | PublicDataTreeLeaf): SerializedLeafValue {
+  if (leaf instanceof BlockHash) {
+    return leaf.toBuffer();
+  } else if (leaf instanceof Fr) {
     return leaf.toBuffer();
   } else if (leaf instanceof NullifierLeaf) {
     return { nullifier: leaf.nullifier.toBuffer() };