1pub 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;
39pub 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 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 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 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 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 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 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}