3 files changed, 136 insertions, 2 deletions
diff --git a/src/disk/allocation.rs b/src/disk/allocation.rs
new file mode 100644
index 0000000..291a136
--- /dev/null
+++ b/src/disk/allocation.rs
@@ -0,0 +1,132 @@
+use crate::memory::cached_block::convert;
+use crate::disk::data_block::{DataBlock, DoubleIndirectBlock, IndirectBlock, TripleIndirectBlock};
+use crate::disk::inode::Inode;
+use crate::AyaFS;
+
+impl AyaFS {
+    /// 为 Inode 分配新 block, 返回 block 的编号
+    pub(crate) fn allocate_block(&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;
+                *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;
+        }
+        // 在 indirect block 里尝试分配
+        if let Some(block_index) = self.allocate_in_indirect(inode.single_indirect) {
+            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;
+        }
+        // 在 double indirect block 里尝试分配
+        if let Some(block_index) = self.alloc_in_double_indirect(inode.double_indirect) {
+            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;
+        }
+        // 在 double indirect block 里尝试分配
+        if let Some(block_index) = self.alloc_in_triple_indirect(inode.triple_indirect) {
+            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;
+
+        if let Some(block) = self
+            .get_block(indirect_entry)
+            .map(convert::<DataBlock, IndirectBlock>)
+        {
+            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
+    }
+}
diff --git a/src/disk/data_block.rs b/src/disk/data_block.rs
index 9c8cc26..d287ee6 100644
--- a/src/disk/data_block.rs
+++ b/src/disk/data_block.rs
@@ -1,5 +1,4 @@
 use crate::disk::inode::Inode;
-use libc::pathconf;
 
 pub trait Block: Default + Clone {}
 
diff --git a/src/disk/mod.rs b/src/disk/mod.rs
index 404c6ab..65f313c 100644
--- a/src/disk/mod.rs
+++ b/src/disk/mod.rs
@@ -1,4 +1,7 @@
+/// On-disk data structures and logic.
+/// Including bitmaps, inodes and blocks.
+
 pub mod bitmap;
 pub mod inode;
-
 pub mod data_block;
+pub mod allocation;