btrfs_explorer/src/btrfs_lookup.rs

use std::convert::identity;
use std::rc::Rc;
use std::sync::Arc;
use std::ops::{Deref, RangeBounds, Bound};

use crate::btrfs_structs::{Leaf, Key, Item, InteriorNode, Node, ParseError, ParseBin, Value, Superblock, ItemType, ZERO_KEY};
use crate::addrmap::{node_at_log, AddressMap};

/// 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],
	pub addr_map: Arc<AddressMap>,
	pub root_addr_log: u64,
}

impl<'a> Tree<'a> {
	pub fn new<T: Into<u64>>(image: &'a [u8], tree_id: T) -> Result<Tree<'a>, ParseError> {
		let addr_map = Arc::new(AddressMap::new(image)?);
		let superblock = Superblock::parse(&image[0x10000..])?;

		let root_tree = Tree {
			image,
			addr_map: Arc::clone(&addr_map),
			root_addr_log: superblock.root
		};
		let tree_root_item = root_tree.find_key(Key::new(tree_id.into(), ItemType::Root, 0))?;

		let root_addr_log = match tree_root_item.value {
			Value::Root(root) => root.bytenr,
			_ => return Err("root item invalid".into())
		};

		Ok(Tree { image, addr_map, root_addr_log })
	}

	pub fn root(image: &'a [u8]) -> Result<Tree<'a>, ParseError> {
		let addr_map = Arc::new(AddressMap::new(image)?);
		let superblock = Superblock::parse(&image[0x10000..])?;

		Ok(Tree { image, addr_map, root_addr_log: superblock.root })
	}

	pub fn chunk(image: &'a [u8]) -> Result<Tree<'a>, ParseError> {
		let addr_map = Arc::new(AddressMap::new(image)?);
		let superblock = Superblock::parse(&image[0x10000..])?;

		Ok(Tree { image, addr_map, root_addr_log: superblock.chunk_root })
	}

	/// Read a node given its logical address
	pub fn get_node(&self, addr: u64) -> Result<Arc<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(Arc::new(node))
	}
}

/***** looking up keys *****/

impl Leaf {
	pub fn find_key(&self, key: Key) -> Option<Item> {
		self.items
			.iter()
			.find(|x|x.key == key)
			.map(|x|x.clone())
	}

	pub fn find_key_or_previous(&self, key: Key) -> Option<usize> {
		self.items
			.iter()
			.take_while(|x|x.key <= key)
			.enumerate()
			.last()
			.map(|x|x.0)
	}

}

impl InteriorNode {
	/// Return the index of the last child which has key at most `key`. This is the
	/// branch which contains `key` if it exists. Returns `None` if all children are greater than
	/// `key`, which guarantees that `key` is not among the descendants of `self`.
	pub fn find_key_or_previous(&self, key: Key) -> Option<usize> {
		self.children
			.iter()
			.take_while(|x|x.key <= key)
			.enumerate()
			.last()
			.map(|x|x.0)
	}
}


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;
				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<Item, ParseError> {
		self.find_key_in_node(self.root_addr_log, key)
	}
}

/***** iterator *****/

pub struct RangeIter<'a, 'b> {
	tree: &'b Tree<'a>,

	start: Bound<Key>,
	end: Bound<Key>,
	forward_skip_fn: Box<dyn Fn(Key) -> Key>,
	backward_skip_fn: Box<dyn Fn(Key) -> Key>,
}

impl<'a> Tree<'a> {
	/// Given a tree, a range of indices, and two "skip functions", produces a double
	/// ended iterator which iterates through the keys contained in the range, in ascending
	/// or descending order.
    ///
	/// The skip functions are ignored for now, but are intended as an optimization:
	/// after a key `k` was returned by the iterator (or the reverse iterator), all keys
	/// strictly lower than `forward_skip_fn(k)` are skipped (resp. all keys strictly above
	/// `backward_skip_fn(k)` are skipped.
    ///
	/// If `forward_skip_fn` and `backward_skip_fn` are the identity, nothing is skipped
	pub fn range_with_skip<'b, R, F1, F2>(&'b self, range: R, forward_skip_fn: F1, backward_skip_fn: F2) -> RangeIter<'a, 'b>
	where
		R: RangeBounds<Key>,
		F1: Fn(Key) -> Key + 'static,
		F2: Fn(Key) -> Key + 'static {
		RangeIter {
			tree: self,
			start: range.start_bound().cloned(),
			end: range.end_bound().cloned(),
			forward_skip_fn: Box::new(forward_skip_fn),
			backward_skip_fn: Box::new(backward_skip_fn),
		}
	}

	pub fn range<'b, R: RangeBounds<Key>>(&'b self, range: R) -> RangeIter<'a, 'b> {
		RangeIter {
			tree: self,
			start: range.start_bound().cloned(),
			end: range.end_bound().cloned(),
			forward_skip_fn: Box::new(identity),
			backward_skip_fn: Box::new(identity),
		}
	}


	pub fn iter<'b>(&'b self) -> RangeIter<'a, 'b> {
		RangeIter {
			tree: self,
			start: Bound::Unbounded,
			end: Bound::Unbounded,
			forward_skip_fn: Box::new(identity),
			backward_skip_fn: Box::new(identity),
		}
	}
}


