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::{BlockDevice, Index};
use crate::crc32;
use crate::structs::Superblock;

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

/// # 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 inode: Index,
}

/// # 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 Directory {
    /// 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, inode: Index) {
        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(), inode });
            }
        } else {
            // if it doesn't, create a new list
            let entries = vec![DirectoryListing { name: name.to_string(), inode }];
            self.btree.insert(crc32, BTreeEntry { crc32, entries });
        }
    }

    /// Returns an inode for the given entry name
    pub fn find(&mut self, name: &str) -> Option<Index> {
        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().inode);
            }
        }
        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() {
                }
            }
        }
    }

    /// 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());
        bytes.extend_from_slice(&self.inode.to_be_bytes());
        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();
        let inode = Index::from_be_bytes(bytes[4 + name_len..12 + name_len].try_into().unwrap());
        Self { name, inode }
    }
}

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
    }
}