1 files changed, 39 insertions, 1 deletions

diff --git a/ayafs-core/src/disk/bitmap.rs b/ayafs-core/src/disk/bitmap.rs
index b68c341..cf7eae5 100644
--- a/ayafs-core/src/disk/bitmap.rs
+++ b/ayafs-core/src/disk/bitmap.rs
@@ -6,22 +6,59 @@ pub struct Bitmap {
     pub length: usize,
     pub device: Arc<dyn BlockDevice>,
     pub data: Vec<u8>,
+    pub count: u64,
 }
 
 impl Bitmap {
-    pub(crate) fn new(starting_block: usize, length: usize, device: Arc<dyn BlockDevice>) -> Self {
+    pub(crate) fn new(starting_block: usize, length: usize, device: Arc<dyn BlockDevice>, count: u64) -> Self {
         Self {
             starting_block,
             length,
             device,
             data: vec![0u8; length * BLOCK_SIZE],
+            count,
         }
     }
+
+    pub(crate) fn load(starting_block: usize, length: usize, device: Arc<dyn BlockDevice>, count: u64) -> Self {
+        let mut data = vec![0u8; length * BLOCK_SIZE];
+        for id in 0 .. length {
+            let block_id = starting_block + id;
+            device.read(block_id, &mut data[id * BLOCK_SIZE .. (id + 1) * BLOCK_SIZE]);
+        }
+        Self {
+            starting_block,
+            length,
+            device,
+            data,
+            count,
+        }
+    }
+
+    pub(crate) fn write_back(&self) {
+        for id in 0 .. self.length {
+            let block_id = self.starting_block + id;
+            self.device.write(block_id, &self.data[id * BLOCK_SIZE .. (id + 1) * BLOCK_SIZE]);
+        }
+    }
+
+    // Allocate, but doesn't modify self.
+    pub(crate) fn peek(&self) -> Option<usize> {
+        for (i, byte) in self.data.iter().enumerate() {
+            let leading_ones = byte.leading_ones();
+            if leading_ones != 8 {
+                return Some(i * 8 + leading_ones as usize);
+            }
+        }
+        None
+    }
+
     pub(crate) fn allocate(&mut self) -> Option<usize> {
         for (i, byte) in self.data.iter_mut().enumerate() {
             let leading_ones = byte.leading_ones();
             if leading_ones != 8 {
                 *byte |= (1 << (7 - leading_ones)) as u8;
+                self.count += 1;
                 return Some(i * 8 + leading_ones as usize);
             }
         }
@@ -39,6 +76,7 @@ impl Bitmap {
     pub(crate) fn deallocate(&mut self, index: usize) -> bool {
         if self.query(index) {
             let mask = !(1u8 << (7 - index % 8));
+            self.count -= 1;
             self.data[index / 8] &= mask;
             true
         } else {