bootloader/btrfs/src/v2/tree.rs

use alloc::boxed::Box;
use core::cell::RefCell;
use core::fmt::Display;
use core::marker::PhantomData;
use core::mem::size_of;
use core::ops::Deref;
use core::ptr::NonNull;

use crate::structs::{Header, Item, Key, KeyPtr, KnownObjectId, ObjectType, TreeItem};
use crate::{Error, Result};
use alloc::rc::Rc;
use alloc::vec::Vec;
use derivative::Derivative;
use scroll::Pread;
use zerocopy::FromBytes;

use super::volume::Volume;

/// A leaf node in a btrfs tree, containing different items
#[derive(Debug, Clone)]
pub struct BTreeLeafNode {
    pub header: Header,
    /// actual leaf data
    pub items: Vec<Item>,
}

#[derive(Debug)]
pub enum NodePtr {
    Unvisited(KeyPtr),
    Visited { key: KeyPtr, node: BoxedNode }, // TODO: this doesnt need to be an Rc, can just be a NonNull with manual memory management
}

impl Clone for NodePtr {
    fn clone(&self) -> Self {
        match self {
            Self::Unvisited(arg0) => Self::Unvisited(arg0.clone()),
            Self::Visited { key, node } => Self::Visited {
                key: key.clone(),
                node: Node::clone_from_nonnull(node),
            },
        }
    }
}

impl NodePtr {
    pub fn key_ptr(&self) -> &KeyPtr {
        match self {
            NodePtr::Unvisited(key) => key,
            NodePtr::Visited { key, .. } => key,
        }
    }

    pub fn node(&self) -> Option<&BoxedNode> {
        match self {
            NodePtr::Unvisited(_) => None,
            NodePtr::Visited { node, .. } => Some(&node),
        }
    }

    pub fn key(&self) -> &Key {
        &self.key_ptr().key
    }

    pub fn destroy(self) {
        match self {
            Self::Visited { node, .. } => {
                _ = node;
                // TODO: from box drop
            }
            _ => {}
        }
    }
}

/// An internal node in a btrfs tree, containing `KeyPtr`s to other internal nodes or leaf nodes.
#[derive(Derivative, Clone)]
#[derivative(Debug)]
pub struct BTreeInternalNode {
    pub header: Header,
    #[derivative(Debug = "ignore")]
    children: Vec<RefCell<NodePtr>>,
}

impl BTreeInternalNode {
    pub fn visit_child<R: super::Read>(
        &self,
        idx: usize,
        volume: &super::volume::Volume<R>,
    ) -> Result<BoxedNode> {
        match self.children.get(idx) {
            Some(child) => self.visit_child_inner(child, volume),
            None => Err(Error::OutOfBounds {
                range: 0..self.children.len(),
                index: idx,
            }),
        }
    }

    fn visit_child_inner<R: super::Read>(
        &self,
        child: &RefCell<NodePtr>,
        volume: &super::volume::Volume<R>,
    ) -> Result<BoxedNode> {
        match unsafe { &*child.as_ptr() } {
            NodePtr::Unvisited(keyptr) => {
                let node = volume
                    .read_keyptr(keyptr)
                    .and_then(|bytes| Node::boxed_from_bytes(bytes))?;

                child.replace(NodePtr::Visited {
                    key: *keyptr,
                    node: node.clone(),
                });
                Ok(node)
            }
            NodePtr::Visited { node, .. } => Ok(node.clone()),
        }
    }

    pub fn visit_children_keys(
        &self,
    ) -> impl Iterator<Item = (usize, Key)> + DoubleEndedIterator + '_ {
        self.children
            .iter()
            .enumerate()
            .map(|(i, child)| (i, unsafe { *(&*child.as_ptr()).key() }))
    }

    pub fn visit_children<'a, 'b, R: super::Read>(
        &'a self,
        volume: &'b super::volume::Volume<R>,
    ) -> impl Iterator<Item = (usize, Result<BoxedNode>)> + 'a
    where
        'b: 'a,
    {
        self.children
            .iter()
            .enumerate()
            .map(|(i, child)| (i, self.visit_child_inner(child, volume)))
    }
}

impl PartialEq for BTreeInternalNode {
    fn eq(&self, other: &Self) -> bool {
        self.header == other.header
    }
}

