diff --git a/full-node/db/sov-db/src/native_db.rs b/full-node/db/sov-db/src/native_db.rs index 97e627b23..14b41a6cb 100644 --- a/full-node/db/sov-db/src/native_db.rs +++ b/full-node/db/sov-db/src/native_db.rs @@ -5,7 +5,10 @@ use sov_schema_db::{SchemaBatch, DB}; use crate::rocks_db_config::gen_rocksdb_options; use crate::schema::tables::{ModuleAccessoryState, NATIVE_TABLES}; -use crate::schema::types::StateKey; +use crate::schema::types::AccessoryKey; + +/// Specifies a particular version of the Accessory state. +pub type Version = u64; /// A typed wrapper around RocksDB for storing native-only accessory state. /// Internally, this is roughly just an [`Arc`]. @@ -37,20 +40,37 @@ impl NativeDB { } /// Queries for a value in the [`NativeDB`], given a key. - pub fn get_value_option(&self, key: &StateKey) -> anyhow::Result>> { - self.db - .get::(key) - .map(Option::flatten) + pub fn get_value_option( + &self, + key: &AccessoryKey, + version: Version, + ) -> anyhow::Result>> { + let mut iter = self.db.iter::()?; + iter.seek_for_prev(&(key.to_vec(), version))?; + let found = iter.next(); + match found { + Some(result) => { + let ((found_key, found_version), value) = result?; + if &found_key == key { + anyhow::ensure!(found_version <= version, "Bug! iterator isn't returning expected values. expected a version <= {version:} but found {found_version:}"); + Ok(value) + } else { + Ok(None) + } + } + None => Ok(None), + } } /// Sets a sequence of key-value pairs in the [`NativeDB`]. The write is atomic. pub fn set_values( &self, key_value_pairs: impl IntoIterator, Option>)>, + version: Version, ) -> anyhow::Result<()> { let mut batch = SchemaBatch::default(); for (key, value) in key_value_pairs { - batch.put::(&key, &value)?; + batch.put::(&(key, version), &value)?; } self.db.write_schemas(batch) } @@ -144,9 +164,13 @@ mod tests { let key = b"foo".to_vec(); let value = b"bar".to_vec(); - db.set_values(vec![(key.clone(), Some(value.clone()))]) + db.set_values(vec![(key.clone(), Some(value.clone()))], 0) + .unwrap(); + assert_eq!(db.get_value_option(&key, 0).unwrap(), Some(value.clone())); + let value2 = b"bar2".to_vec(); + db.set_values(vec![(key.clone(), Some(value2.clone()))], 1) .unwrap(); - assert_eq!(db.get_value_option(&key).unwrap(), Some(value)); + assert_eq!(db.get_value_option(&key, 0).unwrap(), Some(value)); } #[test] @@ -155,8 +179,8 @@ mod tests { let db = NativeDB::with_path(tmpdir.path()).unwrap(); let key = b"deleted".to_vec(); - db.set_values(vec![(key.clone(), None)]).unwrap(); - assert_eq!(db.get_value_option(&key).unwrap(), None); + db.set_values(vec![(key.clone(), None)], 0).unwrap(); + assert_eq!(db.get_value_option(&key, 0).unwrap(), None); } #[test] @@ -165,6 +189,6 @@ mod tests { let db = NativeDB::with_path(tmpdir.path()).unwrap(); let key = b"spam".to_vec(); - assert_eq!(db.get_value_option(&key).unwrap(), None); + assert_eq!(db.get_value_option(&key, 0).unwrap(), None); } } diff --git a/full-node/db/sov-db/src/schema/tables.rs b/full-node/db/sov-db/src/schema/tables.rs index 726ae09f5..2ed806a19 100644 --- a/full-node/db/sov-db/src/schema/tables.rs +++ b/full-node/db/sov-db/src/schema/tables.rs @@ -218,11 +218,6 @@ define_table_with_default_codec!( (SlotByHash) DbHash => SlotNumber ); -define_table_with_default_codec!( - /// Non-JMT state stored by a module for JSON-RPC use. - (ModuleAccessoryState) AccessoryKey => AccessoryStateValue -); - define_table_with_seek_key_codec!( /// The primary source for batch data (BatchByNumber) BatchNumber => StoredBatch @@ -347,3 +342,47 @@ define_table_with_default_codec!( /// which requires the ability to fetch values by hash. (KeyHashToKey) [u8;32] => StateKey ); + +define_table_without_codec!( + /// Non-JMT state stored by a module for JSON-RPC use. + (ModuleAccessoryState) (AccessoryKey, Version) => AccessoryStateValue +); + +impl KeyEncoder for (AccessoryKey, Version) { + fn encode_key(&self) -> sov_schema_db::schema::Result> { + let mut out = Vec::with_capacity(self.0.len() + std::mem::size_of::() + 8); + self.0 + .as_slice() + .serialize(&mut out) + .map_err(CodecError::from)?; + // Write the version in big-endian order so that sorting order is based on the most-significant bytes of the key + out.write_u64::(self.1) + .expect("serialization to vec is infallible"); + Ok(out) + } +} + +impl SeekKeyEncoder for (AccessoryKey, Version) { + fn encode_seek_key(&self) -> sov_schema_db::schema::Result> { + <(Vec, u64) as KeyEncoder>::encode_key(self) + } +} + +impl KeyDecoder for (AccessoryKey, Version) { + fn decode_key(data: &[u8]) -> sov_schema_db::schema::Result { + let mut cursor = maybestd::io::Cursor::new(data); + let key = Vec::::deserialize_reader(&mut cursor)?; + let version = cursor.read_u64::()?; + Ok((key, version)) + } +} + +impl ValueCodec for AccessoryStateValue { + fn encode_value(&self) -> sov_schema_db::schema::Result> { + self.try_to_vec().map_err(CodecError::from) + } + + fn decode_value(data: &[u8]) -> sov_schema_db::schema::Result { + Ok(Self::deserialize_reader(&mut &data[..])?) + } +} diff --git a/module-system/module-implementations/sov-bank/src/token.rs b/module-system/module-implementations/sov-bank/src/token.rs index e2f837ba5..f43e1886d 100644 --- a/module-system/module-implementations/sov-bank/src/token.rs +++ b/module-system/module-implementations/sov-bank/src/token.rs @@ -145,7 +145,6 @@ impl Token { self.balances.set(from, &from_balance, working_set); self.balances.set(to, &to_balance, working_set); - Ok(()) } /// Burns a specified `amount` of token from the address `from`. First check that the address has enough token to burn, diff --git a/module-system/module-implementations/sov-bank/tests/archival_query_test.rs b/module-system/module-implementations/sov-bank/tests/archival_query_test.rs new file mode 100644 index 000000000..aa8f78f17 --- /dev/null +++ b/module-system/module-implementations/sov-bank/tests/archival_query_test.rs @@ -0,0 +1,279 @@ +mod helpers; + +use helpers::*; +use sov_bank::{get_genesis_token_address, Amount, Bank, CallMessage, Coins}; +use sov_modules_api::default_context::DefaultContext; +use sov_modules_api::{Address, Context, Module, StateReaderAndWriter, WorkingSet}; +use sov_state::storage::{StorageKey, StorageValue}; +use sov_state::{DefaultStorageSpec, ProverStorage, Storage}; + +#[test] +fn transfer_initial_token() { + let initial_balance = 100; + let bank_config = create_bank_config_with_token(3, initial_balance); + let tmpdir = tempfile::tempdir().unwrap(); + let prover_storage = ProverStorage::with_path(tmpdir.path()).unwrap(); + let mut working_set = WorkingSet::new(prover_storage.clone()); + let bank = Bank::default(); + bank.genesis(&bank_config, &mut working_set).unwrap(); + + let token_address = get_genesis_token_address::( + &bank_config.tokens[0].token_name, + bank_config.tokens[0].salt, + ); + let sender_address = bank_config.tokens[0].address_and_balances[0].0; + let receiver_address = bank_config.tokens[0].address_and_balances[1].0; + assert_ne!(sender_address, receiver_address); + + let (sender_balance, receiver_balance) = query_sender_receiver_balances( + &bank, + token_address, + sender_address, + receiver_address, + &mut working_set, + ); + assert_eq!((sender_balance, receiver_balance), (100, 100)); + commit(working_set, prover_storage.clone()); + + let mut working_set: WorkingSet = WorkingSet::new(prover_storage.clone()); + + transfer( + &bank, + token_address, + sender_address, + receiver_address, + 10, + &mut working_set, + ); + let (sender_balance, receiver_balance) = query_sender_receiver_balances( + &bank, + token_address, + sender_address, + receiver_address, + &mut working_set, + ); + assert_eq!((sender_balance, receiver_balance), (90, 110)); + + commit(working_set, prover_storage.clone()); + + let mut working_set: WorkingSet = WorkingSet::new(prover_storage.clone()); + + transfer( + &bank, + token_address, + sender_address, + receiver_address, + 10, + &mut working_set, + ); + let (sender_balance, receiver_balance) = query_sender_receiver_balances( + &bank, + token_address, + sender_address, + receiver_address, + &mut working_set, + ); + assert_eq!((sender_balance, receiver_balance), (80, 120)); + commit(working_set, prover_storage.clone()); + + // Archival tests + + let archival_slot: u64 = 2; + let mut working_set: WorkingSet = WorkingSet::new(prover_storage.clone()); + working_set.set_archival_version(archival_slot); + + let (sender_balance, receiver_balance) = query_sender_receiver_balances( + &bank, + token_address, + sender_address, + receiver_address, + &mut working_set, + ); + assert_eq!((sender_balance, receiver_balance), (90, 110)); + + // modify in archival + transfer( + &bank, + token_address, + sender_address, + receiver_address, + 5, + &mut working_set, + ); + + let (sender_balance, receiver_balance) = query_sender_receiver_balances( + &bank, + token_address, + sender_address, + receiver_address, + &mut working_set, + ); + assert_eq!((sender_balance, receiver_balance), (85, 115)); + + let archival_slot: u64 = 1; + let mut working_set: WorkingSet = WorkingSet::new(prover_storage.clone()); + working_set.set_archival_version(archival_slot); + let (sender_balance, receiver_balance) = query_sender_receiver_balances( + &bank, + token_address, + sender_address, + receiver_address, + &mut working_set, + ); + assert_eq!((sender_balance, receiver_balance), (100, 100)); + + transfer( + &bank, + token_address, + sender_address, + receiver_address, + 45, + &mut working_set, + ); + + let (sender_balance, receiver_balance) = query_sender_receiver_balances( + &bank, + token_address, + sender_address, + receiver_address, + &mut working_set, + ); + assert_eq!((sender_balance, receiver_balance), (55, 145)); + + working_set.unset_archival_version(); + let (sender_balance, receiver_balance) = query_sender_receiver_balances( + &bank, + token_address, + sender_address, + receiver_address, + &mut working_set, + ); + assert_eq!((sender_balance, receiver_balance), (80, 120)); + + // Accessory tests + + transfer( + &bank, + token_address, + sender_address, + receiver_address, + 10, + &mut working_set, + ); + let (sender_balance, receiver_balance) = query_sender_receiver_balances( + &bank, + token_address, + sender_address, + receiver_address, + &mut working_set, + ); + assert_eq!((sender_balance, receiver_balance), (70, 130)); + + let mut accessory_state = working_set.accessory_state(); + accessory_state.set(&StorageKey::from("k"), StorageValue::from(b"v1".to_vec())); + let val = accessory_state.get(&StorageKey::from("k")).unwrap(); + assert_eq!("v1", String::from_utf8(val.value().to_vec()).unwrap()); + + commit(working_set, prover_storage.clone()); + + // next block + + let mut working_set: WorkingSet = WorkingSet::new(prover_storage.clone()); + transfer( + &bank, + token_address, + sender_address, + receiver_address, + 10, + &mut working_set, + ); + let (sender_balance, receiver_balance) = query_sender_receiver_balances( + &bank, + token_address, + sender_address, + receiver_address, + &mut working_set, + ); + assert_eq!((sender_balance, receiver_balance), (60, 140)); + let mut accessory_state = working_set.accessory_state(); + accessory_state.set(&StorageKey::from("k"), StorageValue::from(b"v2".to_vec())); + let val = accessory_state.get(&StorageKey::from("k")).unwrap(); + assert_eq!("v2", String::from_utf8(val.value().to_vec()).unwrap()); + + commit(working_set, prover_storage.clone()); + + // archival versioned state query + + let archival_slot = 3; + let mut working_set: WorkingSet = WorkingSet::new(prover_storage.clone()); + working_set.set_archival_version(archival_slot); + let mut accessory_state = working_set.accessory_state(); + let val = accessory_state.get(&StorageKey::from("k")).unwrap(); + assert_eq!("v1", String::from_utf8(val.value().to_vec()).unwrap()); + + // archival accessory set + + accessory_state.set(&StorageKey::from("k"), StorageValue::from(b"v3".to_vec())); + let val = accessory_state.get(&StorageKey::from("k")).unwrap(); + assert_eq!("v3", String::from_utf8(val.value().to_vec()).unwrap()); + + working_set.unset_archival_version(); + let mut accessory_state = working_set.accessory_state(); + let val = accessory_state.get(&StorageKey::from("k")).unwrap(); + assert_eq!("v2", String::from_utf8(val.value().to_vec()).unwrap()); +} + +fn query_sender_receiver_balances( + bank: &Bank, + token_address: Address, + sender_address: Address, + receiver_address: Address, + working_set: &mut WorkingSet, +) -> (u64, u64) { + let sender_balance = bank + .get_balance_of(sender_address, token_address, working_set) + .unwrap(); + let receiver_balance = bank + .get_balance_of(receiver_address, token_address, working_set) + .unwrap(); + (sender_balance, receiver_balance) +} + +fn transfer( + bank: &Bank, + token_address: Address, + sender_address: Address, + receiver_address: Address, + transfer_amount: Amount, + working_set: &mut WorkingSet, +) { + let transfer_message = CallMessage::Transfer { + to: receiver_address, + coins: Coins { + amount: transfer_amount, + token_address, + }, + }; + + let sender_context = C::new(sender_address, 1); + + bank.call(transfer_message, &sender_context, working_set) + .expect("Transfer call failed"); +} + +fn commit(working_set: WorkingSet, storage: ProverStorage) { + // Save checkpoint + let mut checkpoint = working_set.checkpoint(); + + let (cache_log, witness) = checkpoint.freeze(); + + let (_, authenticated_node_batch) = storage + .compute_state_update(cache_log, &witness) + .expect("jellyfish merkle tree update must succeed"); + + let working_set = checkpoint.to_revertable(); + + let accessory_log = working_set.checkpoint().freeze_non_provable(); + + storage.commit(&authenticated_node_batch, &accessory_log); +} diff --git a/module-system/sov-modules-api/src/containers/map.rs b/module-system/sov-modules-api/src/containers/map.rs index 8bd63344d..8f1d564a5 100644 --- a/module-system/sov-modules-api/src/containers/map.rs +++ b/module-system/sov-modules-api/src/containers/map.rs @@ -4,7 +4,6 @@ use sov_modules_core::{Context, Prefix, StateCodec, StateKeyCodec, StateValueCod use sov_state::codec::BorshCodec; use super::traits::StateMapAccessor; - /// A container that maps keys to values. /// /// # Type parameters diff --git a/module-system/sov-modules-api/src/containers/mod.rs b/module-system/sov-modules-api/src/containers/mod.rs index 83a7a9a90..84645cfbc 100644 --- a/module-system/sov-modules-api/src/containers/mod.rs +++ b/module-system/sov-modules-api/src/containers/mod.rs @@ -81,7 +81,10 @@ mod test { prover_storage .validate_and_commit(cache, &witness) .expect("storage is valid"); - assert_eq!(test.value, prover_storage.get(&test.key, &witness).unwrap()); + assert_eq!( + test.value, + prover_storage.get(&test.key, None, &witness).unwrap() + ); } let version_after = get_state_db_version(path); assert_eq!(version_after, test.version + 1) @@ -95,7 +98,7 @@ mod test { for test in tests { assert_eq!( test.value, - storage.get(&test.key, &Default::default()).unwrap() + storage.get(&test.key, None, &Default::default()).unwrap() ); } } @@ -130,7 +133,7 @@ mod test { assert!(!prover_storage.is_empty()); assert_eq!( value, - prover_storage.get(&key, &Default::default()).unwrap() + prover_storage.get(&key, None, &Default::default()).unwrap() ); } } diff --git a/module-system/sov-modules-api/src/lib.rs b/module-system/sov-modules-api/src/lib.rs index e9b153d3f..5fbf27166 100644 --- a/module-system/sov-modules-api/src/lib.rs +++ b/module-system/sov-modules-api/src/lib.rs @@ -34,9 +34,10 @@ use serde::{Deserialize, Serialize}; #[cfg(feature = "native")] pub use sov_modules_core::PrivateKey; pub use sov_modules_core::{ - runtime, AccessoryWorkingSet, Address, AddressBech32, CallResponse, Context, DispatchCall, - EncodeCall, GasUnit, Genesis, Module, ModuleCallJsonSchema, ModuleError, ModuleError as Error, - ModuleInfo, ModulePrefix, PublicKey, Signature, Spec, StateCheckpoint, WorkingSet, + archival_state, runtime, AccessoryWorkingSet, Address, AddressBech32, CallResponse, Context, + DispatchCall, EncodeCall, GasUnit, Genesis, Module, ModuleCallJsonSchema, ModuleError, + ModuleError as Error, ModuleInfo, ModulePrefix, PublicKey, Signature, Spec, StateCheckpoint, + StateReaderAndWriter, WorkingSet, }; pub use sov_rollup_interface::da::{BlobReaderTrait, DaSpec}; pub use sov_rollup_interface::services::da::SlotData; diff --git a/module-system/sov-modules-core/src/lib.rs b/module-system/sov-modules-core/src/lib.rs index 5a6bd8445..a8dfe54e1 100644 --- a/module-system/sov-modules-core/src/lib.rs +++ b/module-system/sov-modules-core/src/lib.rs @@ -3,6 +3,7 @@ #![doc = include_str!("../README.md")] extern crate alloc; +extern crate core; #[cfg(all(feature = "sync", not(target_has_atomic = "ptr")))] compile_error!("The `sync` feature is not supported on this architecture."); diff --git a/module-system/sov-modules-core/src/storage/cache.rs b/module-system/sov-modules-core/src/storage/cache.rs index 262eece19..67978193c 100644 --- a/module-system/sov-modules-core/src/storage/cache.rs +++ b/module-system/sov-modules-core/src/storage/cache.rs @@ -367,9 +367,19 @@ pub struct StorageInternalCache { pub tx_cache: CacheLog, /// Ordered reads and writes. pub ordered_db_reads: Vec<(CacheKey, Option)>, + /// Version for versioned usage with cache + pub version: Option, } impl StorageInternalCache { + /// Wrapper around default that can create the cache with knowledge of the version + pub fn new_with_version(version: u64) -> Self { + StorageInternalCache { + version: Some(version), + ..Default::default() + } + } + /// Gets a value from the cache or reads it from the provided `ValueReader`. pub fn get_or_fetch( &mut self, @@ -377,14 +387,14 @@ impl StorageInternalCache { value_reader: &S, witness: &S::Witness, ) -> Option { - let cache_key = key.to_cache_key(); + let cache_key = key.to_cache_key_version(self.version); let cache_value = self.get_value_from_cache(&cache_key); match cache_value { ValueExists::Yes(cache_value_exists) => cache_value_exists.map(Into::into), // If the value does not exist in the cache, then fetch it from an external source. ValueExists::No => { - let storage_value = value_reader.get(key, witness); + let storage_value = value_reader.get(key, self.version, witness); let cache_value = storage_value.as_ref().map(|v| v.clone().into_cache_value()); self.add_read(cache_key, cache_value); @@ -395,20 +405,20 @@ impl StorageInternalCache { /// Gets a keyed value from the cache, returning a wrapper on whether it exists. pub fn try_get(&self, key: &StorageKey) -> ValueExists { - let cache_key = key.to_cache_key(); + let cache_key = key.to_cache_key_version(self.version); self.get_value_from_cache(&cache_key) } /// Replaces the keyed value on the storage. pub fn set(&mut self, key: &StorageKey, value: StorageValue) { - let cache_key = key.to_cache_key(); + let cache_key = key.to_cache_key_version(self.version); let cache_value = value.into_cache_value(); self.tx_cache.add_write(cache_key, Some(cache_value)); } /// Deletes a keyed value from the cache. pub fn delete(&mut self, key: &StorageKey) { - let cache_key = key.to_cache_key(); + let cache_key = key.to_cache_key_version(self.version); self.tx_cache.add_write(cache_key, None); } diff --git a/module-system/sov-modules-core/src/storage/mod.rs b/module-system/sov-modules-core/src/storage/mod.rs index 049f4d5e6..1cb6399b0 100644 --- a/module-system/sov-modules-core/src/storage/mod.rs +++ b/module-system/sov-modules-core/src/storage/mod.rs @@ -48,6 +48,22 @@ impl StorageKey { } } + /// Converts this key into a [`CacheKey`] via cloning. + pub fn to_cache_key_version(&self, version: Option) -> CacheKey { + match version { + None => CacheKey { + key: self.key.clone(), + }, + Some(v) => { + let mut bytes = v.to_be_bytes().to_vec(); + bytes.extend((*self.key).clone()); + CacheKey { + key: RefCount::new(bytes), + } + } + } + } + /// Converts this key into a [`CacheKey`]. pub fn into_cache_key(self) -> CacheKey { CacheKey { key: self.key } @@ -197,7 +213,12 @@ pub trait Storage: Clone { fn with_config(config: Self::RuntimeConfig) -> anyhow::Result; /// Returns the value corresponding to the key or None if key is absent. - fn get(&self, key: &StorageKey, witness: &Self::Witness) -> Option; + fn get( + &self, + key: &StorageKey, + version: Option, + witness: &Self::Witness, + ) -> Option; /// Returns the value corresponding to the key or None if key is absent. /// @@ -206,7 +227,7 @@ pub trait Storage: Clone { /// execution environments** (i.e. outside of the zmVM) **SHOULD** override /// this method to return a value. This is because accessory state **MUST /// NOT** be readable from within the zmVM. - fn get_accessory(&self, _key: &StorageKey) -> Option { + fn get_accessory(&self, _key: &StorageKey, _version: Option) -> Option { None } diff --git a/module-system/sov-modules-core/src/storage/scratchpad.rs b/module-system/sov-modules-core/src/storage/scratchpad.rs index 0e9ab3674..f324ae6a7 100644 --- a/module-system/sov-modules-core/src/storage/scratchpad.rs +++ b/module-system/sov-modules-core/src/storage/scratchpad.rs @@ -7,12 +7,15 @@ pub use kernel_state::{KernelWorkingSet, VersionedWorkingSet}; use sov_rollup_interface::maybestd::collections::HashMap; use sov_rollup_interface::stf::Event; +use crate::archival_state::{ArchivalAccessoryWorkingSet, ArchivalJmtWorkingSet}; use crate::common::{GasMeter, Prefix}; use crate::module::{Context, Spec}; use crate::storage::{ CacheKey, CacheValue, EncodeKeyLike, NativeStorage, OrderedReadsAndWrites, StateCodec, StateValueCodec, Storage, StorageInternalCache, StorageKey, StorageProof, StorageValue, }; +use crate::Version; + /// A storage reader and writer pub trait StateReaderAndWriter { /// Get a value from the storage. @@ -154,15 +157,18 @@ pub struct Delta { } impl Delta { - fn new(inner: S) -> Self { - Self::with_witness(inner, Default::default()) + fn new(inner: S, version: Option) -> Self { + Self::with_witness(inner, Default::default(), version) } - fn with_witness(inner: S, witness: S::Witness) -> Self { + fn with_witness(inner: S, witness: S::Witness, version: Option) -> Self { Self { inner, witness, - cache: Default::default(), + cache: match version { + None => Default::default(), + Some(v) => StorageInternalCache::new_with_version(v), + }, } } @@ -194,7 +200,13 @@ impl StateReaderAndWriter for Delta { } } -type RevertableWrites = HashMap>; +// type RevertableWrites = HashMap>; + +#[derive(Default)] +struct RevertableWrites { + pub cache: HashMap>, + pub version: Option, +} struct AccessoryDelta { // This inner storage is never accessed inside the zkVM because reads are @@ -205,18 +217,22 @@ struct AccessoryDelta { } impl AccessoryDelta { - fn new(storage: S) -> Self { - Self { - storage, - writes: Default::default(), - } + fn new(storage: S, version: Option) -> Self { + let writes = match version { + None => Default::default(), + Some(v) => RevertableWrites { + cache: Default::default(), + version: Some(v), + }, + }; + Self { storage, writes } } fn freeze(&mut self) -> OrderedReadsAndWrites { let mut reads_and_writes = OrderedReadsAndWrites::default(); let writes = mem::take(&mut self.writes); - for write in writes { + for write in writes.cache { reads_and_writes.ordered_writes.push((write.0, write.1)); } @@ -226,20 +242,24 @@ impl AccessoryDelta { impl StateReaderAndWriter for AccessoryDelta { fn get(&mut self, key: &StorageKey) -> Option { - let cache_key = key.to_cache_key(); - if let Some(value) = self.writes.get(&cache_key) { + let cache_key = key.to_cache_key_version(self.writes.version); + if let Some(value) = self.writes.cache.get(&cache_key) { return value.clone().map(Into::into); } - self.storage.get_accessory(key) + self.storage.get_accessory(key, self.writes.version) } fn set(&mut self, key: &StorageKey, value: StorageValue) { - self.writes - .insert(key.to_cache_key(), Some(value.into_cache_value())); + self.writes.cache.insert( + key.to_cache_key_version(self.writes.version), + Some(value.into_cache_value()), + ); } fn delete(&mut self, key: &StorageKey) { - self.writes.insert(key.to_cache_key(), None); + self.writes + .cache + .insert(key.to_cache_key_version(self.writes.version), None); } } @@ -258,8 +278,8 @@ impl StateCheckpoint { /// by the given [`Storage`]. pub fn new(inner: ::Storage) -> Self { Self { - delta: Delta::new(inner.clone()), - accessory_delta: AccessoryDelta::new(inner), + delta: Delta::new(inner.clone(), None), + accessory_delta: AccessoryDelta::new(inner, None), } } @@ -270,18 +290,20 @@ impl StateCheckpoint { witness: <::Storage as Storage>::Witness, ) -> Self { Self { - delta: Delta::with_witness(inner.clone(), witness), - accessory_delta: AccessoryDelta::new(inner), + delta: Delta::with_witness(inner.clone(), witness, None), + accessory_delta: AccessoryDelta::new(inner, None), } } /// Transforms this [`StateCheckpoint`] back into a [`WorkingSet`]. pub fn to_revertable(self) -> WorkingSet { WorkingSet { - delta: RevertableWriter::new(self.delta), - accessory_delta: RevertableWriter::new(self.accessory_delta), + delta: RevertableWriter::new(self.delta, None), + accessory_delta: RevertableWriter::new(self.accessory_delta, None), events: Default::default(), gas_meter: GasMeter::default(), + archival_working_set: None, + archival_accessory_working_set: None, } } @@ -319,6 +341,8 @@ pub struct WorkingSet { accessory_delta: RevertableWriter>, events: Vec, gas_meter: GasMeter, + archival_working_set: Option>, + archival_accessory_working_set: Option>, } impl WorkingSet { @@ -338,6 +362,28 @@ impl WorkingSet { AccessoryWorkingSet { ws: self } } + /// Returns a handler for the archival state (JMT state). + fn archival_state(&mut self, version: Version) -> ArchivalJmtWorkingSet { + ArchivalJmtWorkingSet::new(&self.delta.inner.inner, version) + } + + /// Returns a handler for the archival accessory state (non-JMT state). + fn archival_accessory_state(&mut self, version: Version) -> ArchivalAccessoryWorkingSet { + ArchivalAccessoryWorkingSet::new(&self.accessory_delta.inner.storage, version) + } + + /// Sets archival version for a working set + pub fn set_archival_version(&mut self, version: Version) { + self.archival_working_set = Some(self.archival_state(version)); + self.archival_accessory_working_set = Some(self.archival_accessory_state(version)); + } + + /// Unset archival version + pub fn unset_archival_version(&mut self) { + self.archival_working_set = None; + self.archival_accessory_working_set = None; + } + /// Returns a handler for the kernel state (priveleged jmt state) /// /// You can use this method when calling getters and setters on accessory @@ -424,15 +470,24 @@ impl WorkingSet { impl StateReaderAndWriter for WorkingSet { fn get(&mut self, key: &StorageKey) -> Option { - self.delta.get(key) + match &mut self.archival_working_set { + None => self.delta.get(key), + Some(ref mut archival_working_set) => archival_working_set.get(key), + } } fn set(&mut self, key: &StorageKey, value: StorageValue) { - self.delta.set(key, value) + match &mut self.archival_working_set { + None => self.delta.set(key, value), + Some(ref mut archival_working_set) => archival_working_set.set(key, value), + } } fn delete(&mut self, key: &StorageKey) { - self.delta.delete(key) + match &mut self.archival_working_set { + None => self.delta.delete(key), + Some(ref mut archival_working_set) => archival_working_set.delete(key), + } } } @@ -447,16 +502,96 @@ impl<'a, C: Context> StateReaderAndWriter for AccessoryWorkingSet<'a, C> { if !cfg!(feature = "native") { None } else { - self.ws.accessory_delta.get(key) + match &mut self.ws.archival_accessory_working_set { + None => self.ws.accessory_delta.get(key), + Some(ref mut archival_working_set) => archival_working_set.get(key), + } } } fn set(&mut self, key: &StorageKey, value: StorageValue) { - self.ws.accessory_delta.set(key, value) + match &mut self.ws.archival_accessory_working_set { + None => self.ws.accessory_delta.set(key, value), + Some(ref mut archival_working_set) => archival_working_set.set(key, value), + } } fn delete(&mut self, key: &StorageKey) { - self.ws.accessory_delta.delete(key) + match &mut self.ws.archival_accessory_working_set { + None => self.ws.accessory_delta.delete(key), + Some(ref mut archival_working_set) => archival_working_set.delete(key), + } + } +} + +/// Module for archival state +pub mod archival_state { + use super::*; + + /// Archival JMT + pub struct ArchivalJmtWorkingSet { + delta: RevertableWriter>, + } + + impl ArchivalJmtWorkingSet { + /// create a new instance of ArchivalJmtWorkingSet + pub fn new(inner: &::Storage, version: Version) -> Self { + Self { + delta: RevertableWriter::new( + Delta::new(inner.clone(), Some(version)), + Some(version), + ), + } + } + } + + /// Archival Accessory + pub struct ArchivalAccessoryWorkingSet { + delta: RevertableWriter>, + } + + impl ArchivalAccessoryWorkingSet { + /// create a new instance of ArchivalAccessoryWorkingSet + pub fn new(inner: &::Storage, version: Version) -> Self { + Self { + delta: RevertableWriter::new( + AccessoryDelta::new(inner.clone(), Some(version)), + Some(version), + ), + } + } + } + + impl StateReaderAndWriter for ArchivalJmtWorkingSet { + fn get(&mut self, key: &StorageKey) -> Option { + self.delta.get(key) + } + + fn set(&mut self, key: &StorageKey, value: StorageValue) { + self.delta.set(key, value) + } + + fn delete(&mut self, key: &StorageKey) { + self.delta.delete(key) + } + } + + impl StateReaderAndWriter for ArchivalAccessoryWorkingSet { + fn get(&mut self, key: &StorageKey) -> Option { + if !cfg!(feature = "native") { + None + } else { + self.delta.get(key) + } + } + + fn set(&mut self, key: &StorageKey, value: StorageValue) { + self.delta.set(key, value) + } + + fn delete(&mut self, key: &StorageKey) { + self.delta.delete(key) + } } } @@ -545,6 +680,7 @@ pub mod kernel_state { struct RevertableWriter { inner: T, writes: HashMap>, + version: Option, } impl fmt::Debug for RevertableWriter { @@ -559,10 +695,11 @@ impl RevertableWriter where T: StateReaderAndWriter, { - fn new(inner: T) -> Self { + fn new(inner: T, version: Option) -> Self { Self { inner, writes: Default::default(), + version, } } @@ -585,7 +722,7 @@ where impl StateReaderAndWriter for RevertableWriter { fn get(&mut self, key: &StorageKey) -> Option { - if let Some(value) = self.writes.get(&key.to_cache_key()) { + if let Some(value) = self.writes.get(&key.to_cache_key_version(self.version)) { value.as_ref().cloned().map(Into::into) } else { self.inner.get(key) @@ -593,11 +730,14 @@ impl StateReaderAndWriter for RevertableWriter { } fn set(&mut self, key: &StorageKey, value: StorageValue) { - self.writes - .insert(key.to_cache_key(), Some(value.into_cache_value())); + self.writes.insert( + key.to_cache_key_version(self.version), + Some(value.into_cache_value()), + ); } fn delete(&mut self, key: &StorageKey) { - self.writes.insert(key.to_cache_key(), None); + self.writes + .insert(key.to_cache_key_version(self.version), None); } } diff --git a/module-system/sov-state/src/prover_storage.rs b/module-system/sov-state/src/prover_storage.rs index 6259d3e0f..63cd7f731 100644 --- a/module-system/sov-state/src/prover_storage.rs +++ b/module-system/sov-state/src/prover_storage.rs @@ -54,10 +54,11 @@ impl ProverStorage { } } - fn read_value(&self, key: &StorageKey) -> Option { + fn read_value(&self, key: &StorageKey, version: Option) -> Option { + let version_to_use = version.unwrap_or_else(|| self.db.get_next_version()); match self .db - .get_value_option_by_key(self.db.get_next_version(), key.as_ref()) + .get_value_option_by_key(version_to_use, key.as_ref()) { Ok(value) => value.map(Into::into), // It is ok to panic here, we assume the db is available and consistent. @@ -69,7 +70,6 @@ impl ProverStorage { pub struct ProverStateUpdate { pub(crate) node_batch: NodeBatch, pub key_preimages: Vec<(KeyHash, CacheKey)>, - // pub accessory_update: OrderedReadsAndWrites, } impl Storage for ProverStorage { @@ -83,16 +83,22 @@ impl Storage for ProverStorage { Self::with_path(config.path.as_path()) } - fn get(&self, key: &StorageKey, witness: &Self::Witness) -> Option { - let val = self.read_value(key); + fn get( + &self, + key: &StorageKey, + version: Option, + witness: &Self::Witness, + ) -> Option { + let val = self.read_value(key, version); witness.add_hint(val.clone()); val } #[cfg(feature = "native")] - fn get_accessory(&self, key: &StorageKey) -> Option { + fn get_accessory(&self, key: &StorageKey, version: Option) -> Option { + let version_to_use = version.unwrap_or_else(|| self.db.get_next_version() - 1); self.native_db - .get_value_option(key.as_ref()) + .get_value_option(key.as_ref(), version_to_use) .unwrap() .map(Into::into) } @@ -167,6 +173,7 @@ impl Storage for ProverStorage { } fn commit(&self, state_update: &Self::StateUpdate, accessory_writes: &OrderedReadsAndWrites) { + let latest_version = self.db.get_next_version() - 1; self.db .put_preimages( state_update @@ -182,6 +189,7 @@ impl Storage for ProverStorage { .ordered_writes .iter() .map(|(k, v_opt)| (k.key.to_vec(), v_opt.as_ref().map(|v| v.value.to_vec()))), + latest_version, ) .expect("native db write must succeed"); diff --git a/module-system/sov-state/src/zk_storage.rs b/module-system/sov-state/src/zk_storage.rs index 01d6876b2..5e2c2e457 100644 --- a/module-system/sov-state/src/zk_storage.rs +++ b/module-system/sov-state/src/zk_storage.rs @@ -47,7 +47,12 @@ impl Storage for ZkStorage { Ok(Self::new()) } - fn get(&self, _key: &StorageKey, witness: &Self::Witness) -> Option { + fn get( + &self, + _key: &StorageKey, + _version: Option, + witness: &Self::Witness, + ) -> Option { witness.get_hint() }