vapfs_ufs/src/btree.rs

use alloc::string::{String, ToString};
use alloc::vec;
use alloc::vec::Vec;
use ondisk_btree::{BTree, FromBytes, SizeOf, ToBytes};
use vapfs::{Index};
use crate::crc32;

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

/// # EntryType
/// the type of an entry in a directory listing
#[derive(Debug, Clone)]
pub enum EntryType {
    /// a file or directory, stored in an inode
    Inode(Index),
    /// a hard link (i.e. pointer to inode, cannot be directory)
    HardLink(Index),
    /// a soft link (i.e. path string, can be directory)
    SoftLink(String),
    /// this entry is corrupt
    Corrupt,
}

/// # DirectoryListing
/// an entry in a list of all files and directories sharing a common crc32 hash collision (unlikely to happen but just in case!)
#[derive(Debug, Clone)]
pub struct DirectoryListing {
    pub name: String,
    pub data: EntryType,
}

/// # BTreeEntry
/// the aforementioned list, stored in a BTree
#[derive(Debug, Clone, Default)]
pub struct BTreeEntry {
    pub crc32: u32,
    pub entries: Vec<DirectoryListing>,
}

/// # Directory
/// a directory, containing a list of all files and directories in it
#[derive(Clone)]
pub struct Directory {
    /// used in the case of b-tree corruption, to restore the b-tree
    pub backup_entries: Vec<DirectoryListing>,
    /// a crc32 hash of the b-tree
    pub crc32: u32,
    /// the b-tree
    pub btree: BTree<BTreeEntry>,
}

impl Default for Directory {
    fn default() -> Self {
        Self::new()
    }
}

impl Directory {
    /// creates a new, empty directory
    pub fn new() -> Self {
        let mut new = Self { backup_entries: Vec::new(), crc32: 0, btree: BTree::new(BTD) };
        // calculate crc32 hash of empty b-tree
        new.crc32 = crc32::crc32(&new.btree.to_bytes());
        new
    }
    /// Reads the directory into memory from the given bytes
    pub fn open(bytes: &[u8]) -> Self {
        Self::from_bytes(bytes)
    }

    /// Converts the directory back into bytes for writing to disk
    pub fn write(&self) -> Vec<u8> {
        self.to_bytes()
    }

    /// Adds a new entry to the directory listing
    pub fn new_entry(&mut self, name: &str, data: EntryType) {
        let crc32 = crc32::crc32(name.as_bytes());
        // check if this crc32 hash already exists
        let index = self.btree.search(crc32);
        if let Some(node) = index {
            if let Some(key) = node.keys.iter_mut().find(|key| key.0 == crc32) {
                // if it does, add the new entry to the existing list
                key.1.entries.push(DirectoryListing { name: name.to_string(), data });
            }
        } else {
            // if it doesn't, create a new list
            let entries = vec![DirectoryListing { name: name.to_string(), data }];
            self.btree.insert(crc32, BTreeEntry { crc32, entries });
        }
    }

    /// Returns an inode for the given entry name
    pub fn find(&mut self, name: &str) -> Option<EntryType> {
        let crc32 = crc32::crc32(name.as_bytes());
        let index = self.btree.search(crc32);
        if let Some(node) = index {
            if let Some(key) = node.keys.iter().find(|key| key.0 == crc32) {
                return Some(key.1.entries.iter().find(|entry| entry.name == name).unwrap().data.clone());
            }
        }
        None
    }

    /// Removes an entry from the directory listing
    pub fn remove_entry(&mut self, name: &str) {
        let crc32 = crc32::crc32(name.as_bytes());
        let index = self.btree.search(crc32);
        if let Some(node) = index {
            if let Some(key) = node.keys.iter_mut().find(|key| key.0 == crc32) {
                key.1.entries.retain(|entry| entry.name != name);
                if key.1.entries.is_empty() {
                    self.btree.remove(crc32);
                }
            }
        }
    }

    /// Returns a list of all entries in the directory
    pub fn list(&mut self) -> Vec<DirectoryListing> {
        let mut entries = Vec::new();
        if let Some(root) = self.btree.root {
            for entry in self.btree.traverse_in_order_values(root) {
                entries.extend_from_slice(&entry.entries);
            }
        }
        entries
    }

    /// Rebuilds the b-tree from the backup entries
    pub fn rebuild(&mut self) {
        self.btree = BTree::new(BTD);
        for entry in &self.backup_entries {
            let crc32 = crc32::crc32(entry.name.as_bytes());
            let index = self.btree.search(crc32);
            if let Some(node) = index {
                if let Some(key) = node.keys.iter_mut().find(|key| key.0 == crc32) {
                    key.1.entries.push(entry.clone());
                }
            } else {
                let entries = vec![entry.clone()];
                self.btree.insert(crc32, BTreeEntry { crc32, entries });
            }
        }
    }
}

