trie_test/
iter.rs

1use database::DatabaseGuard;
2use merkle::{MerkleNode, MerkleValue};
3use merkle::nibble::{into_key, NibbleVec};
4use rlp::{self, Rlp};
5
6use std::ops::Deref;
7use std::marker::PhantomData;
8
9pub struct FixedMerkleIterator<'a, D: DatabaseGuard + 'a, K, V> {
10    iter: MerkleIterator<'a, D>,
11    _marker: PhantomData<(K, V)>
12}
13
14impl<'a, D: DatabaseGuard + 'a, K: rlp::Decodable, V: rlp::Decodable> FixedMerkleIterator<'a, D, K, V> {
15    pub fn new(iter: MerkleIterator<'a, D>) -> Self {
16        FixedMerkleIterator {
17            iter, _marker: PhantomData
18        }
19    }
20}
21
22impl<'a, D: DatabaseGuard + 'a, K: rlp::Decodable, V: rlp::Decodable> Iterator for FixedMerkleIterator<'a, D, K, V> {
23    type Item = (K, V);
24
25    fn next(&mut self) -> Option<(K, V)> {
26        if let Some((raw_key, raw_value)) = self.iter.next() {
27            let key = rlp::decode(&raw_key);
28            let value = rlp::decode(&raw_value);
29            Some((key, value))
30        } else {
31            None
32        }
33    }
34}
35
36pub struct MerkleIterator<'a, D: DatabaseGuard + 'a> {
37    database: &'a D,
38    prefix: NibbleVec,
39    value: Vec<u8>,
40    index: usize,
41    child: Option<Box<MerkleIterator<'a, D>>>,
42    is_empty: bool,
43}
44
45impl<'a, D: DatabaseGuard + 'a> MerkleIterator<'a, D> {
46    pub fn new(database: &'a D, value: Vec<u8>) -> Self {
47        Self {
48            database, value,
49            index: 0, child: None, prefix: NibbleVec::new(),
50            is_empty: false,
51        }
52    }
53
54    pub fn empty(database: &'a D) -> Self {
55        Self {
56            database,
57            value: Vec::new(), index: 0, child: None, prefix: NibbleVec::new(),
58            is_empty: true,
59        }
60    }
61}
62
63fn prepare_value_as_child<'a, D: DatabaseGuard + 'a>(
64    database: &'a D, prefix: NibbleVec, subnibble: NibbleVec, value: MerkleValue
65) -> Option<Box<MerkleIterator<'a, D>>> {
66    let mut nibble = prefix.clone();
67    nibble.extend(subnibble);
68
69    match value {
70        MerkleValue::Empty => None,
71        MerkleValue::Full(sub_node) => {
72            let value = rlp::encode(sub_node.deref()).to_vec();
73            Some(Box::new(
74                MerkleIterator {
75                    database: database,
76                    prefix: nibble,
77                    value, index: 0, child: None,
78                    is_empty: false,
79                }
80            ))
81        },
82        MerkleValue::Hash(h) => {
83            let value = database.get(h).unwrap();
84            Some(Box::new(
85                MerkleIterator {
86                    database: database,
87                    prefix: nibble,
88                    value, index: 0, child: None,
89                    is_empty: false,
90                }
91            ))
92        }
93    }
94}
95
96impl<'a, D: DatabaseGuard + 'a> Iterator for MerkleIterator<'a, D> {
97    type Item = (Vec<u8>, Vec<u8>);
98
99    fn next(&mut self) -> Option<(Vec<u8>, Vec<u8>)> {
100        if self.is_empty {
101            return None;
102        }
103
104        let node = MerkleNode::decode(&Rlp::new(&self.value));
105
106        match node {
107            MerkleNode::Leaf(node_nibble, node_value) => {
108                debug_assert!(self.child.is_none());
109
110                if self.index == 0 {
111                    self.index += 1;
112
113                    let mut nibble = self.prefix.clone();
114                    nibble.extend(node_nibble);
115
116                    Some((into_key(&nibble), node_value.into()))
117                } else {
118                    None
119                }
120            },
121            MerkleNode::Extension(node_nibble, node_value) => {
122                if self.index == 0 {
123                    debug_assert!(self.child.is_none());
124
125                    self.child = prepare_value_as_child(
126                        self.database, self.prefix.clone(), node_nibble, node_value);
127
128                    if self.child.is_some() {
129                        self.index += 1;
130                        self.child.as_mut().unwrap().next()
131                    } else {
132                        None
133                    }
134                } else {
135                    debug_assert!(self.child.is_some());
136
137                    self.child.as_mut().unwrap().next()
138                }
139            },
140            MerkleNode::Branch(nodes, additional) => {
141                debug_assert!(self.index <= 17);
142                while self.index <= 16 {
143                    if self.index < 16 {
144                        if self.child.is_some() {
145                            match self.child.as_mut().unwrap().next() {
146                                Some(val) => return Some(val),
147                                None => {
148                                    self.child = None;
149                                },
150                            }
151                        } else {
152                            let value = nodes[self.index].clone();
153                            let subnibble = vec![self.index.into()];
154                            self.index += 1;
155                            self.child = prepare_value_as_child(
156                                self.database, self.prefix.clone(), subnibble, value);
157                        }
158                    } else {
159                        self.index += 1;
160                        match additional {
161                            Some(val) => return Some((into_key(&self.prefix), val.into())),
162                            None => (),
163                        }
164                    }
165                }
166                None
167            },
168        }
169    }
170}
trie_test/iter.rs

trie_test/
iter.rs