impl PartialEq for BTreeLeafNode {
    fn eq(&self, other: &Self) -> bool {
        self.header == other.header
    }
}

impl Eq for BTreeLeafNode {}
impl Eq for BTreeInternalNode {}

#[derive(Debug, PartialEq, Eq, Clone)]
pub enum BTreeNode {
    Internal(BTreeInternalNode),
    Leaf(BTreeLeafNode),
}

#[derive(Derivative, Eq, Clone)]
#[derivative(Debug, PartialEq)]
pub struct Node {
    inner: BTreeNode,
    #[derivative(Debug = "ignore")]
    #[derivative(PartialEq = "ignore")]
    bytes: Vec<u8>,
}

#[repr(transparent)]
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct BoxedNode(NonNull<Node>);

impl Deref for BoxedNode {
    type Target = Node;

    fn deref(&self) -> &Self::Target {
        self.as_ref()
    }
}

impl From<NonNull<Node>> for BoxedNode {
    fn from(value: NonNull<Node>) -> Self {
        Self(value)
    }
}

impl From<Node> for BoxedNode {
    fn from(value: Node) -> Self {
        Self(unsafe { NonNull::new_unchecked(Box::leak(Box::new(value))) })
    }
}

impl BoxedNode {
    pub fn as_ref(&self) -> &Node {
        unsafe { self.0.as_ref() }
    }

    pub fn as_ptr(self) -> *mut Node {
        self.0.as_ptr()
    }
}

//type BoxedNode = NonNull<Node>;

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct NodeHandle {
    parents: Vec<(BoxedNode, u32)>,
    node: BoxedNode,
    idx: u32,
}

impl Display for NodeHandle {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        write!(
            f,
            "NodeHandle {{\n\tindex: {},\n\tnode: {:#?},\n\t ..\n}}\n",
            self.idx,
            self.node.inner.header()
        )
    }
}

