btrfs_explorer/src/btrfs_lookup.rs

use std::rc::Rc;
use std::ops::Deref;

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

pub struct Tree<'a> {
	pub image: &'a [u8],
	pub addr_map: Rc<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 = Rc::new(AddressMap::new(image)?);
		let superblock = Superblock::parse(&image[0x10000..])?;

		let root_tree = Tree {
			image: image,
			addr_map: Rc::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 = Rc::new(AddressMap::new(image)?);
		let superblock = Superblock::parse(&image[0x10000..])?;

		Ok(Tree { image, addr_map, root_addr_log: superblock.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<Item> {
		self.items
			.iter()
			.take_while(|x|x.key <= key)
			.last()
			.map(|x|x.clone())
	}

}

impl InteriorNode {
	pub fn find_key_or_previous(&self, key: Key) -> Option<u64> {
		self.children
			.iter()
			.take_while(|x|x.key <= key)
			.last()
			.map(|x|x.ptr)
	}
}

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 {
		Node::Interior(interior_node) => {
			let next_node_log = interior_node.find_key_or_previous(key).unwrap();
			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<'_> {
	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)
	}
}

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

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

	// path to the last returned item
	nodes: Vec<InteriorNode>,
	leaf: Option<Box<Leaf>>,
	indices: Vec<usize>,

	lower_limit: Option<Key>,
	upper_limit: Option<Key>,
}

impl Tree<'_> {
	pub fn iter<'a>(&'a self) -> Iter<'a> {
		self.range(None, None)
	}

	pub fn range<'a>(&'a self, lower: Option<Key>, upper: Option<Key>) -> Iter<'a> {
		Iter {
			tree: self,
			nodes: Vec::new(),
			leaf: None,
			indices: Vec::new(), // in nodes and leaf
			lower_limit: lower,
			upper_limit: upper,
		}
	}

	pub fn range_id<'a>(&'a self, id: u64) -> Iter<'a> {
		if id == u64::MAX {
			self.range(
				Some(Key::new(id, ItemType::Invalid, 0)),
				None
			)
		} else {
			self.range(
				Some(Key::new(id, ItemType::Invalid, 0)),
				Some(Key::new(id+1, ItemType::Invalid, 0))
			)
		}
	}
}

impl Iter<'_> {
	fn move_down_and_get_first_item(&mut self, mut node_addr: u64) -> Option<Item> {
		loop {
			let node = Node::parse(node_at_log(self.tree.image, self.tree.addr_map.deref(), node_addr).ok()?).ok()?;
			match node {
				Node::Interior(int_node) => {
					node_addr = int_node.children.first()?.ptr;
					self.nodes.push(int_node);
					self.indices.push(0);
				},
				Node::Leaf(leaf_node) => {
					let result = leaf_node.items.first()?.clone();
					self.leaf = Some(Box::new(leaf_node));
					self.indices.push(0);
					return Some(result);
				},
			}
		}
	}

	fn move_down_and_get_item_or_previous(&mut self, mut node_addr: u64, key: Key) -> Option<Item> {
		loop {
			let node = Node::parse(node_at_log(self.tree.image, self.tree.addr_map.deref(), node_addr).ok()?).ok()?;

			match node {
				Node::Interior(int_node) => {
					let (i, new_node_ptr) = int_node
						.children
						.iter()
						.enumerate()
						.take_while(|(_,bp)|bp.key <= key)
						.last()?;

					node_addr = new_node_ptr.ptr;
					self.nodes.push(int_node);
					self.indices.push(i);
				},
				Node::Leaf(leaf_node) => {
					let (i, result) = leaf_node
						.items
						.iter()
						.enumerate()
						.take_while(|(_,item)|item.key <= key)
						.last()?;

					let result_cloned = result.clone();
					self.leaf = Some(Box::new(leaf_node));
					self.indices.push(i);
					return Some(result_cloned);
				},
			}
		}
	}

}

impl Iterator for Iter<'_> {
	type Item = Item;

	// for now we just silently stop when we encounter an error, maybe that isn't the best solution
	fn next(&mut self) -> Option<Item> {
		if self.leaf.is_none() && self.nodes.len() == 0 {
			// first item
			// finding the first item is a bit tricky
			// if there is a lower limit, the B+ tree only allows us to either find the item
			// or the previous one if there is no exact match; in the latter case, go one further

			let result = if let Some(lim) = self.lower_limit {
				let first_res = self.move_down_and_get_item_or_previous(self.tree.root_addr_log, lim);
				if let Some(item) = first_res {
					if item.key == lim {
						// found exactly the limit, that's the easy case
						Some(item)
					} else {
						// found a previous item; so we want the next one
						self.next()
					}
				} else {
					// did not find an item, so everything must come after lower limit
					// just get the first
					self.move_down_and_get_first_item(self.tree.root_addr_log)
				}
			} else {
				// there is no lower limit, so also just get the first
				self.move_down_and_get_first_item(self.tree.root_addr_log)
			};

			result.filter(|item|self.upper_limit.is_none() || item.key < self.upper_limit.unwrap())
		} else if self.leaf.is_none() {
			// already through the iterator
			return None;
		} else {
			let height = self.indices.len(); // must be at least 1
			let leaf = self.leaf.as_ref().unwrap();

			self.indices[height-1] += 1;
			if let Some(item) = leaf.items.get(self.indices[height-1]) {
				// there's a next item in the same leaf
				if self.upper_limit.is_none() || item.key < self.upper_limit.unwrap() {
					return Some(item.clone());
				} else {
					return None;
				}
			} else if height == 1 {
				// the tree has height 1 and we're through the (only) leaf, there's nothing left
				return None;
			} else {
				// try to advance in one of the higher nodes
				self.leaf = None;
				self.indices.pop();
				let mut level = height - 2;

				// go up until we can move forward in a node
				let node_addr = loop {
					let node = &self.nodes[level];

					self.indices[level] += 1;
					if let Some(blockptr) = node.children.get(self.indices[level]) {
						break blockptr.ptr;
					} else {
						if level == 0 {
							return None;
						}
						self.indices.pop();
						self.nodes.pop();
						level -= 1;
					}
				};

				// first first item under this node
				self.move_down_and_get_first_item(node_addr)
					.filter(|item|self.upper_limit.is_none() || item.key < self.upper_limit.unwrap())
			}
		}
	}
}