Remove anys in getSnapshot

AztecProtocol · Mar 26, 2024 · 0060129 · 0060129
1 parent 93ce80a
commit 0060129
Show file tree

Hide file tree

Showing 4 changed files with 21 additions and 13 deletions.
diff --git a/yarn-project/end-to-end/src/e2e_dapp_subscription.test.ts b/yarn-project/end-to-end/src/e2e_dapp_subscription.test.ts
@@ -28,7 +28,7 @@ import {
   setup,
 } from './fixtures/utils.js';
 
-jest.setTimeout(1_000_000);
+jest.setTimeout(100_000);
 
 const TOKEN_NAME = 'BananaCoin';
 const TOKEN_SYMBOL = 'BC';

diff --git a/yarn-project/merkle-tree/src/interfaces/indexed_tree.ts b/yarn-project/merkle-tree/src/interfaces/indexed_tree.ts
@@ -1,7 +1,9 @@
 import { SiblingPath } from '@aztec/circuit-types';
 import { IndexedTreeLeaf, IndexedTreeLeafPreimage } from '@aztec/foundation/trees';
 
+import { IndexedTreeSnapshot, TreeSnapshot, TreeSnapshotBuilder } from '../snapshots/snapshot_builder.js';
 import { AppendOnlyTree } from './append_only_tree.js';
+import { MerkleTree } from './merkle_tree.js';
 
 /**
  * Factory for creating leaf preimages.
@@ -73,7 +75,10 @@ export interface BatchInsertionResult<TreeHeight extends number, SubtreeSiblingP
 /**
  * Indexed merkle tree.
  */
-export interface IndexedTree extends AppendOnlyTree<Buffer> {
+export interface IndexedTree
+  extends MerkleTree<Buffer>,
+    TreeSnapshotBuilder<IndexedTreeSnapshot>,
+    Omit<AppendOnlyTree<Buffer>, keyof TreeSnapshotBuilder<TreeSnapshot<Buffer>>> {
   /**
    * 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.

diff --git a/yarn-project/merkle-tree/src/snapshots/snapshot_builder_test_suite.ts b/yarn-project/merkle-tree/src/snapshots/snapshot_builder_test_suite.ts
@@ -6,7 +6,7 @@ import { jest } from '@jest/globals';
 import { TreeBase } from '../tree_base.js';
 import { TreeSnapshot, TreeSnapshotBuilder } from './snapshot_builder.js';
 
-jest.setTimeout(1_000_000);
+jest.setTimeout(50_000);
 
 /** Creates a test suit for snapshots */
 export function describeSnapshotBuilderTestSuite<

diff --git a/yarn-project/world-state/src/world-state-db/merkle_trees.ts b/yarn-project/world-state/src/world-state-db/merkle_trees.ts
@@ -524,17 +524,20 @@ export class MerkleTrees implements MerkleTreeDb {
   }
 
   public async getSnapshot(blockNumber: number): Promise<TreeSnapshots> {
-    // TODO #5448 fix "as any"
+    const snapshots = await Promise.all([
+      this.trees[MerkleTreeId.NULLIFIER_TREE].getSnapshot(blockNumber),
+      this.trees[MerkleTreeId.NOTE_HASH_TREE].getSnapshot(blockNumber),
+      this.trees[MerkleTreeId.PUBLIC_DATA_TREE].getSnapshot(blockNumber),
+      this.trees[MerkleTreeId.L1_TO_L2_MESSAGE_TREE].getSnapshot(blockNumber),
+      this.trees[MerkleTreeId.ARCHIVE].getSnapshot(blockNumber),
+    ]);
+
     return {
-      [MerkleTreeId.NULLIFIER_TREE]: (await this.trees[MerkleTreeId.NULLIFIER_TREE].getSnapshot(blockNumber)) as any,
-      [MerkleTreeId.NOTE_HASH_TREE]: await this.trees[MerkleTreeId.NOTE_HASH_TREE].getSnapshot(blockNumber),
-      [MerkleTreeId.PUBLIC_DATA_TREE]: (await this.trees[MerkleTreeId.PUBLIC_DATA_TREE].getSnapshot(
-        blockNumber,
-      )) as any,
-      [MerkleTreeId.L1_TO_L2_MESSAGE_TREE]: await this.trees[MerkleTreeId.L1_TO_L2_MESSAGE_TREE].getSnapshot(
-        blockNumber,
-      ),
-      [MerkleTreeId.ARCHIVE]: await this.trees[MerkleTreeId.ARCHIVE].getSnapshot(blockNumber),
+      [MerkleTreeId.NULLIFIER_TREE]: snapshots[0],
+      [MerkleTreeId.NOTE_HASH_TREE]: snapshots[1],
+      [MerkleTreeId.PUBLIC_DATA_TREE]: snapshots[2],
+      [MerkleTreeId.L1_TO_L2_MESSAGE_TREE]: snapshots[3],
+      [MerkleTreeId.ARCHIVE]: snapshots[4],
     };
   }