/// 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>>,
    pub(crate) start: RootOrEdge,
    pub(crate) end: RootOrEdge,
    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_and_value(&self) -> Result<(Item, TreeItem)> {
            self.node.parse_item()
        }

        pub fn value(&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> {
    #[derivative(Debug = "ignore")]
    volume: Rc<Volume<R>>,
    root: BoxedNode,
}

impl<R: super::Read> Clone for Tree<R> {
    fn clone(&self) -> Self {
        Self {
            volume: self.volume.clone(),
            root: Node::clone_from_nonnull(&self.root),
        }
    }
}

impl<R: super::Read> Drop for Tree<R> {
    fn drop(&mut self) {
        log::debug!("======= cleaning up tree =======");
        Node::destroy(self.root.clone());
        log::debug!("======= [done] =======");
    }
}

impl<R: super::Read> Tree<R> {
    pub fn from_logical_offset(volume: Rc<Volume<R>>, logical: u64) -> Result<Self> {
        // TODO: this might read a very big range, far more than needed
        let bytes = volume
            .read_range_from_logical(logical)?
            .ok_or(Error::BadLogicalAddress)?; // TODO: make this a better error
        let root = Node::boxed_from_bytes(bytes)?;

        Ok(Self { volume, root })
    }

    fn find_node_inner<K, F>(&self, key: &K, find: F) -> Result<Option<NodeHandle>>
    where
        K: PartialEq<Key> + PartialOrd<Key>,
        F: Fn(NodeHandle, &K) -> SearchResult,
    {
        let mut node = NodeHandle::start(self.root.clone());

        loop {
            let search = find(node, key);
            match search {
                SearchResult::Leaf(a) => {
                    return Ok(Some(a));
                }
                SearchResult::Edge(mut edge) => match &edge.node.inner {
                    BTreeNode::Internal(internal) => {
                        let child = internal
                            .visit_child(edge.idx as usize, &self.volume)
                            .expect("child");
                        edge.parents.push((edge.node, edge.idx));
                        node = NodeHandle {
                            parents: edge.parents,
                            node: child,
                            idx: 0,
                        };
                        // recurse
                    }
                    BTreeNode::Leaf(_) => {
                        // leaf node returning and edge means key is not present
                        return Ok(None);
                    }
                },
            }
        }
    }

    fn find_node_rev<K>(&self, key: &K) -> Result<Option<NodeHandle>>
    where
        K: PartialEq<Key> + PartialOrd<Key>,
    {
        self.find_node_inner(key, NodeHandle::find_key_rev)
    }

    fn find_node<K>(&self, key: &K) -> Result<Option<NodeHandle>>
    where
        K: PartialEq<Key> + PartialOrd<Key>,
    {
        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_range<K>(&self, key: &K) -> Result<Range<R>>
    where
        K: PartialEq<Key> + PartialOrd<Key>,
    {
        let start = self.find_node(key)?;
        let end = self.find_node_rev(key)?;

        Ok(Range::from_handles(
            self.volume.clone(),
            start.unwrap_or(NodeHandle::start(self.root.clone())),
            end.unwrap_or(NodeHandle::end(self.root.clone())),
        ))
    }

    pub fn iter(&self) -> Range<'_, R> {
        Range::new(self.volume.clone(), self.root.clone(), self.root.clone())
    }
}

impl BTreeLeafNode {
    pub fn parse(header: Header, bytes: &[u8]) -> Result<Self> {
        log::debug!("leaf:");

        let offset = &mut 0;
        let items = core::iter::from_fn(|| {
            if *offset as usize + size_of::<Item>() < bytes.len() {
                let item = Item::read_from(&bytes[*offset..*offset + size_of::<Item>()]);
                *offset += size_of::<Item>();

                if let Some(item) = item.as_ref() {
                    log::debug!("\t{item:?}");
                }

                item
            } else {
                None
            }
        })
        .take(header.nritems.get() as usize)
        .collect::<Vec<_>>();
        Ok(Self { header, items })
    }
}

impl BTreeInternalNode {
    pub fn parse(header: Header, bytes: &[u8]) -> Result<Self> {
        log::debug!("internal lvl: {}", header.level);

        let offset = &mut 0;
        let size = size_of::<KeyPtr>();
        let children = core::iter::from_fn(|| {
            if *offset as usize + size < bytes.len() {
                let item = KeyPtr::read_from(&bytes[*offset..*offset + size]);
                *offset += size;

                if let Some(item) = item.as_ref() {
                    log::debug!(
                        "\tchild gen: {} offset: {}",
                        item.generation.get(),
                        item.key.offset.get()
                    );
                }

                item
            } else {
                None
            }
        })
        .take(header.nritems.get() as usize)
        .map(|ptr| RefCell::new(NodePtr::Unvisited(ptr)))
        .collect::<Vec<_>>();

        Ok(Self { header, children })
    }
}

impl BTreeNode {
    pub fn parse(bytes: &[u8]) -> Result<Self> {
        let offset = &mut 0;
        let header = bytes.gread::<Header>(offset)?;

        if header.level == 0 {
            Ok(Self::Leaf(BTreeLeafNode::parse(header, &bytes[*offset..])?))
        } else {
            Ok(Self::Internal(BTreeInternalNode::parse(
                header,
                &bytes[*offset..],
            )?))
        }
    }

    pub fn header(&self) -> &Header {
        match self {
            BTreeNode::Internal(node) => &node.header,
            BTreeNode::Leaf(node) => &node.header,
        }
    }

    /// Returns `true` if the btree node is [`Internal`].
    ///
    /// [`Internal`]: BTreeNode::Internal
    #[must_use]
    pub fn is_internal(&self) -> bool {
        matches!(self, Self::Internal(..))
    }

    /// Returns `true` if the btree node is [`Leaf`].
    ///
    /// [`Leaf`]: BTreeNode::Leaf
    #[must_use]
    pub fn is_leaf(&self) -> bool {
        matches!(self, Self::Leaf(..))
    }

    pub fn as_internal(&self) -> Option<&BTreeInternalNode> {
        if let Self::Internal(v) = self {
            Some(v)
        } else {
            None
        }
    }

    pub fn as_leaf(&self) -> Option<&BTreeLeafNode> {
        if let Self::Leaf(v) = self {
            Some(v)
        } else {
            None
        }
    }
}

pub enum NodeHandleAdvanceResult {
    Decend {
        parent: NodeHandle,
        child_ptr: KeyPtr,
    },
    Next(NodeHandle),
    Ascend,
}

pub enum SearchResult {
    Leaf(NodeHandle),
    Edge(NodeHandle),
}

#[derive(Debug, Clone, Eq)]
pub(crate) enum RootOrEdge {
    Root(NodeHandle),
    Edge(NodeHandle),
}

impl RootOrEdge {
    pub fn into_handle(&self) -> NodeHandle {
        match self {
            RootOrEdge::Root(handle) => handle,
            RootOrEdge::Edge(handle) => handle,
        }
        .clone()
    }

    fn into_next_node_any<R, F>(
        self,
        volume: &super::volume::Volume<R>,
        f: F,
    ) -> core::result::Result<Self, Self>
    where
        F: Fn(
            NodeHandle,
            &super::volume::Volume<R>,
        ) -> core::result::Result<NodeHandle, NodeHandle>,
        R: super::Read,
    {
        let node = match self {
            RootOrEdge::Root(root) => Ok(root),
            RootOrEdge::Edge(edge) => f(edge, volume),
        };

        match node {
            Ok(node) => Ok(Self::Edge(node)),
            Err(node) => Ok(Self::Edge(node)),
        }
    }

    pub fn into_next_node<R>(
        self,
        volume: &super::volume::Volume<R>,
    ) -> core::result::Result<Self, Self>
    where
        R: super::Read,
    {
        self.into_next_node_any(volume, NodeHandle::into_next)
    }

    pub fn into_next_back_node<R>(
        self,
        volume: &super::volume::Volume<R>,
    ) -> core::result::Result<Self, Self>
    where
        R: super::Read,
    {
        self.into_next_node_any(volume, NodeHandle::into_next_back)
    }

    fn into_next_leaf_any<R, F>(
        self,
        volume: &super::volume::Volume<R>,
        f: F,
    ) -> core::result::Result<Self, Self>
    where
        F: Fn(
            NodeHandle,
            &super::volume::Volume<R>,
        ) -> core::result::Result<NodeHandle, NodeHandle>,
        R: super::Read,
    {
        // this will either give us the next immediate node, or self if self was a Root.
        let mut this = self.into_next_node_any(volume, &f)?;

        loop {
            // we are only interested in leaf nodes
            if this.node.inner.is_leaf() {
                break Ok(this);
            } else {
                // this will eventually return Err(self) so will always return
                this = this.into_next_node_any(volume, &f)?;
            }
        }
    }

    /// returns the next leaf, or the end of the tree as Err
    pub fn into_next_leaf<R: super::Read>(
        self,
        volume: &super::volume::Volume<R>,
    ) -> core::result::Result<Self, Self> {
        self.into_next_leaf_any(volume, NodeHandle::into_next)
    }

    /// returns the next RootOrEdge, or the end of the tree as Err, in reverse direction
    pub fn into_next_back_leaf<R: super::Read>(
        self,
        volume: &super::volume::Volume<R>,
    ) -> core::result::Result<Self, Self> {
        self.into_next_leaf_any(volume, NodeHandle::into_next_back)
    }

    pub fn as_handle(&self) -> &NodeHandle {
        match self {
            RootOrEdge::Root(handle) => handle,
            RootOrEdge::Edge(handle) => handle,
        }
    }
}

impl Deref for RootOrEdge {
    type Target = NodeHandle;

    fn deref(&self) -> &Self::Target {
        match self {
            RootOrEdge::Root(node) => node,
            RootOrEdge::Edge(node) => node,
        }
    }
}

impl PartialEq for RootOrEdge {
    fn eq(&self, other: &Self) -> bool {
        match self {
            RootOrEdge::Root(root) => match other {
                RootOrEdge::Root(_) => false,
                RootOrEdge::Edge(other) => root.eq(other),
            },
            RootOrEdge::Edge(edge) => match other {
                RootOrEdge::Edge(_) => false,
                RootOrEdge::Root(other) => edge.eq(other),
            },
        }
    }
}

impl Node {
    pub fn clone_from_nonnull(this: &BoxedNode) -> BoxedNode {
        (*this.as_ref()).clone().into()
    }
    pub fn boxed_from_bytes(bytes: Vec<u8>) -> Result<BoxedNode> {
        Ok(Self::from_bytes(bytes)?.into())
    }

    pub fn from_bytes(bytes: Vec<u8>) -> Result<Self> {
        let inner = BTreeNode::parse(&bytes)?;

        Ok(Self { inner, bytes })
    }

    /// must not be called if any outstanding live references to `this` exist
    pub fn destroy(this: BoxedNode) {
        log::debug!("Node::destroy");
        {
            log::debug!("free: {:?}", this);
            let inner = &this.as_ref().inner;
            match inner {
                BTreeNode::Internal(node) => {
                    log::debug!("destroying children..");
                    for child in node.children.iter() {
                        match unsafe { &*child.as_ptr() } {
                            NodePtr::Visited { node, .. } => {
                                Self::destroy(node.clone());
                            }
                            _ => {}
                        }
                    }
                    log::debug!("destroying children [end]");
                }
                BTreeNode::Leaf(_) => {}
            }
        }

        log::debug!("dropping: {:?}", this);
        unsafe {
            drop(Box::from_raw(this.as_ptr()));
        }
    }

    /// 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 = 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)?;
                    Some((*item, value))
                } else {
                    None
                };

                Ok(item)
            }
        }
    }
}

