From 07e8093ee0007475e9ff1651ac0e4740a1e7cd97 Mon Sep 17 00:00:00 2001
From: Jeremy Soller <jackpot51@gmail.com>
Date: Wed, 14 Feb 2024 12:53:06 -0700
Subject: [PATCH] Allow allocating in different sizes

---
 src/allocator.rs   | 12 ++++++------
 src/filesystem.rs  |  2 +-
 src/transaction.rs |  6 +++---
 3 files changed, 10 insertions(+), 10 deletions(-)

diff --git a/src/allocator.rs b/src/allocator.rs
index dbe23efe80..42cfa372dc 100644
--- a/src/allocator.rs
+++ b/src/allocator.rs
@@ -26,10 +26,10 @@ impl Allocator {
         free
     }
 
-    pub fn allocate(&mut self) -> Option<u64> {
+    pub fn allocate(&mut self, addr_level: usize) -> Option<u64> {
         // First, find the lowest level with a free block
         let mut addr_opt = None;
-        let mut level = 0;
+        let mut level = addr_level;
         while level < self.levels.len() {
             if !self.levels[level].is_empty() {
                 addr_opt = self.levels[level].pop();
@@ -38,9 +38,9 @@ impl Allocator {
             level += 1;
         }
 
-        // Next, if a free block was found, split it up until you have a usable level 0 block
+        // Next, if a free block was found, split it up until you have a usable block of the right level
         let addr = addr_opt?;
-        while level > 0 {
+        while level > addr_level {
             level -= 1;
             let level_size = 1 << level;
             self.levels[level].push(addr + level_size);
@@ -79,10 +79,10 @@ impl Allocator {
         addr_opt
     }
 
-    pub fn deallocate(&mut self, mut addr: u64) {
+    pub fn deallocate(&mut self, mut addr: u64, addr_level: usize) {
         // See if block matches with a sibling - if so, join them into a larger block, and populate
         // this all the way to the top level
-        let mut level = 0;
+        let mut level = addr_level;
         loop {
             while level >= self.levels.len() {
                 self.levels.push(Vec::new());
diff --git a/src/filesystem.rs b/src/filesystem.rs
index 491eba2b89..81dc25099f 100644
--- a/src/filesystem.rs
+++ b/src/filesystem.rs
@@ -253,7 +253,7 @@ impl<D: Disk> FileSystem<D> {
                     }
                 } else {
                     for i in 0..count {
-                        self.allocator.deallocate(addr + i as u64);
+                        self.allocator.deallocate(addr + i as u64, 0);
                     }
                 }
             }
diff --git a/src/transaction.rs b/src/transaction.rs
index fdec114091..7cf87645cd 100644
--- a/src/transaction.rs
+++ b/src/transaction.rs
@@ -53,7 +53,7 @@ impl<'a, D: Disk> Transaction<'a, D> {
 
     // Unsafe because order must be done carefully and changes must be flushed to disk
     unsafe fn allocate(&mut self) -> Result<u64> {
-        match self.allocator.allocate() {
+        match self.allocator.allocate(0) {
             Some(addr) => {
                 self.allocator_log.push_back(AllocEntry::new(addr, -1));
                 Ok(addr)
@@ -83,7 +83,7 @@ impl<'a, D: Disk> Transaction<'a, D> {
 
         if found {
             // Deallocate immediately since it is an allocation that was not needed
-            self.allocator.deallocate(addr);
+            self.allocator.deallocate(addr, 0);
         } else {
             // Deallocate later when syncing filesystem, to avoid re-use
             self.deallocate.push(addr);
@@ -164,7 +164,7 @@ impl<'a, D: Disk> Transaction<'a, D> {
         // De-allocate old blocks (after allocation to prevent re-use)
         //TODO: optimize allocator log in memory
         while let Some(addr) = self.deallocate.pop() {
-            self.allocator.deallocate(addr);
+            self.allocator.deallocate(addr, 0);
             self.allocator_log.push_back(AllocEntry::new(addr, 1));
         }