AkurAI-Framework / crates/storage/src/node.rs

Latest change 1d2c3fa4f1f97676ca091dd5b541e69abf5a431f - feat(storage): B+tree node encode/decode (leaf + interior) by Ólafur Búi Ólafsson

//! B+tree node encoding: the byte layout of a single page-sized node.
//!
//! Two shapes share one page format, distinguished by a leading tag byte:
//! - **Leaf**: sorted `(key, value)` pairs — all data lives here.
//! - **Interior**: `N` separator keys + `N+1` child page pointers; routing only.
//!
//! Decode trusts its input: pages come from our own writes (and meta pages are
//! checksum-guarded), so we don't defend against arbitrary bytes here.

use crate::pager::{Page, PageId, PAGE_SIZE};

const TAG_LEAF: u8 = 0;
const TAG_INTERIOR: u8 = 1;

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Node {
    /// Sorted key/value pairs.
    Leaf(Vec<(Vec<u8>, Vec<u8>)>),
    /// `children.len() == keys.len() + 1`. To route key `s`: pick `children[i]`
    /// for the first `i` with `s < keys[i]`, else the last child.
    Interior {
        keys: Vec<Vec<u8>>,
        children: Vec<PageId>,
    },
}

impl Node {
    /// Bytes this node needs on a page. Callers check against [`Node::fits`]
    /// before [`Node::encode`], which panics on overflow.
    pub fn encoded_len(&self) -> usize {
        match self {
            Node::Leaf(entries) => {
                1 + 2
                    + entries
                        .iter()
                        .map(|(k, v)| 2 + k.len() + 4 + v.len())
                        .sum::<usize>()
            }
            Node::Interior { keys, .. } => {
                1 + 2 + 8 + keys.iter().map(|k| 2 + k.len() + 8).sum::<usize>()
            }
        }
    }

    /// Does this node fit in a single page?
    pub fn fits(&self) -> bool {
        self.encoded_len() <= PAGE_SIZE
    }

    /// Encode into a zeroed page. Panics if the node exceeds [`PAGE_SIZE`];
    /// guard with [`Node::fits`] (the tree splits before it gets here).
    pub fn encode(&self) -> Page {
        let mut buf = Vec::with_capacity(self.encoded_len());
        match self {
            Node::Leaf(entries) => {
                buf.push(TAG_LEAF);
                buf.extend_from_slice(&(entries.len() as u16).to_le_bytes());
                for (k, v) in entries {
                    buf.extend_from_slice(&(k.len() as u16).to_le_bytes());
                    buf.extend_from_slice(k);
                    buf.extend_from_slice(&(v.len() as u32).to_le_bytes());
                    buf.extend_from_slice(v);
                }
            }
            Node::Interior { keys, children } => {
                buf.push(TAG_INTERIOR);
                buf.extend_from_slice(&(keys.len() as u16).to_le_bytes());
                buf.extend_from_slice(&children[0].to_le_bytes());
                for (k, child) in keys.iter().zip(children.iter().skip(1)) {
                    buf.extend_from_slice(&(k.len() as u16).to_le_bytes());
                    buf.extend_from_slice(k);
                    buf.extend_from_slice(&child.to_le_bytes());
                }
            }
        }
        let mut page = [0u8; PAGE_SIZE];
        page[..buf.len()].copy_from_slice(&buf);
        page
    }

    /// Decode a node from a page. Assumes the page was produced by [`encode`].
    pub fn decode(page: &Page) -> Node {
        let mut c = Cursor { buf: page, pos: 0 };
        let tag = c.u8();
        let count = c.u16() as usize;
        match tag {
            TAG_LEAF => {
                let mut entries = Vec::with_capacity(count);
                for _ in 0..count {
                    let klen = c.u16() as usize;
                    let k = c.bytes(klen);
                    let vlen = c.u32() as usize;
                    let v = c.bytes(vlen);
                    entries.push((k, v));
                }
                Node::Leaf(entries)
            }
            TAG_INTERIOR => {
                let mut children = Vec::with_capacity(count + 1);
                let mut keys = Vec::with_capacity(count);
                children.push(c.u64());
                for _ in 0..count {
                    let klen = c.u16() as usize;
                    keys.push(c.bytes(klen));
                    children.push(c.u64());
                }
                Node::Interior { keys, children }
            }
            other => panic!("invalid node tag {other}"),
        }
    }
}

/// A tiny forward-only reader over a page's bytes.
struct Cursor<'a> {
    buf: &'a [u8],
    pos: usize,
}

impl Cursor<'_> {
    fn u8(&mut self) -> u8 {
        let b = self.buf[self.pos];
        self.pos += 1;
        b
    }
    fn u16(&mut self) -> u16 {
        let v = u16::from_le_bytes(self.buf[self.pos..self.pos + 2].try_into().unwrap());
        self.pos += 2;
        v
    }
    fn u32(&mut self) -> u32 {
        let v = u32::from_le_bytes(self.buf[self.pos..self.pos + 4].try_into().unwrap());
        self.pos += 4;
        v
    }
    fn u64(&mut self) -> u64 {
        let v = u64::from_le_bytes(self.buf[self.pos..self.pos + 8].try_into().unwrap());
        self.pos += 8;
        v
    }
    fn bytes(&mut self, n: usize) -> Vec<u8> {
        let s = self.buf[self.pos..self.pos + n].to_vec();
        self.pos += n;
        s
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn leaf_round_trips() {
        let node = Node::Leaf(vec![
            (b"alpha".to_vec(), b"1".to_vec()),
            (b"beta".to_vec(), b"two".to_vec()),
            (b"".to_vec(), b"".to_vec()), // empty key and value
        ]);
        assert_eq!(Node::decode(&node.encode()), node);
    }

    #[test]
    fn interior_round_trips() {
        let node = Node::Interior {
            keys: vec![b"m".to_vec(), b"t".to_vec()],
            children: vec![10, 20, 30],
        };
        assert_eq!(Node::decode(&node.encode()), node);
    }

    #[test]
    fn empty_leaf_round_trips() {
        let node = Node::Leaf(vec![]);
        assert_eq!(Node::decode(&node.encode()), node);
    }

    #[test]
    fn encoded_len_is_exact() {
        let node = Node::Leaf(vec![(b"key".to_vec(), b"value".to_vec())]);
        // tag(1) + count(2) + klen(2)+3 + vlen(4)+5 = 17
        assert_eq!(node.encoded_len(), 17);
        assert!(node.fits());
    }

    #[test]
    fn oversized_leaf_does_not_fit() {
        let big = vec![0u8; PAGE_SIZE];
        let node = Node::Leaf(vec![(b"k".to_vec(), big)]);
        assert!(!node.fits());
    }
}