impl<'tree, R> Range<'tree, R>
where
    R: super::Read,
{
    pub fn from_handles(volume: Rc<Volume<R>>, start: NodeHandle, end: NodeHandle) -> Self {
        Self {
            volume,
            start: RootOrEdge::Root(start),
            end: RootOrEdge::Root(end),
            phantom: PhantomData,
        }
    }

    pub fn new(volume: Rc<Volume<R>>, start: BoxedNode, end: BoxedNode) -> Self {
        Self::from_handles(volume, NodeHandle::start(start), NodeHandle::end(end))
    }

    pub fn is_empty(&self) -> bool {
        return self.start == self.end;
    }
}

impl<'tree, R> Iterator for Range<'tree, R>
where
    R: super::Read,
{
    type Item = (Item, TreeItem);

    fn next(&mut self) -> Option<Self::Item> {
        loop {
            if !self.is_empty() {
                replace_with::replace_with_or_abort(&mut self.start, |start| {
                    match start.into_next_node(&self.volume) {
                        Ok(next) => next,
                        Err(next) => next,
                    }
                });

                if self.start.node.inner.is_leaf() {
                    break Some(self.start.as_handle().parse_item().expect("range item"));
                }
                // else repeat
            } else {
                break None;
            }
        }
    }
}

impl<'tree, R> DoubleEndedIterator for Range<'tree, R>
where
    R: super::Read,
{
    fn next_back(&mut self) -> Option<Self::Item> {
        loop {
            if !self.is_empty() {
                replace_with::replace_with_or_abort(&mut self.end, |start| {
                    match start.into_next_back_node(&self.volume) {
                        Ok(next) => next,
                        Err(next) => next,
                    }
                });

                if self.end.node.inner.is_leaf() {
                    break Some(self.end.as_handle().parse_item().expect("range item"));
                }
                // else repeat
            } else {
                break None;
            }
        }
    }
}

