merkletree_mintlayer/merkle/tree/
padding.rs1use std::iter::FusedIterator;
17
18pub struct IncrementalPaddingIterator<T, I: Iterator<Item = T> + FusedIterator, F: Fn(&T) -> T> {
23 leaves: I,
24 padding_function: F,
25 last_value: Option<T>,
26 current_index: usize,
27}
28
29impl<T, I: Iterator<Item = T> + FusedIterator, F: Fn(&T) -> T> IncrementalPaddingIterator<T, I, F> {
30 pub fn new(leaves: I, padding_function: F) -> Self {
31 IncrementalPaddingIterator {
32 leaves,
33 padding_function,
34 last_value: None,
35 current_index: 0,
36 }
37 }
38}
39
40impl<T: Clone, I: Iterator<Item = T> + FusedIterator, F: Fn(&T) -> T> Iterator
41 for IncrementalPaddingIterator<T, I, F>
42{
43 type Item = T;
44
45 fn next(&mut self) -> Option<T> {
46 match self.leaves.next() {
47 None => {
48 if self.current_index == self.current_index.next_power_of_two()
51 || self.current_index == 0
52 {
53 None
54 } else {
55 let res =
56 (self.padding_function)(self.last_value.as_ref().expect("Never at zero"));
57 self.current_index += 1;
58 self.last_value = Some(res.clone());
59 Some(res)
60 }
61 }
62 Some(leaf) => {
63 self.current_index += 1;
64 self.last_value = Some(leaf.clone());
65 Some(leaf)
66 }
67 }
68 }
69}
70
71#[cfg(test)]
72mod tests {
73
74 use crate::internal::{hash_data, HashedData};
75
76 use super::*;
77
78 fn leaves_with_inc_padding(n: usize) -> Vec<HashedData> {
79 let mut leaves = Vec::new();
80 for i in 0..n {
81 leaves.push(HashedData::from_low_u64_be(i as u64));
82 }
83 for _ in n..n.next_power_of_two() {
84 leaves.push(hash_data(*leaves.last().unwrap()));
85 }
86 leaves
87 }
88
89 #[test]
90 fn non_zero_size() {
91 let f = |i: &HashedData| hash_data(i);
92
93 for i in 1..130 {
94 let all_leaves = leaves_with_inc_padding(i);
95 let leaves = &leaves_with_inc_padding(i)[0..i];
96
97 let vec =
98 IncrementalPaddingIterator::new(leaves.iter().copied(), f).collect::<Vec<_>>();
99 assert_eq!(vec, all_leaves);
100 }
101 }
102
103 #[test]
104 fn zero_size() {
105 let f = |i: &HashedData| hash_data(i);
106
107 let vec = IncrementalPaddingIterator::new(Vec::new().into_iter(), f).collect::<Vec<_>>();
108 assert_eq!(vec, Vec::new());
109 }
110}