path: root/src/disk/allocation.rs

use crate::disk::block::{DataBlock, DoubleIndirectBlock, IndirectBlock, TripleIndirectBlock};
use crate::disk::inode::Inode;
use crate::memory::cached_block::convert;
use crate::AyaFS;

impl AyaFS {
    /// 为 Inode 分配新 block, 返回 block 的编号
    pub(crate) fn allocate_block_for(&mut self, inode: &mut Inode) -> Option<u32> {
        // 先看这个 inode 的 direct block 有没有空闲
        for index in inode.direct.iter_mut() {
            if !self.data_bitmap.query(*index as usize) {
                let block_index = self.data_bitmap.allocate().unwrap() as u32;
                // println!("allocating {} for direct", block_index);
                *index = block_index;
                inode.n_blocks += 1;
                // 当调用 get_inode_mut 拿出 &mut Inode 的时候对应的 block 在 cache 里已经脏了
                return Some(block_index);
            }
        }

        // direct block 全部分配完了, 先检查 indirect block 有没有分配, 没有就分配一个
        if !self.data_bitmap.query(inode.single_indirect as usize) {
            inode.single_indirect = self
                .data_bitmap
                .allocate()
                .expect("No free space for new block") as u32;
            // println!("allocating {} for indirect", inode.single_indirect);
        }
        // 在 indirect block 里尝试分配
        if let Some(block_index) = self.allocate_in_indirect(inode.single_indirect) {
            // println!("allocating {} in indirect", block_index);
            inode.n_blocks += 1;
            return Some(block_index);
        }

        // direct & indirect block 全部分配完了, 先检查 double indirect block 有没有分配, 没有就分配一个
        if !self.data_bitmap.query(inode.double_indirect as usize) {
            inode.double_indirect = self
                .data_bitmap
                .allocate()
                .expect("No free space for new block") as u32;
            // println!("allocating {} for double indirect", inode.double_indirect);
        }
        // 在 double indirect block 里尝试分配
        if let Some(block_index) = self.alloc_in_double_indirect(inode.double_indirect) {
            // println!("allocating {} in double indirect", block_index);
            inode.n_blocks += 1;
            return Some(block_index);
        }

        // direct, indirect & double indirect block 全部分配完了, 先检查 triple indirect block 有没有分配, 没有就分配一个
        if !self.data_bitmap.query(inode.triple_indirect as usize) {
            inode.triple_indirect = self
                .data_bitmap
                .allocate()
                .expect("No free space for new block") as u32;
            // println!("allocating {} for triple indirect", inode.triple_indirect);
        }
        // 在 double indirect block 里尝试分配
        if let Some(block_index) = self.alloc_in_triple_indirect(inode.triple_indirect) {
            // println!("allocating {} in triple indirect", block_index);
            inode.n_blocks += 1;
            return Some(block_index);
        }
        None
    }

    fn allocate_in_indirect(&mut self, indirect_entry: u32) -> Option<u32> {
        // 取出 single indirect block, 尝试在里面分配
        let indirect_entry = indirect_entry as usize;
        // println!("finding indirect block with number {}, bitmap says {}", indirect_entry, self.data_bitmap.query(indirect_entry));

        if let Some(block) = self
            .get_block(indirect_entry)
            .map(convert::<DataBlock, IndirectBlock>)
        {
            // println!("found indirect block with number {}", indirect_entry);
            let mut indirect_block = block.clone();
            for entry in indirect_block.block.entries.iter_mut() {
                if self.data_bitmap.query(*entry as usize) == false {
                    indirect_block.dirty = true; // 把这个块标记为 dirty
                    let block_index = self.data_bitmap.allocate().expect("No free space") as u32;
                    *entry = block_index;
                    self.update_block(indirect_block);
                    return Some(block_index);
                }
            }
        }
        None
    }

    fn alloc_in_double_indirect(&mut self, double_indirect_entry: u32) -> Option<u32> {
        let double_indirect_entry = double_indirect_entry as usize;

        if let Some(block) = self
            .get_block(double_indirect_entry)
            .map(convert::<DataBlock, DoubleIndirectBlock>)
        {
            let mut double_indirect_block = block.clone();
            for indirect_entry in double_indirect_block.block.indirect.iter_mut() {
                if self.data_bitmap.query(*indirect_entry as usize) == false {
                    double_indirect_block.dirty = true;
                    *indirect_entry = self.data_bitmap.allocate().expect("No free space") as u32;
                }

                if let Some(block_index) = self.allocate_in_indirect(*indirect_entry) {
                    if double_indirect_block.dirty {
                        self.update_block(double_indirect_block);
                    }
                    return Some(block_index);
                }
            }
        }
        None
    }

    fn alloc_in_triple_indirect(&mut self, triple_indirect_entry: u32) -> Option<u32> {
        let triple_indirect_entry = triple_indirect_entry as usize;

        if let Some(block) = self
            .get_block(triple_indirect_entry)
            .map(convert::<DataBlock, TripleIndirectBlock>)
        {
            let mut triple_indirect_block = block.clone();
            for double_indirect_entry in triple_indirect_block.block.double_indirect.iter_mut() {
                if self.data_bitmap.query(*double_indirect_entry as usize) == false {
                    triple_indirect_block.dirty = true;
                    *double_indirect_entry =
                        self.data_bitmap.allocate().expect("No free space") as u32;
                }
                if let Some(block_index) = self.alloc_in_double_indirect(*double_indirect_entry) {
                    if triple_indirect_block.dirty {
                        self.update_block(triple_indirect_block);
                    }
                    return Some(block_index);
                }
            }
        }
        None
    }
}