/// 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> {
	match tree.get_node(addr)?.deref() {
		Node::Interior(intnode) => {
			get_first_item(tree, intnode.children[0].ptr)
		},
		Node::Leaf(leafnode) => {
			Ok(leafnode.items[0].clone())
		},
	}
}

/// 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> {
	match tree.get_node(addr)?.deref() {
		Node::Interior(intnode) => {
			get_last_item(tree, intnode.children.last().unwrap().ptr)
		},
		Node::Leaf(leafnode) => {
			Ok(leafnode.items.last().unwrap().clone())
		},
	}
}

#[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
/// key less than / greater than / greater or equal to / less or equal to `key`.
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; 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;

	loop {
		let node = tree.get_node(current)?;
		match node.deref() {
			Node::Interior(intnode) => {
				match intnode.find_key_or_previous(key) {
					Some(idx) => {
						if let Some(kp) = (idx > 0).then(|| intnode.children.get(idx-1)).flatten() {
							prev = Some(kp.ptr);
						}
						if let Some(kp) = intnode.children.get(idx+1) {
							next = Some(kp.ptr);
						}

						current = intnode.children[idx].ptr;
					},
					None => {
						// this can only happen if every key in the current node is `> key`
						// which really should only happen if we're in the root node, as otherwise
						// we wouldn't have descended into this branch; so assume every key in the
						// tree is above `> key`.
						if mode == FindKeyMode::LT || mode == FindKeyMode::LE {
							return Ok(None);
						} else {
							// return the first item in tree; we are an interior node so we really should have
							// at least one child
							let addr = intnode.children[0].ptr;
							return Ok(Some(get_first_item(tree, addr)?));
						}
					}
				}
			},
			Node::Leaf(leafnode) => {
				match leafnode.find_key_or_previous(key) {
					Some(idx) => {
						// the standard case, we found a key `k` with the guarantee that `k <= key`
						let Item {key: k, value: v} = leafnode.items[idx].clone();

						if mode == FindKeyMode::LE || mode == FindKeyMode::LT && k < key || mode == FindKeyMode::GE && k == key {
							return Ok(Some(Item {key: k, value: v}))
						} else if mode == FindKeyMode::LT && k == key {
							// prev
							if idx > 0 {
								return Ok(Some(leafnode.items[idx-1].clone()));
							} else {
								// use prev
								if let Some(addr) = prev {
									return Ok(Some(get_last_item(tree, addr)?));
								} else {
									return Ok(None);
								}
							}
						} else {
							// next
							if let Some(item) = leafnode.items.get(idx+1) {
								return Ok(Some(item.clone()));
							} else {
								// use next
								if let Some(addr) = next {
									return Ok(Some(get_first_item(tree, addr)?));
								} else {
									return Ok(None);
								}
							}
						}
					},
					None => {
						// same as above, but this can only happen if the root node is a leaf
						if mode == FindKeyMode::LT || mode == FindKeyMode::LE {
							return Ok(None);
						} else {
							// return the first item in tree if it exists
							return Ok(leafnode.items.get(0).map(|x|x.clone()));
						}
					},
				}
			},
		}
	}
}

fn range_valid<T: Ord>(start: Bound<T>, end: Bound<T>) -> bool {
	match (start, end) {
		(Bound::Included(x), Bound::Included(y)) => x <= y,
		(Bound::Excluded(x), Bound::Included(y)) => x < y,
		(Bound::Included(x), Bound::Excluded(y)) => x < y,
		(Bound::Excluded(x), Bound::Excluded(y)) => x < y, // could technically be empty if "y = x+1", but we can't check
		(_, _) => true, // one of them is unbounded
	}
}

impl<'a, 'b> Iterator for RangeIter<'a, 'b> {
	type Item = Item;

	fn next(&mut self) -> Option<Item> {
		if !range_valid(self.start.as_ref(), self.end.as_ref()) {
			return None;
		}

		let (start_key, mode) : (Key, FindKeyMode) = match &self.start {
			&Bound::Included(x) => (x, FindKeyMode::GE),
			&Bound::Excluded(x) => (x, FindKeyMode::GT),
			&Bound::Unbounded => (ZERO_KEY, FindKeyMode::GE),
		};

		// FIX: proper error handling
		let result = find_closest_key(self.tree, start_key, mode)
			.expect("file system should be consistent (or this is a bug)");

		if let Some(item) = &result {
			self.start = Bound::Excluded((self.forward_skip_fn)(item.key));
		}

		let end_filter = |item : &Item| {
			match &self.end {
				&Bound::Included(x) => item.key <= x,
				&Bound::Excluded(x) => item.key < x,
				&Bound::Unbounded => true,
			}
		};

		result
			.filter(end_filter)
			.map(|item|item.clone())
	}
}