1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78
use crate::collections::BTreeMap; use crate::{vec::Vec, MMRStore, Merge, MerkleProof, Result, MMR}; use core::cell::RefCell; use core::fmt::Debug; use core::marker::PhantomData; #[derive(Clone)] pub struct MemStore<T>(RefCell<BTreeMap<u64, T>>); impl<T> Default for MemStore<T> { fn default() -> Self { Self::new() } } impl<T> MemStore<T> { fn new() -> Self { MemStore(RefCell::new(Default::default())) } } impl<T: Clone> MMRStore<T> for &MemStore<T> { fn get_elem(&self, pos: u64) -> Result<Option<T>> { Ok(self.0.borrow().get(&pos).cloned()) } fn append(&mut self, pos: u64, elems: Vec<T>) -> Result<()> { let mut store = self.0.borrow_mut(); for (i, elem) in elems.into_iter().enumerate() { store.insert(pos + i as u64, elem); } Ok(()) } } pub struct MemMMR<T, M> { store: MemStore<T>, mmr_size: u64, merge: PhantomData<M>, } impl<T: Clone + Debug + PartialEq, M: Merge<Item = T>> Default for MemMMR<T, M> { fn default() -> Self { Self::new(0, Default::default()) } } impl<T: Clone + Debug + PartialEq, M: Merge<Item = T>> MemMMR<T, M> { pub fn new(mmr_size: u64, store: MemStore<T>) -> Self { MemMMR { mmr_size, store, merge: PhantomData, } } pub fn store(&self) -> &MemStore<T> { &self.store } pub fn get_root(&self) -> Result<T> { let mmr = MMR::<T, M, &MemStore<T>>::new(self.mmr_size, &self.store); mmr.get_root() } pub fn push(&mut self, elem: T) -> Result<u64> { let mut mmr = MMR::<T, M, &MemStore<T>>::new(self.mmr_size, &self.store); let pos = mmr.push(elem)?; self.mmr_size = mmr.mmr_size(); mmr.commit()?; Ok(pos) } pub fn gen_proof(&self, pos_list: Vec<u64>) -> Result<MerkleProof<T, M>> { let mmr = MMR::<T, M, &MemStore<T>>::new(self.mmr_size, &self.store); mmr.gen_proof(pos_list) } }