impl AyaFS {
    /// 从 inode 中删去最后一个 block
    pub(crate) fn deallocate_block_for(&mut self, inode: &mut Inode) -> Option<u32> {
        // 如果 triple indirect 块存在, 则尝试从中销毁一个块
        if self.data_bitmap.query(inode.triple_indirect as usize) {
            if let Some(block_index) = self.deallocate_from_triple_indirect(inode.triple_indirect) {
                inode.n_blocks -= 1;
                return Some(block_index); // 销毁成功, 直接返回
            } else {
                // 销毁失败, 说明 triple indirect 空了, 把它也销毁.
                self.data_bitmap.deallocate(inode.triple_indirect as usize);
                inode.triple_indirect = 0; // 这个地方理论上应该不用?
            }
        }
        // 如果 double indirect 块存在, 则从其中销毁
        if self.data_bitmap.query(inode.double_indirect as usize) {
            if let Some(block_index) = self.deallocate_from_double_indirect(inode.double_indirect) {
                inode.n_blocks -= 1;
                return Some(block_index);
            } else {
                self.data_bitmap.deallocate(inode.double_indirect as usize);
                inode.double_indirect = 0; // 这个地方理论上应该不用?
            }
        }
        // 如果 indirect 块存在, 则从其中销毁
        if self.data_bitmap.query(inode.single_indirect as usize) {
            if let Some(block_index) = self.deallocate_from_indirect(inode.single_indirect) {
                inode.n_blocks -= 1;
                return Some(block_index);
            } else {
                self.data_bitmap.deallocate(inode.single_indirect as usize);
                inode.single_indirect = 0; // 这个地方理论上应该不用?
            }
        }
        // 都没有，直接从 direct 块中销毁
        for entry in inode.direct.iter_mut().rev() {
            if self.data_bitmap.query(*entry as usize) {
                let index = std::mem::replace(entry, 0); // let index = *entry; *entry = 0;
                inode.n_blocks -= 1;
                self.data_bitmap.deallocate(index as usize);
                return Some(index);
            }
        }
        None
    }

    fn deallocate_from_triple_indirect(&mut self, triple_indirect_entry: u32) -> Option<u32> {
        let triple_indirect_entry = triple_indirect_entry as usize;
        if let Some(triple_indirect_block) = self
            .get_block(triple_indirect_entry)
            .map(convert::<DataBlock, TripleIndirectBlock>)
        {
            let mut triple_indirect_block = triple_indirect_block.clone();
            let mut block_modified = false;
            for double_indirect_entry in
                triple_indirect_block.block.double_indirect.iter_mut().rev()
            {
                // 如果这个位置的 double indirect 存在
                if self.data_bitmap.query(*double_indirect_entry as usize) {
                    // 尝试从中销毁一个块
                    if let Some(block_index) =
                        self.deallocate_from_double_indirect(*double_indirect_entry)
                    {
                        if block_modified {
                            self.update_block(triple_indirect_block);
                        }
                        return Some(block_index); // 成功则直接返回
                    } else {
                        // 失败则把这个 double indirect 销毁
                        let double_indirect_entry_to_deallocate =
                            std::mem::replace(double_indirect_entry, 0);
                        self.data_bitmap
                            .deallocate(double_indirect_entry_to_deallocate as usize);
                        triple_indirect_block.dirty = true;
                        block_modified = true;
                    }
                }
            }
            if block_modified {
                self.update_block(triple_indirect_block);
            }
        }
        None
    }

    fn deallocate_from_double_indirect(&mut self, double_indirect_entry: u32) -> Option<u32> {
        let double_indirect_entry = double_indirect_entry as usize;
        if let Some(double_indirect_block) = self
            .get_block(double_indirect_entry)
            .map(convert::<DataBlock, DoubleIndirectBlock>)
        {
            let mut double_indirect_block = double_indirect_block.clone();
            let mut block_modified = false;
            for indirect_entry in double_indirect_block.block.indirect.iter_mut().rev() {
                // 如果这个位置的 indirect 存在
                if self.data_bitmap.query(*indirect_entry as usize) {
                    // 尝试从中销毁一个块
                    if let Some(block_index) = self.deallocate_from_indirect(*indirect_entry) {
                        if block_modified {
                            self.update_block(double_indirect_block);
                        }
                        return Some(block_index); // 成功则直接返回
                    } else {
                        // 失败则把这个 indirect 销毁
                        let indirect_entry_to_deallocate = std::mem::replace(indirect_entry, 0);
                        self.data_bitmap
                            .deallocate(indirect_entry_to_deallocate as usize);
                        double_indirect_block.dirty = true;
                        block_modified = true;
                    }
                }
            }
            if block_modified {
                self.update_block(double_indirect_block);
            }
        }
        None
    }

    fn deallocate_from_indirect(&mut self, indirect_entry: u32) -> Option<u32> {
        let indirect_entry = indirect_entry as usize;
        if let Some(indirect_block) = self
            .get_block(indirect_entry)
            .map(convert::<DataBlock, IndirectBlock>)
        {
            let mut indirect_block = indirect_block.clone();
            // 遍历 indirect block 里的每个 block
            for entry in indirect_block.block.entries.iter_mut().rev() {
                // 如果这个 block 存在, 销毁它
                if self.data_bitmap.query(*entry as usize) {
                    let entry_to_deallocate = std::mem::replace(entry, 0);

                    self.data_bitmap.deallocate(entry_to_deallocate as usize);
                    indirect_block.dirty = true;
                    self.update_block(indirect_block);

                    return Some(entry_to_deallocate);
                }
            }
        }
        None
    }
}