impl ToBytes for Directory {
    fn to_bytes(&self) -> Vec<u8> {
        let mut bytes = Vec::new();
        bytes.extend_from_slice(&(self.backup_entries.len() as u32).to_be_bytes());
        for entry in &self.backup_entries {
            bytes.extend_from_slice(&(entry.size_of()).to_be_bytes());
            bytes.extend_from_slice(&entry.to_bytes());
        }
        bytes.extend_from_slice(&self.crc32.to_be_bytes());
        bytes.extend_from_slice(&self.btree.to_bytes());
        bytes
    }
}

impl FromBytes for Directory {
    fn from_bytes(bytes: &[u8]) -> Self {
        let backup_entry_count = u32::from_be_bytes(bytes[0..4].try_into().unwrap()) as usize;
        let mut backup_entries = Vec::new();
        let mut offset = 4;
        for _ in 0..backup_entry_count {
            let entry_size = u32::from_be_bytes(bytes[offset..offset + 4].try_into().unwrap()) as usize;
            backup_entries.push(DirectoryListing::from_bytes(&bytes[offset + 4..offset + 4 + entry_size]));
            offset += 4 + entry_size;
        }
        let crc32 = u32::from_be_bytes(bytes[offset..offset + 4].try_into().unwrap());
        let btree = BTree::from_bytes(&bytes[offset + 4..]);
        Self { backup_entries, crc32, btree }
    }
}

impl ToBytes for DirectoryListing {
    fn to_bytes(&self) -> Vec<u8> {
        let mut bytes = Vec::new();
        let name_len: u32 = self.name.len() as u32;
        bytes.extend_from_slice(&name_len.to_be_bytes());
        bytes.extend_from_slice(self.name.as_bytes());
        match &self.data {
            EntryType::Inode(index) => {
                // 0 for inode
                bytes.push(0);
                bytes.extend_from_slice(&index.to_be_bytes());
            }
            EntryType::HardLink(index) => {
                // 1 for hard link
                bytes.push(1);
                bytes.extend_from_slice(&index.to_be_bytes());
            }
            EntryType::SoftLink(str) => {
                // 2 for soft link
                bytes.push(2);
                let str_len: u32 = str.len() as u32;
                bytes.extend_from_slice(&str_len.to_be_bytes());
                bytes.extend_from_slice(str.as_bytes());
            }
            EntryType::Corrupt => {
                // 3 for corrupt
                bytes.push(3);
            }
        }
        bytes
    }
}

impl FromBytes for DirectoryListing {
    fn from_bytes(bytes: &[u8]) -> Self {
        let name_len = u32::from_be_bytes(bytes[0..4].try_into().unwrap()) as usize;
        let name = String::from_utf8(bytes[4..4 + name_len].to_vec()).unwrap();
        match bytes[4 + name_len] {
            0 => {
                let index = u64::from_be_bytes(bytes[5 + name_len..13 + name_len].try_into().unwrap());
                Self { name, data: EntryType::Inode(index) }
            }
            1 => {
                let index = u64::from_be_bytes(bytes[5 + name_len..13 + name_len].try_into().unwrap());
                Self { name, data: EntryType::HardLink(index) }
            }
            2 => {
                let str_len = u32::from_be_bytes(bytes[5 + name_len..9 + name_len].try_into().unwrap()) as usize;
                let str = String::from_utf8(bytes[9 + name_len..9 + name_len + str_len].to_vec()).unwrap();
                Self { name, data: EntryType::SoftLink(str) }
            }
            _ => {
                Self { name, data: EntryType::Corrupt }
            }
        }
    }
}

impl SizeOf for DirectoryListing {
    fn size_of(&self) -> u32 {
        4 + self.name.len() as u32 + 8
    }
}

impl ToBytes for BTreeEntry {
    fn to_bytes(&self) -> Vec<u8> {
        let mut bytes = Vec::new();
        bytes.extend_from_slice(&self.crc32.to_be_bytes());
        bytes.extend_from_slice(&(self.entries.len() as u32).to_be_bytes());
        for entry in &self.entries {
            bytes.extend_from_slice(&(entry.size_of()).to_be_bytes());
            bytes.extend_from_slice(&entry.to_bytes());
        }
        bytes
    }
}

impl FromBytes for BTreeEntry {
    fn from_bytes(bytes: &[u8]) -> Self {
        let crc32 = u32::from_be_bytes(bytes[0..4].try_into().unwrap());
        let entry_count = u32::from_be_bytes(bytes[4..8].try_into().unwrap()) as usize;
        let mut entries = Vec::new();
        let mut offset = 8;
        for _ in 0..entry_count {
            let entry_size = u32::from_be_bytes(bytes[offset..offset + 4].try_into().unwrap()) as usize;
            entries.push(DirectoryListing::from_bytes(&bytes[offset + 4..offset + 4 + entry_size]));
            offset += 4 + entry_size;
        }
        Self { crc32, entries }
    }
}

impl SizeOf for BTreeEntry {
    fn size_of(&self) -> u32 {
        let mut size = 8;
        for entry in &self.entries {
            size += 4 + entry.size_of();
        }
        size
    }
}