impl NodeHandle {
    pub fn start(node: BoxedNode) -> Self {
        let parents = Vec::with_capacity(node.inner.header().level as usize);
        Self {
            node,
            parents,
            idx: 0,
        }
    }

    /// 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 {
        let header = self.node.inner.header();
        if header.nritems.get() >= self.idx + 1 {
            NodeHandleAdvanceResult::Ascend
        } else {
            match &self.node.inner {
                BTreeNode::Internal(_) => NodeHandleAdvanceResult::Next(Self {
                    idx: self.idx + 1,
                    ..self
                }),
                _ => unreachable!(),
            }
        }
    }

    pub fn find_key_rev<K: PartialEq<Key> + PartialOrd<Key>>(self, key: &K) -> SearchResult {
        match &self.node.inner {
            BTreeNode::Internal(node) => {
                let idx = node
                    .visit_children_keys()
                    .rev()
                    .find_map(|(i, child)| match key.partial_cmp(&child) {
                        Some(core::cmp::Ordering::Greater)
                        | Some(core::cmp::Ordering::Equal)
                        | None => Some(i as u32),
                        _ => None,
                    })
                    .unwrap_or(0);

                SearchResult::Edge(NodeHandle { idx, ..self })
            }
            BTreeNode::Leaf(node) => {
                for (i, child) in node.items.iter().enumerate().rev() {
                    if key.eq(&child.key) {
                        return SearchResult::Leaf(NodeHandle {
                            idx: i as u32,
                            ..self
                        });
                    }
                }

                log::debug!("key definitely not found!");
                SearchResult::Edge(NodeHandle { idx: 0, ..self })
            }
        }
    }

    pub fn find_key<K: PartialEq<Key> + PartialOrd<Key>>(self, key: &K) -> SearchResult {
        match &self.node.inner {
            BTreeNode::Internal(node) => {
                let idx = node
                    .visit_children_keys()
                    .find_map(|(i, child)| match key.partial_cmp(&child) {
                        Some(core::cmp::Ordering::Less) => {
                            Some(if i == 0 { 0 } else { i as u32 - 1 })
                        }
                        Some(core::cmp::Ordering::Equal) | None => Some(i as u32),
                        _ => None,
                    })
                    .unwrap_or(node.children.len() as u32 - 1);

                SearchResult::Edge(NodeHandle { idx, ..self })
            }
            BTreeNode::Leaf(node) => {
                for (i, child) in node.items.iter().enumerate() {
                    // if key < &child.key {
                    //     return SearchResult::Leaf(NodeHandle {
                    //         node: self.clone(),
                    //         idx: if i == 0 { 0 } else { i as u32 - 1 },
                    //     });
                    // } else
                    if key.eq(&child.key) {
                        return SearchResult::Leaf(NodeHandle {
                            idx: i as u32,
                            ..self
                        });
                    }
                }

                log::debug!("key definitely not found!");
                SearchResult::Edge(NodeHandle {
                    idx: node.items.len() as u32 - 1,
                    ..self
                })
            }
        }
    }

    // needs reference to volume to be able to read children key_ptrs
    // this is a live reference so if in the future I want to have a &mut to something in Volume
    // greppable: volume ref
    pub fn into_next_back<R: super::Read>(
        self,
        volume: &super::volume::Volume<R>,
    ) -> core::result::Result<Self, Self> {
        let Self {
            mut parents,
            node,
            idx,
        } = self;

        if idx < 1 {
            // go up
            match parents.pop() {
                Some((node, idx)) => Self { node, idx, parents }.into_next_back(volume),
                None => Err(Self { node, idx, parents }),
            }
        } else {
            match &node.inner {
                BTreeNode::Internal(internal) => {
                    let node = match internal.visit_child(idx as usize, volume) {
                        Ok(child) => {
                            parents.push((node, idx));
                            Ok(Self {
                                parents,
                                idx: child.inner.header().nritems.get() - 1,
                                node: child,
                            })
                        }
                        Err(_) => Err(Self { parents, node, idx }),
                    };

                    // TODO: better error or panic here? this would return self, indicating the end of the tree, even though we simply failed to retrieve the next node
                    node
                }

                BTreeNode::Leaf(_) => Ok(Self {
                    idx: idx - 1,
                    parents,
                    node,
                }),
            }
        }
    }

    // needs reference to volume to be able to read children key_ptrs
    // this is a live reference so if in the future I want to have a &mut to something in Volume
    // greppable: volume ref
    /// returns Ok(next) or Err(self) if self is already the last node
    pub fn into_next<R: super::Read>(
        self,
        volume: &super::volume::Volume<R>,
    ) -> core::result::Result<Self, Self> {
        let Self {
            mut parents,
            node,
            idx,
        } = self;

        let header = node.inner.header();

        if idx + 1 >= header.nritems.get() {
            // go up
            match parents.pop() {
                Some((node, idx)) => Self {
                    idx: (idx + 1).min(node.inner.header().nritems.get()),
                    node,
                    parents,
                }
                .into_next(volume),
                None => Err(Self { node, idx, parents }),
            }
        } else {
            match &node.inner {
                BTreeNode::Internal(internal) => {
                    let node = match internal.visit_child(idx as usize, volume) {
                        Ok(child) => {
                            parents.push((node, idx));
                            Ok(Self {
                                parents,
                                idx: 0,
                                node: child,
                            })
                        }
                        Err(_) => {
                            log::error!("failed to read child node!");
                            Err(Self { parents, node, idx })
                        }
                    };

                    // TODO: better error or panic here? this would return self, indicating the end of the tree, even though we simply failed to retrieve the next node
                    node
                }
                BTreeNode::Leaf(_) => Ok(Self {
                    idx: idx + 1,
                    parents,
                    node,
                }),
            }
        }
    }

    pub fn end(node: BoxedNode) -> Self {
        let parents = Vec::with_capacity(node.inner.header().level as usize);
        Self {
            parents,
            idx: node.inner.header().nritems.get() - 1,
            node,
        }
    }
}

