refactor reading nodes to prepare caching

src/btrfs_lookup.rs

@@ -5,7 +5,7 @@ use std::ops::{Deref, RangeBounds, Bound};
 use crate::btrfs_structs::{Leaf, Key, Item, InteriorNode, Node, ParseError, ParseBin, Value, Superblock, ItemType, ZERO_KEY, DirItem};
 use crate::addrmap::{node_at_log, LogToPhys, AddressMap};
 
-/// represents a B-Tree inside a filesystem image. Can be used to look up keys,
+/// Represents a B-Tree inside a filesystem image. Can be used to look up keys,
 /// and handles the tree traversal and the virtual address translation.
 pub struct Tree<'a> {
 	pub image: &'a [u8],
@@ -39,6 +39,15 @@ impl<'a> Tree<'a> {
 
 		Ok(Tree { image, addr_map, root_addr_log: superblock.root })
 	}
+
+	/// Read a node given its logical address
+	pub fn get_node(&self, addr: u64) -> Result<Rc<Node>, ParseError> {
+		eprintln!("Reading node at {:x}", addr);
+
+		let node_data = node_at_log(self.image, self.addr_map.as_ref(), addr)?;
+		let node = Node::parse(node_data)?;
+		Ok(Rc::new(node))
+	}
 }
 
 /***** looking up keys *****/
@@ -76,30 +85,31 @@ impl InteriorNode {
 	}
 }
 
-/// Recursively traverse a tree to find a key, given they key and logical address
-/// of the tree root. Internal function, `Tree::find_key` is the public interface.
-fn find_key_in_node<T: LogToPhys>(image: &[u8], addr: &T, root_addr_log: u64, key: Key) -> Result<Item, ParseError> {
-	let node = Node::parse(node_at_log(image, addr, root_addr_log)?)?;
 
-	match node {
+impl Tree<'_> {
+	/// Recursively traverse a tree to find a key, given they key and logical address
+	/// of the tree root. Internal function, `Tree::find_key` is the public interface.
+	fn find_key_in_node(&self, addr: u64, key: Key) -> Result<Item, ParseError> {
+		let node = self.get_node(self.root_addr_log)?;
+
+		match node.deref() {
 			Node::Interior(interior_node) => {
 				let next_node_index = interior_node.find_key_or_previous(key).unwrap();
 				let next_node_log = interior_node.children[next_node_index].ptr;
-			find_key_in_node(image, addr, next_node_log, key)
+				self.find_key_in_node(next_node_log, key)
 			},
 			Node::Leaf(leaf) => {
 				leaf.find_key(key).ok_or(
 					error!(
 						"Item with key ({},{:?},{}) was not found in the leaf at logical address 0x{:x}",
-					key.key_id, key.key_type, key.key_offset, root_addr_log)
+						key.key_id, key.key_type, key.key_offset, addr)
 				)
 			}
 		}
-}
+	}
 
-impl Tree<'_> {
 	pub fn find_key(&self, key: Key) -> Result<Item, ParseError> {
-		find_key_in_node(self.image, self.addr_map.deref(), self.root_addr_log, key)
+		self.find_key_in_node(self.root_addr_log, key)
 	}
 }
 
@@ -161,12 +171,11 @@ impl<'a> Tree<'a> {
 	}
 }
 
-#[derive(Debug,PartialEq,Eq,Clone,Copy)]
-enum FindKeyMode {LT, GT, GE, LE}
 
+/// Get the first item under the node at logical address `addr`.
+/// This function panics if there are no items
 fn get_first_item(tree: &Tree, addr: u64) -> Result<Item, ParseError> {
-	let node_data = node_at_log(tree.image, tree.addr_map.deref(), addr)?;
+	match tree.get_node(addr)?.deref() {
-	match Node::parse(node_data)? {
 		Node::Interior(intnode) => {
 			get_first_item(tree, intnode.children[0].ptr)
 		},
@@ -176,9 +185,10 @@ fn get_first_item(tree: &Tree, addr: u64) -> Result<Item, ParseError> {
 	}
 }
 
+/// Get the last item under the node at logical address `addr`.
+/// This function panics if there are no items
 fn get_last_item(tree: &Tree, addr: u64) -> Result<Item, ParseError> {
-	let node_data = node_at_log(tree.image, tree.addr_map.deref(), addr)?;
+	match tree.get_node(addr)?.deref() {
-	match Node::parse(node_data)? {
 		Node::Interior(intnode) => {
 			get_last_item(tree, intnode.children.last().unwrap().ptr)
 		},
@@ -188,6 +198,9 @@ fn get_last_item(tree: &Tree, addr: u64) -> Result<Item, ParseError> {
 	}
 }
 
+#[derive(Debug,PartialEq,Eq,Clone,Copy)]
+enum FindKeyMode {LT, GT, GE, LE}
+
 /// Try to find the item with key `key` if it exists in the tree, and return
 /// the "closest" match. The exact meaning of "closest" is given by the `mode` argument:
 /// If `mode` is `LT`/`GT`/`GE`/`LE`, return the item with the greatest / least / greatest / least
@@ -195,14 +208,14 @@ fn get_last_item(tree: &Tree, addr: u64) -> Result<Item, ParseError> {
 fn find_closest_key(tree: &Tree, key: Key, mode: FindKeyMode) -> Result<Option<Item>, ParseError> {
 
 	// in some cases, this task can't be accomplished by a single traversal
-	// but we might have to go back up the tree; this state allows to quickly go back to the right node
+	// but we might have to go back up the tree; prev/next allows to quickly go back to the right node
+	let mut current: u64 = tree.root_addr_log;
 	let mut prev: Option<u64> = None;
 	let mut next: Option<u64> = None;
 
-	let mut node_data = node_at_log(tree.image, tree.addr_map.deref(), tree.root_addr_log)?;
-
 	loop {
-		match Node::parse(node_data)? {
+		let node = tree.get_node(current)?;
+		match node.deref() {
 			Node::Interior(intnode) => {
 				match intnode.find_key_or_previous(key) {
 					Some(idx) => {
@@ -213,7 +226,7 @@ fn find_closest_key(tree: &Tree, key: Key, mode: FindKeyMode) -> Result<Option<I
 							next = Some(kp.ptr);
 						}
 
-						node_data = node_at_log(tree.image, tree.addr_map.deref(), intnode.children[idx].ptr)?;
+						current = intnode.children[idx].ptr;
 					},
 					None => {
 						// this can only happen if every key in the current node is `> key`
@@ -309,7 +322,7 @@ impl<'a, 'b> Iterator for RangeIter<'a, 'b> {
 			.expect("file system should be consistent (or this is a bug)");
 
 		if let Some(item) = &result {
-			self.start = Bound::Excluded(item.key);
+			self.start = Bound::Excluded((self.forward_skip_fn)(item.key));
 		}
 
 		let end_filter = |item : &Item| {