From 61307e04cfe446b3d2398dd8441c6e16c9c97dd3 Mon Sep 17 00:00:00 2001 From: Richard Pringle Date: Mon, 11 Dec 2023 15:32:31 -0500 Subject: [PATCH] Stop prefetching next node in iterator --- firewood/src/merkle/stream.rs | 179 +++++++++++----------------------- 1 file changed, 59 insertions(+), 120 deletions(-) diff --git a/firewood/src/merkle/stream.rs b/firewood/src/merkle/stream.rs index ddd0bffd1..b0e6916f5 100644 --- a/firewood/src/merkle/stream.rs +++ b/firewood/src/merkle/stream.rs @@ -16,11 +16,7 @@ pub(super) enum IteratorState<'a> { /// Start iterating at the specified key StartAtKey(Vec), /// Continue iterating after the given `next_node` and parents - Iterating { - visit_last: bool, - next_result: Option<(Vec, Vec)>, - parents: Vec<(ObjRef<'a>, u8)>, - }, + Iterating { parents: Vec<(ObjRef<'a>, u8)> }, } impl IteratorState<'_> { pub(super) fn new>(starting: Option) -> Self { @@ -64,130 +60,79 @@ impl<'a, S: ShaleStore + Send + Sync, T> Stream for MerkleKeyValueStream<' .get_node(*merkle_root) .map_err(|e| api::Error::InternalError(Box::new(e)))?; + match root.inner() { + NodeType::Branch(branch) => { + if let Some(value) = branch.value.as_ref() { + let key = vec![]; + let value = value.to_vec(); + + self.key_state = IteratorState::Iterating { + parents: vec![(root, 0)], + }; + + return Poll::Ready(Some(Ok((key, value)))); + } + } + _ => (), + } + // always put the sentinal node in parents let mut parents = vec![(root, 0)]; - let next_result = find_next_result(merkle, &mut parents, true) - .map_err(|e| api::Error::InternalError(Box::new(e)))?; - - let (next_result, visit_last) = next_result - .map(|(next, visit)| (Some(next), visit)) - .unwrap_or_default(); + let next_result = find_next_result(merkle, &mut parents) + .map_err(|e| api::Error::InternalError(Box::new(e))) + .transpose(); - self.key_state = IteratorState::Iterating { - visit_last, - next_result, - parents, - }; + self.key_state = IteratorState::Iterating { parents }; - self.poll_next(_cx) + Poll::Ready(next_result) } IteratorState::StartAtKey(key) => { + if key.is_empty() { + self.key_state = IteratorState::StartAtBeginning; + return self.poll_next(_cx); + } + let root_node = merkle .get_node(*merkle_root) .map_err(|e| api::Error::InternalError(Box::new(e)))?; - let (next_result, parents) = { - let (found_node, parents) = { - let mut parents = vec![]; - - let found_node = merkle - .get_node_by_key_with_callbacks( - root_node, - &key, - |addr, nib| { - parents.push((addr, nib)); - }, - |_, _| {}, - ) - .map_err(|e| api::Error::InternalError(Box::new(e)))?; - - (found_node, parents) - }; - - let mut parents = parents - .into_iter() - .map(|(addr, nib)| merkle.get_node(addr).map(|node| (node, nib))) - .collect::, _>>() - .map_err(|e| api::Error::InternalError(Box::new(e)))?; - - // means the keys definitely match - if let Some(found_node) = found_node { - let (value, visit_last) = match found_node.inner() { - NodeType::Extension(_) => unreachable!(), - - NodeType::Branch(branch) => { - let value = branch.value.as_ref().map(|v| v.to_vec()); - let visit_last = true; - parents.push((found_node, 0)); - - (value, visit_last) - } - - NodeType::Leaf(leaf) => { - let value = leaf.data.to_vec(); - let visit_last = false; + let (found_node, mut parents) = merkle + .get_node_and_parents_by_key(root_node, &key) + .map_err(|e| api::Error::InternalError(Box::new(e)))?; - (Some(value), visit_last) - } - }; + if let Some(found_node) = found_node { + let value = match found_node.inner() { + NodeType::Branch(branch) => branch.value.as_ref(), + NodeType::Leaf(leaf) => Some(&leaf.data), + NodeType::Extension(_) => None, + }; - let next_result = value.map(|value| (std::mem::take(key), value)); + let next_result = value.map(|value| { + let value = value.to_vec(); - self.key_state = IteratorState::Iterating { - visit_last, - next_result, - parents, - }; + Ok(dbg!(std::mem::take(key), value)) + }); - return self.poll_next(_cx); - } + parents.push((found_node, 0)); - let key_from_parents = key_from_nibble_iter(nibble_iter_from_parents(&parents)); - let mut visit_last = key == &key_from_parents; - - if key.as_slice() < key_from_parents.as_slice() { - let _ = parents.pop(); - visit_last = true; - } + self.key_state = IteratorState::Iterating { parents }; - let next_result = find_next_result(merkle, &mut parents, visit_last) - .map_err(|e| api::Error::InternalError(Box::new(e)))?; - - (next_result, parents) - }; + return Poll::Ready(next_result); + } - let (next_result, visit_last) = next_result - .map(|(next, visit)| (Some(next), visit)) - .unwrap_or_default(); - - self.key_state = IteratorState::Iterating { - visit_last, - next_result, - parents, - }; + self.key_state = IteratorState::Iterating { parents }; self.poll_next(_cx) } - IteratorState::Iterating { - visit_last, - next_result, - parents, - } => { - let Some((next_key, next_value)) = next_result.take() else { - return Poll::Ready(None); - }; - - let next = find_next_result(merkle, parents, *visit_last) - .map_err(|e| api::Error::InternalError(Box::new(e)))?; - - (*next_result, *visit_last) = next - .map(|(next, visit)| (Some(next), visit)) - .unwrap_or_default(); + IteratorState::Iterating { parents } => { + let next = find_next_result(merkle, parents) + .map_err(|e| api::Error::InternalError(Box::new(e))) + .transpose(); - Poll::Ready(Some(Ok((next_key, next_value)))) + Poll::Ready(next) } } } @@ -220,28 +165,23 @@ impl<'a> ParentNode<'a> { } } -type NextResult = Option<((Vec, Vec), bool)>; +type NextResult = Option<(Vec, Vec)>; fn find_next_result<'a, S: ShaleStore, T>( merkle: &'a Merkle, parents: &mut Vec<(ObjRef<'a>, u8)>, - visit_last: bool, ) -> Result { - let next = find_next_node_with_data(merkle, parents, visit_last)?.map(|next| { - let (next_node, value) = next; + let next = find_next_node_with_data(merkle, parents)?.map(|(next_node, value)| { let node_path_iter = match next_node.inner() { NodeType::Leaf(leaf) => leaf.path.iter().copied(), _ => [].iter().copied(), }; let key = key_from_nibble_iter(nibble_iter_from_parents(parents).chain(node_path_iter)); - let visit_last = matches!(next_node.inner(), NodeType::Branch(_)); - if visit_last { - parents.push((next_node, 0)); - } + parents.push((next_node, 0)); - ((key, value), visit_last) + (key, value) }); Ok(next) @@ -250,7 +190,6 @@ fn find_next_result<'a, S: ShaleStore, T>( fn find_next_node_with_data<'a, S: ShaleStore, T>( merkle: &'a Merkle, visited_parents: &mut Vec<(ObjRef<'a>, u8)>, - visit_last: bool, ) -> Result, Vec)>, super::MerkleError> { use InnerNode::*; @@ -260,18 +199,16 @@ fn find_next_node_with_data<'a, S: ShaleStore, T>( let mut node = ParentNode::Visited(visited_parent); let mut pos = visited_pos; + let mut first_loop = true; loop { match node.inner() { - // TODO: find a better way to handle this impossible case - Visited(NodeType::Leaf(_)) => unreachable!(), - New(NodeType::Leaf(leaf)) => { let value = leaf.data.to_vec(); return Ok(Some((node.into_node(), value))); } - Visited(NodeType::Extension(_)) => { + Visited(NodeType::Leaf(_)) | Visited(NodeType::Extension(_)) => { let Some((next_parent, next_pos)) = visited_parents.pop() else { return Ok(None); }; @@ -290,7 +227,7 @@ fn find_next_node_with_data<'a, S: ShaleStore, T>( } Visited(NodeType::Branch(branch)) => { - let compare_op = if visit_last { + let compare_op = if first_loop { ::ge } else { ::gt @@ -323,6 +260,8 @@ fn find_next_node_with_data<'a, S: ShaleStore, T>( } } } + + first_loop = false; } }