From 577abca2db42dffc8f404267ae31a105d72de1b2 Mon Sep 17 00:00:00 2001
From: Janis <janis@nirgendwo.xyz>
Date: Wed, 5 Apr 2023 17:32:33 +0200
Subject: [PATCH] entry api

---
 btrfs/src/v2/tree.rs | 144 +++++++++++++++++++++++++++++++++++++++----
 1 file changed, 132 insertions(+), 12 deletions(-)

diff --git a/btrfs/src/v2/tree.rs b/btrfs/src/v2/tree.rs
index b87cc66..0fab4f0 100644
--- a/btrfs/src/v2/tree.rs
+++ b/btrfs/src/v2/tree.rs
@@ -166,6 +166,7 @@ impl Display for NodeHandle {
     }
 }
 
+/// range of nodes that iterates thru all leaf nodes from start to end, inclusively.
 #[derive(Debug)]
 pub struct Range<'tree, R: super::Read> {
     volume: Rc<Volume<R>>,
@@ -174,6 +175,66 @@ pub struct Range<'tree, R: super::Read> {
     phantom: PhantomData<&'tree ()>,
 }
 
+pub mod entry {
+    use super::*;
+
+    #[derive(Debug, PartialEq, Eq)]
+    pub enum Entry<'tree> {
+        Occupied(OccupiedEntry<'tree>),
+        Vacant(VacantEntry<'tree>),
+    }
+
+    impl<'tree> From<VacantEntry<'tree>> for Entry<'tree> {
+        fn from(v: VacantEntry<'tree>) -> Self {
+            Self::Vacant(v)
+        }
+    }
+
+    impl<'tree> From<OccupiedEntry<'tree>> for Entry<'tree> {
+        fn from(v: OccupiedEntry<'tree>) -> Self {
+            Self::Occupied(v)
+        }
+    }
+
+    #[derive(Debug, PartialEq, Eq)]
+    pub struct OccupiedEntry<'tree> {
+        key: Key,
+        node: NodeHandle,
+        phantom: PhantomData<&'tree ()>,
+    }
+
+    impl<'tree> OccupiedEntry<'tree> {
+        pub fn new(key: Key, node: NodeHandle) -> Self {
+            Self {
+                key,
+                node,
+                phantom: PhantomData,
+            }
+        }
+
+        pub fn key(&self) -> Key {
+            self.key
+        }
+
+        pub fn item(&self) -> Result<TreeItem> {
+            Ok(self.node.parse_item()?.1)
+        }
+    }
+
+    #[derive(Debug, PartialEq, Eq)]
+    pub struct VacantEntry<'tree> {
+        phantom: PhantomData<&'tree ()>,
+    }
+
+    impl<'tree> VacantEntry<'tree> {
+        pub fn new() -> Self {
+            Self {
+                phantom: PhantomData,
+            }
+        }
+    }
+}
+
 #[derive(Derivative)]
 #[derivative(Debug)]
 pub struct Tree<R: super::Read> {
@@ -250,12 +311,36 @@ impl<R: super::Read> Tree<R> {
         self.find_node_inner(key, NodeHandle::find_key)
     }
 
+    pub fn entry<K>(&self, key: &K) -> Result<entry::Entry>
+    where
+        K: PartialEq<Key> + PartialOrd<Key>,
+    {
+        let entry: entry::Entry = match self.find_node(key)? {
+            Some(node) => entry::OccupiedEntry::new(node.parse_key(), node).into(),
+            None => entry::VacantEntry::new().into(),
+        };
+
+        Ok(entry)
+    }
+
+    pub fn entry_rev<K>(&self, key: &K) -> Result<entry::Entry>
+    where
+        K: PartialEq<Key> + PartialOrd<Key>,
+    {
+        let entry: entry::Entry = match self.find_node_rev(key)? {
+            Some(node) => entry::OccupiedEntry::new(node.parse_key(), node).into(),
+            None => entry::VacantEntry::new().into(),
+        };
+
+        Ok(entry)
+    }
+
     pub fn find_key<K>(&self, key: &K) -> Result<Option<(Item, TreeItem)>>
     where
         K: PartialEq<Key> + PartialOrd<Key>,
     {
         match self.find_node(key)? {
-            Some(node) => node.parse_item(),
+            Some(node) => node.parse_item().map(|item| Some(item)),
             None => Ok(None),
         }
     }
@@ -265,7 +350,7 @@ impl<R: super::Read> Tree<R> {
         K: PartialEq<Key> + PartialOrd<Key>,
     {
         match self.find_node_rev(key)? {
-            Some(node) => node.parse_item(),
+            Some(node) => node.parse_item().map(|item| Some(item)),
             None => Ok(None),
         }
     }
@@ -556,19 +641,43 @@ impl Node {
         Ok(Self { inner, bytes })
     }
 
+    /// returns Ok(None) if `i` is out of bounds
+    fn read_nth_key(&self, i: usize) -> Option<Key> {
+        match &self.inner {
+            BTreeNode::Internal(internal) => {
+                let item = internal
+                    .children
+                    .get(i)
+                    .map(|child| *unsafe { &*child.as_ptr() }.key());
+                item
+            }
+            BTreeNode::Leaf(leaf) => {
+                let key = leaf.items.get(i).map(|item| item.key);
+
+                key
+            }
+        }
+    }
+
+    /// returns None if pointing at an internal node or `i` is out of bounds.
+    /// returns an Error if parsing the item failed.
     pub fn read_nth_item(&self, i: usize) -> Result<Option<(Item, TreeItem)>> {
         match &self.inner {
             BTreeNode::Internal(_) => Ok(None),
             BTreeNode::Leaf(leaf) => {
                 // TODO: better error to indicate that it was out of bounds
-                let item = leaf.items.get(i).ok_or(Error::ReadFailed)?;
-                let start = size_of::<Header>() + item.offset.get() as usize;
-                let size = item.size.get() as usize;
-                let bytes = &self.bytes[start..start + size];
+                let item = if let Some(item) = leaf.items.get(i) {
+                    let start = size_of::<Header>() + item.offset.get() as usize;
+                    let size = item.size.get() as usize;
+                    let bytes = &self.bytes[start..start + size];
 
-                let value = TreeItem::parse(item, bytes)?;
+                    let value = TreeItem::parse(item, bytes)?;
+                    Some((*item, value))
+                } else {
+                    None
+                };
 
-                Ok(Some((*item, value)))
+                Ok(item)
             }
         }
     }
@@ -613,7 +722,7 @@ where
                 });
 
                 if self.start.node.inner.is_leaf() {
-                    break self.start.as_handle().parse_item().expect("range item");
+                    break Some(self.start.as_handle().parse_item().expect("range item"));
                 }
                 // else repeat
             } else {
@@ -638,7 +747,7 @@ where
                 });
 
                 if self.end.node.inner.is_leaf() {
-                    break self.end.as_handle().parse_item().expect("range item");
+                    break Some(self.end.as_handle().parse_item().expect("range item"));
                 }
                 // else repeat
             } else {
@@ -658,8 +767,19 @@ impl NodeHandle {
         }
     }
 
-    pub fn parse_item(&self) -> Result<Option<(Item, TreeItem)>> {
-        self.node.read_nth_item(self.idx as usize)
+    /// returns None if pointing at an internal node or `i` is out of bounds.
+    /// returns an Error if parsing the item failed.
+    pub fn parse_item(&self) -> Result<(Item, TreeItem)> {
+        self.node
+            .read_nth_item(self.idx as usize)
+            .map(|result| result.unwrap())
+    }
+
+    /// returns an Error if the key read fails
+    pub fn parse_key(&self) -> Key {
+        self.node
+            .read_nth_key(self.idx as usize)
+            .expect("idx out of bounds")
     }
 
     pub fn advance_sideways(self) -> NodeHandleAdvanceResult {