// PARTIAL KEY

/// key lookup that will find the first key that matches the present items in this partial key
pub struct PartialKey {
    pub id: Option<KnownObjectId>,
    pub ty: Option<ObjectType>,
    pub offset: Option<u64>,
}

impl PartialKey {
    pub fn new(id: Option<KnownObjectId>, ty: Option<ObjectType>, offset: Option<u64>) -> Self {
        Self { id, ty, offset }
    }
}

impl PartialEq<Key> for PartialKey {
    fn eq(&self, other: &Key) -> bool {
        self.id.map(|id| id == other.id()).unwrap_or(true)
            && self.ty.map(|ty| ty == other.ty()).unwrap_or(true)
            && self
                .offset
                .map(|offset| offset == other.offset.get())
                .unwrap_or(true)
    }
}

/// compares Self to a key, by comparing each item with the element in key if present, and skipping to the next item if missing.
impl PartialOrd<Key> for PartialKey {
    fn partial_cmp(&self, other: &Key) -> Option<core::cmp::Ordering> {
        let id = self.id.and_then(|id| id.partial_cmp(&other.id()));
        let ty = self.ty.and_then(|ty| ty.partial_cmp(&other.ty()));
        let offset = self
            .offset
            .and_then(|offset| offset.partial_cmp(&other.offset.get()));

        match id {
            Some(core::cmp::Ordering::Equal) | None => match ty {
                Some(core::cmp::Ordering::Equal) | None => offset,
                ord => ord,
            },
            ord => ord,
        }
    }
}

