1use std::cmp::Ordering;
2use std::fmt;
3use std::sync::Arc;
4
5use crate::store::NodeStore;
6
7use super::node::{Hash, Node, NodeError};
8use super::source_error::TreeOperationError;
9
10#[derive(Debug, Clone, PartialEq, Eq)]
11pub enum TreeError {
12 MissingNode { hash: Hash },
13 InvalidNode,
14}
15
16impl fmt::Display for TreeError {
17 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
18 match self {
19 Self::MissingNode { hash } => write!(formatter, "missing tree node {hash}"),
20 Self::InvalidNode => write!(formatter, "invalid tree node"),
21 }
22 }
23}
24
25impl std::error::Error for TreeError {}
26
27impl From<NodeError> for TreeError {
28 fn from(_: NodeError) -> Self {
29 Self::InvalidNode
30 }
31}
32
33#[derive(Debug)]
34pub struct Cursor<'a, S: NodeStore + ?Sized> {
35 store: &'a S,
36 root_hash: Hash,
37}
38
39impl<'a, S: NodeStore + ?Sized> Cursor<'a, S> {
40 pub const fn new(store: &'a S, root_hash: Hash) -> Self {
41 Self { store, root_hash }
42 }
43
44 pub const fn root_hash(&self) -> Hash {
45 self.root_hash
46 }
47
48 pub fn get(&self, key: &[u8]) -> Result<Option<Vec<u8>>, TreeError> {
49 self.get_with_source(key)
50 .map_err(TreeOperationError::into_tree_error)
51 }
52
53 pub(crate) fn get_with_source(
54 &self,
55 key: &[u8],
56 ) -> Result<Option<Vec<u8>>, TreeOperationError<S::Error>> {
57 let mut next_hash = self.root_hash;
58
59 loop {
60 match &*load_node_with_source(self.store, next_hash)? {
61 Node::Leaf(leaf) => return leaf_value(leaf, key).map_err(TreeOperationError::Tree),
62 Node::Internal(internal) => {
63 let children = internal.children();
64 let index =
65 child_index_for_key(children, key).map_err(TreeOperationError::Tree)?;
66 let Some((_separator, child_hash)) = children.get(index) else {
67 return Err(TreeOperationError::Tree(TreeError::InvalidNode));
68 };
69 next_hash = *child_hash;
70 }
71 }
72 }
73 }
74
75 pub fn range(&self, from: &[u8], to: &[u8]) -> RangeIter<'_, S> {
76 RangeIter::new(self.store, self.root_hash, from, to)
77 }
78}
79
80#[derive(Debug)]
81pub struct RangeIter<'a, S: NodeStore + ?Sized> {
82 store: &'a S,
83 root_hash: Hash,
84 from: Vec<u8>,
85 to: Vec<u8>,
86 stack: Vec<RangeFrame>,
87 leaf: Option<Arc<Node>>,
88 leaf_index: usize,
89 started: bool,
90 finished: bool,
91}
92
93impl<'a, S: NodeStore + ?Sized> RangeIter<'a, S> {
94 fn new(store: &'a S, root_hash: Hash, from: &[u8], to: &[u8]) -> Self {
95 Self {
96 store,
97 root_hash,
98 from: from.to_vec(),
99 to: to.to_vec(),
100 stack: Vec::new(),
101 leaf: None,
102 leaf_index: 0,
103 started: false,
104 finished: to <= from,
105 }
106 }
107
108 fn start(&mut self) -> Result<(), TreeError> {
109 self.started = true;
110 self.descend_for_key(self.root_hash)
111 }
112
113 fn descend_for_key(&mut self, hash: Hash) -> Result<(), TreeError> {
114 let mut next_hash = hash;
115
116 loop {
117 let node = load_node(self.store, next_hash)?;
118 match &*node {
119 Node::Leaf(leaf) => {
120 self.leaf_index = lower_bound(leaf.entries(), self.from.as_slice());
121 self.leaf = Some(Arc::clone(&node));
122 return Ok(());
123 }
124 Node::Internal(internal) => {
125 let children = internal.children();
126 let index = child_index_for_key(children, self.from.as_slice())?;
127 let Some((_separator, child_hash)) = children.get(index) else {
128 return Err(TreeError::InvalidNode);
129 };
130 next_hash = *child_hash;
131 self.stack.push(RangeFrame {
132 node: Arc::clone(&node),
133 next_index: index.saturating_add(1),
134 });
135 }
136 }
137 }
138 }
139
140 fn descend_leftmost(&mut self, hash: Hash) -> Result<(), TreeError> {
141 let mut next_hash = hash;
142
143 loop {
144 let node = load_node(self.store, next_hash)?;
145 match &*node {
146 Node::Leaf(_leaf) => {
147 self.leaf_index = 0;
148 self.leaf = Some(Arc::clone(&node));
149 return Ok(());
150 }
151 Node::Internal(internal) => {
152 let Some((_separator, child_hash)) = internal.children().first() else {
153 return Err(TreeError::InvalidNode);
154 };
155 next_hash = *child_hash;
156 self.stack.push(RangeFrame {
157 node: Arc::clone(&node),
158 next_index: 1,
159 });
160 }
161 }
162 }
163 }
164
165 fn advance_to_next_leaf(&mut self) -> Result<(), TreeError> {
166 while let Some(mut frame) = self.stack.pop() {
167 let children = frame_children(&frame.node)?;
168 if frame.next_index < children.len() {
169 let Some((separator, child_hash)) = children.get(frame.next_index) else {
170 return Err(TreeError::InvalidNode);
171 };
172 let separator_at_or_past_end = separator.as_slice() >= self.to.as_slice();
173 let child_hash = *child_hash;
174
175 frame.next_index = frame.next_index.saturating_add(1);
176 self.stack.push(frame);
177
178 if separator_at_or_past_end {
179 self.finished = true;
180 return Ok(());
181 }
182
183 return self.descend_leftmost(child_hash);
184 }
185 }
186
187 self.finished = true;
188 Ok(())
189 }
190}
191
192impl<S: NodeStore + ?Sized> Iterator for RangeIter<'_, S> {
193 type Item = Result<(Vec<u8>, Vec<u8>), TreeError>;
194
195 fn next(&mut self) -> Option<Self::Item> {
196 if self.finished {
197 return None;
198 }
199
200 if !self.started {
201 match self.start() {
202 Ok(()) => {}
203 Err(error) => {
204 self.finished = true;
205 return Some(Err(error));
206 }
207 }
208 }
209
210 loop {
211 if let Some(node) = self.leaf.clone()
212 && let Node::Leaf(leaf) = &*node
213 {
214 let entries = leaf.entries();
215 while let Some((key, value)) = entries.get(self.leaf_index) {
216 self.leaf_index = self.leaf_index.saturating_add(1);
217
218 if key.as_slice() < self.from.as_slice() {
219 continue;
220 }
221
222 if key.as_slice() >= self.to.as_slice() {
223 self.finished = true;
224 return None;
225 }
226
227 return Some(Ok((key.clone(), value.clone())));
228 }
229 }
230
231 if let Err(error) = self.advance_to_next_leaf() {
232 self.finished = true;
233 return Some(Err(error));
234 }
235
236 if self.finished {
237 return None;
238 }
239 }
240 }
241}
242
243#[derive(Debug, Clone)]
244struct RangeFrame {
245 node: Arc<Node>,
246 next_index: usize,
247}
248
249fn frame_children(node: &Node) -> Result<&[(Vec<u8>, Hash)], TreeError> {
250 match node {
251 Node::Internal(internal) => Ok(internal.children()),
252 Node::Leaf(_leaf) => Err(TreeError::InvalidNode),
253 }
254}
255
256pub(crate) fn load_node<S: NodeStore + ?Sized>(
257 store: &S,
258 hash: Hash,
259) -> Result<Arc<Node>, TreeError> {
260 load_node_with_source(store, hash).map_err(TreeOperationError::into_tree_error)
261}
262
263pub(crate) fn load_node_with_source<S: NodeStore + ?Sized>(
264 store: &S,
265 hash: Hash,
266) -> Result<Arc<Node>, TreeOperationError<S::Error>> {
267 store
268 .get(&hash)
269 .map_err(|source| TreeOperationError::NodeGet { hash, source })?
270 .ok_or(TreeOperationError::Tree(TreeError::MissingNode { hash }))
271}
272
273pub(crate) fn child_index_for_key(
274 children: &[(Vec<u8>, Hash)],
275 key: &[u8],
276) -> Result<usize, TreeError> {
277 if children.is_empty() {
278 return Err(TreeError::InvalidNode);
279 }
280
281 Ok(children
282 .partition_point(|(separator, _hash)| separator.as_slice() <= key)
283 .saturating_sub(1))
284}
285
286fn leaf_value(leaf: &super::node::LeafNode, key: &[u8]) -> Result<Option<Vec<u8>>, TreeError> {
287 match leaf
288 .entries()
289 .binary_search_by(|(entry_key, _value)| entry_key.as_slice().cmp(key))
290 {
291 Ok(index) => leaf
292 .entries()
293 .get(index)
294 .map(|(_key, value)| Some(value.clone()))
295 .ok_or(TreeError::InvalidNode),
296 Err(_index) => Ok(None),
297 }
298}
299
300fn lower_bound(entries: &[(Vec<u8>, Vec<u8>)], key: &[u8]) -> usize {
301 entries.partition_point(|(entry_key, _value)| {
302 matches!(entry_key.as_slice().cmp(key), Ordering::Less)
303 })
304}