last somewhat working version before reimplementing tree traversal

src/btrfs_lookup.rs

@@ -1,9 +1,11 @@
 use std::rc::Rc;
-use std::ops::Deref;
+use std::ops::{Deref, RangeBounds};
 
 use crate::btrfs_structs::{Leaf, Key, Item, InteriorNode, Node, ParseError, ParseBin, Value, Superblock, ItemType};
 use crate::addrmap::{node_at_log, LogToPhys, 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: Rc<AddressMap>,
@@ -16,7 +18,7 @@ impl<'a> Tree<'a> {
 		let superblock = Superblock::parse(&image[0x10000..])?;
 
 		let root_tree = Tree {
-			image: image,
+			image,
 			addr_map: Rc::clone(&addr_map),
 			root_addr_log: superblock.root
 		};
@@ -94,7 +96,7 @@ impl Tree<'_> {
 
 /***** iterator *****/
 
-pub struct Iter<'a> {
+pub struct RangeIter<'a, R: RangeBounds<Key>, F: Fn(Key) -> Key = fn(Key) -> Key> {
 	tree: &'a Tree<'a>,
 
 	// path to the last returned item
@@ -102,27 +104,28 @@ pub struct Iter<'a> {
 	leaf: Option<Box<Leaf>>,
 	indices: Vec<usize>,
 
-	lower_limit: Option<Key>,
-	upper_limit: Option<Key>,
+	bounds: R,
+	skip_fn: F,
 }
 
 impl Tree<'_> {
-	pub fn iter<'a>(&'a self) -> Iter<'a> {
+	pub fn iter<'a>(&'a self) -> RangeIter<'a> {
 		self.range(None, None)
 	}
 
-	pub fn range<'a>(&'a self, lower: Option<Key>, upper: Option<Key>) -> Iter<'a> {
-		Iter {
+	pub fn range<'a>(&'a self, lower: Option<Key>, upper: Option<Key>) -> RangeIter<'a> {
+		RangeIter {
 			tree: self,
 			nodes: Vec::new(),
 			leaf: None,
 			indices: Vec::new(), // in nodes and leaf
 			lower_limit: lower,
 			upper_limit: upper,
+			skip_fn: |x|x
 		}
 	}
 
-	pub fn range_id<'a>(&'a self, id: u64) -> Iter<'a> {
+	pub fn range_id<'a>(&'a self, id: u64) -> RangeIter<'a> {
 		if id == u64::MAX {
 			self.range(
 				Some(Key::new(id, ItemType::Invalid, 0)),
@@ -135,9 +138,29 @@ impl Tree<'_> {
 			)
 		}
 	}
+
+	/// 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` are skipped.
+	pub fn range_with_skip<'a, R, F>(&'a self, range: R, forward_skip_fn: F, backward_skip_fn: F) -> RangeIter<'a, F>
+	where
+		R: RangeBounds<Key>,
+		F: Fn(Key) -> Key {
+		RangeIter {
+			tree: self,
+			nodes: Vec::new(),
+			leaf: None,
+			indices: Vec::new(),
+		}
+	}
 }
 
-impl Iter<'_> {
+impl<F: Fn(Key) -> Key> RangeIter<'_, F> {
 	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()?;
@@ -190,10 +213,9 @@ impl Iter<'_> {
 			}
 		}
 	}
-
 }
 
-impl Iterator for Iter<'_> {
+impl<F: Fn(Key) -> Key> Iterator for RangeIter<'_, F> {
 	type Item = Item;
 
 	// for now we just silently stop when we encounter an error, maybe that isn't the best solution
@@ -267,7 +289,7 @@ impl Iterator for Iter<'_> {
 				};
 
 				// first first item under this node
-				self.move_down_and_get_first_item(node_addr)
+				return self.move_down_and_get_first_item(node_addr)
 					.filter(|item|self.upper_limit.is_none() || item.key < self.upper_limit.unwrap())
 			}
 		}

src/btrfs_structs.rs

