From 98514c495ebec7229c13a90764f9398c78c38db0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Fran=C3=A7ois=20Garillot?= Date: Tue, 21 Apr 2026 06:23:33 -0400 Subject: [PATCH] Harden LargeSmt reader snapshots --- miden-crypto/src/boxed_storage.rs | 4 +- miden-crypto/src/merkle/smt/large/mod.rs | 19 ++-- .../src/merkle/smt/large/storage/memory.rs | 86 +++++++++++++++++-- .../src/merkle/smt/large/storage/mod.rs | 9 +- .../src/merkle/smt/large/storage/rocksdb.rs | 46 +++++++--- miden-crypto/src/merkle/smt/mod.rs | 5 +- miden-crypto/tests/rocksdb_large_smt.rs | 26 +++++- 7 files changed, 158 insertions(+), 37 deletions(-) diff --git a/miden-crypto/src/boxed_storage.rs b/miden-crypto/src/boxed_storage.rs index b375736e43..52bde48635 100644 --- a/miden-crypto/src/boxed_storage.rs +++ b/miden-crypto/src/boxed_storage.rs @@ -69,8 +69,8 @@ impl SmtStorageReader for BoxedStorage { impl SmtStorage for BoxedStorage { type Reader = Box; - fn reader(&self) -> Self::Reader { - Box::new(BoxedStorage(self.0.reader())) + fn reader(&self) -> Result { + Ok(Box::new(BoxedStorage(self.0.reader()?))) } fn insert_value( &mut self, diff --git a/miden-crypto/src/merkle/smt/large/mod.rs b/miden-crypto/src/merkle/smt/large/mod.rs index ef64c5ffb4..94c271472d 100644 --- a/miden-crypto/src/merkle/smt/large/mod.rs +++ b/miden-crypto/src/merkle/smt/large/mod.rs @@ -261,8 +261,8 @@ pub use subtree::{Subtree, SubtreeError}; mod storage; pub use storage::{ - MemoryStorage, SmtStorage, SmtStorageReader, StorageError, StorageUpdateParts, StorageUpdates, - SubtreeUpdate, + CloneableSmtStorageReader, MemoryStorage, SmtStorage, SmtStorageReader, SmtStorageSnapshot, + StorageError, StorageUpdateParts, StorageUpdates, SubtreeUpdate, }; #[cfg(feature = "rocksdb")] pub use storage::{RocksDbConfig, RocksDbStorage}; @@ -334,9 +334,8 @@ 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: Clone`. Read-only storage types (e.g. in-memory -/// snapshots) may implement `Clone`; mutable backends (e.g. database handles) typically do not, -/// which prevents accidental duplication of a shared mutable store. +/// `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. #[derive(Debug)] pub struct LargeSmt { storage: S, @@ -352,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(), @@ -515,13 +514,13 @@ impl LargeSmt { /// 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. - pub fn reader(&self) -> LargeSmt { - LargeSmt { - storage: self.storage.reader(), + pub fn reader(&self) -> Result, LargeSmtError> { + Ok(LargeSmt { + storage: self.storage.reader()?, in_memory_nodes: self.in_memory_nodes.clone(), leaf_count: self.leaf_count, entry_count: self.entry_count, - } + }) } } diff --git a/miden-crypto/src/merkle/smt/large/storage/memory.rs b/miden-crypto/src/merkle/smt/large/storage/memory.rs index c67b70c773..0364779a8a 100644 --- a/miden-crypto/src/merkle/smt/large/storage/memory.rs +++ b/miden-crypto/src/merkle/smt/large/storage/memory.rs @@ -1,7 +1,8 @@ use alloc::{boxed::Box, vec::Vec}; use super::{ - SmtStorage, SmtStorageReader, StorageError, StorageUpdateParts, StorageUpdates, SubtreeUpdate, + CloneableSmtStorageReader, SmtStorage, SmtStorageReader, StorageError, StorageUpdateParts, + StorageUpdates, SubtreeUpdate, }; use crate::{ EMPTY_WORD, Map, MapEntry, Word, @@ -30,6 +31,14 @@ pub struct MemoryStorage { pub subtrees: Map, } +/// Read-only snapshot of SMT storage data. +/// +/// This type intentionally implements [`SmtStorageReader`] only. It is used as the reader view for +/// 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); + impl MemoryStorage { /// Creates a new, empty in-memory storage for a Sparse Merkle Tree. /// @@ -37,6 +46,11 @@ impl MemoryStorage { pub fn new() -> Self { Self { leaves: Map::new(), subtrees: Map::new() } } + + /// Converts this storage into a read-only snapshot. + pub fn into_snapshot(self) -> SmtStorageSnapshot { + SmtStorageSnapshot(self) + } } impl Default for MemoryStorage { @@ -45,6 +59,62 @@ impl Default for MemoryStorage { } } +impl SmtStorageReader for SmtStorageSnapshot { + fn leaf_count(&self) -> Result { + self.0.leaf_count() + } + + fn entry_count(&self) -> Result { + self.0.entry_count() + } + + fn get_leaf(&self, index: u64) -> Result, StorageError> { + self.0.get_leaf(index) + } + + fn get_leaves(&self, indices: &[u64]) -> Result>, StorageError> { + self.0.get_leaves(indices) + } + + fn has_leaves(&self) -> Result { + self.0.has_leaves() + } + + fn get_subtree(&self, index: NodeIndex) -> Result, StorageError> { + self.0.get_subtree(index) + } + + fn get_subtrees(&self, indices: &[NodeIndex]) -> Result>, StorageError> { + self.0.get_subtrees(indices) + } + + fn get_leaf_and_subtrees( + &self, + leaf_index: u64, + subtree_indices: &[NodeIndex], + ) -> Result<(Option, Vec>), StorageError> { + self.0.get_leaf_and_subtrees(leaf_index, subtree_indices) + } + + fn get_inner_node(&self, index: NodeIndex) -> Result, StorageError> { + self.0.get_inner_node(index) + } + + fn iter_leaves(&self) -> Result + '_>, StorageError> { + self.0.iter_leaves() + } + + fn iter_subtrees(&self) -> Result + '_>, StorageError> { + self.0.iter_subtrees() + } + + fn get_depth24(&self) -> Result, StorageError> { + self.0.get_depth24() + } +} + +impl CloneableSmtStorageReader for SmtStorageSnapshot {} + impl SmtStorageReader for MemoryStorage { /// Gets the total number of non-empty leaves currently stored. fn leaf_count(&self) -> Result { @@ -131,16 +201,14 @@ impl SmtStorageReader for MemoryStorage { } } +impl CloneableSmtStorageReader for MemoryStorage {} + impl SmtStorage for MemoryStorage { - // The reader for in-memory storage is not a read-only view; it's a clone of the storage. - type Reader = MemoryStorage; + type Reader = SmtStorageSnapshot; - /// Returns a snapshot of this in-memory storage by cloning it. - /// - /// The reader is provided as a clone because we don't support a read-only type for in-memory - /// storage. - fn reader(&self) -> Self::Reader { - self.clone() + /// Returns a read-only snapshot of this in-memory storage by cloning it. + fn reader(&self) -> Result { + Ok(self.clone().into_snapshot()) } /// Inserts a key-value pair into the leaf at the given index. diff --git a/miden-crypto/src/merkle/smt/large/storage/mod.rs b/miden-crypto/src/merkle/smt/large/storage/mod.rs index 6f25a97a9a..dcf8706c83 100644 --- a/miden-crypto/src/merkle/smt/large/storage/mod.rs +++ b/miden-crypto/src/merkle/smt/large/storage/mod.rs @@ -21,7 +21,7 @@ mod rocksdb; pub use rocksdb::{RocksDbConfig, RocksDbStorage}; mod memory; -pub use memory::MemoryStorage; +pub use memory::{MemoryStorage, SmtStorageSnapshot}; mod updates; pub use updates::{StorageUpdateParts, StorageUpdates, SubtreeUpdate}; @@ -137,6 +137,9 @@ 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 { @@ -226,7 +229,7 @@ pub trait SmtStorage: SmtStorageReader { /// 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. - fn reader(&self) -> Self::Reader; + fn reader(&self) -> Result; /// Inserts a key-value pair into the SMT leaf at the specified logical `index`. /// @@ -339,7 +342,7 @@ impl SmtStorage for Box { type Reader = T::Reader; #[inline] - fn reader(&self) -> Self::Reader { + fn reader(&self) -> Result { self.deref().reader() } diff --git a/miden-crypto/src/merkle/smt/large/storage/rocksdb.rs b/miden-crypto/src/merkle/smt/large/storage/rocksdb.rs index 9178992683..14e003061a 100644 --- a/miden-crypto/src/merkle/smt/large/storage/rocksdb.rs +++ b/miden-crypto/src/merkle/smt/large/storage/rocksdb.rs @@ -7,7 +7,8 @@ use rocksdb::{ }; use super::{ - SmtStorage, SmtStorageReader, StorageError, StorageUpdateParts, StorageUpdates, SubtreeUpdate, + MemoryStorage, SmtStorage, SmtStorageReader, SmtStorageSnapshot, StorageError, + StorageUpdateParts, StorageUpdates, SubtreeUpdate, }; use crate::{ EMPTY_WORD, Word, @@ -534,16 +535,41 @@ impl SmtStorageReader for RocksDbStorage { } impl SmtStorage for RocksDbStorage { - type Reader = RocksDbStorage; + type Reader = SmtStorageSnapshot; - /// Returns a read-only handle to the same RocksDB database. - /// - /// Because `RocksDbStorage` wraps an `Arc`, the returned value shares the underlying - /// database with `self`. The returned value is intentionally typed as a reader, so the caller - /// cannot perform writes through it, but it is not isolated from subsequent writes made - /// through `self`. - fn reader(&self) -> Self::Reader { - self.clone() + /// Returns a detached read-only snapshot of the current RocksDB-backed storage. + fn reader(&self) -> Result { + let snapshot = self.db.snapshot(); + + let leaves_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); + } + + 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); + } + } + + Ok(MemoryStorage { leaves, subtrees }.into_snapshot()) } /// Inserts a key-value pair into the SMT leaf at the specified logical `index`. diff --git a/miden-crypto/src/merkle/smt/mod.rs b/miden-crypto/src/merkle/smt/mod.rs index 6e8d2084bf..46ae9b1653 100644 --- a/miden-crypto/src/merkle/smt/mod.rs +++ b/miden-crypto/src/merkle/smt/mod.rs @@ -22,8 +22,9 @@ mod large; pub use full::concurrent::{SubtreeLeaf, build_subtree_for_bench}; #[cfg(feature = "concurrent")] pub use large::{ - LargeSmt, LargeSmtError, MemoryStorage, SmtStorage, SmtStorageReader, StorageError, - StorageUpdateParts, StorageUpdates, Subtree, SubtreeError, SubtreeUpdate, + CloneableSmtStorageReader, LargeSmt, LargeSmtError, MemoryStorage, SmtStorage, + SmtStorageReader, SmtStorageSnapshot, StorageError, StorageUpdateParts, StorageUpdates, + Subtree, SubtreeError, SubtreeUpdate, }; #[cfg(feature = "rocksdb")] pub use large::{RocksDbConfig, RocksDbStorage}; diff --git a/miden-crypto/tests/rocksdb_large_smt.rs b/miden-crypto/tests/rocksdb_large_smt.rs index c42dff8b61..4f250ba4ab 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}, + smt::{LargeSmt, LargeSmtError, RocksDbConfig, RocksDbStorage, SmtStorageSnapshot}, }, }; use tempfile::TempDir; @@ -48,6 +48,30 @@ fn rocksdb_sanity_insert_and_get() { assert_eq!(smt.get_value(&key), val); } +#[test] +fn rocksdb_reader_is_detached_snapshot() { + fn assert_snapshot_reader(_: &LargeSmt) {} + + let entries = generate_entries(1000); + let (storage, _tmp) = setup_storage(); + let mut smt = LargeSmt::::with_entries(storage, entries).unwrap(); + + let reader = smt.reader().unwrap(); + assert_snapshot_reader(&reader); + + let reader_root = reader.root(); + 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)]); + + assert_eq!(reader.get_value(&key), EMPTY_WORD); + + smt.insert(key, value).unwrap(); + + assert_ne!(smt.root(), reader_root); + assert_eq!(reader.root(), reader_root); + assert_eq!(reader.get_value(&key), EMPTY_WORD); +} + #[test] fn rocksdb_persistence_reopen() { let entries = generate_entries(1000);