btrfs_explorer/src/btrfs_lookup.rs

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

use crate::btrfs_structs::{Leaf, Key, Item, InteriorNode, Node, ParseError, ParseBin, Value, Superblock, ItemType, ZERO_KEY, LAST_KEY};
use crate::nodereader::NodeReader;

/// 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 reader: Rc<NodeReader<'a>>,
	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 superblock = Superblock::parse(&image[0x10000..])?;
		let reader = Rc::new(NodeReader::new(image)?);

		let root_tree = Tree {
			image,
			reader: Rc::clone(&reader),
			root_addr_log: superblock.root
		};

		//		let tree_root_item = root_tree.find_key(Key::new(tree_id.into(), ItemType::Root, 0))?;
		let tree_id = tree_id.into();
		let root_item_key = Key::new(tree_id, ItemType::Root, 0);
		let tree_root_item = root_tree.range(root_item_key..)
			.next()
			.filter(|x| x.key.key_id == tree_id && x.key.key_type == ItemType::Root);

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

		Ok(Tree { image, reader: Rc::clone(&reader), root_addr_log })
	}

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

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

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

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

/***** 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> {
		// if the key is not exactly matched, binary_search returns the next index, but we want the previous one
		match self.children.binary_search_by_key(&key, |x|x.key) {
			Ok(idx) => Some(idx),
			Err(idx) if idx == 0 => None,
			Err(idx) => Some(idx-1),
		}
	}
}


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.reader.get_node(addr)?;

		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.reader.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.reader.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.reader.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())
	}
}

impl<'a, 'b> DoubleEndedIterator for RangeIter<'a, 'b> {
	fn next_back(&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.end {
			&Bound::Included(x) => (x, FindKeyMode::LE),
			&Bound::Excluded(x) => (x, FindKeyMode::LT),
			&Bound::Unbounded => (LAST_KEY, FindKeyMode::LE),
		};

		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.end = Bound::Excluded((self.backward_skip_fn)(item.key));
		}

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

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