@@ -63,6 +63,16 @@ pub struct Key {
 	pub key_offset: u64,
 }
 
+impl Key {
+	pub fn new(key_id: u64, key_type: ItemType, key_offset: u64) -> Key {
+		Key { key_id, key_type, key_offset }
+	}
+
+	pub fn id(key_id: u64) -> Key {
+		Key { key_id, key_type: ItemType::Invalid, key_offset: 0 }
+	}
+}
+
 #[allow(unused)]
 #[derive(Debug,Clone)]
 pub enum Value {
@@ -709,8 +719,16 @@ impl fmt::Debug for Checksum {
 	}
 }
 
-impl Key {
-	pub fn new(key_id: u64, key_type: ItemType, key_offset: u64) -> Key {
-		Key { key_id, key_type, key_offset }
-	}
+macro_rules! key {
+	($arg1:expr) => {
+		btrfs_structs::Key { key_id: $arg1, key_type: btrfs_structs::ItemType::Invalid, key_offset: 0 }
+	};
+	($arg1:expr, $arg2:expr) => {
+		btrfs_structs::Key { key_id: $arg1, key_type: $arg2, key_offset: 0 }
+	};
+	($arg1:expr, $arg2:expr, $arg3:expr) => {
+		btrfs_structs::Key { key_id: $arg1, key_type: $arg2, key_offset: $arg3 }
+	};
 }
+
+pub(crate) use key;

src/main.rs

@@ -1,17 +1,15 @@
-use memmap2::{Mmap, MmapOptions};
-use std::fs::File;
-use rouille::Request;
-use rouille::Response;
-use rouille::router;
-use std::iter;
-use std::env;
-use std::{fs::OpenOptions, os::unix::fs::OpenOptionsExt};
-use std::collections::HashMap;
-
-use parsebtrfs::btrfs_structs::{TreeID, Value::Extent, Value::BlockGroup, ParseError, NODE_SIZE, ItemType, Node, ParseBin};
-use parsebtrfs::btrfs_lookup::Tree;
-
-use parsebtrfs::addrmap::{AddressMap, LogToPhys};
+use std::{
+	iter,
+	env,
+	fs::OpenOptions,
+	collections::HashMap,
+};
+use memmap2::Mmap;
+use rouille::{Request, Response, router};
+use parsebtrfs::{
+	btrfs_structs::{TreeID, Value::Extent, Value::BlockGroup, ParseError, NODE_SIZE, ItemType},
+	btrfs_lookup::Tree,
+};
 
 const COLORS: &[&str] = &["#e6194b", "#3cb44b", "#ffe119", "#4363d8", "#f58231", "#911eb4", "#46f0f0", "#f032e6", "#bcf60c", "#fabebe", "#008080", "#e6beff", "#9a6324", "#fffac8", "#800000", "#aaffc3", "#808000", "#ffd8b1", "#000075", "#808080", "#000000"];
 

src/test.rs

@@ -1,5 +1,9 @@
 use super::*;
-use btrfs_structs::{Key, ItemType};
+use btrfs_structs::key;
+use btrfs_structs::{Key, ItemType, TreeID};
+use btrfs_lookup::Tree;
+use memmap2::Mmap;
+use std::fs::OpenOptions;
 
 #[test]
 fn test_key_new() {
@@ -8,3 +12,17 @@ fn test_key_new() {
 		Key { key_id: 1, key_type: ItemType::Root, key_offset: 2 }
 	);
 }
+
+#[test]
+fn image_test() {
+	let file = OpenOptions::new().read(true).open("../image").unwrap();
+	let image = unsafe { Mmap::map(&file).unwrap() };
+
+	let extent_tree = Tree::new(&image, TreeID::Extent).unwrap();
+
+	let skip = |Key { key_id, key_type, key_offset }| Key { key_id: key_id + 0x10000, key_type, key_offset };
+
+	for item in extent_tree.range_with_skip(Some(key!(0x1d70000)), Some(key!(0x2000000)), skip) {
+		println!("{item:x?}");
+	}
+}