refactor reading nodes to prepare caching

This commit is contained in:
Florian Stecker 2024-02-07 09:31:44 -05:00
parent 4397a02c4e
commit e2fb0cbb47

View File

@ -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::btrfs_structs::{Leaf, Key, Item, InteriorNode, Node, ParseError, ParseBin, Value, Superblock, ItemType, ZERO_KEY, DirItem};
use crate::addrmap::{node_at_log, LogToPhys, AddressMap}; 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. /// and handles the tree traversal and the virtual address translation.
pub struct Tree<'a> { pub struct Tree<'a> {
pub image: &'a [u8], pub image: &'a [u8],
@ -39,6 +39,15 @@ impl<'a> Tree<'a> {
Ok(Tree { image, addr_map, root_addr_log: superblock.root }) 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 *****/ /***** looking up keys *****/
@ -76,30 +85,31 @@ impl InteriorNode {
} }
} }
impl Tree<'_> {
/// Recursively traverse a tree to find a key, given they key and logical address /// 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. /// 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> { fn find_key_in_node(&self, addr: u64, key: Key) -> Result<Item, ParseError> {
let node = Node::parse(node_at_log(image, addr, root_addr_log)?)?; let node = self.get_node(self.root_addr_log)?;
match node { match node.deref() {
Node::Interior(interior_node) => { Node::Interior(interior_node) => {
let next_node_index = interior_node.find_key_or_previous(key).unwrap(); let next_node_index = interior_node.find_key_or_previous(key).unwrap();
let next_node_log = interior_node.children[next_node_index].ptr; 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) => { Node::Leaf(leaf) => {
leaf.find_key(key).ok_or( leaf.find_key(key).ok_or(
error!( error!(
"Item with key ({},{:?},{}) was not found in the leaf at logical address 0x{:x}", "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> { 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> { 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) => { Node::Interior(intnode) => {
get_first_item(tree, intnode.children[0].ptr) 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> { 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) => { Node::Interior(intnode) => {
get_last_item(tree, intnode.children.last().unwrap().ptr) 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 /// 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: /// 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 /// 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> { 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 // 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 prev: Option<u64> = None;
let mut next: 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 { loop {
match Node::parse(node_data)? { let node = tree.get_node(current)?;
match node.deref() {
Node::Interior(intnode) => { Node::Interior(intnode) => {
match intnode.find_key_or_previous(key) { match intnode.find_key_or_previous(key) {
Some(idx) => { Some(idx) => {
@ -213,7 +226,7 @@ fn find_closest_key(tree: &Tree, key: Key, mode: FindKeyMode) -> Result<Option<I
next = Some(kp.ptr); 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 => { None => {
// this can only happen if every key in the current node is `> key` // 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)"); .expect("file system should be consistent (or this is a bug)");
if let Some(item) = &result { 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| { let end_filter = |item : &Item| {