diff --git a/CHANGELOG.md b/CHANGELOG.md index 0e4268300..39f07d2fa 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,7 +5,7 @@ - [BREAKING] Changed the serialization format of `PartialSmt` to be more compact on the wire. If you depended on the serialization format directly (or if it is stored), this will be breaking. This ser/de is now covered by fuzz tests. - [BREAKING] Changed `SmtLeaf::hash` to perform domain-separated hashing, reducing the risk of a collision with the hash of an inner node. Miden VM **must** be updated to comply with this. - Fixed `SimpleSmt::set_subtree()` to clear stale leaves and inner nodes in the replaced subtree region ([#981](https://github.com/0xMiden/crypto/pull/981)). -- [BREAKING] Extracted `SmtStorageReader` and `SparseMerkleTreeReader`, allowing `LargeSmt` to work with read-only storage backends ([#958](https://github.com/0xMiden/crypto/pull/958)). +- [BREAKING] Extracted `SmtStorageReader` and `SparseMerkleTreeReader`, allowing `LargeSmt` to work with read-only storage backends ([#967](https://github.com/0xMiden/crypto/pull/967)). ## 0.24.0 (2026-04-19) diff --git a/miden-crypto/src/boxed_storage.rs b/miden-crypto/src/boxed_storage.rs index 52bde4863..fc17bdb80 100644 --- a/miden-crypto/src/boxed_storage.rs +++ b/miden-crypto/src/boxed_storage.rs @@ -1,14 +1,17 @@ +//! Adapter used only by the `miden-crypto` executable benchmark. +//! +//! The executable chooses memory or RocksDB storage at runtime. This wrapper erases the concrete +//! storage type while preserving the associated reader type needed by `SmtStorage`. + use miden_crypto::{ Map, Word, merkle::smt::{InnerNode, SmtLeaf, SmtStorage, SmtStorageReader, StorageError, StorageUpdates}, }; -/// Type alias for boxed storage with boxed reader -pub type BoxedSmtStorage = Box>>; +pub(crate) type BoxedSmtStorage = Box>>; -/// Generic wrapper that boxes the Reader type for any SmtStorage implementation #[derive(Debug)] -pub struct BoxedStorage(pub T); +pub(crate) struct BoxedStorage(pub(crate) T); impl SmtStorageReader for BoxedStorage { fn leaf_count(&self) -> Result { diff --git a/miden-crypto/src/main.rs b/miden-crypto/src/main.rs index 635135572..37dfa3e80 100644 --- a/miden-crypto/src/main.rs +++ b/miden-crypto/src/main.rs @@ -11,6 +11,7 @@ use miden_crypto::{ }; use rand::{Rng, prelude::IteratorRandom, rng}; +#[cfg(feature = "executable")] mod boxed_storage; use boxed_storage::{BoxedSmtStorage as Storage, BoxedStorage}; diff --git a/miden-crypto/src/merkle/smt/large/batch_ops.rs b/miden-crypto/src/merkle/smt/large/batch_ops.rs index 81f0c2b70..05d1b4f5a 100644 --- a/miden-crypto/src/merkle/smt/large/batch_ops.rs +++ b/miden-crypto/src/merkle/smt/large/batch_ops.rs @@ -449,7 +449,7 @@ impl LargeSmt { } let new_root = leaves[0][0].hash; - self.in_memory_nodes[ROOT_MEMORY_INDEX] = new_root; + self.in_memory_nodes_mut()[ROOT_MEMORY_INDEX] = new_root; // Build leaf updates for storage (convert Empty to None for deletion) let mut leaf_update_map = leaf_map; @@ -578,7 +578,7 @@ impl LargeSmt { } = prepared; // Update the root in memory - self.in_memory_nodes[ROOT_MEMORY_INDEX] = new_root; + self.in_memory_nodes_mut()[ROOT_MEMORY_INDEX] = new_root; // Process node mutations - group by subtree and apply in batch // Since mutations are sorted by subtree root, we can process them in groups diff --git a/miden-crypto/src/merkle/smt/large/construction.rs b/miden-crypto/src/merkle/smt/large/construction.rs index 4268420e6..c6c5b3e02 100644 --- a/miden-crypto/src/merkle/smt/large/construction.rs +++ b/miden-crypto/src/merkle/smt/large/construction.rs @@ -137,7 +137,7 @@ impl LargeSmt { if is_empty { return Ok(Self { storage, - in_memory_nodes, + in_memory_nodes: in_memory_nodes.into(), leaf_count: 0, entry_count: 0, }); @@ -192,7 +192,7 @@ impl LargeSmt { Ok(Self { storage, - in_memory_nodes, + in_memory_nodes: in_memory_nodes.into(), leaf_count, entry_count, }) @@ -340,9 +340,10 @@ impl LargeSmt { for (index, node) in nodes { if index.depth() < IN_MEMORY_DEPTH { let memory_index = to_memory_index(&index); + let in_memory_nodes = self.in_memory_nodes_mut(); // Store in flat layout: left at 2*i, right at 2*i+1 - self.in_memory_nodes[memory_index * 2] = node.left; - self.in_memory_nodes[memory_index * 2 + 1] = node.right; + in_memory_nodes[memory_index * 2] = node.left; + in_memory_nodes[memory_index * 2 + 1] = node.right; } } } diff --git a/miden-crypto/src/merkle/smt/large/iter.rs b/miden-crypto/src/merkle/smt/large/iter.rs index 64640f3f6..18a24accb 100644 --- a/miden-crypto/src/merkle/smt/large/iter.rs +++ b/miden-crypto/src/merkle/smt/large/iter.rs @@ -13,7 +13,7 @@ use crate::{ enum InnerNodeIteratorState<'a> { InMemory { current_index: usize, - large_smt_in_memory_nodes: &'a Vec, + large_smt_in_memory_nodes: &'a [Word], }, Subtree { subtree_iter: Box + 'a>, diff --git a/miden-crypto/src/merkle/smt/large/mod.rs b/miden-crypto/src/merkle/smt/large/mod.rs index 94c271472..cda428cc1 100644 --- a/miden-crypto/src/merkle/smt/large/mod.rs +++ b/miden-crypto/src/merkle/smt/large/mod.rs @@ -237,7 +237,7 @@ //! To optimize memory and I/O: group updates by key locality so that keys sharing //! high-order bits are processed together. -use alloc::vec::Vec; +use alloc::{sync::Arc, vec::Vec}; use super::{ EmptySubtreeRoots, InnerNode, InnerNodeInfo, LeafIndex, MerkleError, NodeIndex, NodeMutation, @@ -261,11 +261,11 @@ pub use subtree::{Subtree, SubtreeError}; mod storage; pub use storage::{ - CloneableSmtStorageReader, MemoryStorage, SmtStorage, SmtStorageReader, SmtStorageSnapshot, - StorageError, StorageUpdateParts, StorageUpdates, SubtreeUpdate, + MemoryStorage, MemoryStorageSnapshot, SmtStorage, SmtStorageReader, StorageError, + StorageUpdateParts, StorageUpdates, SubtreeUpdate, }; #[cfg(feature = "rocksdb")] -pub use storage::{RocksDbConfig, RocksDbStorage}; +pub use storage::{RocksDbConfig, RocksDbSnapshotStorage, RocksDbStorage}; mod iter; pub use iter::LargeSmtInnerNodeIterator; @@ -334,15 +334,15 @@ type MutatedLeaves = (MutatedSubtreeLeaves, Map, Map, /// - Depths 0-23: Stored in memory as a flat array for fast access /// - Depths 24-64: Stored in external storage organized as subtrees for efficient batch operations /// -/// `LargeSmt` implements [`Clone`] only when `S` is a cloneable reader storage type. This prevents -/// accidental duplication of writable storage backends while keeping read-only snapshots cloneable. +/// `LargeSmt` implements [`Clone`] when its storage is cloneable. The in-memory top is shared and +/// detaches on mutation. #[derive(Debug)] pub struct LargeSmt { storage: S, - /// Flat vector representation of in-memory nodes. + /// Shared flat array representation of in-memory nodes. /// Index 0 is unused; index 1 is root. /// For node at index i: left child at 2*i, right child at 2*i+1. - in_memory_nodes: Vec, + in_memory_nodes: Arc<[Word]>, /// Cached count of non-empty leaves. Initialized from storage on load, /// updated after each mutation. leaf_count: usize, @@ -351,7 +351,7 @@ pub struct LargeSmt { entry_count: usize, } -impl Clone for LargeSmt { +impl Clone for LargeSmt { fn clone(&self) -> Self { Self { storage: self.storage.clone(), @@ -478,6 +478,10 @@ impl LargeSmt { >::get_inner_node(self, index) } + pub(crate) fn in_memory_nodes_mut(&mut self) -> &mut [Word] { + Arc::make_mut(&mut self.in_memory_nodes) + } + /// Helper to get an in-memory node if not empty. /// /// # Panics @@ -503,7 +507,7 @@ impl LargeSmt { // -------------------------------------------------------------------------------------------- #[cfg(test)] - pub(crate) fn in_memory_nodes(&self) -> &Vec { + pub(crate) fn in_memory_nodes(&self) -> &[Word] { &self.in_memory_nodes } } @@ -512,8 +516,8 @@ impl LargeSmt { /// Returns a read-only `LargeSmt` backed by a reader view of this tree's storage. /// /// The new tree shares the same root, leaf count, and entry count as `self`, and its storage - /// is produced by [`SmtStorage::reader`]. The returned tree's storage type is - /// `S::Reader: SmtStorageReader`, so it cannot be used for mutations. + /// is a point-in-time snapshot produced by [`SmtStorage::reader`]. The returned tree's storage + /// type is `S::Reader: SmtStorageReader`, so it cannot be used for mutations. pub fn reader(&self) -> Result, LargeSmtError> { Ok(LargeSmt { storage: self.storage.reader()?, diff --git a/miden-crypto/src/merkle/smt/large/smt_trait.rs b/miden-crypto/src/merkle/smt/large/smt_trait.rs index 4323ac349..b0b547273 100644 --- a/miden-crypto/src/merkle/smt/large/smt_trait.rs +++ b/miden-crypto/src/merkle/smt/large/smt_trait.rs @@ -157,19 +157,20 @@ impl SparseMerkleTreeReader for LargeSmt { impl SparseMerkleTree for LargeSmt { fn set_root(&mut self, root: Word) { - self.in_memory_nodes[ROOT_MEMORY_INDEX] = root; + self.in_memory_nodes_mut()[ROOT_MEMORY_INDEX] = root; } fn insert_inner_node(&mut self, index: NodeIndex, inner_node: InnerNode) -> Option { if index.depth() < IN_MEMORY_DEPTH { let i = to_memory_index(&index); + let nodes = self.in_memory_nodes_mut(); // Get the old node before replacing - let old_left = self.in_memory_nodes[i * 2]; - let old_right = self.in_memory_nodes[i * 2 + 1]; + let old_left = nodes[i * 2]; + let old_right = nodes[i * 2 + 1]; // Store new node in flat layout - self.in_memory_nodes[i * 2] = inner_node.left; - self.in_memory_nodes[i * 2 + 1] = inner_node.right; + nodes[i * 2] = inner_node.left; + nodes[i * 2 + 1] = inner_node.right; // Check if the old node was empty if is_empty_parent(old_left, old_right, index.depth() + 1) { @@ -186,15 +187,16 @@ impl SparseMerkleTree for LargeSmt { fn remove_inner_node(&mut self, index: NodeIndex) -> Option { if index.depth() < IN_MEMORY_DEPTH { let memory_index = to_memory_index(&index); + let nodes = self.in_memory_nodes_mut(); // Get the old node before replacing with empty hashes - let old_left = self.in_memory_nodes[memory_index * 2]; - let old_right = self.in_memory_nodes[memory_index * 2 + 1]; + let old_left = nodes[memory_index * 2]; + let old_right = nodes[memory_index * 2 + 1]; // Replace with empty hashes let child_depth = index.depth() + 1; let empty_hash = *EmptySubtreeRoots::entry(SMT_DEPTH, child_depth); - self.in_memory_nodes[memory_index * 2] = empty_hash; - self.in_memory_nodes[memory_index * 2 + 1] = empty_hash; + nodes[memory_index * 2] = empty_hash; + nodes[memory_index * 2 + 1] = empty_hash; // Return the old node if it wasn't already empty if is_empty_parent(old_left, old_right, child_depth) { diff --git a/miden-crypto/src/merkle/smt/large/storage/memory.rs b/miden-crypto/src/merkle/smt/large/storage/memory.rs index c3bb7a261..88b472b7c 100644 --- a/miden-crypto/src/merkle/smt/large/storage/memory.rs +++ b/miden-crypto/src/merkle/smt/large/storage/memory.rs @@ -1,8 +1,7 @@ use alloc::{boxed::Box, vec::Vec}; use super::{ - CloneableSmtStorageReader, SmtStorage, SmtStorageReader, StorageError, StorageUpdateParts, - StorageUpdates, SubtreeUpdate, + SmtStorage, SmtStorageReader, StorageError, StorageUpdateParts, StorageUpdates, SubtreeUpdate, }; use crate::{ EMPTY_WORD, Map, MapEntry, Word, @@ -37,7 +36,7 @@ pub struct MemoryStorage { /// storage backends that need to hand out a detached point-in-time copy without also exposing /// mutation methods through [`SmtStorage`]. #[derive(Debug, Clone)] -pub struct SmtStorageSnapshot(MemoryStorage); +pub struct MemoryStorageSnapshot(MemoryStorage); impl MemoryStorage { /// Creates a new, empty in-memory storage for a Sparse Merkle Tree. @@ -48,8 +47,8 @@ impl MemoryStorage { } /// Converts this storage into a read-only snapshot. - pub fn into_snapshot(self) -> SmtStorageSnapshot { - SmtStorageSnapshot(self) + pub fn into_snapshot(self) -> MemoryStorageSnapshot { + MemoryStorageSnapshot(self) } } @@ -59,7 +58,7 @@ impl Default for MemoryStorage { } } -impl SmtStorageReader for SmtStorageSnapshot { +impl SmtStorageReader for MemoryStorageSnapshot { fn leaf_count(&self) -> Result { self.0.leaf_count() } @@ -113,8 +112,6 @@ impl SmtStorageReader for SmtStorageSnapshot { } } -impl CloneableSmtStorageReader for SmtStorageSnapshot {} - impl SmtStorageReader for MemoryStorage { /// Gets the total number of non-empty leaves currently stored. fn leaf_count(&self) -> Result { @@ -212,10 +209,8 @@ impl SmtStorageReader for MemoryStorage { } } -impl CloneableSmtStorageReader for MemoryStorage {} - impl SmtStorage for MemoryStorage { - type Reader = SmtStorageSnapshot; + type Reader = MemoryStorageSnapshot; /// Returns a read-only snapshot of this in-memory storage by cloning it. fn reader(&self) -> Result { diff --git a/miden-crypto/src/merkle/smt/large/storage/mod.rs b/miden-crypto/src/merkle/smt/large/storage/mod.rs index dcf8706c8..6a7cab0c3 100644 --- a/miden-crypto/src/merkle/smt/large/storage/mod.rs +++ b/miden-crypto/src/merkle/smt/large/storage/mod.rs @@ -18,10 +18,10 @@ pub use error::StorageError; #[cfg(feature = "rocksdb")] mod rocksdb; #[cfg(feature = "rocksdb")] -pub use rocksdb::{RocksDbConfig, RocksDbStorage}; +pub use rocksdb::{RocksDbConfig, RocksDbSnapshotStorage, RocksDbStorage}; mod memory; -pub use memory::{MemoryStorage, SmtStorageSnapshot}; +pub use memory::{MemoryStorage, MemoryStorageSnapshot}; mod updates; pub use updates::{StorageUpdateParts, StorageUpdates, SubtreeUpdate}; @@ -38,7 +38,9 @@ pub use updates::{StorageUpdateParts, StorageUpdates, SubtreeUpdate}; /// All methods are expected to handle potential storage errors by returning a /// `Result<_, StorageError>`. /// -/// Implementations may return either a point-in-time snapshot or a live view. +/// Implementations used as [`SmtStorage::Reader`] must be point-in-time snapshots. This is required +/// because `LargeSmt::reader()` copies the in-memory portion of the tree and pairs it with the +/// returned storage reader. pub trait SmtStorageReader: 'static + fmt::Debug + Send + Sync { /// Retrieves the total number of leaf nodes currently stored. /// @@ -137,9 +139,6 @@ pub trait SmtStorageReader: 'static + fmt::Debug + Send + Sync { fn get_depth24(&self) -> Result, StorageError>; } -/// Marker trait for reader storage types that can be cloned without duplicating mutable storage. -pub trait CloneableSmtStorageReader: SmtStorageReader + Clone {} - impl SmtStorageReader for Box { #[inline] fn leaf_count(&self) -> Result { @@ -220,15 +219,14 @@ pub trait SmtStorage: SmtStorageReader { /// The read-only view type returned by [`Self::reader`]. type Reader: SmtStorageReader; - /// Returns a read-only view of this storage that observes its current state. + /// Returns a read-only snapshot of this storage at its current committed state. /// /// The returned value is used to construct a read-only `LargeSmt` (via /// [`super::LargeSmt::reader`]) from a writable one. Implementations are responsible for - /// ensuring that the returned reader is consistent with `self` at the time of the call. + /// ensuring that the returned reader remains consistent with `self` at the time of the call. /// - /// Implementations may return either a point-in-time snapshot or a live view. Either way, the - /// view must be of consistent / committed state (not partial). Holding the reader must not - /// block writes in any way. + /// Implementations must return a point-in-time snapshot. Later writes through `self` must not + /// affect the returned reader. Holding the reader must not block writes in any way. fn reader(&self) -> Result; /// Inserts a key-value pair into the SMT leaf at the specified logical `index`. diff --git a/miden-crypto/src/merkle/smt/large/storage/rocksdb.rs b/miden-crypto/src/merkle/smt/large/storage/rocksdb.rs index 14e003061..ae639f3d9 100644 --- a/miden-crypto/src/merkle/smt/large/storage/rocksdb.rs +++ b/miden-crypto/src/merkle/smt/large/storage/rocksdb.rs @@ -1,5 +1,6 @@ use alloc::{boxed::Box, vec::Vec}; -use std::{path::PathBuf, sync::Arc}; +use core::fmt; +use std::{mem::ManuallyDrop, path::PathBuf, sync::Arc}; use rocksdb::{ BlockBasedOptions, Cache, ColumnFamilyDescriptor, DB, DBCompactionStyle, DBCompressionType, @@ -7,8 +8,7 @@ use rocksdb::{ }; use super::{ - MemoryStorage, SmtStorage, SmtStorageReader, SmtStorageSnapshot, StorageError, - StorageUpdateParts, StorageUpdates, SubtreeUpdate, + SmtStorage, SmtStorageReader, StorageError, StorageUpdateParts, StorageUpdates, SubtreeUpdate, }; use crate::{ EMPTY_WORD, Word, @@ -64,6 +64,77 @@ pub struct RocksDbStorage { db: Arc, } +/// Read-only, cloneable SMT storage backed by a native RocksDB point-in-time snapshot. +#[derive(Clone)] +pub struct RocksDbSnapshotStorage { + inner: Arc, +} + +impl fmt::Debug for RocksDbSnapshotStorage { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("RocksDbSnapshotStorage").finish_non_exhaustive() + } +} + +/// Owns a RocksDB snapshot together with the database it borrows from. +/// +/// `rocksdb::Snapshot<'a>` borrows the database used to create it. This type stores an `Arc` +/// beside the snapshot and releases the snapshot before the `Arc` is dropped, so the borrowed +/// database remains alive for the full lifetime of the snapshot. +struct RocksDbSnapshotInner { + snapshot: ManuallyDrop>, + db: Arc, +} + +impl RocksDbSnapshotInner { + fn new(db: Arc) -> Self { + let snapshot = db.snapshot(); + // SAFETY: The snapshot internally stores a reference to the same `DB` allocation owned by + // `db`. `RocksDbSnapshotInner` keeps that `Arc` alive and its `Drop` implementation + // manually releases the snapshot before the `Arc` field is dropped. + let snapshot = unsafe { + core::mem::transmute::, rocksdb::Snapshot<'static>>(snapshot) + }; + Self { + snapshot: ManuallyDrop::new(snapshot), + db, + } + } +} + +impl Drop for RocksDbSnapshotInner { + fn drop(&mut self) { + // SAFETY: `snapshot` was placed in `ManuallyDrop` only to control field drop order. It is + // dropped exactly once here, before `db` is dropped by Rust's normal field cleanup. + unsafe { + ManuallyDrop::drop(&mut self.snapshot); + } + } +} + +impl RocksDbSnapshotStorage { + /// Creates a snapshot-backed storage reader from a shared RocksDB handle. + pub fn new(db: Arc) -> Self { + Self { + inner: Arc::new(RocksDbSnapshotInner::new(db)), + } + } + + /// Retrieves a handle to a RocksDB column family by its name. + fn cf_handle(&self, name: &str) -> Result<&rocksdb::ColumnFamily, StorageError> { + self.inner + .db + .cf_handle(name) + .ok_or_else(|| StorageError::Unsupported(format!("unknown column family `{name}`"))) + } + + #[inline(always)] + fn subtree_cf(&self, index: NodeIndex) -> &rocksdb::ColumnFamily { + let name = cf_for_depth(index.depth()); + self.cf_handle(name).expect("CF handle missing") + } +} + impl RocksDbStorage { /// Opens or creates a RocksDB database at the specified `path` and configures it for SMT /// storage. @@ -534,42 +605,190 @@ impl SmtStorageReader for RocksDbStorage { } } -impl SmtStorage for RocksDbStorage { - type Reader = SmtStorageSnapshot; +impl SmtStorageReader for RocksDbSnapshotStorage { + /// Retrieves the total count of non-empty leaves from the snapshot. + fn leaf_count(&self) -> Result { + let cf = self.cf_handle(METADATA_CF)?; + self.inner + .snapshot + .get_cf(cf, LEAF_COUNT_KEY)? + .map_or(Ok(0), |bytes| read_count("leaf count", &bytes)) + } - /// Returns a detached read-only snapshot of the current RocksDB-backed storage. - fn reader(&self) -> Result { - let snapshot = self.db.snapshot(); + /// Retrieves the total count of key-value entries from the snapshot. + fn entry_count(&self) -> Result { + let cf = self.cf_handle(METADATA_CF)?; + self.inner + .snapshot + .get_cf(cf, ENTRY_COUNT_KEY)? + .map_or(Ok(0), |bytes| read_count("entry count", &bytes)) + } - let leaves_cf = self.cf_handle(LEAVES_CF)?; + /// Retrieves a single SMT leaf node by its logical `index` from the snapshot. + fn get_leaf(&self, index: u64) -> Result, StorageError> { + let cf = self.cf_handle(LEAVES_CF)?; + let key = RocksDbStorage::index_db_key(index); + match self.inner.snapshot.get_cf(cf, key)? { + Some(bytes) => { + let leaf = SmtLeaf::read_from_bytes_with_budget(&bytes, bytes.len())?; + Ok(Some(leaf)) + }, + None => Ok(None), + } + } + + /// Retrieves multiple SMT leaf nodes by their logical `indices` from the snapshot. + fn get_leaves(&self, indices: &[u64]) -> Result>, StorageError> { + let cf = self.cf_handle(LEAVES_CF)?; + let db_keys: Vec<[u8; 8]> = + indices.iter().map(|&idx| RocksDbStorage::index_db_key(idx)).collect(); + let results = self.inner.snapshot.multi_get_cf(db_keys.iter().map(|k| (cf, k.as_ref()))); + + results + .into_iter() + .map(|result| match result { + Ok(Some(bytes)) => { + Ok(Some(SmtLeaf::read_from_bytes_with_budget(&bytes, bytes.len())?)) + }, + Ok(None) => Ok(None), + Err(e) => Err(e.into()), + }) + .collect() + } + + /// Returns true if the snapshot has any leaves. + fn has_leaves(&self) -> Result { + Ok(self.leaf_count()? > 0) + } + + /// Retrieves a single SMT Subtree by its root `NodeIndex` from the snapshot. + fn get_subtree(&self, index: NodeIndex) -> Result, StorageError> { + let cf = self.subtree_cf(index); + let key = RocksDbStorage::subtree_db_key(index); + match self.inner.snapshot.get_cf(cf, key)? { + Some(bytes) => { + let subtree = Subtree::from_vec(index, &bytes)?; + Ok(Some(subtree)) + }, + None => Ok(None), + } + } + + /// Retrieves multiple subtrees from the snapshot. + fn get_subtrees(&self, indices: &[NodeIndex]) -> Result>, StorageError> { + use p3_maybe_rayon::prelude::*; + + let mut depth_buckets: [Vec<(usize, NodeIndex)>; 5] = Default::default(); + + for (original_index, &node_index) in indices.iter().enumerate() { + let depth = node_index.depth(); + let bucket_index = match depth { + 56 => 0, + 48 => 1, + 40 => 2, + 32 => 3, + 24 => 4, + _ => { + return Err(StorageError::Unsupported(format!( + "unsupported subtree depth {depth}" + ))); + }, + }; + depth_buckets[bucket_index].push((original_index, node_index)); + } + let mut results = vec![None; indices.len()]; + + let bucket_results: Result, StorageError> = depth_buckets + .into_par_iter() + .enumerate() + .filter(|(_, bucket)| !bucket.is_empty()) + .map( + |(bucket_index, bucket)| -> Result)>, StorageError> { + let depth = LargeSmt::::SUBTREE_DEPTHS[bucket_index]; + let cf = self.cf_handle(cf_for_depth(depth))?; + let keys: Vec<_> = bucket + .iter() + .map(|(_, idx)| RocksDbStorage::subtree_db_key(*idx)) + .collect(); + + let db_results = + self.inner.snapshot.multi_get_cf(keys.iter().map(|k| (cf, k.as_ref()))); + + bucket + .into_iter() + .zip(db_results) + .map(|((original_index, node_index), db_result)| { + let subtree = match db_result { + Ok(Some(bytes)) => Some(Subtree::from_vec(node_index, &bytes)?), + Ok(None) => None, + Err(e) => return Err(e.into()), + }; + Ok((original_index, subtree)) + }) + .collect() + }, + ) + .collect(); + + for bucket_result in bucket_results? { + for (original_index, subtree) in bucket_result { + results[original_index] = subtree; + } + } + + Ok(results) + } + + /// Retrieves a single inner node from within a snapshot subtree. + fn get_inner_node(&self, index: NodeIndex) -> Result, StorageError> { + if index.depth() < IN_MEMORY_DEPTH { + return Err(StorageError::Unsupported( + "Cannot get inner node from upper part of the tree".into(), + )); + } + let subtree_root_index = Subtree::find_subtree_root(index); + Ok(self + .get_subtree(subtree_root_index)? + .and_then(|subtree| subtree.get_inner_node(index))) + } + + /// Returns an iterator over all leaves in this snapshot. + fn iter_leaves(&self) -> Result + '_>, StorageError> { + let cf = self.cf_handle(LEAVES_CF)?; let mut read_opts = ReadOptions::default(); read_opts.set_total_order_seek(true); - let mut leaves = Map::new(); - for item in snapshot.iterator_cf_opt(leaves_cf, read_opts, IteratorMode::Start) { - let (key_bytes, value_bytes) = item?; - let leaf_idx = index_from_key_bytes(&key_bytes)?; - let leaf = SmtLeaf::read_from_bytes_with_budget(&value_bytes, value_bytes.len())?; - leaves.insert(leaf_idx, leaf); - } + let db_iter = self.inner.snapshot.iterator_cf_opt(cf, read_opts, IteratorMode::Start); + Ok(Box::new(RocksDbDirectLeafIterator { iter: db_iter })) + } + + /// Returns an iterator over all subtrees in this snapshot. + fn iter_subtrees(&self) -> Result + '_>, StorageError> { const SUBTREE_CFS: [&str; 5] = [SUBTREE_24_CF, SUBTREE_32_CF, SUBTREE_40_CF, SUBTREE_48_CF, SUBTREE_56_CF]; - let mut subtrees = Map::new(); - for (cf_index, cf_name) in SUBTREE_CFS.into_iter().enumerate() { - let cf = self.cf_handle(cf_name)?; - let depth = IN_MEMORY_DEPTH + (cf_index as u8 * 8); - let mut read_opts = ReadOptions::default(); - read_opts.set_total_order_seek(true); - for item in snapshot.iterator_cf_opt(cf, read_opts, IteratorMode::Start) { - let (key_bytes, value_bytes) = item?; - let node_idx = subtree_root_from_key_bytes(&key_bytes, depth)?; - let subtree = Subtree::from_vec(node_idx, &value_bytes)?; - subtrees.insert(subtree.root_index(), subtree); - } + let mut cf_handles = Vec::new(); + for cf_name in SUBTREE_CFS { + cf_handles.push(self.cf_handle(cf_name)?); } - Ok(MemoryStorage { leaves, subtrees }.into_snapshot()) + Ok(Box::new(RocksDbSnapshotSubtreeIterator::new(&self.inner.snapshot, cf_handles))) + } + + /// Retrieves all depth 24 hashes from this snapshot. + fn get_depth24(&self) -> Result, StorageError> { + let cf = self.cf_handle(DEPTH_24_CF)?; + let iter = self.inner.snapshot.iterator_cf(cf, IteratorMode::Start); + collect_depth24(iter) + } +} + +impl SmtStorage for RocksDbStorage { + type Reader = RocksDbSnapshotStorage; + + /// Returns a detached read-only snapshot of the current RocksDB-backed storage. + fn reader(&self) -> Result { + Ok(RocksDbSnapshotStorage::new(Arc::clone(&self.db))) } /// Inserts a key-value pair into the SMT leaf at the specified logical `index`. @@ -1061,6 +1280,63 @@ struct RocksDbSubtreeIterator<'a> { current_iter: Option>, } +/// An iterator over subtrees from multiple RocksDB column families in a single snapshot. +struct RocksDbSnapshotSubtreeIterator<'a> { + snapshot: &'a rocksdb::Snapshot<'static>, + cf_handles: Vec<&'a rocksdb::ColumnFamily>, + current_cf_index: usize, + current_iter: Option>, +} + +impl<'a> RocksDbSnapshotSubtreeIterator<'a> { + fn new( + snapshot: &'a rocksdb::Snapshot<'static>, + cf_handles: Vec<&'a rocksdb::ColumnFamily>, + ) -> Self { + let mut iterator = Self { + snapshot, + cf_handles, + current_cf_index: 0, + current_iter: None, + }; + iterator.advance_to_next_cf(); + iterator + } + + fn advance_to_next_cf(&mut self) { + if self.current_cf_index < self.cf_handles.len() { + let cf = self.cf_handles[self.current_cf_index]; + let mut read_opts = ReadOptions::default(); + read_opts.set_total_order_seek(true); + self.current_iter = + Some(self.snapshot.iterator_cf_opt(cf, read_opts, IteratorMode::Start)); + } else { + self.current_iter = None; + } + } +} + +impl Iterator for RocksDbSnapshotSubtreeIterator<'_> { + type Item = Subtree; + + fn next(&mut self) -> Option { + loop { + let iter = self.current_iter.as_mut()?; + + if let Some(subtree) = + RocksDbSubtreeIterator::try_next_from_iter(iter, self.current_cf_index) + { + return Some(subtree); + } + + self.current_cf_index += 1; + self.advance_to_next_cf(); + + self.current_iter.as_ref()?; + } + } +} + impl<'a> RocksDbSubtreeIterator<'a> { fn new(db: &'a DB, cf_handles: Vec<&'a rocksdb::ColumnFamily>) -> Self { let mut iterator = Self { @@ -1273,6 +1549,32 @@ fn index_from_key_bytes(key_bytes: &[u8]) -> Result { Ok(u64::from_be_bytes(arr)) } +fn read_count(what: &'static str, bytes: &[u8]) -> Result { + let arr: [u8; 8] = bytes.try_into().map_err(|_| StorageError::BadValueLen { + what, + expected: 8, + found: bytes.len(), + })?; + Ok(usize::from_be_bytes(arr)) +} + +fn collect_depth24( + iter: DBIteratorWithThreadMode<'_, DB>, +) -> Result, StorageError> { + let mut hashes = Vec::new(); + + for item in iter { + let (key_bytes, value_bytes) = item?; + + let index = index_from_key_bytes(&key_bytes)?; + let hash = Word::read_from_bytes_with_budget(&value_bytes, value_bytes.len())?; + + hashes.push((index, hash)); + } + + Ok(hashes) +} + /// Reconstructs a `NodeIndex` from the variable-length subtree key stored in RocksDB. /// /// * `key_bytes` is the big-endian tail of the 64-bit value: diff --git a/miden-crypto/src/merkle/smt/large/tests.rs b/miden-crypto/src/merkle/smt/large/tests.rs index b5ee9bd94..442416d29 100644 --- a/miden-crypto/src/merkle/smt/large/tests.rs +++ b/miden-crypto/src/merkle/smt/large/tests.rs @@ -767,3 +767,41 @@ fn test_memory_storage_snapshot_depth24() { "root reconstructed from snapshot must match the original" ); } + +#[test] +fn reader_shares_in_memory_top_until_writer_mutates() { + let entries = generate_entries(1000); + let storage = MemoryStorage::new(); + let mut smt = LargeSmt::::with_entries(storage, entries).unwrap(); + + let reader = smt.reader().unwrap(); + let reader_root = reader.root(); + assert_eq!(smt.in_memory_nodes().as_ptr(), reader.in_memory_nodes().as_ptr()); + + let key = Word::new([ONE, ONE, Felt::new_unchecked(10_000), Felt::new_unchecked(10_000)]); + let value = Word::new([ONE, ONE, ONE, Felt::new_unchecked(10_000)]); + smt.insert(key, value).unwrap(); + + assert_ne!(smt.in_memory_nodes().as_ptr(), reader.in_memory_nodes().as_ptr()); + assert_ne!(smt.root(), reader_root); + assert_eq!(reader.root(), reader_root); +} + +#[test] +fn clone_shares_in_memory_top_until_mutation() { + let entries = generate_entries(1000); + let storage = MemoryStorage::new(); + let smt = LargeSmt::::with_entries(storage, entries).unwrap(); + + let mut clone = smt.clone(); + let original_root = smt.root(); + assert_eq!(smt.in_memory_nodes().as_ptr(), clone.in_memory_nodes().as_ptr()); + + let key = Word::new([ONE, ONE, Felt::new_unchecked(10_001), Felt::new_unchecked(10_001)]); + let value = Word::new([ONE, ONE, ONE, Felt::new_unchecked(10_001)]); + clone.insert(key, value).unwrap(); + + assert_ne!(smt.in_memory_nodes().as_ptr(), clone.in_memory_nodes().as_ptr()); + assert_eq!(smt.root(), original_root); + assert_ne!(clone.root(), original_root); +} diff --git a/miden-crypto/src/merkle/smt/mod.rs b/miden-crypto/src/merkle/smt/mod.rs index 9121ca1fb..6a0f6988c 100644 --- a/miden-crypto/src/merkle/smt/mod.rs +++ b/miden-crypto/src/merkle/smt/mod.rs @@ -22,12 +22,11 @@ mod large; pub use full::concurrent::{SubtreeLeaf, build_subtree_for_bench}; #[cfg(feature = "concurrent")] pub use large::{ - CloneableSmtStorageReader, LargeSmt, LargeSmtError, MemoryStorage, SmtStorage, - SmtStorageReader, SmtStorageSnapshot, StorageError, StorageUpdateParts, StorageUpdates, - Subtree, SubtreeError, SubtreeUpdate, + LargeSmt, LargeSmtError, MemoryStorage, MemoryStorageSnapshot, SmtStorage, SmtStorageReader, + StorageError, StorageUpdateParts, StorageUpdates, Subtree, SubtreeError, SubtreeUpdate, }; #[cfg(feature = "rocksdb")] -pub use large::{RocksDbConfig, RocksDbStorage}; +pub use large::{RocksDbConfig, RocksDbSnapshotStorage, RocksDbStorage}; mod large_forest; pub use large_forest::{ diff --git a/miden-crypto/tests/rocksdb_large_smt.rs b/miden-crypto/tests/rocksdb_large_smt.rs index 4f250ba4a..781f08983 100644 --- a/miden-crypto/tests/rocksdb_large_smt.rs +++ b/miden-crypto/tests/rocksdb_large_smt.rs @@ -2,7 +2,7 @@ use miden_crypto::{ EMPTY_WORD, Felt, ONE, Word, ZERO, merkle::{ InnerNodeInfo, - smt::{LargeSmt, LargeSmtError, RocksDbConfig, RocksDbStorage, SmtStorageSnapshot}, + smt::{LargeSmt, LargeSmtError, RocksDbConfig, RocksDbSnapshotStorage, RocksDbStorage}, }, }; use tempfile::TempDir; @@ -50,9 +50,11 @@ fn rocksdb_sanity_insert_and_get() { #[test] fn rocksdb_reader_is_detached_snapshot() { - fn assert_snapshot_reader(_: &LargeSmt) {} + fn assert_snapshot_reader(_: &LargeSmt) {} let entries = generate_entries(1000); + let existing_key = entries[10].0; + let existing_value = entries[10].1; let (storage, _tmp) = setup_storage(); let mut smt = LargeSmt::::with_entries(storage, entries).unwrap(); @@ -70,6 +72,13 @@ fn rocksdb_reader_is_detached_snapshot() { assert_ne!(smt.root(), reader_root); assert_eq!(reader.root(), reader_root); assert_eq!(reader.get_value(&key), EMPTY_WORD); + assert_eq!(reader.get_value(&existing_key), existing_value); + + drop(smt); + + assert_eq!(reader.root(), reader_root); + assert_eq!(reader.get_value(&key), EMPTY_WORD); + assert_eq!(reader.get_value(&existing_key), existing_value); } #[test]