ZkStorage implementation (Sovereign-Labs#51)

first-read-last-write-cache/src/cache.rs

@@ -1,6 +1,7 @@
 use crate::access::{Access, MergeError};
 use crate::{CacheKey, CacheValue};
 use std::collections::{hash_map::Entry, HashMap};
+use std::sync::Arc;
 use thiserror::Error;
 
 #[derive(Error, Debug, Eq, PartialEq)]
@@ -14,12 +15,12 @@ pub enum ReadError {
 
 /// Cache entry can be in three states:
 /// - Does not exists, a given key was never inserted in the cache:             
-///     ExistsInCache::No
+///     ValueExists::No
 /// - Exists but the value is empty.
-///      ExistsInCache::Yes(None)
+///      ValueExists::Yes(None)
 /// - Exists and contains a value:
-///     ExistsInCache::Yes(Some(value))
-pub enum ExistsInCache {
+///     ValueExists::Yes(Some(value))
+pub enum ValueExists {
     Yes(CacheValue),
     No,
 }
@@ -30,6 +31,28 @@ pub struct CacheLog {
     log: HashMap<CacheKey, Access>,
 }
 
+/// Represents all reads from a CacheLog.
+#[derive(Default, Clone, Debug)]
+pub struct FirstReads {
+    reads: Arc<HashMap<CacheKey, CacheValue>>,
+}
+
+impl FirstReads {
+    pub fn new(reads: HashMap<CacheKey, CacheValue>) -> Self {
+        Self {
+            reads: Arc::new(reads),
+        }
+    }
+
+    /// Returns a value corresponding to the key.
+    pub fn get(&self, key: &CacheKey) -> ValueExists {
+        match self.reads.get(key) {
+            Some(read) => ValueExists::Yes(read.clone()),
+            None => ValueExists::No,
+        }
+    }
+}
+
 impl CacheLog {
     pub fn with_capacity(capacity: usize) -> Self {
         Self {
@@ -39,11 +62,22 @@ impl CacheLog {
 }
 
 impl CacheLog {
-    /// Gets value form the cache.
-    pub fn get_value(&self, key: &CacheKey) -> ExistsInCache {
+    /// Returns all reads from the CacheLog.
+    pub fn get_first_reads(&self) -> FirstReads {
+        let reads = self
+            .log
+            .iter()
+            .filter_map(|(k, v)| filter_first_reads(k.clone(), v.clone()))
+            .collect::<HashMap<_, _>>();
+
+        FirstReads::new(reads)
+    }
+
+    /// Returns a value corresponding to the key.
+    pub fn get_value(&self, key: &CacheKey) -> ValueExists {
         match self.log.get(key) {
-            Some(value) => ExistsInCache::Yes(value.last_value().clone()),
-            None => ExistsInCache::No,
+            Some(value) => ValueExists::Yes(value.last_value().clone()),
+            None => ValueExists::No,
         }
     }
 
@@ -117,16 +151,24 @@ impl CacheLog {
     }
 }
 
+fn filter_first_reads(k: CacheKey, access: Access) -> Option<(CacheKey, CacheValue)> {
+    match access {
+        Access::Read(read) => Some((k, read)),
+        Access::ReadThenWrite { original, .. } => Some((k, original)),
+        Access::Write(_) => None,
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
     use crate::utils::test_util::{create_key, create_value};
 
-    impl ExistsInCache {
+    impl ValueExists {
         fn get(self) -> CacheValue {
             match self {
-                ExistsInCache::Yes(value) => value,
-                ExistsInCache::No => unreachable!(),
+                ValueExists::Yes(value) => value,
+                ValueExists::No => unreachable!(),
             }
         }
     }
@@ -362,4 +404,27 @@ mod tests {
             assert!(result.is_err());
         }
     }
+
+    #[test]
+    fn test_first_reads() {
+        let mut cache = CacheLog::default();
+        let entries = vec![
+            new_cache_entry(1, 11),
+            new_cache_entry(2, 22),
+            new_cache_entry(3, 33),
+        ];
+
+        for entry in entries.clone() {
+            cache.add_read(entry.key, entry.value).unwrap();
+        }
+
+        let first_reads = cache.get_first_reads();
+
+        for entry in entries {
+            match first_reads.get(&entry.key) {
+                ValueExists::Yes(value) => assert_eq!(entry.value, value),
+                ValueExists::No => unreachable!(),
+            }
+        }
+    }
 }

sov-modules/sov-modules-api/src/mocks.rs

@@ -1,5 +1,5 @@
 use crate::Context;
-use sov_state::JmtStorage;
+use sov_state::{JmtStorage, ZkStorage};
 
 /// Mock for Context::PublicKey, useful for testing.
 #[derive(borsh::BorshDeserialize, borsh::BorshSerialize, PartialEq, Eq)]
@@ -41,3 +41,19 @@ impl Context for MockContext {
         &self.sender
     }
 }
+
+pub struct ZkMockContext {
+    sender: MockPublicKey,
+}
+
+impl Context for ZkMockContext {
+    type Storage = ZkStorage;
+
+    type Signature = MockSignature;
+
+    type PublicKey = MockPublicKey;
+
+    fn sender(&self) -> &Self::PublicKey {
+        &self.sender
+    }
+}

sov-modules/sov-modules-impl/tests/tests.rs

@@ -1,8 +1,8 @@
-use sov_modules_api::mocks::MockContext;
+use sov_modules_api::mocks::{MockContext, ZkMockContext};
 use sov_modules_api::{Context, ModuleInfo, Prefix};
 use sov_modules_macros::ModuleInfo;
 use sov_state::storage::{StorageKey, StorageValue};
-use sov_state::{JmtStorage, StateMap, StateValue, Storage};
+use sov_state::{JmtStorage, StateMap, StateValue, Storage, ZkStorage};
 
 pub mod module_a {
     use super::*;
@@ -57,7 +57,7 @@ mod module_c {
     }
 
     impl<C: Context> ModuleC<C> {
-        pub fn update(&mut self, key: &str, value: &str) {
+        pub fn execute(&mut self, key: &str, value: &str) {
             self.mod_1_a.update(key, value);
             self.mod_1_b.update(key, value);
             self.mod_1_a.update(key, value);
@@ -67,34 +67,52 @@ mod module_c {
 
 #[test]
 fn nested_module_call_test() {
-    type C = MockContext;
-    let test_storage = JmtStorage::default();
+    let native_storage = JmtStorage::default();
+
+    // Test the `native` execution.
+    {
+        execute_module_logic::<MockContext>(native_storage.clone());
+        test_state_update::<MockContext>(native_storage.clone());
+    }
+
+    // Test the `zk` execution.
+    {
+        let zk_storage = ZkStorage::new(native_storage.get_first_reads());
+        execute_module_logic::<ZkMockContext>(zk_storage.clone());
+        test_state_update::<ZkMockContext>(zk_storage);
+    }
+}
+
+fn execute_module_logic<C: Context>(storage: C::Storage) {
+    let module = &mut module_c::ModuleC::<C>::new(storage);
+    module.execute("some_key", "some_value");
+}
 
-    let module = &mut <module_c::ModuleC<C> as ModuleInfo<C>>::new(test_storage.clone());
-    module.update("some_key", "some_value");
+fn test_state_update<C: Context>(storage: C::Storage) {
+    let module = <module_c::ModuleC<C> as ModuleInfo<C>>::new(storage.clone());
 
     let expected_value = StorageValue::new("some_value");
 
     {
         let prefix = Prefix::new("tests::module_a", "ModuleA", "state_1_a");
         let key = StorageKey::new(&prefix.into(), "some_key");
-        let value = test_storage.get(key).unwrap();
+        let value = storage.get(key).unwrap();
 
         assert_eq!(expected_value, value);
     }
 
     {
         let prefix = Prefix::new("tests::module_b", "ModuleB", "state_1_b");
         let key = StorageKey::new(&prefix.into(), "some_key");
-        let value = test_storage.get(key).unwrap();
+        let value = storage.get(key).unwrap();
 
         assert_eq!(expected_value, value);
     }
 
     {
         let prefix = Prefix::new("tests::module_a", "ModuleA", "state_1_a");
         let key = StorageKey::new(&prefix.into(), "key_from_b");
-        let value = test_storage.get(key).unwrap();
+        let value = storage.get(key).unwrap();
 
         assert_eq!(expected_value, value);
     }

sov-modules/sov-state/src/internal_cache.rs

@@ -0,0 +1,58 @@
+use crate::storage::{StorageKey, StorageValue};
+use first_read_last_write_cache::{
+    cache::{self, CacheLog},
+    CacheValue,
+};
+use std::{cell::RefCell, rc::Rc};
+
+/// `ValueReader` Reads a value from an external data source.
+pub trait ValueReader {
+    fn read_value(&self, key: StorageKey) -> Option<StorageValue>;
+}
+
+/// Caches reads and writes for a (key, value) pair. On the first read the value is fetched
+/// from an external source represented by the `ValueReader` trait. On following reads,
+/// the cache checks if the value we read was inserted before.
+#[derive(Default, Clone)]
+pub(crate) struct StorageInternalCache {
+    pub(crate) cache: Rc<RefCell<CacheLog>>,
+}
+
+impl StorageInternalCache {
+    /// Gets a value from the cache or reads it from the provided `ValueReader`.
+    pub(crate) fn get_or_fetch<VR: ValueReader>(
+        &self,
+        key: StorageKey,
+        value_reader: &VR,
+    ) -> Option<StorageValue> {
+        let cache_key = key.clone().as_cache_key();
+        let cache_value = self.cache.borrow().get_value(&cache_key);
+
+        match cache_value {
+            cache::ValueExists::Yes(cache_value_exists) => {
+                self.cache
+                    .borrow_mut()
+                    .add_read(cache_key, cache_value_exists.clone())
+                    // It is ok to panic here, we must guarantee that the cache is consistent.
+                    .unwrap_or_else(|e| panic!("Inconsistent read from the cache: {e:?}"));
+
+                cache_value_exists.value.map(|value| StorageValue { value })
+            }
+            // If the value does not exist in the cache, then fetch it from an external source.
+            cache::ValueExists::No => value_reader.read_value(key),
+        }
+    }
+
+    pub(crate) fn set(&mut self, key: StorageKey, value: StorageValue) {
+        let cache_key = key.as_cache_key();
+        let cache_value = value.as_cache_value();
+        self.cache.borrow_mut().add_write(cache_key, cache_value);
+    }
+
+    pub(crate) fn delete(&mut self, key: StorageKey) {
+        let cache_key = key.as_cache_key();
+        self.cache
+            .borrow_mut()
+            .add_write(cache_key, CacheValue::empty());
+    }
+}

sov-modules/sov-state/src/jmt_storage.rs

@@ -1,50 +1,36 @@
-use std::{cell::RefCell, rc::Rc};
-
-use crate::storage::{Storage, StorageKey, StorageValue};
-use first_read_last_write_cache::{
-    cache::{self, CacheLog},
-    CacheValue,
+use crate::{
+    internal_cache::{StorageInternalCache, ValueReader},
+    storage::{GenericStorage, StorageKey, StorageValue},
 };
+use first_read_last_write_cache::cache::{CacheLog, FirstReads};
 use jellyfish_merkle_generic::Version;
 
-// Storage backed by JMT.
 #[derive(Default, Clone)]
-pub struct JmtStorage {
-    // Caches first read and last write for a particular key.
-    cache: Rc<RefCell<CacheLog>>,
+pub struct JmtDb {
     _version: Version,
 }
 
-impl Storage for JmtStorage {
-    fn get(&self, key: StorageKey) -> Option<StorageValue> {
-        let cache_key = key.as_cache_key();
-        let cache_value = self.cache.borrow().get_value(&cache_key);
+impl ValueReader for JmtDb {
+    fn read_value(&self, _key: StorageKey) -> Option<StorageValue> {
+        todo!()
+    }
+}
 
-        match cache_value {
-            cache::ExistsInCache::Yes(cache_value_exists) => {
-                self.cache
-                    .borrow_mut()
-                    .add_read(cache_key, cache_value_exists.clone())
-                    // It is ok to panic here, we must guarantee that the cache is consistent.
-                    .unwrap_or_else(|e| panic!("Inconsistent read from the cache: {e:?}"));
+/// Storage backed by JmtDb.
+pub type JmtStorage = GenericStorage<JmtDb>;
 
-                cache_value_exists.value.map(|value| StorageValue { value })
-            }
-            // TODO If the value does not exist in the cache, then fetch it from the JMT.
-            cache::ExistsInCache::No => todo!(),
+impl JmtStorage {
+    /// Creates a new JmtStorage.
+    pub fn new(jmt: JmtDb) -> Self {
+        Self {
+            internal_cache: StorageInternalCache::default(),
+            value_reader: jmt,
         }
     }
 
-    fn set(&mut self, key: StorageKey, value: StorageValue) {
-        let cache_key = key.as_cache_key();
-        let cache_value = value.as_cache_value();
-        self.cache.borrow_mut().add_write(cache_key, cache_value);
-    }
-
-    fn delete(&mut self, key: StorageKey) {
-        let cache_key = key.as_cache_key();
-        self.cache
-            .borrow_mut()
-            .add_write(cache_key, CacheValue::empty());
+    /// Gets the first reads from the JmtStorage.
+    pub fn get_first_reads(&self) -> FirstReads {
+        let cache: &CacheLog = &self.internal_cache.cache.borrow();
+        cache.get_first_reads()
     }
 }

sov-modules/sov-state/src/lib.rs

@@ -1,13 +1,20 @@
 mod backend;
+mod internal_cache;
 mod jmt_storage;
 mod map;
 pub mod storage;
 mod utils;
 mod value;
+mod zk_storage;
+
+#[cfg(test)]
+mod storage_test;
+
 pub use jmt_storage::JmtStorage;
 pub use map::StateMap;
 pub use storage::Storage;
 use utils::AlignedVec;
+pub use zk_storage::ZkStorage;
 
 pub use value::StateValue;