aigc_core/core/pmmr/
vec_backend.rs

1// Copyright 2021 The Aigc Developers
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::collections::HashSet;
16use std::convert::TryFrom;
17
18use croaring::Bitmap;
19
20use crate::core::hash::Hash;
21use crate::core::pmmr::{self, Backend};
22use crate::core::BlockHeader;
23use crate::ser::PMMRable;
24
25/// Simple/minimal/naive MMR backend implementation backed by Vec<T> and Vec<Hash>.
26/// Removed pos are maintained in a HashSet<u64>.
27#[derive(Clone, Debug)]
28pub struct VecBackend<T: PMMRable> {
29	/// Backend elements (optional, possible to just store hashes).
30	pub data: Option<Vec<T>>,
31	/// Vec of hashes for the PMMR (both leaves and parents).
32	pub hashes: Vec<Hash>,
33	/// Positions of removed elements (is this applicable if we do not store data?)
34	pub removed: HashSet<u64>,
35}
36
37impl<T: PMMRable> Backend<T> for VecBackend<T> {
38	fn append(&mut self, elmt: &T, hashes: &[Hash]) -> Result<(), String> {
39		if let Some(data) = &mut self.data {
40			data.push(elmt.clone());
41		}
42		self.hashes.extend_from_slice(hashes);
43		Ok(())
44	}
45
46	fn append_pruned_subtree(&mut self, _hash: Hash, _pos0: u64) -> Result<(), String> {
47		unimplemented!()
48	}
49
50	fn get_hash(&self, pos0: u64) -> Option<Hash> {
51		if self.removed.contains(&pos0) {
52			None
53		} else {
54			self.get_from_file(pos0)
55		}
56	}
57
58	fn get_data(&self, pos0: u64) -> Option<T::E> {
59		if self.removed.contains(&pos0) {
60			None
61		} else {
62			self.get_data_from_file(pos0)
63		}
64	}
65
66	fn get_from_file(&self, pos0: u64) -> Option<Hash> {
67		let idx = usize::try_from(pos0).expect("usize from u64");
68		self.hashes.get(idx).cloned()
69	}
70
71	fn get_peak_from_file(&self, pos0: u64) -> Option<Hash> {
72		self.get_from_file(pos0)
73	}
74
75	fn get_data_from_file(&self, pos0: u64) -> Option<T::E> {
76		if let Some(data) = &self.data {
77			let idx = usize::try_from(pmmr::n_leaves(1 + pos0) - 1).expect("usize from u64");
78			data.get(idx).map(|x| x.as_elmt())
79		} else {
80			None
81		}
82	}
83
84	/// Number of leaves in the MMR
85	fn n_unpruned_leaves(&self) -> u64 {
86		unimplemented!()
87	}
88
89	fn leaf_pos_iter(&self) -> Box<dyn Iterator<Item = u64> + '_> {
90		Box::new(
91			self.hashes
92				.iter()
93				.enumerate()
94				.map(|(x, _)| x as u64)
95				.filter(move |x| pmmr::is_leaf(*x) && !self.removed.contains(x)),
96		)
97	}
98
99	/// NOTE this function is needlessly inefficient with repeated calls to n_leaves()
100	fn leaf_idx_iter(&self, from_idx: u64) -> Box<dyn Iterator<Item = u64> + '_> {
101		let from_pos = pmmr::insertion_to_pmmr_index(from_idx);
102		Box::new(
103			self.leaf_pos_iter()
104				.skip_while(move |x| *x < from_pos)
105				.map(|x| pmmr::n_leaves(x + 1) - 1),
106		)
107	}
108
109	fn remove(&mut self, pos0: u64) -> Result<(), String> {
110		self.removed.insert(pos0);
111		Ok(())
112	}
113
114	fn rewind(&mut self, position: u64, _rewind_rm_pos: &Bitmap) -> Result<(), String> {
115		if let Some(data) = &mut self.data {
116			let idx = pmmr::n_leaves(position);
117			data.truncate(usize::try_from(idx).expect("usize from u64"));
118		}
119		self.hashes
120			.truncate(usize::try_from(position).expect("usize from u64"));
121		Ok(())
122	}
123
124	fn snapshot(&self, _header: &BlockHeader) -> Result<(), String> {
125		Ok(())
126	}
127
128	fn release_files(&mut self) {}
129
130	fn dump_stats(&self) {}
131}
132
133impl<T: PMMRable> VecBackend<T> {
134	/// Instantiates a new empty vec backend.
135	pub fn new() -> VecBackend<T> {
136		VecBackend {
137			data: Some(vec![]),
138			hashes: vec![],
139			removed: HashSet::new(),
140		}
141	}
142
143	/// Instantiate a new empty "hash only" vec backend.
144	pub fn new_hash_only() -> VecBackend<T> {
145		VecBackend {
146			data: None,
147			hashes: vec![],
148			removed: HashSet::new(),
149		}
150	}
151
152	/// Size of this vec backend in hashes.
153	pub fn size(&self) -> u64 {
154		self.hashes.len() as u64
155	}
156}