merkletree_mintlayer/merkle/tree/
padding.rs

1// Copyright (c) 2021-2024 RBB S.r.l
2// opensource@mintlayer.org
3// SPDX-License-Identifier: MIT
4// Licensed under the MIT License;
5// you may not use this file except in compliance with the License.
6// You may obtain a copy of the License at
7//
8// https://github.com/mintlayer/merkletree-mintlayer/blob/master/LICENSE
9//
10// Unless required by applicable law or agreed to in writing, software
11// distributed under the License is distributed on an "AS IS" BASIS,
12// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13// See the License for the specific language governing permissions and
14// limitations under the License.
15
16use std::iter::FusedIterator;
17
18/// An iterator that pads the leaves of a Merkle tree with incremental padding,
19/// i.e. the padding function is applied to the last value of the iterator,
20/// iteratively, until the next power of two is reached.
21/// An empty iterator will return no values.
22pub 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                // index == 0 means that we have no leaves at all;
49                // otherwise, we have to check if we have reached the next power of two to complete the padding.
50                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}