basecoin_store/impls/
in_memory.rs

1use core::convert::Infallible;
2
3use ics23::CommitmentProof;
4use tendermint::hash::Algorithm;
5use tendermint::Hash;
6use tracing::trace;
7
8use crate::avl::{AsBytes, AvlTree};
9use crate::context::{ProvableStore, Store};
10use crate::types::{Height, Path, RawHeight, State};
11
12/// A wrapper type around [`Vec`] that more easily facilitates the pruning of
13/// its elements at a particular height / index. Keeps track of the latest
14/// height at which its elements were pruned.
15///
16/// This type is used by [`InMemoryStore`] in order to prune old store entries.
17#[derive(Debug, Clone, Default)]
18pub struct PrunedVec<T> {
19    vec: Vec<T>,
20    /// The latest index at which elements were pruned. In other words,
21    /// elements that exist at and before this index are no longer accessible.
22    pruned: usize,
23}
24
25impl<T> PrunedVec<T> {
26    pub fn push(&mut self, value: T) {
27        self.vec.push(value);
28    }
29
30    pub fn get(&self, index: usize) -> Option<&T> {
31        self.vec.get(index.checked_sub(self.pruned)?)
32    }
33
34    pub fn last(&self) -> Option<&T> {
35        self.vec.last()
36    }
37
38    /// Returns the number of elements currently in the `PrunedVec`,
39    /// i.e., the total number of elements minus the pruned elements.
40    pub fn current_length(&self) -> usize {
41        self.vec.len()
42    }
43
44    /// Returns the number of elements that have been pruned over the
45    /// lifetime of the instance of this type.
46    pub fn pruned_length(&self) -> usize {
47        self.pruned
48    }
49
50    /// Returns the total number of elements that have been added to
51    /// the `PrunedVec` over the lifetime of the instance of this type.
52    /// This includes the number of pruned elements in its count.
53    pub fn original_length(&self) -> usize {
54        self.current_length() + self.pruned_length()
55    }
56
57    /// Removes all elements from the `PrunedVec` up to the specified
58    /// index, inclusive. Note that `index` needs to be strictly greater
59    /// than the current `self.pruned` index, otherwise this method is
60    /// a no-op.
61    pub fn prune(&mut self, index: usize) {
62        trace!("pruning at index = {}", index);
63        if index > self.pruned {
64            self.vec.drain(0..index - self.pruned);
65            self.pruned = index;
66        }
67    }
68}
69
70/// An in-memory store backed by an AvlTree.
71///
72/// [`InMemoryStore`] has two copies of the current working store - `staged` and `pending`.
73///
74/// Each transaction works on the `pending` copy. When a transaction returns:
75/// - If it succeeded, the store _applies_ the transaction changes by copying `pending` to `staged`.
76/// - If it failed, the store _reverts_ the transaction changes by copying `staged` to `pending`.
77///
78/// When a block is committed, the staged copy is copied into the committed store.
79///
80/// Note that this store implementation is not production-friendly. After each transaction,
81/// the entire store is copied from `pending` to `staged`, or from `staged` to `pending`.
82#[derive(Clone, Debug)]
83pub struct InMemoryStore {
84    /// A collection of states corresponding to every committed block height.
85    store: PrunedVec<State>,
86    /// The changes made as a result of successful transactions that are staged
87    /// and waiting to be committed.
88    staged: State,
89    /// The dirty changes resulting from transactions that have not yet completed.
90    pending: State,
91}
92
93impl InMemoryStore {
94    #[inline]
95    fn get_state(&self, height: Height) -> Option<&State> {
96        match height {
97            Height::Pending => Some(&self.pending),
98            Height::Latest => self.store.last(),
99            Height::Stable(height) => {
100                if height == 0 {
101                    None
102                } else {
103                    let h = height as usize;
104                    self.store.get(h - 1)
105                }
106            }
107        }
108    }
109}
110
111impl Default for InMemoryStore {
112    /// The store starts out with an empty state. We also initialize the pending location as empty.
113    fn default() -> Self {
114        let genesis_state = AvlTree::new();
115
116        let store = PrunedVec::default();
117        let staged = genesis_state.clone();
118        let pending = genesis_state.clone();
119
120        Self {
121            store,
122            staged,
123            pending,
124        }
125    }
126}
127
128impl Store for InMemoryStore {
129    type Error = Infallible;
130
131    fn set(&mut self, path: Path, value: Vec<u8>) -> Result<Option<Vec<u8>>, Self::Error> {
132        trace!("set at path = {}", path.to_string());
133        Ok(self.pending.insert(path, value))
134    }
135
136    fn get(&self, height: Height, path: &Path) -> Option<Vec<u8>> {
137        trace!(
138            "get at path = {} at height = {:?}",
139            path.to_string(),
140            height
141        );
142        self.get_state(height).and_then(|v| v.get(path).cloned())
143    }
144
145    fn delete(&mut self, path: &Path) {
146        self.pending.remove(path.clone());
147    }
148
149    fn commit(&mut self) -> Result<Vec<u8>, Self::Error> {
150        self.apply()?;
151        trace!("committing height: {}", self.current_height());
152        self.store.push(self.staged.clone());
153        Ok(self.root_hash())
154    }
155
156    fn apply(&mut self) -> Result<(), Self::Error> {
157        trace!("applying height: {}", self.current_height());
158        self.staged = self.pending.clone();
159        Ok(())
160    }
161
162    fn reset(&mut self) {
163        trace!("resetting height: {}", self.current_height());
164        self.pending = self.staged.clone();
165    }
166
167    fn prune(&mut self, height: RawHeight) -> Result<RawHeight, Self::Error> {
168        let h = height as usize;
169        self.store.prune(h);
170        Ok(height)
171    }
172
173    fn current_height(&self) -> u64 {
174        self.store.original_length() as u64
175    }
176
177    fn get_keys(&self, key_prefix: &Path) -> Vec<Path> {
178        let key_prefix = key_prefix.as_bytes();
179        self.pending
180            .get_keys()
181            .into_iter()
182            .filter(|&key| key.as_bytes().as_ref().starts_with(key_prefix.as_ref()))
183            .cloned()
184            .collect()
185    }
186}
187
188impl ProvableStore for InMemoryStore {
189    fn root_hash(&self) -> Vec<u8> {
190        self.get_state(Height::Latest)
191            .and_then(|s| s.root_hash())
192            .unwrap_or(&Hash::from_bytes(Algorithm::Sha256, &[0u8; 32]).unwrap())
193            .as_bytes()
194            .to_vec()
195    }
196
197    fn get_proof(&self, height: Height, key: &Path) -> Option<CommitmentProof> {
198        trace!(
199            "get proof at path = {} at height = {:?}",
200            key.to_string(),
201            height
202        );
203        self.get_state(height).map(|v| v.get_proof(key))
204    }
205}
206
207// TODO(hu55a1n1): import tests
208
209#[cfg(test)]
210mod tests {
211    use super::*;
212
213    #[test]
214    fn test_pruned_vec() {
215        let mut pv = PrunedVec::default();
216        pv.push(1);
217        pv.push(2);
218        pv.push(3);
219        pv.push(4);
220        pv.push(5);
221        assert_eq!(pv.original_length(), 5);
222        pv.prune(2);
223        assert_eq!(pv.original_length(), 5);
224        assert_eq!(pv.pruned_length(), 2);
225        assert_eq!(pv.current_length(), 3);
226        assert_eq!(pv.get(0), None);
227        assert_eq!(pv.get(1), None);
228        assert_eq!(pv.get(2), Some(&3));
229        assert_eq!(pv.get(3), Some(&4));
230        assert_eq!(pv.get(4), Some(&5));
231        assert_eq!(pv.get(5), None);
232        assert_eq!(pv.last(), Some(&5));
233    }
234
235    #[test]
236    fn test_in_memory_store() {
237        let mut store = InMemoryStore::default();
238        assert!(!store.root_hash().is_empty());
239        assert_eq!(store.current_height(), 0);
240
241        let path = Path::from("a".to_owned());
242        let value1 = vec![1, 2, 3];
243        let value2 = vec![4, 5, 6];
244
245        store.set(path.clone(), value1.clone()).unwrap();
246        assert_eq!(store.get(Height::Pending, &path), Some(value1.clone()));
247        assert_eq!(store.get(Height::Latest, &path), None);
248        assert_eq!(store.get(Height::Stable(1), &path), None);
249
250        store.apply().unwrap();
251        store.commit().unwrap();
252
253        assert_eq!(store.get(Height::Pending, &path), Some(value1.clone()));
254        assert_eq!(store.get(Height::Latest, &path), Some(value1.clone()));
255        assert_eq!(store.get(Height::Stable(1), &path), Some(value1.clone()));
256        assert_eq!(store.get(Height::Stable(2), &path), None);
257        assert_eq!(store.current_height(), 1);
258        assert!(!store.root_hash().is_empty());
259
260        store.set(path.clone(), value2.clone()).unwrap();
261        assert_eq!(store.get(Height::Pending, &path), Some(value2.clone()));
262        assert_eq!(store.get(Height::Latest, &path), Some(value1.clone()));
263        assert_eq!(store.get(Height::Stable(1), &path), Some(value1.clone()));
264
265        store.apply().unwrap();
266        store.commit().unwrap();
267
268        assert_eq!(store.get(Height::Pending, &path), Some(value2.clone()));
269        assert_eq!(store.get(Height::Latest, &path), Some(value2.clone()));
270        assert_eq!(store.get(Height::Stable(1), &path), Some(value1.clone()));
271        assert_eq!(store.get(Height::Stable(2), &path), Some(value2.clone()));
272        assert_eq!(store.get(Height::Stable(3), &path), None);
273        assert_eq!(store.current_height(), 2);
274        assert!(!store.root_hash().is_empty());
275    }
276}
basecoin_store/impls/in_memory.rs

basecoin_store/impls/
in_memory.rs
