feat: Add FPI support to NtxDataStore #1521
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,4 +1,5 @@ | ||
| use std::collections::{BTreeMap, BTreeSet}; | ||
| use std::sync::Arc; | ||
|
|
||
| use miden_node_proto::clients::ValidatorClient; | ||
| use miden_node_proto::generated::{self as proto}; | ||
|
|
@@ -8,9 +9,11 @@ use miden_protocol::Word; | |
| use miden_protocol::account::{ | ||
| Account, | ||
| AccountId, | ||
| AccountStorageHeader, | ||
| PartialAccount, | ||
| StorageMapWitness, | ||
| StorageSlotName, | ||
| StorageSlotType, | ||
| }; | ||
| use miden_protocol::asset::{AssetVaultKey, AssetWitness}; | ||
| use miden_protocol::block::{BlockHeader, BlockNumber}; | ||
|
|
@@ -45,6 +48,7 @@ use miden_tx::{ | |
| TransactionMastStore, | ||
| TransactionProverError, | ||
| }; | ||
| use tokio::sync::Mutex; | ||
| use tokio::task::JoinError; | ||
| use tracing::{Instrument, instrument}; | ||
|
|
||
|
|
@@ -222,7 +226,7 @@ impl NtxContext { | |
|
|
||
| match Box::pin(checker.check_notes_consumability( | ||
| data_store.account.id(), | ||
| data_store.reference_block.block_num(), | ||
| notes, | ||
| TransactionArgs::default(), | ||
| )) | ||
|
|
@@ -256,7 +260,7 @@ impl NtxContext { | |
|
|
||
| Box::pin(executor.execute_transaction( | ||
| data_store.account.id(), | ||
| data_store.reference_block.block_num(), | ||
| notes, | ||
| TransactionArgs::default(), | ||
| )) | ||
|
|
@@ -322,20 +326,42 @@ impl NtxContext { | |
| /// This is sufficient for executing a network transaction. | ||
| struct NtxDataStore { | ||
| account: Account, | ||
| reference_block: BlockHeader, | ||
| chain_mmr: PartialBlockchain, | ||
| mast_store: TransactionMastStore, | ||
| /// Store client for retrieving note scripts. | ||
| store: StoreClient, | ||
| /// LRU cache for storing retrieved note scripts to avoid repeated store calls. | ||
| script_cache: LruCache<Word, NoteScript>, | ||
| /// Mapping of storage map roots to storage slot names observed during various calls. | ||
| /// | ||
| /// The registered slot names are subsequently used to retrieve storage map witnesses from the | ||
| /// store. We need this because the store interface (and the underling SMT forest) use storage | ||
| /// slot names, but the `DataStore` interface works with tree roots. To get around this problem | ||
| /// we populate this map when: | ||
| /// - The the native account is loaded (in `get_transaction_inputs()`). | ||
| /// - When a foreign account is loaded (in `get_foreign_account_inputs`). | ||
| /// | ||
| /// The assumption here are: | ||
| /// - Once an account is loaded, the mapping between `(account_id, map_root)` and slot names do | ||
| /// not change. This is always the case. | ||
| /// - New storage slots created during transaction execution will not be accesses in the same | ||
| /// transaction. The mechanism for adding new storage slots is not implemented yet, but the | ||
| /// plan for it is consistent with this assumption. | ||
| /// | ||
| /// One nuance worth mentioning: it is possible that there could be a root collision where an | ||
| /// account has two storage maps with the same root. In this case, the map will contain only a | ||
| /// single entry with the storage slot name that was added last. Thus, technically, requests | ||
| /// to the store could be "wrong", but given that two identical maps have identical witnesses | ||
| /// this does not cause issues in practice. | ||
| storage_slots: Arc<Mutex<BTreeMap<(AccountId, Word), StorageSlotName>>>, | ||
sergerad marked this conversation as resolved.
Show resolved
Hide resolved
Comment on lines
336
to
357
|
||
| } | ||
|
|
||
| impl NtxDataStore { | ||
| /// Creates a new `NtxDataStore` with default cache size. | ||
| fn new( | ||
| account: Account, | ||
| reference_block: BlockHeader, | ||
| chain_mmr: PartialBlockchain, | ||
| store: StoreClient, | ||
| script_cache: LruCache<Word, NoteScript>, | ||
|
|
@@ -345,11 +371,28 @@ impl NtxDataStore { | |
|
|
||
| Self { | ||
| account, | ||
| reference_block, | ||
| chain_mmr, | ||
| mast_store, | ||
| store, | ||
| script_cache, | ||
| storage_slots: Arc::new(Mutex::new(BTreeMap::default())), | ||
| } | ||
| } | ||
|
|
||
| /// Registers storage map slot names for the given account ID and storage header. | ||
| /// | ||
| /// These slot names are subsequently used to query for storage map witnesses against the store. | ||
| async fn register_storage_map_slots( | ||
| &self, | ||
| account_id: AccountId, | ||
| storage_header: &AccountStorageHeader, | ||
| ) { | ||
| let mut storage_slots = self.storage_slots.lock().await; | ||
| for slot_header in storage_header.slots() { | ||
| if let StorageSlotType::Map = slot_header.slot_type() { | ||
| storage_slots.insert((account_id, slot_header.value()), slot_header.name().clone()); | ||
| } | ||
| } | ||
sergerad marked this conversation as resolved.
Show resolved
Hide resolved
|
||
| } | ||
| } | ||
|
|
@@ -366,53 +409,63 @@ impl DataStore for NtxDataStore { | |
| return Err(DataStoreError::AccountNotFound(account_id)); | ||
| } | ||
|
|
||
| // The latest supplied reference block must match the current reference block. | ||
| match ref_blocks.last().copied() { | ||
| Some(reference) if reference == self.reference_block.block_num() => {}, | ||
| Some(other) => return Err(DataStoreError::BlockNotFound(other)), | ||
| None => return Err(DataStoreError::other("no reference block requested")), | ||
| } | ||
|
|
||
| // Register slot names from the native account for later use. | ||
| self.register_storage_map_slots(account_id, &self.account.storage().to_header()) | ||
| .await; | ||
|
|
||
| let partial_account = PartialAccount::from(&self.account); | ||
| Ok((partial_account, self.reference_block.clone(), self.chain_mmr.clone())) | ||
| } | ||
| } | ||
|
|
||
| fn get_foreign_account_inputs( | ||
| &self, | ||
| foreign_account_id: AccountId, | ||
| ref_block: BlockNumber, | ||
| ) -> impl FutureMaybeSend<Result<AccountInputs, DataStoreError>> { | ||
| async move { | ||
| debug_assert_eq!(ref_block, self.reference_block.block_num()); | ||
|
|
||
| // Get foreign account inputs from store. | ||
| let account_inputs = | ||
| self.store.get_account_inputs(foreign_account_id, ref_block).await.map_err( | ||
| |err| DataStoreError::other_with_source("failed to get account inputs", err), | ||
| )?; | ||
|
|
||
| // Register slot names from the foreign account for later use. | ||
| self.register_storage_map_slots(foreign_account_id, account_inputs.storage().header()) | ||
| .await; | ||
|
|
||
| Ok(account_inputs) | ||
|
Comment on lines
+433
to
+446
Contributor
There was a problem hiding this comment. Following up on the conversation we had during the call, I actually don't think we need to add foreign account code to the MAST store in this function (i.e., nothing needs to change). The reason for this is that transaction host does this after calling Technically, I think we didn't have to do this in the validator either - but in the validator, we are planning to change how re-execution works and so, I'd leave it as is. cc @igamigo to confirm.
Collaborator
There was a problem hiding this comment. Looking at the code, I think this is right. But also, not sure why the validator was erroring during client integration FPI tests then (which prompted #1561).
Collaborator
There was a problem hiding this comment. Oh actually it's because it never even calls Although then I wonder how you plan on doing away with adding the MAST forest in the validator. |
||
| } | ||
| } | ||
|
|
||
| fn get_vault_asset_witnesses( | ||
| &self, | ||
| account_id: AccountId, | ||
| _vault_root: Word, | ||
| vault_keys: BTreeSet<AssetVaultKey>, | ||
| ) -> impl FutureMaybeSend<Result<Vec<AssetWitness>, DataStoreError>> { | ||
| async move { | ||
| let ref_block = self.reference_block.block_num(); | ||
|
|
||
| // Get vault asset witnesses from the store. | ||
| let witnesses = self | ||
| .store | ||
| .get_vault_asset_witnesses(account_id, vault_keys, Some(ref_block)) | ||
| .await | ||
| .map_err(|err| { | ||
| DataStoreError::other_with_source("failed to get vault asset witnesses", err) | ||
| })?; | ||
|
|
||
| Ok(witnesses) | ||
| } | ||
| } | ||
|
|
||
|
|
@@ -423,27 +476,27 @@ impl DataStore for NtxDataStore { | |
| map_key: Word, | ||
| ) -> impl FutureMaybeSend<Result<StorageMapWitness, DataStoreError>> { | ||
| async move { | ||
| // The slot name that corresponds to the given account ID and map root must have been | ||
| // registered during previous calls of this data store. | ||
| let storage_slots = self.storage_slots.lock().await; | ||
| let Some(slot_name) = storage_slots.get(&(account_id, map_root)) else { | ||
| return Err(DataStoreError::other( | ||
| "requested storage slot has not been registered", | ||
| )); | ||
| }; | ||
|
Comment on lines
+479
to
+486
Collaborator
There was a problem hiding this comment. Are we essentially relying on knowing the order of operations of the caller of data store?
Collaborator
There was a problem hiding this comment. nit: Seems like you can drop the lock by this point |
||
|
|
||
| let ref_block = self.reference_block.block_num(); | ||
|
|
||
| // Get storage map witness from the store. | ||
| let witness = self | ||
| .store | ||
| .get_storage_map_witness(account_id, slot_name.clone(), map_key, Some(ref_block)) | ||
| .await | ||
| .map_err(|err| { | ||
| DataStoreError::other_with_source("failed to get storage map witness", err) | ||
| })?; | ||
|
|
||
| Ok(witness) | ||
| } | ||
| } | ||
|
|
||
|
|
||
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Because we know this lock is low-contention and also not held across awaits, I think an
std::sync::Mutexwould be lower overhead and overall faster