From 7f40fa1631d408f5f96186b57f706dc22d3f8cbc Mon Sep 17 00:00:00 2001
From: Christopher Berner <me@cberner.com>
Date: Sat, 7 Sep 2024 11:24:26 -0700
Subject: [PATCH] Resolve TODO

---
 src/db.rs                                 | 15 ++++++++-------
 src/tree_store/page_store/page_manager.rs | 23 ++++-------------------
 2 files changed, 12 insertions(+), 26 deletions(-)

diff --git a/src/db.rs b/src/db.rs
index 49c31e9d..f09fdcbd 100644
--- a/src/db.rs
+++ b/src/db.rs
@@ -518,7 +518,9 @@ impl Database {
                 FreedPageList::fixed_width(),
                 mem.clone(),
             )?;
-            mem.mark_pages_allocated(freed_pages_iter, false)?;
+            for page in freed_pages_iter {
+                mem.mark_page_allocated(page?);
+            }
         }
 
         let freed_table: ReadOnlyTable<FreedTableKey, FreedPageList<'static>> = ReadOnlyTable::new(
@@ -530,11 +532,9 @@ impl Database {
         )?;
         for result in freed_table.range::<FreedTableKey>(..)? {
             let (_, freed_page_list) = result?;
-            let mut freed_page_list_as_vec = vec![];
             for i in 0..freed_page_list.value().len() {
-                freed_page_list_as_vec.push(Ok(freed_page_list.value().get(i)));
+                mem.mark_page_allocated(freed_page_list.value().get(i));
             }
-            mem.mark_pages_allocated(freed_page_list_as_vec.into_iter(), false)?;
         }
 
         Ok(())
@@ -569,7 +569,9 @@ impl Database {
         // Repair the allocator state
         // All pages in the master table
         let master_pages_iter = AllPageNumbersBtreeIter::new(root, None, None, mem.clone())?;
-        mem.mark_pages_allocated(master_pages_iter, false)?;
+        for page in master_pages_iter {
+            mem.mark_page_allocated(page?);
+        }
 
         // Iterate over all other tables
         let iter: BtreeRangeIter<&str, InternalTableDefinition> =
@@ -579,8 +581,7 @@ impl Database {
         for entry in iter {
             let definition = entry?.value();
             definition.visit_all_pages(mem.clone(), |path| {
-                // TODO: simplify mark_pages_allocated()
-                mem.mark_pages_allocated([Ok(path.page_number())].into_iter(), false)?;
+                mem.mark_page_allocated(path.page_number());
                 Ok(())
             })?;
         }
diff --git a/src/tree_store/page_store/page_manager.rs b/src/tree_store/page_store/page_manager.rs
index 88c4475a..35214bc6 100644
--- a/src/tree_store/page_store/page_manager.rs
+++ b/src/tree_store/page_store/page_manager.rs
@@ -356,26 +356,11 @@ impl TransactionalMemory {
         Ok(())
     }
 
-    pub(crate) fn mark_pages_allocated(
-        &self,
-        allocated_pages: impl Iterator<Item = Result<PageNumber>>,
-        allow_duplicates: bool,
-    ) -> Result<()> {
+    pub(crate) fn mark_page_allocated(&self, page_number: PageNumber) {
         let mut state = self.state.lock().unwrap();
-
-        for page_number in allocated_pages {
-            let page_number = page_number?;
-            let region_index = page_number.region;
-            let allocator = state.get_region_mut(region_index);
-            if allow_duplicates
-                && allocator.is_allocated(page_number.page_index, page_number.page_order)
-            {
-                continue;
-            }
-            allocator.record_alloc(page_number.page_index, page_number.page_order);
-        }
-
-        Ok(())
+        let region_index = page_number.region;
+        let allocator = state.get_region_mut(region_index);
+        allocator.record_alloc(page_number.page_index, page_number.page_order);
     }
 
     fn write_header(&self, header: &DatabaseHeader, swap_primary: bool) -> Result {