merkle_heapless/
compactable.rs

1//! ```rust
2//! use merkle_heapless::compactable::{DefaultCompactable};
3//!
4//! const ARITY: usize = 4;
5//! const HEIGHT: usize = 3;
6//!
7//! let mut cmt = DefaultCompactable::<ARITY, HEIGHT, StdHash>::try_from::<&[u8]>(&[
8//!     "apple", "apricot", "banana", "cherry",
9//! ]).unwrap();
10//!
11//! cmt.try_remove(0).unwrap();
12//! cmt.compact();
13//! // will try to create a smaller tree from the compacted tree
14//! let mut reduced = cmt.try_reduce().unwrap();
15//! ```
16//!
17
18/// Compactable tree with [Proof] of tree's height
19pub type DefaultCompactable<
20    const ARITY: usize,
21    const HEIGHT: usize,
22    H,
23    const MAX_INPUT_LEN: usize,
24> = CompactableHeaplessTree<
25    ARITY,
26    HEIGHT,
27    H,
28    MAX_INPUT_LEN,
29    Proof<ARITY, HEIGHT, H, MAX_INPUT_LEN>,
30>;
31
32use crate::traits::CanRemove;
33use crate::{
34    is_pow2, layer_size, location_in_prefixed, max_leaves, num_of_prefixed, Assert, Error, HashT,
35    IsTrue, Prefixed, Proof, ProofBuilder, StaticTree, StaticTreeTrait,
36};
37use core::fmt::Debug;
38use core::ops::Deref;
39/// Tree that can be compacted after leaf removal and reduced to a smaller tree
40pub struct CompactableHeaplessTree<
41    const ARITY: usize,
42    const HEIGHT: usize,
43    H,
44    const MAX_INPUT_LEN: usize,
45    PB = Proof<ARITY, HEIGHT, H, MAX_INPUT_LEN>,
46> where
47    [(); num_of_prefixed!(ARITY, HEIGHT)]: Sized,
48    [(); max_leaves!(ARITY, HEIGHT)]: Sized,
49    Assert<{ is_pow2!(ARITY) }>: IsTrue,
50    H: HashT,
51    PB: ProofBuilder<ARITY, H>,
52{
53    tree: StaticTree<ARITY, HEIGHT, H, MAX_INPUT_LEN, PB>,
54    num_of_leaves: usize,
55    leaves_present: [bool; max_leaves!(ARITY, HEIGHT)],
56}
57
58impl<const ARITY: usize, const HEIGHT: usize, H, const MAX_INPUT_LEN: usize, PB>
59    CompactableHeaplessTree<ARITY, HEIGHT, H, MAX_INPUT_LEN, PB>
60where
61    [(); num_of_prefixed!(ARITY, HEIGHT)]: Sized,
62    [(); max_leaves!(ARITY, HEIGHT)]: Sized,
63    Assert<{ is_pow2!(ARITY) }>: IsTrue,
64    H: HashT,
65    PB: ProofBuilder<ARITY, H>,
66{
67    /// creates a tree from an input if possible
68    pub fn try_from<T: AsRef<[u8]> + Deref<Target = [u8]>>(input: &[T]) -> Result<Self, Error> {
69        let mut this = Self {
70            tree: StaticTree::try_from(input)?,
71            num_of_leaves: input.len(),
72            leaves_present: [false; max_leaves!(ARITY, HEIGHT)],
73        };
74        for i in 0..input.len() {
75            this.leaves_present[i] = true;
76        }
77        Ok(this)
78    }
79
80    /// creates a tree from an input if possible
81    pub fn from<T: AsRef<[u8]> + Deref<Target = [u8]>>(input: &[T]) -> Self {
82        let mut this = Self {
83            tree: StaticTree::from(input),
84            num_of_leaves: input.len(),
85            leaves_present: [false; max_leaves!(ARITY, HEIGHT)],
86        };
87        for i in 0..input.len() {
88            this.leaves_present[i] = true;
89        }
90        this
91    }
92
93    /// creates a tree from hashed leaves (of another tree)
94    pub fn try_from_leaves(prefixed: &[Prefixed<ARITY, H>]) -> Result<Self, Error> {
95        let mut leaves_present = [false; max_leaves!(ARITY, HEIGHT)];
96        let mut num_of_leaves = 0;
97        let default_hash = Prefixed::<ARITY, H>::default_hash();
98        let mut j = 0;
99        for leaf in prefixed {
100            num_of_leaves += leaf
101                .hashes
102                .iter()
103                .enumerate()
104                .filter(|&(_, h)| (h != &default_hash))
105                .map(|(i, _)| {
106                    leaves_present[i + j] = true;
107                })
108                .count();
109            j += ARITY;
110        }
111        Ok(Self {
112            tree: StaticTree::try_from_leaves(prefixed)?,
113            num_of_leaves,
114            leaves_present,
115        })
116    }
117
118    fn try_from_compacted(
119        leaves: &[Prefixed<ARITY, H>; layer_size!(ARITY, HEIGHT, 0)],
120        num_of_leaves: usize,
121    ) -> Result<Self, Error> {
122        (num_of_leaves <= max_leaves!(ARITY, HEIGHT))
123            .then_some(())
124            .ok_or(Error::Create)
125            .and_then(|()| {
126                let mut leaves_present = [false; max_leaves!(ARITY, HEIGHT)];
127                for present in leaves_present.iter_mut().take(num_of_leaves) {
128                    *present = true;
129                }
130                Ok(Self {
131                    tree: StaticTree::try_from_leaves(leaves)?,
132                    num_of_leaves,
133                    leaves_present,
134                })
135            })
136    }
137
138    fn compacted_leaves<const EXPECTED_HEIGHT: usize>(
139        &self,
140    ) -> Result<[Prefixed<ARITY, H>; layer_size!(ARITY, EXPECTED_HEIGHT, 0)], Error> {
141        if self.num_of_leaves > max_leaves!(ARITY, EXPECTED_HEIGHT) {
142            return Err(Error::Create);
143        }
144
145        let mut prefixed =
146            [Prefixed::<ARITY, H>::default(); layer_size!(ARITY, EXPECTED_HEIGHT, 0)];
147
148        let mut j = 0;
149        for i in 0..self.leaves_present.len() {
150            if self.leaves_present[i] {
151                let (old_index, old_offset) = location_in_prefixed::<ARITY>(i);
152                let (new_index, new_offset) = location_in_prefixed::<ARITY>(j);
153
154                prefixed[new_index].hashes[new_offset] =
155                    self.tree.leaves()[old_index].hashes[old_offset];
156                j += 1;
157            }
158        }
159        assert_eq!(self.num_of_leaves(), j);
160        Ok(prefixed)
161    }
162    /// move all existing leaves leftwards
163    pub fn compact(&mut self)
164    where
165        [(); num_of_prefixed!(ARITY, HEIGHT)]: Sized,
166        [(); max_leaves!(ARITY, HEIGHT)]: Sized,
167        H: HashT,
168        PB: ProofBuilder<ARITY, H>,
169    {
170        *self = self
171            .compacted_leaves::<HEIGHT>()
172            .and_then(|leaves| Self::try_from_leaves(&leaves))
173            .expect("can create from compacted leaves. qed");
174    }
175    /// tries to compact this tree to a size of a tree with height-1 and create an instance of the new tree
176    /// Note: takes ownership, but as it implements Copy trait may need explicit dropping
177    /// to prevent being any longer available
178    pub fn try_reduce(
179        self,
180    ) -> Result<CompactableHeaplessTree<ARITY, { HEIGHT - 1 }, H, MAX_INPUT_LEN, PB>, Self>
181    where
182        [(); num_of_prefixed!(ARITY, { HEIGHT - 1 })]: Sized,
183        [(); max_leaves!(ARITY, { HEIGHT - 1 })]: Sized,
184        H: HashT,
185        PB: ProofBuilder<ARITY, H>,
186    {
187        self.compacted_leaves::<{ HEIGHT - 1 }>()
188            .and_then(|leaves| {
189                CompactableHeaplessTree::<ARITY, { HEIGHT - 1 }, H, MAX_INPUT_LEN, PB>::try_from_compacted(
190                    &leaves,
191                    self.num_of_leaves(),
192                )
193            })
194            .map_err(|_| self)
195    }
196}
197
198impl<const ARITY: usize, const HEIGHT: usize, H, const MAX_INPUT_LEN: usize, PB>
199    StaticTreeTrait<ARITY, H, PB> for CompactableHeaplessTree<ARITY, HEIGHT, H, MAX_INPUT_LEN, PB>
200where
201    [(); num_of_prefixed!(ARITY, HEIGHT)]: Sized,
202    [(); max_leaves!(ARITY, HEIGHT)]: Sized,
203    Assert<{ is_pow2!(ARITY) }>: IsTrue,
204    H: HashT,
205    PB: ProofBuilder<ARITY, H>,
206{
207    fn generate_proof(&self, index: usize) -> PB {
208        self.tree.generate_proof(index)
209    }
210    fn replace(&mut self, index: usize, input: &[u8]) {
211        self.tree.replace(index, input);
212
213        if !self.leaves_present[index] {
214            self.num_of_leaves += 1;
215        }
216        self.leaves_present[index] = true;
217    }
218    fn replace_leaf(&mut self, index: usize, leaf: H::Output) {
219        self.tree.replace_leaf(index, leaf);
220
221        if !self.leaves_present[index] {
222            self.num_of_leaves += 1;
223        }
224        self.leaves_present[index] = true;
225    }
226    fn root(&self) -> H::Output {
227        self.tree.root()
228    }
229    fn leaves(&self) -> &[Prefixed<ARITY, H>] {
230        &self.tree.prefixed[..layer_size!(ARITY, HEIGHT, 0)]
231    }
232    fn base_layer_size(&self) -> usize {
233        layer_size!(ARITY, HEIGHT, 0)
234    }
235    fn arity(&self) -> usize {
236        ARITY
237    }
238    fn height(&self) -> usize {
239        HEIGHT
240    }
241}
242
243impl<const ARITY: usize, const HEIGHT: usize, H, const MAX_INPUT_LEN: usize, PB> CanRemove
244    for CompactableHeaplessTree<ARITY, HEIGHT, H, MAX_INPUT_LEN, PB>
245where
246    [(); num_of_prefixed!(ARITY, HEIGHT)]: Sized,
247    [(); max_leaves!(ARITY, HEIGHT)]: Sized,
248    Assert<{ is_pow2!(ARITY) }>: IsTrue,
249    H: HashT,
250    PB: ProofBuilder<ARITY, H>,
251{
252    // remove element by replacing with nothing
253    // panics if index is out of leaf layer bound
254    fn remove(&mut self, index: usize) {
255        self.tree.replace(index, &[]);
256
257        if self.leaves_present[index] {
258            self.num_of_leaves -= 1;
259        }
260        self.leaves_present[index] = false;
261    }
262
263    fn num_of_leaves(&self) -> usize {
264        self.num_of_leaves
265    }
266}
267
268impl<const ARITY: usize, const HEIGHT: usize, H, const MAX_INPUT_LEN: usize, PB> Clone
269    for CompactableHeaplessTree<ARITY, HEIGHT, H, MAX_INPUT_LEN, PB>
270where
271    [(); num_of_prefixed!(ARITY, HEIGHT)]: Sized,
272    [(); max_leaves!(ARITY, HEIGHT)]: Sized,
273    Assert<{ is_pow2!(ARITY) }>: IsTrue,
274    H: HashT,
275    PB: ProofBuilder<ARITY, H>,
276{
277    fn clone(&self) -> Self {
278        Self {
279            tree: self.tree,
280            num_of_leaves: self.num_of_leaves,
281            leaves_present: self.leaves_present,
282        }
283    }
284}
285
286impl<const ARITY: usize, const HEIGHT: usize, H, const MAX_INPUT_LEN: usize, PB> Default
287    for CompactableHeaplessTree<ARITY, HEIGHT, H, MAX_INPUT_LEN, PB>
288where
289    [(); num_of_prefixed!(ARITY, HEIGHT)]: Sized,
290    [(); max_leaves!(ARITY, HEIGHT)]: Sized,
291    Assert<{ is_pow2!(ARITY) }>: IsTrue,
292    H: HashT,
293    PB: ProofBuilder<ARITY, H>,
294{
295    fn default() -> Self {
296        Self {
297            tree: Default::default(),
298            num_of_leaves: 0,
299            leaves_present: [false; max_leaves!(ARITY, HEIGHT)],
300        }
301    }
302}
303
304impl<const ARITY: usize, const HEIGHT: usize, H, const MAX_INPUT_LEN: usize, PB> Debug
305    for CompactableHeaplessTree<ARITY, HEIGHT, H, MAX_INPUT_LEN, PB>
306where
307    [(); num_of_prefixed!(ARITY, HEIGHT)]: Sized,
308    [(); max_leaves!(ARITY, HEIGHT)]: Sized,
309    Assert<{ is_pow2!(ARITY) }>: IsTrue,
310    H: HashT,
311    PB: ProofBuilder<ARITY, H>,
312{
313    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> Result<(), core::fmt::Error> {
314        writeln!(f, "{:?}", self.tree)?;
315        writeln!(f, "[num of leaves]: {:?}", self.num_of_leaves())
316    }
317}