#[cfg(test)]
mod partial_key_tests {
    use test_log::test;

    use super::*;

    #[test]
    fn test_partial_key_ord() {
        let key = Key::new(
            KnownObjectId::ChunkTree,
            ObjectType::DirItem,
            0x8dbfc2d2, // crc of "default"
        );

        let pkey = PartialKey::new(
            Some(KnownObjectId::ChunkTree),
            Some(ObjectType::DirItem),
            None,
        );
        assert_eq!(pkey.partial_cmp(&key), None);

        let pkey = PartialKey::new(
            Some(KnownObjectId::ChunkTree),
            Some(ObjectType::DirItem),
            Some(0xdeadbeef),
        );
        assert_ne!(pkey.partial_cmp(&key), Some(core::cmp::Ordering::Equal));

        let pkey = PartialKey::new(None, Some(ObjectType::DirItem), Some(0xdeadbeef));
        assert_ne!(pkey.partial_cmp(&key), None);

        let pkey = PartialKey::new(
            Some(KnownObjectId::ChunkTree),
            Some(ObjectType::DirItem),
            Some(0x8dbfc2d2),
        );
        assert_eq!(pkey.partial_cmp(&key), Some(core::cmp::Ordering::Equal));
    }

    #[test]
    fn test_partial_eq_partial_key() {
        let key = Key::new(
            KnownObjectId::ChunkTree,
            ObjectType::DirItem,
            0x8dbfc2d2, // crc of "default"
        );

        let pkey = PartialKey::new(
            Some(KnownObjectId::ChunkTree),
            Some(ObjectType::DirItem),
            None,
        );
        assert!(pkey.eq(&key));

        let pkey = PartialKey::new(
            Some(KnownObjectId::ChunkTree),
            Some(ObjectType::DirItem),
            Some(0xdeadbeef),
        );
        assert!(!pkey.eq(&key));

        let pkey = PartialKey::new(None, Some(ObjectType::DirItem), Some(0xdeadbeef));
        assert!(!pkey.eq(&key));

        let pkey = PartialKey::new(
            Some(KnownObjectId::ChunkTree),
            Some(ObjectType::DirItem),
            Some(0x8dbfc2d2),
        );
        assert!(pkey.eq(&key));
    }
}