vapfs_ufs/src/btree.rs

use alloc::string::{String, ToString};
use alloc::vec::Vec;
use vapfs::{BlockDevice, Index};

/// # BTd
/// min-degree of all BTrees
pub const BTd: u32 = 3;

/// # BTKeyType
/// key used to index a BTree
pub type BTKeyType = u32;

/// # BTKey
/// whole value of a key in a BTree
#[repr(C)]
#[derive(Copy, Clone)]
pub struct BTKey {
    pub key: BTKeyType,
    /// length of this entry in bytes
    pub length: u32,
    /// amount of files represented by this key, used in the case of hash collisions
    pub count: u32,
    /// if count is 1, this is the inode number of the file represented by this key
    /// if count is >1, offset from first BTNode in bytes to an HCList struct
    pub value: Index,
    /// length of filename
    /// zero if count is >1
    pub filename_length: u32,
    // rest of the struct is the filename
}

/// # HCList
/// list of files represented by a single key in a BTree, used in the case of hash collisions
#[repr(C)]
#[derive(Copy, Clone)]
pub struct HCList {
    /// offset from first HCList in bytes to the next HCList
    pub next: Index,
    /// inode number of the file represented by this HCList
    pub inode: Index,
    /// length of filename
    pub filename_length: u32,
    // rest of the struct is the filename
}

#[repr(C)]
#[derive(Copy, Clone)]
pub struct BTNode {
    /// offset from first BTNode in bytes, each entry is a DirectoryKey
    pub key_array: Index,
    pub key_count: u32,
    /// offset from first BTNode in bytes, each entry is another BTNode
    pub child_array: Index,
    pub child_count: u32,
}

impl BTKey {
    /// in-place conversion from the storage representation (big endian) to the native representation
    pub fn convert_big_endian_to_native(&mut self) {
        #[cfg(target_endian = "little")]
        {
            self.key = u32::from_be(self.key);
            self.length = u32::from_be(self.length);
            self.count = u32::from_be(self.count);
            self.value = u64::from_be(self.value);
            self.filename_length = u32::from_be(self.filename_length);
        }
    }

    /// in-place conversion from the native representation to the storage representation (big endian)
    pub fn convert_native_to_big_endian(&mut self) {
        #[cfg(target_endian = "little")]
        {
            self.key = u32::to_be(self.key);
            self.length = u32::to_be(self.length);
            self.count = u32::to_be(self.count);
            self.value = u64::to_be(self.value);
            self.filename_length = u32::to_be(self.filename_length);
        }
    }

    /// offset is from 0x0 in the block device to the beginning of the struct
    pub fn read(offset: Index, bd: &mut dyn BlockDevice) -> Option<(Self, String)> {
        let mut buf: [u8; core::mem::size_of::<Self>()] = [0; core::mem::size_of::<Self>()];
        bd.seek(offset);
        let read_count = bd.read_blocks(&mut buf);
        if read_count < core::mem::size_of::<Self>() {
            return None;
        }
        let mut key = unsafe { core::ptr::read(buf.as_ptr() as *const Self) };
        key.convert_big_endian_to_native();
        let filename = key.read_filename(offset, bd);
        Some((key, filename))
    }

    /// offset is from 0x0 in the block device to the beginning of the struct
    pub fn read_filename(&self, offset: Index, bd: &mut dyn BlockDevice) -> String {
        let mut filename_buf: Vec<u8> = Vec::new();
        filename_buf.resize(self.filename_length as usize, 0);
        bd.seek(offset);
        bd.read_blocks(&mut filename_buf);
        String::from_utf8_lossy(unsafe { core::slice::from_raw_parts(filename_buf.as_ptr(), filename_buf.len()) }).to_string()
    }
}

impl BTNode {
    pub const fn new(key_array: Index, key_count: u32, child_array: Index, child_count: u32) -> Self {
        Self {
            key_array,
            key_count,
            child_array,
            child_count,
        }
    }

    /// in-place conversion from the storage representation (big endian) to the native representation
    pub fn convert_big_endian_to_native(&mut self) {
        #[cfg(target_endian = "little")]
        {
            self.key_array = u64::from_be(self.key_array);
            self.key_count = u32::from_be(self.key_count);
            self.child_array = u64::from_be(self.child_array);
            self.child_count = u32::from_be(self.child_count);
        }
    }

    /// in-place conversion from the native representation to the storage representation (big endian)
    pub fn convert_native_to_big_endian(&mut self) {
        #[cfg(target_endian = "little")]
        {
            self.key_array = u64::to_be(self.key_array);
            self.key_count = u32::to_be(self.key_count);
            self.child_array = u64::to_be(self.child_array);
            self.child_count = u32::to_be(self.child_count);
        }
    }

    /// offset is from root in bytes
    pub fn read(root: Index, offset: Index, bd: &mut dyn BlockDevice) -> Option<Self> {
        let mut buf: [u8; core::mem::size_of::<Self>()] = [0; core::mem::size_of::<Self>()];
        bd.seek(root + offset);
        let read_count = bd.read_blocks(&mut buf);
        if read_count < core::mem::size_of::<Self>() {
            return None;
        }
        let mut node = unsafe { core::ptr::read(buf.as_ptr() as *const Self) };
        node.convert_big_endian_to_native();
        Some(node)
    }

    pub fn search(self, root: Index, key_to_find: BTKeyType, bd: &mut dyn BlockDevice) -> Option<(BTKey, String)> {
        let mut node = self;
        loop {
            let mut i = 0;
            let mut last_key: Option<BTKey> = None;
            while i < node.key_count {
                let key_offset = if i == 0 {
                    node.key_array
                } else {
                    last_key.as_ref().unwrap().length as Index + last_key.as_ref().unwrap().value
                };
                let key = BTKey::read(key_offset, bd)?;
                if key.0.key == key_to_find {
                    return Some(key);
                } else if key.0.key > key_to_find {
                    break;
                }
                i += 1;
                last_key = Some(key.0);
            }
            if node.child_count == 0 {
                return None;
            }
            let child_offset = if i == 0 {
                node.child_array
            } else {
                last_key.unwrap().value
            };
            node = BTNode::read(root, child_offset, bd)?;
        }
    }
}