Revertible in-memory store

basecoin/app/src/builder.rs

@@ -5,7 +5,7 @@ use basecoin_modules::context::Module;
 use basecoin_modules::error::Error;
 use basecoin_modules::types::{IdentifiedModule, ModuleList, ModuleStore};
 use basecoin_store::context::ProvableStore;
-use basecoin_store::impls::{RevertibleStore, SharedStore};
+use basecoin_store::impls::SharedStore;
 use basecoin_store::types::{Identifier, MainStore};
 use basecoin_store::utils::{SharedRw, SharedRwExt};
 use cosmrs::AccountId;
@@ -22,7 +22,7 @@ impl<S: Default + ProvableStore> Builder<S> {
     /// Constructor.
     pub fn new(store: S) -> Self {
         Self {
-            store: SharedStore::new(RevertibleStore::new(store)),
+            store: SharedStore::new(store),
             modules: Arc::new(RwLock::new(vec![])),
         }
     }
@@ -35,7 +35,7 @@ impl<S: Default + ProvableStore> Builder<S> {
             .iter()
             .find(|m| &m.id == prefix)
             .map(|IdentifiedModule { module, .. }| module.store().share())
-            .unwrap_or_else(|| SharedStore::new(ModuleStore::new(S::default())))
+            .unwrap_or_else(|| SharedStore::new(S::default()))
     }
 
     #[inline]

basecoin/modules/src/types.rs

@@ -1,11 +1,10 @@
-use basecoin_store::impls::RevertibleStore;
 use basecoin_store::types::Identifier;
 use tendermint::merkle::proof::ProofOp;
 
 use crate::context::Module;
 
 pub type ModuleList<S> = Vec<IdentifiedModule<S>>;
-pub type ModuleStore<S> = RevertibleStore<S>;
+pub type ModuleStore<S> = S;
 
 pub struct IdentifiedModule<S> {
     pub id: Identifier,

basecoin/src/helper.rs

@@ -6,11 +6,10 @@ use basecoin_modules::bank::Bank;
 use basecoin_modules::context::{prefix, Identifiable};
 use basecoin_modules::ibc::Ibc;
 use basecoin_store::context::ProvableStore;
-use basecoin_store::impls::RevertibleStore;
 use basecoin_store::utils::SharedRwExt;
 
 /// Gives access to the IBC module.
-pub fn ibc<S>(app: BaseCoinApp<S>) -> Ibc<RevertibleStore<S>>
+pub fn ibc<S>(app: BaseCoinApp<S>) -> Ibc<S>
 where
     S: ProvableStore + Default + Debug,
 {
@@ -19,23 +18,12 @@ where
     modules
         .iter()
         .find(|m| m.id == prefix::Ibc {}.identifier())
-        .and_then(|m| {
-            m.module
-                .as_any()
-                .downcast_ref::<Ibc<RevertibleStore<S>>>()
-                .cloned()
-        })
+        .and_then(|m| m.module.as_any().downcast_ref::<Ibc<S>>().cloned())
         .expect("IBC module not found")
 }
 
 /// Gives access to the Bank module.
-pub fn bank<S>(
-    app: BaseCoinApp<S>,
-) -> Bank<
-    RevertibleStore<S>,
-    AuthAccountReader<RevertibleStore<S>>,
-    AuthAccountKeeper<RevertibleStore<S>>,
->
+pub fn bank<S>(app: BaseCoinApp<S>) -> Bank<S, AuthAccountReader<S>, AuthAccountKeeper<S>>
 where
     S: ProvableStore + Default + Debug,
 {
@@ -47,11 +35,7 @@ where
         .and_then(|m| {
             m.module
                 .as_any()
-                .downcast_ref::<Bank<
-                    RevertibleStore<S>,
-                    AuthAccountReader<RevertibleStore<S>>,
-                    AuthAccountKeeper<RevertibleStore<S>>,
-                >>()
+                .downcast_ref::<Bank<S, AuthAccountReader<S>, AuthAccountKeeper<S>>>()
                 .cloned()
         })
         .expect("Bank module not found")

basecoin/store/src/impls/in_memory.rs

@@ -1,18 +1,92 @@
+use core::convert::Infallible;
+
 use ics23::CommitmentProof;
 use tendermint::hash::Algorithm;
 use tendermint::Hash;
 use tracing::trace;
 
 use crate::avl::{AsBytes, AvlTree};
 use crate::context::{ProvableStore, Store};
-use crate::types::{Height, Path, State};
+use crate::types::{Height, Path, RawHeight, State};
+
+/// A wrapper type around [`Vec`] that more easily facilitates the pruning of
+/// its elements at a particular height / index. Keeps track of the latest
+/// height at which its elements were pruned.
+///
+/// This type is used by [`InMemoryStore`] in order to prune old store entries.
+#[derive(Debug, Clone, Default)]
+pub struct PrunedVec<T> {
+    vec: Vec<T>,
+    /// The latest index at which elements were pruned. In other words,
+    /// elements that exist at and before this index are no longer accessible.
+    pruned: usize,
+}
+
+impl<T> PrunedVec<T> {
+    pub fn push(&mut self, value: T) {
+        self.vec.push(value);
+    }
+
+    pub fn get(&self, index: usize) -> Option<&T> {
+        self.vec.get(index.checked_sub(self.pruned)?)
+    }
+
+    pub fn last(&self) -> Option<&T> {
+        self.vec.last()
+    }
+
+    /// Returns the number of elements currently in the `PrunedVec`,
+    /// i.e., the total number of elements minus the pruned elements.
+    pub fn current_length(&self) -> usize {
+        self.vec.len()
+    }
+
+    /// Returns the number of elements that have been pruned over the
+    /// lifetime of the instance of this type.
+    pub fn pruned_length(&self) -> usize {
+        self.pruned
+    }
+
+    /// Returns the total number of elements that have been added to
+    /// the `PrunedVec` over the lifetime of the instance of this type.
+    /// This includes the number of pruned elements in its count.
+    pub fn original_length(&self) -> usize {
+        self.current_length() + self.pruned_length()
+    }
+
+    /// Removes all elements from the `PrunedVec` up to the specified
+    /// index, inclusive. Note that `index` needs to be strictly greater
+    /// than the current `self.pruned` index, otherwise this method is
+    /// a no-op.
+    pub fn prune(&mut self, index: usize) {
+        trace!("pruning at index = {}", index);
+        if index > self.pruned {
+            self.vec.drain(0..index - self.pruned);
+            self.pruned = index;
+        }
+    }
+}
 
 /// An in-memory store backed by an AvlTree.
+///
+/// [`InMemoryStore`] has two copies of the current working store - `staged` and `pending`.
+///
+/// Each transaction works on the `pending` copy. When a transaction returns:
+/// - If it succeeded, the store _applies_ the transaction changes by copying `pending` to `staged`.
+/// - If it failed, the store _reverts_ the transaction changes by copying `staged` to `pending`.
+///
+/// When a block is committed, the staged copy is copied into the committed store.
+///
+/// Note that this store implementation is not production-friendly. After each transaction,
+/// the entire store is copied from `pending` to `staged`, or from `staged` to `pending`.
 #[derive(Clone, Debug)]
 pub struct InMemoryStore {
-    /// collection of states corresponding to every committed block height
-    store: Vec<State>,
-    /// pending block state
+    /// A collection of states corresponding to every committed block height.
+    store: PrunedVec<State>,
+    /// The changes made as a result of successful transactions that are staged
+    /// and waiting to be committed.
+    staged: State,
+    /// The dirty changes resulting from transactions that have not yet completed.
     pending: State,
 }
 
@@ -23,11 +97,11 @@ impl InMemoryStore {
             Height::Pending => Some(&self.pending),
             Height::Latest => self.store.last(),
             Height::Stable(height) => {
-                let h = height as usize;
-                if h <= self.store.len() {
-                    self.store.get(h - 1)
-                } else {
+                if height == 0 {
                     None
+                } else {
+                    let h = height as usize;
+                    self.store.get(h - 1)
                 }
             }
         }
@@ -37,15 +111,22 @@ impl InMemoryStore {
 impl Default for InMemoryStore {
     /// The store starts out with an empty state. We also initialize the pending location as empty.
     fn default() -> Self {
+        let genesis_state = AvlTree::new();
+
+        let store = PrunedVec::default();
+        let staged = genesis_state.clone();
+        let pending = genesis_state.clone();
+
         Self {
-            store: vec![],
-            pending: AvlTree::new(),
+            store,
+            staged,
+            pending,
         }
     }
 }
 
 impl Store for InMemoryStore {
-    type Error = (); // underlying store ops are infallible
+    type Error = Infallible;
 
     fn set(&mut self, path: Path, value: Vec<u8>) -> Result<Option<Vec<u8>>, Self::Error> {
         trace!("set at path = {}", path.to_string());
@@ -66,13 +147,31 @@ impl Store for InMemoryStore {
     }
 
     fn commit(&mut self) -> Result<Vec<u8>, Self::Error> {
-        trace!("committing height: {}", self.store.len());
-        self.store.push(self.pending.clone());
+        self.apply()?;
+        trace!("committing height: {}", self.current_height());
+        self.store.push(self.staged.clone());
         Ok(self.root_hash())
     }
 
+    fn apply(&mut self) -> Result<(), Self::Error> {
+        trace!("applying height: {}", self.current_height());
+        self.staged = self.pending.clone();
+        Ok(())
+    }
+
+    fn reset(&mut self) {
+        trace!("resetting height: {}", self.current_height());
+        self.pending = self.staged.clone();
+    }
+
+    fn prune(&mut self, height: RawHeight) -> Result<RawHeight, Self::Error> {
+        let h = height as usize;
+        self.store.prune(h);
+        Ok(height)
+    }
+
     fn current_height(&self) -> u64 {
-        self.store.len() as u64
+        self.store.original_length() as u64
     }
 
     fn get_keys(&self, key_prefix: &Path) -> Vec<Path> {
@@ -88,8 +187,8 @@ impl Store for InMemoryStore {
 
 impl ProvableStore for InMemoryStore {
     fn root_hash(&self) -> Vec<u8> {
-        self.pending
-            .root_hash()
+        self.get_state(Height::Latest)
+            .and_then(|s| s.root_hash())
             .unwrap_or(&Hash::from_bytes(Algorithm::Sha256, &[0u8; 32]).unwrap())
             .as_bytes()
             .to_vec()
@@ -106,3 +205,72 @@ impl ProvableStore for InMemoryStore {
 }
 
 // TODO(hu55a1n1): import tests
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_pruned_vec() {
+        let mut pv = PrunedVec::default();
+        pv.push(1);
+        pv.push(2);
+        pv.push(3);
+        pv.push(4);
+        pv.push(5);
+        assert_eq!(pv.original_length(), 5);
+        pv.prune(2);
+        assert_eq!(pv.original_length(), 5);
+        assert_eq!(pv.pruned_length(), 2);
+        assert_eq!(pv.current_length(), 3);
+        assert_eq!(pv.get(0), None);
+        assert_eq!(pv.get(1), None);
+        assert_eq!(pv.get(2), Some(&3));
+        assert_eq!(pv.get(3), Some(&4));
+        assert_eq!(pv.get(4), Some(&5));
+        assert_eq!(pv.get(5), None);
+        assert_eq!(pv.last(), Some(&5));
+    }
+
+    #[test]
+    fn test_in_memory_store() {
+        let mut store = InMemoryStore::default();
+        assert!(!store.root_hash().is_empty());
+        assert_eq!(store.current_height(), 0);
+
+        let path = Path::from("a".to_owned());
+        let value1 = vec![1, 2, 3];
+        let value2 = vec![4, 5, 6];
+
+        store.set(path.clone(), value1.clone()).unwrap();
+        assert_eq!(store.get(Height::Pending, &path), Some(value1.clone()));
+        assert_eq!(store.get(Height::Latest, &path), None);
+        assert_eq!(store.get(Height::Stable(1), &path), None);
+
+        store.apply().unwrap();
+        store.commit().unwrap();
+
+        assert_eq!(store.get(Height::Pending, &path), Some(value1.clone()));
+        assert_eq!(store.get(Height::Latest, &path), Some(value1.clone()));
+        assert_eq!(store.get(Height::Stable(1), &path), Some(value1.clone()));
+        assert_eq!(store.get(Height::Stable(2), &path), None);
+        assert_eq!(store.current_height(), 1);
+        assert!(!store.root_hash().is_empty());
+
+        store.set(path.clone(), value2.clone()).unwrap();
+        assert_eq!(store.get(Height::Pending, &path), Some(value2.clone()));
+        assert_eq!(store.get(Height::Latest, &path), Some(value1.clone()));
+        assert_eq!(store.get(Height::Stable(1), &path), Some(value1.clone()));
+
+        store.apply().unwrap();
+        store.commit().unwrap();
+
+        assert_eq!(store.get(Height::Pending, &path), Some(value2.clone()));
+        assert_eq!(store.get(Height::Latest, &path), Some(value2.clone()));
+        assert_eq!(store.get(Height::Stable(1), &path), Some(value1.clone()));
+        assert_eq!(store.get(Height::Stable(2), &path), Some(value2.clone()));
+        assert_eq!(store.get(Height::Stable(3), &path), None);
+        assert_eq!(store.current_height(), 2);
+        assert!(!store.root_hash().is_empty());
+    }
+}

basecoin/store/src/impls/revertible.rs

@@ -4,7 +4,18 @@ use tracing::trace;
 use crate::context::{ProvableStore, Store};
 use crate::types::{Height, Path};
 
-/// A wrapper store that implements rudimentary `apply()`/`reset()` support for other stores
+/// A wrapper store that implements rudimentary `apply()`/`reset()` support for other stores.
+///
+/// [`RevertibleStore`] relies on maintaining a list of processed operations - `delete` and `set`.
+///
+/// If it has to revert a failed transaction, it _reverts_ the previous operations as so:
+/// - If reverting an overwriting `set` or `delete` operation, it performs a `set` with the old value.
+/// - If reverting a non-overwriting `set`, it `delete`s the current value.
+///
+/// Note that this scheme makes it trickier to maintain deterministic Merkle root hashes:
+/// an overwriting `set` doesn't reorganize a Merkle tree - but non-overwriting `set` and `delete`
+/// operations may reorganize a Merkle tree - which may change the root hash. However, a Merkle
+/// store should have no effect on a failed transaction.
 #[derive(Clone, Debug)]
 pub struct RevertibleStore<S> {
     /// backing store

basecoin/store/src/types/height.rs

@@ -1,5 +1,3 @@
-use core::fmt::{Display, Formatter};
-
 /// Block height
 pub type RawHeight = u64;
 
@@ -11,16 +9,6 @@ pub enum Height {
     Stable(RawHeight), // or equivalently `tendermint::block::Height`
 }
 
-impl Display for Height {
-    fn fmt(&self, f: &mut Formatter<'_>) -> core::fmt::Result {
-        match self {
-            Height::Pending => write!(f, "pending"),
-            Height::Latest => write!(f, "latest"),
-            Height::Stable(height) => write!(f, "{}", height),
-        }
-    }
-}
-
 impl From<RawHeight> for Height {
     fn from(value: u64) -> Self {
         match value {