From e2fb0cbb470bc18f1105338fa1945b63fb94cba4 Mon Sep 17 00:00:00 2001 From: Florian Stecker Date: Wed, 7 Feb 2024 09:31:44 -0500 Subject: [PATCH] refactor reading nodes to prepare caching --- src/btrfs_lookup.rs | 81 ++++++++++++++++++++++++++------------------- 1 file changed, 47 insertions(+), 34 deletions(-) diff --git a/src/btrfs_lookup.rs b/src/btrfs_lookup.rs index 32fc9f1..86f31cb 100644 --- a/src/btrfs_lookup.rs +++ b/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, 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(image: &[u8], addr: &T, root_addr_log: u64, key: Key) -> Result { - let node = Node::parse(node_at_log(image, addr, root_addr_log)?)?; - - match node { - 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) - }, - 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) - ) - } - } -} 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 { + 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; + 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, addr) + ) + } + } + } + pub fn find_key(&self, key: Key) -> Result { - 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 { - let node_data = node_at_log(tree.image, tree.addr_map.deref(), addr)?; - match Node::parse(node_data)? { + match tree.get_node(addr)?.deref() { Node::Interior(intnode) => { get_first_item(tree, intnode.children[0].ptr) }, @@ -176,9 +185,10 @@ fn get_first_item(tree: &Tree, addr: u64) -> Result { } } +/// 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 { - let node_data = node_at_log(tree.image, tree.addr_map.deref(), addr)?; - match Node::parse(node_data)? { + match tree.get_node(addr)?.deref() { 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 { } } +#[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 { fn find_closest_key(tree: &Tree, key: Key, mode: FindKeyMode) -> Result, 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 = None; let mut next: Option = 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 { // 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| {