diff --git a/rust/hg-core/src/revlog/nodemap.rs b/rust/hg-core/src/revlog/nodemap.rs --- a/rust/hg-core/src/revlog/nodemap.rs +++ b/rust/hg-core/src/revlog/nodemap.rs @@ -222,17 +222,58 @@ &self, prefix: NodePrefixRef<'p>, ) -> Result, NodeMapError> { - let mut visit = self.len() - 1; - for i in 0..prefix.len() { - let nybble = prefix.get_nybble(i); - match self[visit].read(nybble) { - Element::None => return Ok(None), - Element::Rev(r) => return Ok(Some(r)), - Element::Block(idx) => visit = idx, + for (leaf, _, _, opt) in self.visit(prefix) { + if leaf { + return Ok(opt); } } Err(NodeMapError::MultipleResults) } + + fn visit<'n, 'p>( + &'n self, + prefix: NodePrefixRef<'p>, + ) -> NodeTreeVisitor<'n, 'p> { + NodeTreeVisitor { + nt: self, + prefix: prefix, + visit: self.len() - 1, + nybble_idx: 0, + done: false, + } + } +} + +struct NodeTreeVisitor<'n, 'p> { + nt: &'n NodeTree, + prefix: NodePrefixRef<'p>, + visit: usize, + nybble_idx: usize, + done: bool, +} + +impl<'n, 'p> Iterator for NodeTreeVisitor<'n, 'p> { + type Item = (bool, usize, u8, Option); + + fn next(&mut self) -> Option { + if self.done || self.nybble_idx >= self.prefix.len() { + return None; + } + + let nybble = self.prefix.get_nybble(self.nybble_idx); + let visit = self.visit; + let (leaf, opt) = match self.nt[visit].read(nybble) { + Element::None => (true, None), + Element::Rev(r) => (true, Some(r)), + Element::Block(idx) => { + self.visit = idx; + (false, None) + } + }; + self.nybble_idx += 1; + self.done = leaf; + Some((leaf, visit, nybble, opt)) + } } impl From> for NodeTree {