feat: [2/4] integrate smtforest, avoid ser/de of full account/vault data in database #1394
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
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
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
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 |
|---|---|---|
| @@ -0,0 +1,273 @@ | ||
| use std::collections::BTreeMap; | ||
|
|
||
| use miden_protocol::account::delta::{AccountDelta, AccountStorageDelta, AccountVaultDelta}; | ||
| use miden_protocol::account::{AccountId, NonFungibleDeltaAction, StorageSlotName}; | ||
| use miden_protocol::asset::{Asset, FungibleAsset}; | ||
| use miden_protocol::block::BlockNumber; | ||
| use miden_protocol::crypto::merkle::EmptySubtreeRoots; | ||
| use miden_protocol::crypto::merkle::smt::{SMT_DEPTH, SmtForest}; | ||
| use miden_protocol::{EMPTY_WORD, Word}; | ||
|
|
||
| #[cfg(test)] | ||
| mod tests; | ||
|
|
||
| // INNER FOREST | ||
| // ================================================================================================ | ||
|
|
||
| /// Container for forest-related state that needs to be updated atomically. | ||
| pub(crate) struct InnerForest { | ||
| /// `SmtForest` for efficient account storage reconstruction. | ||
| /// Populated during block import with storage and vault SMTs. | ||
| forest: SmtForest, | ||
|
|
||
| /// Maps (`account_id`, `slot_name`, `block_num`) to SMT root. | ||
| /// Populated during block import for all storage map slots. | ||
| storage_map_roots: BTreeMap<(AccountId, StorageSlotName, BlockNumber), Word>, | ||
|
|
||
| /// Maps (`account_id`, `block_num`) to vault SMT root. | ||
| /// Tracks asset vault versions across all blocks with structural sharing. | ||
| vault_roots: BTreeMap<(AccountId, BlockNumber), Word>, | ||
| } | ||
drahnr marked this conversation as resolved.
Show resolved
Hide resolved
|
||
|
|
||
| impl InnerForest { | ||
| pub(crate) fn new() -> Self { | ||
| Self { | ||
| forest: SmtForest::new(), | ||
| storage_map_roots: BTreeMap::new(), | ||
| vault_roots: BTreeMap::new(), | ||
| } | ||
| } | ||
|
|
||
| // HELPERS | ||
| // -------------------------------------------------------------------------------------------- | ||
|
|
||
| /// Returns the root of an empty SMT. | ||
| const fn empty_smt_root() -> Word { | ||
| *EmptySubtreeRoots::entry(SMT_DEPTH, 0) | ||
| } | ||
|
|
||
| /// Retrieves the most recent vault SMT root for an account. | ||
| /// | ||
| /// Returns the latest vault root entry regardless of block number. | ||
| /// Used when applying incremental deltas where we always want the previous state. | ||
| fn get_latest_vault_root(&self, account_id: AccountId) -> Word { | ||
| self.vault_roots | ||
| .range((account_id, BlockNumber::GENESIS)..) | ||
| .take_while(|((id, _), _)| *id == account_id) | ||
| .last() | ||
| .map_or_else(Self::empty_smt_root, |(_, root)| *root) | ||
| } | ||
|
Comment on lines
+49
to
+59
Contributor
There was a problem hiding this comment. Let's mention in the doc comments that if a vault root for the specified account cannot be found, we return a root of an empty tree. Also, looking at the implementation, I wonder if a better backing structure would be Getting root for a specific block number should be pretty simple as well (we could use a binary search or even just a linear scan since we know that the number of entries will be pretty small per account ID). |
||
|
|
||
| /// Retrieves the vault SMT root for an account at or before the given block. | ||
| /// | ||
| /// Finds the most recent vault root entry for the account, since vault state persists | ||
| /// across blocks where no changes occur. | ||
| // | ||
| // TODO: a fallback to DB lookup is required once pruning lands. | ||
| // Currently returns empty root which would be incorrect | ||
| #[cfg(test)] | ||
| fn get_vault_root(&self, account_id: AccountId, block_num: BlockNumber) -> Word { | ||
| self.vault_roots | ||
| .range((account_id, BlockNumber::GENESIS)..=(account_id, block_num)) | ||
| .next_back() | ||
| .map_or_else(Self::empty_smt_root, |(_, root)| *root) | ||
| } | ||
|
|
||
| /// Retrieves the most recent storage map SMT root for an account slot. | ||
| /// | ||
| /// Returns the latest storage root entry regardless of block number. | ||
| /// Used when applying incremental deltas where we always want the previous state. | ||
| fn get_latest_storage_map_root( | ||
| &self, | ||
| account_id: AccountId, | ||
| slot_name: &StorageSlotName, | ||
| ) -> Word { | ||
| self.storage_map_roots | ||
| .range((account_id, slot_name.clone(), BlockNumber::GENESIS)..) | ||
| .take_while(|((id, name, _), _)| *id == account_id && name == slot_name) | ||
| .last() | ||
| .map_or_else(Self::empty_smt_root, |(_, root)| *root) | ||
| } | ||
|
|
||
| // PUBLIC INTERFACE | ||
| // -------------------------------------------------------------------------------------------- | ||
|
|
||
| /// Applies account updates from a block to the forest. | ||
| /// | ||
| /// Iterates through account updates and applies each delta to the forest. | ||
| /// Private accounts should be filtered out before calling this method. | ||
| /// | ||
| /// # Arguments | ||
| /// | ||
| /// * `block_num` - Block number for which these updates apply | ||
| /// * `account_updates` - Iterator of `AccountDelta` for public accounts | ||
| pub(crate) fn apply_block_updates( | ||
| &mut self, | ||
| block_num: BlockNumber, | ||
| account_updates: impl IntoIterator<Item = AccountDelta>, | ||
| ) { | ||
| for delta in account_updates { | ||
| self.update_account(block_num, &delta); | ||
|
|
||
| tracing::debug!( | ||
| target: crate::COMPONENT, | ||
| account_id = %delta.id(), | ||
| %block_num, | ||
| is_full_state = delta.is_full_state(), | ||
| "Updated forest with account delta" | ||
| ); | ||
| } | ||
| } | ||
|
|
||
| /// Updates the forest with account vault and storage changes from a delta. | ||
| /// | ||
| /// Unified interface for updating all account state in the forest, handling both full-state | ||
| /// deltas (new accounts or reconstruction from DB) and partial deltas (incremental updates | ||
| /// during block application). | ||
| /// | ||
| /// Full-state deltas (`delta.is_full_state() == true`) populate the forest from scratch using | ||
| /// an empty SMT root. Partial deltas apply changes on top of the previous block's state. | ||
| pub(crate) fn update_account(&mut self, block_num: BlockNumber, delta: &AccountDelta) { | ||
| let account_id = delta.id(); | ||
| let is_full_state = delta.is_full_state(); | ||
|
|
||
| if !delta.vault().is_empty() { | ||
| self.update_account_vault(block_num, account_id, delta.vault(), is_full_state); | ||
| } | ||
|
|
||
| if !delta.storage().is_empty() { | ||
| self.update_account_storage(block_num, account_id, delta.storage(), is_full_state); | ||
| } | ||
| } | ||
|
|
||
| // PRIVATE METHODS | ||
| // -------------------------------------------------------------------------------------------- | ||
|
|
||
| /// Updates the forest with vault changes from a delta. | ||
| /// | ||
| /// Processes both fungible and non-fungible asset changes, building entries for the vault SMT | ||
| /// and tracking the new root. | ||
| fn update_account_vault( | ||
| &mut self, | ||
| block_num: BlockNumber, | ||
| account_id: AccountId, | ||
| vault_delta: &AccountVaultDelta, | ||
| is_full_state: bool, | ||
| ) { | ||
| let prev_root = if is_full_state { | ||
| Self::empty_smt_root() | ||
| } else { | ||
| self.get_latest_vault_root(account_id) | ||
| }; | ||
|
Comment on lines
+157
to
+161
Contributor
There was a problem hiding this comment. This could be just: let prev_root = self.get_latest_vault_root(account_id); |
||
|
|
||
| let mut entries = Vec::new(); | ||
|
|
||
| // Process fungible assets | ||
| for (faucet_id, amount_delta) in vault_delta.fungible().iter() { | ||
| let key: Word = | ||
| FungibleAsset::new(*faucet_id, 0).expect("valid faucet id").vault_key().into(); | ||
|
|
||
| let new_amount = if is_full_state { | ||
| // For full-state deltas, amount is the absolute value | ||
| (*amount_delta).try_into().expect("full-state amount should be non-negative") | ||
| } else { | ||
| // For partial deltas, amount is a change that must be applied to previous balance. | ||
| // | ||
| // TODO: SmtForest only exposes `fn open()` which computes a full Merkle | ||
| // proof. We only need the leaf, so a direct `fn get()` method would be faster. | ||
drahnr marked this conversation as resolved.
Show resolved
Hide resolved
|
||
| let prev_amount = self | ||
| .forest | ||
| .open(prev_root, key) | ||
| .ok() | ||
| .and_then(|proof| proof.get(&key)) | ||
| .and_then(|word| FungibleAsset::try_from(word).ok()) | ||
| .map_or(0, |asset| asset.amount()); | ||
|
|
||
| let new_balance = i128::from(prev_amount) + i128::from(*amount_delta); | ||
| u64::try_from(new_balance.max(0)).expect("balance fits in u64") | ||
|
Contributor
There was a problem hiding this comment. If |
||
| }; | ||
|
|
||
| let value = if new_amount == 0 { | ||
| EMPTY_WORD | ||
| } else { | ||
| let asset: Asset = FungibleAsset::new(*faucet_id, new_amount) | ||
| .expect("valid fungible asset") | ||
| .into(); | ||
| Word::from(asset) | ||
| }; | ||
| entries.push((key, value)); | ||
| } | ||
|
|
||
| // Process non-fungible assets | ||
| for (asset, action) in vault_delta.non_fungible().iter() { | ||
| let value = match action { | ||
| NonFungibleDeltaAction::Add => Word::from(Asset::NonFungible(*asset)), | ||
| NonFungibleDeltaAction::Remove => EMPTY_WORD, | ||
| }; | ||
| entries.push((asset.vault_key().into(), value)); | ||
| } | ||
|
|
||
| if entries.is_empty() { | ||
| return; | ||
| } | ||
|
|
||
| let updated_root = self | ||
| .forest | ||
| .batch_insert(prev_root, entries.iter().copied()) | ||
| .expect("forest insertion should succeed"); | ||
|
|
||
| self.vault_roots.insert((account_id, block_num), updated_root); | ||
|
|
||
| tracing::debug!( | ||
| target: crate::COMPONENT, | ||
| %account_id, | ||
| %block_num, | ||
| vault_entries = entries.len(), | ||
| "Updated vault in forest" | ||
| ); | ||
| } | ||
|
|
||
| /// Updates the forest with storage map changes from a delta. | ||
| /// | ||
| /// Processes storage map slot deltas, building SMTs for each modified slot | ||
| /// and tracking the new roots. | ||
| fn update_account_storage( | ||
| &mut self, | ||
| block_num: BlockNumber, | ||
| account_id: AccountId, | ||
| storage_delta: &AccountStorageDelta, | ||
| is_full_state: bool, | ||
| ) { | ||
| for (slot_name, map_delta) in storage_delta.maps() { | ||
| let prev_root = if is_full_state { | ||
| Self::empty_smt_root() | ||
| } else { | ||
| self.get_latest_storage_map_root(account_id, slot_name) | ||
| }; | ||
|
Comment on lines
+242
to
+246
Contributor
There was a problem hiding this comment. Similar to one of the above comments, this could probably be just: let prev_root = self.get_latest_storage_map_root(account_id, slot_name); |
||
|
|
||
| let entries: Vec<_> = | ||
| map_delta.entries().iter().map(|(key, value)| ((*key).into(), *value)).collect(); | ||
|
|
||
| if entries.is_empty() { | ||
| continue; | ||
| } | ||
|
|
||
| let updated_root = self | ||
| .forest | ||
| .batch_insert(prev_root, entries.iter().copied()) | ||
| .expect("forest insertion should succeed"); | ||
|
|
||
| self.storage_map_roots | ||
| .insert((account_id, slot_name.clone(), block_num), updated_root); | ||
|
|
||
| tracing::debug!( | ||
| target: crate::COMPONENT, | ||
| %account_id, | ||
| %block_num, | ||
| ?slot_name, | ||
| entries = entries.len(), | ||
| "Updated storage map in forest" | ||
| ); | ||
| } | ||
| } | ||
| } | ||
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.
Uh oh!
There was an error while loading. Please reload this page.