From a1821e1208403de158def9b9e53d4c768b4fc997 Mon Sep 17 00:00:00 2001
From: Maksym Arutyunyan <maksym.arutyunyan@dfinity.org>
Date: Wed, 20 Aug 2025 16:15:40 +0200
Subject: [PATCH 01/18] test: add test showing broken buckets order

---
 src/memory_manager.rs | 34 ++++++++++++++++++++++++++++++++++
 1 file changed, 34 insertions(+)

diff --git a/src/memory_manager.rs b/src/memory_manager.rs
index ab9893ff..e06d6628 100644
--- a/src/memory_manager.rs
+++ b/src/memory_manager.rs
@@ -1284,6 +1284,40 @@ mod test {
         );
     }
 
+    #[test]
+    fn bucket_ordering() {
+        let mock_stable_memory = make_memory();
+        let mm = MemoryManager::init_with_bucket_size(mock_stable_memory.clone(), 1);
+
+        let (a, b) = (MemoryId::new(0), MemoryId::new(1));
+
+        mm.get(a).grow(1);
+        mm.get(a).write(0, &[1]);
+
+        mm.get(b).grow(1);
+        mm.get(b).write(0, &[2]);
+
+        let mut out = [0; 1];
+        mm.get(a).read(0, &mut out);
+        assert_eq!(&out, &[1]);
+        mm.get(b).read(0, &mut out);
+        assert_eq!(&out, &[2]);
+
+        mm.reclaim_memory(a);
+
+        mm.get(b).read(0, &mut out);
+        assert_eq!(&out, &[2]);
+
+        mm.get(b).grow(1);
+
+        mm.get(b).read(0, &mut out);
+        assert_eq!(&out, &[2]);
+
+        let mm = MemoryManager::init_with_bucket_size(mock_stable_memory.clone(), 1);
+        mm.get(b).read(0, &mut out);
+        assert_eq!(&out, &[2]); // This assert fails because buckets were loaded in a different order.
+    }
+
     #[test]
     fn free_buckets_reused_in_increasing_order() {
         let mem_mgr = MemoryManager::init(make_memory());

From 2fa4e550c8c198219e8927b711412412e1b5cfc5 Mon Sep 17 00:00:00 2001
From: Maksym Arutyunyan <maksym.arutyunyan@dfinity.org>
Date: Wed, 20 Aug 2025 16:41:18 +0200
Subject: [PATCH 02/18] fix bucket ordering

---
 src/memory_manager.rs | 25 +++++++++++++++++++++----
 1 file changed, 21 insertions(+), 4 deletions(-)

diff --git a/src/memory_manager.rs b/src/memory_manager.rs
index e06d6628..6a252429 100644
--- a/src/memory_manager.rs
+++ b/src/memory_manager.rs
@@ -395,11 +395,27 @@ impl<M: Memory> MemoryManagerInner<M> {
         // Allocate new buckets as needed.
         memory_bucket.reserve(new_buckets_needed as usize);
         for _ in 0..new_buckets_needed {
-            let new_bucket_id = if let Some(free_bucket) = self.free_buckets.pop_front() {
-                // Reuse a bucket from the free pool (lowest bucket ID first)
-                free_bucket
+            let new_bucket_id = if let Some(free_bucket) = self.free_buckets.front() {
+                // Only reuse free bucket if its ID is higher than current max for this memory
+                // This ensures bucket ordering is preserved across memory manager reloads
+                // For fresh memories (no buckets), allow reusing any free bucket
+                let can_reuse = if memory_bucket.is_empty() {
+                    true // Fresh memory can reuse any bucket
+                } else {
+                    let current_max_bucket = memory_bucket.iter().map(|b| b.0).max().unwrap();
+                    free_bucket.0 > current_max_bucket
+                };
+
+                if can_reuse {
+                    self.free_buckets.pop_front().unwrap()
+                } else {
+                    // Can't reuse this bucket, allocate a new one
+                    let bucket = BucketId(self.allocated_buckets);
+                    self.allocated_buckets += 1;
+                    bucket
+                }
             } else {
-                // Allocate a new bucket
+                // No free buckets available, allocate a new one
                 let bucket = BucketId(self.allocated_buckets);
                 self.allocated_buckets += 1;
                 bucket
@@ -1313,6 +1329,7 @@ mod test {
         mm.get(b).read(0, &mut out);
         assert_eq!(&out, &[2]);
 
+        // Reload memory manager, eg. after upgrade.
         let mm = MemoryManager::init_with_bucket_size(mock_stable_memory.clone(), 1);
         mm.get(b).read(0, &mut out);
         assert_eq!(&out, &[2]); // This assert fails because buckets were loaded in a different order.

From c56179cbc8791d7a11e43eeb93c0377d04f9fe7b Mon Sep 17 00:00:00 2001
From: Maksym Arutyunyan <maksym.arutyunyan@dfinity.org>
Date: Wed, 20 Aug 2025 17:17:54 +0200
Subject: [PATCH 03/18] fix reusing virtual memory buckets

---
 debug_bucket_test.rs  |  30 +++++
 src/memory_manager.rs | 265 ++++++++++++++++++++++--------------------
 2 files changed, 166 insertions(+), 129 deletions(-)
 create mode 100644 debug_bucket_test.rs

diff --git a/debug_bucket_test.rs b/debug_bucket_test.rs
new file mode 100644
index 00000000..21ca56a0
--- /dev/null
+++ b/debug_bucket_test.rs
@@ -0,0 +1,30 @@
+use std::cell::RefCell;
+use std::rc::Rc;
+
+// This is a quick debug script to understand the bucket allocation behavior
+
+fn main() {
+    println!("Let me trace through the bucket_ordering_preserved_across_reload test...");
+
+    // This is what happens in the test:
+    // 1. Memory A gets bucket 0, Memory B gets bucket 1
+    // 2. Memory A is reclaimed (bucket 0 → free pool)
+    // 3. Memory B grows again
+
+    println!("Step 1: Memory A = [0], Memory B = [1]");
+    println!("Step 2: Reclaim Memory A → free pool = [0], Memory B = [1]");
+    println!("Step 3: Memory B grows again...");
+
+    // With conservative bucket reuse:
+    // - Memory B current max bucket = 1
+    // - Free bucket 0 is available
+    // - Can reuse bucket 0? free_bucket.0 > current_max_bucket → 0 > 1 → FALSE
+    // - Therefore: allocate new bucket 2
+
+    println!("Conservative reuse check:");
+    println!("  Memory B max bucket ID: 1");
+    println!("  Available free bucket: 0");
+    println!("  Can reuse? 0 > 1 = false");
+    println!("  Result: Allocate new bucket 2");
+    println!("Final: Memory B = [1, 2] NOT [1, 0]");
+}
diff --git a/src/memory_manager.rs b/src/memory_manager.rs
index 6a252429..14ab4bbd 100644
--- a/src/memory_manager.rs
+++ b/src/memory_manager.rs
@@ -1219,184 +1219,191 @@ mod test {
     }
 
     #[test]
-    fn memory_grows_without_manual_reclaim() {
+    fn reclaim_empty_memory_returns_zero() {
+        let mem_mgr = MemoryManager::init(make_memory());
+        let memory = mem_mgr.get(MemoryId::new(0));
+
+        assert_eq!(memory.size(), 0);
+        assert_eq!(mem_mgr.reclaim_memory(MemoryId::new(0)), 0);
+        assert_eq!(memory.size(), 0);
+    }
+
+    #[test]
+    fn reclaim_returns_correct_page_count() {
+        let mem_mgr = MemoryManager::init(make_memory());
+        let memory = mem_mgr.get(MemoryId::new(0));
+
+        let pages_allocated = BUCKET_SIZE_IN_PAGES * 2;
+        memory.grow(pages_allocated);
+
+        let pages_reclaimed = mem_mgr.reclaim_memory(MemoryId::new(0));
+        assert_eq!(pages_reclaimed, pages_allocated);
+    }
+
+    #[test]
+    fn reclaim_resets_memory_size_to_zero() {
         let mem = make_memory();
         let mem_mgr = MemoryManager::init(mem.clone());
-        let initial_size = mem.size();
+        let memory = mem_mgr.get(MemoryId::new(0));
+
+        memory.grow(BUCKET_SIZE_IN_PAGES * 2);
+        assert_eq!(memory.size(), BUCKET_SIZE_IN_PAGES * 2);
 
-        // Create first memory and allocate bucket 0
+        mem_mgr.reclaim_memory(MemoryId::new(0));
+        assert_eq!(memory.size(), 0);
+
+        // Verify state persists after reload
+        let mem_mgr_reloaded = MemoryManager::init(mem);
+        let memory_reloaded = mem_mgr_reloaded.get(MemoryId::new(0));
+        assert_eq!(memory_reloaded.size(), 0);
+    }
+
+    #[test]
+    fn memory_grows_without_reclaim() {
+        let mem = make_memory();
+        let mem_mgr = MemoryManager::init(mem.clone());
+
+        // Allocate two memories without reclaiming - should use separate buckets
         let memory_0 = mem_mgr.get(MemoryId::new(0));
+        let memory_1 = mem_mgr.get(MemoryId::new(1));
+
         memory_0.grow(BUCKET_SIZE_IN_PAGES);
-        memory_0.write(0, b"bucket_id_0");
+        memory_0.write(0, b"data_0");
         let size_after_first = mem.size();
 
-        // Verify first allocation grew memory
-        assert_eq!(size_after_first, initial_size + BUCKET_SIZE_IN_PAGES);
-
-        // Create second memory WITHOUT reclaiming first - should allocate bucket 1
-        let memory_1 = mem_mgr.get(MemoryId::new(1));
         memory_1.grow(BUCKET_SIZE_IN_PAGES);
-        memory_1.write(0, b"bucket_id_1");
+        memory_1.write(0, b"data_1");
 
-        // Verify memory grew again (no bucket reuse - new bucket allocated)
-        assert_eq!(mem.size(), size_after_first + BUCKET_SIZE_IN_PAGES);
+        // Second allocation should grow underlying memory
+        assert!(mem.size() > size_after_first);
 
-        // Verify both memories have their expected content (different buckets)
-        let mut buf = vec![0u8; 11];
+        // Verify distinct data in separate buckets
+        let mut buf = vec![0u8; 6];
         memory_0.read(0, &mut buf);
-        assert_eq!(
-            &buf, b"bucket_id_0",
-            "Memory 0 should have bucket 0 content"
-        );
+        assert_eq!(&buf, b"data_0");
         memory_1.read(0, &mut buf);
-        assert_eq!(
-            &buf, b"bucket_id_1",
-            "Memory 1 should have bucket 1 content"
-        );
+        assert_eq!(&buf, b"data_1");
     }
 
     #[test]
-    fn memory_reuses_buckets_with_manual_reclaim() {
+    fn memory_reuses_buckets_after_reclaim() {
         let mem = make_memory();
         let mem_mgr = MemoryManager::init(mem.clone());
 
-        // Create first memory and allocate bucket 0
+        // Allocate and reclaim first memory
         let memory_0 = mem_mgr.get(MemoryId::new(0));
         memory_0.grow(BUCKET_SIZE_IN_PAGES);
-        memory_0.write(0, b"bucket_id_0");
+        memory_0.write(0, b"original");
         let size_after_allocation = mem.size();
 
-        // Manually reclaim first memory
-        let pages_reclaimed = mem_mgr.reclaim_memory(MemoryId::new(0));
-        assert_eq!(
-            pages_reclaimed, BUCKET_SIZE_IN_PAGES,
-            "Should reclaim exactly {} pages",
-            BUCKET_SIZE_IN_PAGES
-        );
-
-        // CRITICAL: Drop the memory_0 after memory reclamation as it's now invalid
-        drop(memory_0);
+        mem_mgr.reclaim_memory(MemoryId::new(0));
+        drop(memory_0); // Drop invalidated memory
 
-        // Verify memory size didn't change (buckets marked free, not deallocated)
-        assert_eq!(mem.size(), size_after_allocation);
-
-        // Create second memory AFTER manual memory reclamation - should reuse bucket 0
+        // Allocate second memory - should reuse bucket without growing
         let memory_1 = mem_mgr.get(MemoryId::new(1));
         memory_1.grow(BUCKET_SIZE_IN_PAGES);
+        assert_eq!(mem.size(), size_after_allocation);
 
-        // Verify memory did NOT grow (bucket reuse - no new allocation)
-        assert_eq!(
-            mem.size(),
-            size_after_allocation,
-            "Memory should not grow when reusing buckets"
-        );
-
-        // Verify bucket reuse by reading original content from bucket 0
-        let mut buf = vec![0u8; 11];
+        // Verify bucket reuse by reading original data
+        let mut buf = vec![0u8; 8];
         memory_1.read(0, &mut buf);
-        assert_eq!(
-            &buf, b"bucket_id_0",
-            "Memory 1 should see original bucket 0 content, proving reuse"
-        );
+        assert_eq!(&buf, b"original");
+
+        // Verify state persists after reload
+        let mem_mgr_reloaded = MemoryManager::init(mem);
+        let memory_1_reloaded = mem_mgr_reloaded.get(MemoryId::new(1));
+        assert_eq!(memory_1_reloaded.size(), BUCKET_SIZE_IN_PAGES);
+        memory_1_reloaded.read(0, &mut buf);
+        assert_eq!(&buf, b"original");
     }
 
     #[test]
-    fn bucket_ordering() {
-        let mock_stable_memory = make_memory();
-        let mm = MemoryManager::init_with_bucket_size(mock_stable_memory.clone(), 1);
-
-        let (a, b) = (MemoryId::new(0), MemoryId::new(1));
-
-        mm.get(a).grow(1);
-        mm.get(a).write(0, &[1]);
-
-        mm.get(b).grow(1);
-        mm.get(b).write(0, &[2]);
-
-        let mut out = [0; 1];
-        mm.get(a).read(0, &mut out);
-        assert_eq!(&out, &[1]);
-        mm.get(b).read(0, &mut out);
-        assert_eq!(&out, &[2]);
-
-        mm.reclaim_memory(a);
+    fn conservative_bucket_reuse_prevents_corruption() {
+        let mem = make_memory();
+        let mem_mgr = MemoryManager::init(mem.clone());
 
-        mm.get(b).read(0, &mut out);
-        assert_eq!(&out, &[2]);
+        // Create memories with buckets 0 and 1
+        let memory_0 = mem_mgr.get(MemoryId::new(0));
+        let memory_1 = mem_mgr.get(MemoryId::new(1));
+        memory_0.grow(BUCKET_SIZE_IN_PAGES);
+        memory_1.grow(BUCKET_SIZE_IN_PAGES);
 
-        mm.get(b).grow(1);
+        // Reclaim memory 0 (frees bucket 0)
+        mem_mgr.reclaim_memory(MemoryId::new(0));
 
-        mm.get(b).read(0, &mut out);
-        assert_eq!(&out, &[2]);
+        // Memory 1 cannot reuse lower-ID bucket 0 - must allocate bucket 2
+        let initial_size = mem_mgr.inner.borrow().memory.size();
+        memory_1.grow(BUCKET_SIZE_IN_PAGES);
+        assert!(mem_mgr.inner.borrow().memory.size() > initial_size);
 
-        // Reload memory manager, eg. after upgrade.
-        let mm = MemoryManager::init_with_bucket_size(mock_stable_memory.clone(), 1);
-        mm.get(b).read(0, &mut out);
-        assert_eq!(&out, &[2]); // This assert fails because buckets were loaded in a different order.
+        // Verify conservative reuse persists after reload
+        let mem_mgr_reloaded = MemoryManager::init(mem);
+        let memory_1_reloaded = mem_mgr_reloaded.get(MemoryId::new(1));
+        assert_eq!(memory_1_reloaded.size(), BUCKET_SIZE_IN_PAGES * 2);
     }
 
     #[test]
-    fn free_buckets_reused_in_increasing_order() {
+    fn free_buckets_reused_in_order() {
         let mem_mgr = MemoryManager::init(make_memory());
 
-        // Allocate multiple buckets in sequence (0, 1, 2, 3, 4)
+        // Allocate buckets 0, 1, 2
         let memory_0 = mem_mgr.get(MemoryId::new(0));
         let memory_1 = mem_mgr.get(MemoryId::new(1));
+        let memory_2 = mem_mgr.get(MemoryId::new(2));
 
-        memory_0.grow(BUCKET_SIZE_IN_PAGES * 3); // allocates buckets 0, 1, 2
-        memory_1.grow(BUCKET_SIZE_IN_PAGES * 2); // allocates buckets 3, 4
-
-        // Write bucket identifiers to prove allocation order
-        memory_0.write(0, b"bucket_id_0");
-        memory_0.write(BUCKET_SIZE_IN_PAGES * WASM_PAGE_SIZE, b"bucket_id_1");
-        memory_0.write(BUCKET_SIZE_IN_PAGES * WASM_PAGE_SIZE * 2, b"bucket_id_2");
-        memory_1.write(0, b"bucket_id_3");
-        memory_1.write(BUCKET_SIZE_IN_PAGES * WASM_PAGE_SIZE, b"bucket_id_4");
+        memory_0.grow(BUCKET_SIZE_IN_PAGES);
+        memory_0.write(0, b"bucket_0");
+        memory_1.grow(BUCKET_SIZE_IN_PAGES);
+        memory_1.write(0, b"bucket_1");
+        memory_2.grow(BUCKET_SIZE_IN_PAGES);
+        memory_2.write(0, b"bucket_2");
 
-        // Reclaim all buckets in reverse order to test sorting
-        let pages_reclaimed_1 = mem_mgr.reclaim_memory(MemoryId::new(1)); // reclaims 3, 4
-        let pages_reclaimed_0 = mem_mgr.reclaim_memory(MemoryId::new(0)); // reclaims 0, 1, 2
-        assert_eq!(
-            pages_reclaimed_1,
-            BUCKET_SIZE_IN_PAGES * 2,
-            "Should reclaim {} pages from memory 1",
-            BUCKET_SIZE_IN_PAGES * 2
-        );
-        assert_eq!(
-            pages_reclaimed_0,
-            BUCKET_SIZE_IN_PAGES * 3,
-            "Should reclaim {} pages from memory 0",
-            BUCKET_SIZE_IN_PAGES * 3
-        );
+        // Reclaim buckets 2 and 0 (out of order)
+        mem_mgr.reclaim_memory(MemoryId::new(2));
+        mem_mgr.reclaim_memory(MemoryId::new(0));
 
-        // Allocate new memories - should reuse buckets in increasing order (0, 1, then 2)
-        let memory_2 = mem_mgr.get(MemoryId::new(2));
+        // New allocations should reuse buckets 0, 2 in increasing order
         let memory_3 = mem_mgr.get(MemoryId::new(3));
         let memory_4 = mem_mgr.get(MemoryId::new(4));
+        memory_3.grow(BUCKET_SIZE_IN_PAGES);
+        memory_4.grow(BUCKET_SIZE_IN_PAGES);
 
-        memory_2.grow(BUCKET_SIZE_IN_PAGES); // should get bucket 0 (lowest available)
-        memory_3.grow(BUCKET_SIZE_IN_PAGES); // should get bucket 1 (next lowest)
-        memory_4.grow(BUCKET_SIZE_IN_PAGES); // should get bucket 2 (next lowest)
-
-        // Verify buckets were reused in increasing order by reading original content
-        let mut buf = vec![0u8; 11];
-        memory_2.read(0, &mut buf);
-        assert_eq!(
-            &buf, b"bucket_id_0",
-            "Memory 2 should reuse bucket 0 (lowest)"
-        );
-
+        let mut buf = vec![0u8; 8];
         memory_3.read(0, &mut buf);
-        assert_eq!(
-            &buf, b"bucket_id_1",
-            "Memory 3 should reuse bucket 1 (next lowest)"
-        );
-
+        assert_eq!(&buf, b"bucket_0"); // Reused bucket 0 first
         memory_4.read(0, &mut buf);
-        assert_eq!(
-            &buf, b"bucket_id_2",
-            "Memory 4 should reuse bucket 2 (next lowest)"
-        );
+        assert_eq!(&buf, b"bucket_2"); // Reused bucket 2 second
+    }
+
+    #[test]
+    fn bucket_ordering_preserved_across_reload() {
+        let mem = make_memory();
+        let mem_mgr = MemoryManager::init_with_bucket_size(mem.clone(), 1);
+
+        // Create memories with buckets 0 and 1
+        let memory_a = mem_mgr.get(MemoryId::new(0));
+        let memory_b = mem_mgr.get(MemoryId::new(1));
+        memory_a.grow(1);
+        memory_a.write(0, b"a_data");
+        memory_b.grow(1);
+        memory_b.write(0, b"b_data");
+
+        // Reclaim memory A and grow memory B
+        // Note: Conservative reuse prevents Memory B from getting bucket 0 (since 0 < 1)
+        // So Memory B gets a new bucket 2, resulting in buckets [1, 2]
+        mem_mgr.reclaim_memory(MemoryId::new(0));
+        memory_b.grow(1);
+
+        // Verify memory B still reads from its first bucket
+        let mut buf = vec![0u8; 6];
+        memory_b.read(0, &mut buf);
+        assert_eq!(&buf, b"b_data");
+
+        // Critical: After reload, memory B should still read same data
+        let mem_mgr_reloaded = MemoryManager::init_with_bucket_size(mem, 1);
+        let memory_b_reloaded = mem_mgr_reloaded.get(MemoryId::new(1));
+        memory_b_reloaded.read(0, &mut buf);
+        assert_eq!(&buf, b"b_data");
     }
 }
 

From 8321cc491a8b7567ae7f0e29b15a7759a10bfabc Mon Sep 17 00:00:00 2001
From: Maksym Arutyunyan <maksym.arutyunyan@dfinity.org>
Date: Wed, 20 Aug 2025 17:19:18 +0200
Subject: [PATCH 04/18] cleanup

---
 debug_bucket_test.rs  | 30 ------------------------------
 src/memory_manager.rs | 11 ++++++-----
 2 files changed, 6 insertions(+), 35 deletions(-)
 delete mode 100644 debug_bucket_test.rs

diff --git a/debug_bucket_test.rs b/debug_bucket_test.rs
deleted file mode 100644
index 21ca56a0..00000000
--- a/debug_bucket_test.rs
+++ /dev/null
@@ -1,30 +0,0 @@
-use std::cell::RefCell;
-use std::rc::Rc;
-
-// This is a quick debug script to understand the bucket allocation behavior
-
-fn main() {
-    println!("Let me trace through the bucket_ordering_preserved_across_reload test...");
-
-    // This is what happens in the test:
-    // 1. Memory A gets bucket 0, Memory B gets bucket 1
-    // 2. Memory A is reclaimed (bucket 0 → free pool)
-    // 3. Memory B grows again
-
-    println!("Step 1: Memory A = [0], Memory B = [1]");
-    println!("Step 2: Reclaim Memory A → free pool = [0], Memory B = [1]");
-    println!("Step 3: Memory B grows again...");
-
-    // With conservative bucket reuse:
-    // - Memory B current max bucket = 1
-    // - Free bucket 0 is available
-    // - Can reuse bucket 0? free_bucket.0 > current_max_bucket → 0 > 1 → FALSE
-    // - Therefore: allocate new bucket 2
-
-    println!("Conservative reuse check:");
-    println!("  Memory B max bucket ID: 1");
-    println!("  Available free bucket: 0");
-    println!("  Can reuse? 0 > 1 = false");
-    println!("  Result: Allocate new bucket 2");
-    println!("Final: Memory B = [1, 2] NOT [1, 0]");
-}
diff --git a/src/memory_manager.rs b/src/memory_manager.rs
index 14ab4bbd..f43cb0ab 100644
--- a/src/memory_manager.rs
+++ b/src/memory_manager.rs
@@ -1259,11 +1259,12 @@ mod test {
     }
 
     #[test]
-    fn memory_grows_without_reclaim() {
+    fn memory_growth_without_reclaim_uses_new_buckets() {
         let mem = make_memory();
         let mem_mgr = MemoryManager::init(mem.clone());
 
-        // Allocate two memories without reclaiming - should use separate buckets
+        // Baseline test: allocate two memories without reclaiming - should use separate buckets
+        // This contrasts with reclaim tests where buckets can be reused
         let memory_0 = mem_mgr.get(MemoryId::new(0));
         let memory_1 = mem_mgr.get(MemoryId::new(1));
 
@@ -1343,7 +1344,7 @@ mod test {
     }
 
     #[test]
-    fn free_buckets_reused_in_order() {
+    fn reclaim_frees_buckets_for_reuse_in_order() {
         let mem_mgr = MemoryManager::init(make_memory());
 
         // Allocate buckets 0, 1, 2
@@ -1358,11 +1359,11 @@ mod test {
         memory_2.grow(BUCKET_SIZE_IN_PAGES);
         memory_2.write(0, b"bucket_2");
 
-        // Reclaim buckets 2 and 0 (out of order)
+        // Reclaim buckets 2 and 0 (out of order) to test free bucket sorting
         mem_mgr.reclaim_memory(MemoryId::new(2));
         mem_mgr.reclaim_memory(MemoryId::new(0));
 
-        // New allocations should reuse buckets 0, 2 in increasing order
+        // New allocations should reuse reclaimed buckets 0, 2 in increasing order
         let memory_3 = mem_mgr.get(MemoryId::new(3));
         let memory_4 = mem_mgr.get(MemoryId::new(4));
         memory_3.grow(BUCKET_SIZE_IN_PAGES);

From 14aab6f4b4ab0ba7afd85a0563ce7c4f7b974314 Mon Sep 17 00:00:00 2001
From: Maksym Arutyunyan <maksym.arutyunyan@dfinity.org>
Date: Wed, 20 Aug 2025 17:25:18 +0200
Subject: [PATCH 05/18] update readme

---
 README.md | 11 ++++++++---
 1 file changed, 8 insertions(+), 3 deletions(-)

diff --git a/README.md b/README.md
index f9486d84..40fbe520 100644
--- a/README.md
+++ b/README.md
@@ -103,6 +103,11 @@ assert_eq!(map_b.get(&1), Some(b'B')); // ✅ Succeeds: No data corruption
 
 During data migration scenarios, you often need to create a new data structure (B) and populate it with data from an existing structure (A). Without memory reclamation, this process doubles memory usage even after A is no longer needed.
 
+> **⚠️ CRITICAL SAFETY REQUIREMENT:** 
+> - **MUST** drop the original structure object before calling `reclaim_memory`
+> - **NEVER** use the original structure after reclamation - causes data corruption
+> - Only reclaim memory for completely finished structures, not active ones
+
 Consider this migration scenario:
 
 ```rust
@@ -133,15 +138,15 @@ let actual_size_after_migration = mem.size();
 assert!(actual_size_before_migration < actual_size_after_migration);
 
 // ========================================
-// Scenario 2: WITH reclamation
+// Scenario 2: WITH reclamation (SAFE)
 // ========================================
 let mut map_a: BTreeMap<u64, u8, _> = BTreeMap::init(mem_mgr.get(mem_id_a));
 map_a.insert(1, b'A');              // Populate map A with data
 let data = map_a.get(&1);           // Extract data for migration
 map_a.clear_new();                  // A is now empty
-drop(map_a);                        // Drop A completely
+drop(map_a);                        // 1. Drop A FIRST (critical!)
 let actual_size_before_migration = mem.size();
-mem_mgr.reclaim_memory(mem_id_a);   // Free A's memory buckets for reuse
+mem_mgr.reclaim_memory(mem_id_a);   // 2. Reclaim A's memory safely
 
 let mut map_b: BTreeMap<u64, u8, _> = BTreeMap::init(mem_mgr.get(mem_id_b));
 map_b.insert(1, data.unwrap());     // B reuses A's reclaimed memory buckets

From 5be72e5fe1968d27be86a21921c2e3d5954a03f8 Mon Sep 17 00:00:00 2001
From: Maksym Arutyunyan <maksym.arutyunyan@dfinity.org>
Date: Wed, 20 Aug 2025 17:39:01 +0200
Subject: [PATCH 06/18] add tests

---
 src/memory_manager.rs | 69 ++++++++++++++++++++++++++++++++-----------
 1 file changed, 52 insertions(+), 17 deletions(-)

diff --git a/src/memory_manager.rs b/src/memory_manager.rs
index f43cb0ab..0b074853 100644
--- a/src/memory_manager.rs
+++ b/src/memory_manager.rs
@@ -395,30 +395,39 @@ impl<M: Memory> MemoryManagerInner<M> {
         // Allocate new buckets as needed.
         memory_bucket.reserve(new_buckets_needed as usize);
         for _ in 0..new_buckets_needed {
-            let new_bucket_id = if let Some(free_bucket) = self.free_buckets.front() {
-                // Only reuse free bucket if its ID is higher than current max for this memory
-                // This ensures bucket ordering is preserved across memory manager reloads
-                // For fresh memories (no buckets), allow reusing any free bucket
-                let can_reuse = if memory_bucket.is_empty() {
-                    true // Fresh memory can reuse any bucket
+            let new_bucket_id = if memory_bucket.is_empty() {
+                // Fresh memory can reuse any free bucket (prefer lowest ID for locality)
+                if let Some(free_bucket) = self.free_buckets.pop_front() {
+                    free_bucket
                 } else {
-                    let current_max_bucket = memory_bucket.iter().map(|b| b.0).max().unwrap();
-                    free_bucket.0 > current_max_bucket
-                };
+                    // No free buckets available, allocate a new one
+                    let bucket = BucketId(self.allocated_buckets);
+                    self.allocated_buckets += 1;
+                    bucket
+                }
+            } else {
+                // Find the first free bucket with ID higher than current max for this memory
+                // This ensures bucket ordering is preserved across memory manager reloads
+                let current_max_bucket = memory_bucket.iter().map(|b| b.0).max().unwrap();
+
+                // Search through all free buckets to find the first one with ID > current_max_bucket
+                let mut reusable_bucket_index = None;
+                for (i, free_bucket) in self.free_buckets.iter().enumerate() {
+                    if free_bucket.0 > current_max_bucket {
+                        reusable_bucket_index = Some(i);
+                        break;
+                    }
+                }
 
-                if can_reuse {
-                    self.free_buckets.pop_front().unwrap()
+                if let Some(index) = reusable_bucket_index {
+                    // Remove and return the found bucket
+                    self.free_buckets.remove(index).unwrap()
                 } else {
-                    // Can't reuse this bucket, allocate a new one
+                    // No suitable free bucket found, allocate a new one
                     let bucket = BucketId(self.allocated_buckets);
                     self.allocated_buckets += 1;
                     bucket
                 }
-            } else {
-                // No free buckets available, allocate a new one
-                let bucket = BucketId(self.allocated_buckets);
-                self.allocated_buckets += 1;
-                bucket
             };
 
             memory_bucket.push(new_bucket_id);
@@ -1343,6 +1352,32 @@ mod test {
         assert_eq!(memory_1_reloaded.size(), BUCKET_SIZE_IN_PAGES * 2);
     }
 
+    #[test]
+    fn optimal_match_conservative_bucket_reuse_prevents_corruption() {
+        let mem = make_memory();
+        let mem_mgr = MemoryManager::init(mem.clone());
+
+        // Create memories with buckets [0, 2] and [1]
+        let memory_0 = mem_mgr.get(MemoryId::new(0));
+        let memory_1 = mem_mgr.get(MemoryId::new(1));
+        memory_0.grow(BUCKET_SIZE_IN_PAGES); // 0
+        memory_1.grow(BUCKET_SIZE_IN_PAGES); // 1
+        memory_0.grow(BUCKET_SIZE_IN_PAGES); // 2
+
+        // Reclaim memory 0 (frees buckets 0 and 2)
+        mem_mgr.reclaim_memory(MemoryId::new(0));
+
+        // Memory 1 cannot reuse lower-ID bucket 0, but can reuse bucket 2 -- no waste
+        let initial_size = mem_mgr.inner.borrow().memory.size();
+        memory_1.grow(BUCKET_SIZE_IN_PAGES);
+        assert!(mem_mgr.inner.borrow().memory.size() == initial_size);
+
+        // Verify conservative reuse persists after reload
+        let mem_mgr_reloaded = MemoryManager::init(mem);
+        let memory_1_reloaded = mem_mgr_reloaded.get(MemoryId::new(1));
+        assert_eq!(memory_1_reloaded.size(), BUCKET_SIZE_IN_PAGES * 2);
+    }
+
     #[test]
     fn reclaim_frees_buckets_for_reuse_in_order() {
         let mem_mgr = MemoryManager::init(make_memory());

From ac57d309d371272d7f163d1ff732bdcaa709c9c5 Mon Sep 17 00:00:00 2001
From: Maksym Arutyunyan <maksym.arutyunyan@dfinity.org>
Date: Wed, 20 Aug 2025 17:52:00 +0200
Subject: [PATCH 07/18] .

---
 README.md | 7 +++----
 1 file changed, 3 insertions(+), 4 deletions(-)

diff --git a/README.md b/README.md
index 40fbe520..c0df579f 100644
--- a/README.md
+++ b/README.md
@@ -106,7 +106,6 @@ During data migration scenarios, you often need to create a new data structure (
 > **⚠️ CRITICAL SAFETY REQUIREMENT:** 
 > - **MUST** drop the original structure object before calling `reclaim_memory`
 > - **NEVER** use the original structure after reclamation - causes data corruption
-> - Only reclaim memory for completely finished structures, not active ones
 
 Consider this migration scenario:
 
@@ -138,15 +137,15 @@ let actual_size_after_migration = mem.size();
 assert!(actual_size_before_migration < actual_size_after_migration);
 
 // ========================================
-// Scenario 2: WITH reclamation (SAFE)
+// Scenario 2: WITH reclamation
 // ========================================
 let mut map_a: BTreeMap<u64, u8, _> = BTreeMap::init(mem_mgr.get(mem_id_a));
 map_a.insert(1, b'A');              // Populate map A with data
 let data = map_a.get(&1);           // Extract data for migration
 map_a.clear_new();                  // A is now empty
-drop(map_a);                        // 1. Drop A FIRST (critical!)
+drop(map_a);                        // Drop A completely
 let actual_size_before_migration = mem.size();
-mem_mgr.reclaim_memory(mem_id_a);   // 2. Reclaim A's memory safely
+mem_mgr.reclaim_memory(mem_id_a);   // Free A's memory buckets for reuse
 
 let mut map_b: BTreeMap<u64, u8, _> = BTreeMap::init(mem_mgr.get(mem_id_b));
 map_b.insert(1, data.unwrap());     // B reuses A's reclaimed memory buckets

From 19b632dff657bb4964b03611dcc82709dbfe065f Mon Sep 17 00:00:00 2001
From: Maksym Arutyunyan <maksym.arutyunyan@dfinity.org>
Date: Wed, 20 Aug 2025 18:18:56 +0200
Subject: [PATCH 08/18] update readme

---
 README.md | 21 +++++++++++++--------
 1 file changed, 13 insertions(+), 8 deletions(-)

diff --git a/README.md b/README.md
index c0df579f..f8654c2e 100644
--- a/README.md
+++ b/README.md
@@ -101,11 +101,17 @@ assert_eq!(map_b.get(&1), Some(b'B')); // ✅ Succeeds: No data corruption
 
 ### Memory Reclamation
 
-During data migration scenarios, you often need to create a new data structure (B) and populate it with data from an existing structure (A). Without memory reclamation, this process doubles memory usage even after A is no longer needed.
+During migrations you often create a new structure (B) and copy data from an existing one (A). Without reclamation, this can double memory usage even after A is no longer needed.
 
-> **⚠️ CRITICAL SAFETY REQUIREMENT:** 
-> - **MUST** drop the original structure object before calling `reclaim_memory`
-> - **NEVER** use the original structure after reclamation - causes data corruption
+Bucket IDs are an internal implementation detail — hidden and not user-controllable — and each virtual memory must receive bucket IDs in strictly ascending order. Because of this, reuse of freed buckets is guaranteed when allocating into a newly created (empty) structure. For existing structures, reuse may or may not work: it succeeds only if there is a free bucket with an ID greater than the structure’s current maximum; otherwise a new bucket is allocated.
+
+Example: A = `[0, 4, 5]`, B = `[1, 2, 3]`. After releasing A, `free = [0, 4, 5]`. When B grows, it can’t take `0` (must be `> 3`) but can take `4` → `B = [1, 2, 3, 4]`, `free = [0, 5]`.
+
+**Recommendation:** for predictable reuse migrate into a newly created structure rather than relying on reuse with a populated one.
+
+> **⚠️ CRITICAL SAFETY REQUIREMENT:**
+> - **MUST** drop the original structure object before calling `reclaim_memory`.
+> - **NEVER** use the original structure after reclamation — doing so corrupts data.
 
 Consider this migration scenario:
 
@@ -129,7 +135,7 @@ map_a.clear_new();                  // A is now empty
 drop(map_a);                        // Memory stays allocated to mem_id_a
 let actual_size_before_migration = mem.size();
 
-let mut map_b: BTreeMap<u64, u8, _> = BTreeMap::init(mem_mgr.get(mem_id_b));
+let mut map_b: BTreeMap<u64, u8, _> = BTreeMap::new(mem_mgr.get(mem_id_b));
 map_b.insert(1, data.unwrap());     // B allocates NEW memory
 let actual_size_after_migration = mem.size();
                                     // Result: ~2x memory usage
@@ -147,7 +153,8 @@ drop(map_a);                        // Drop A completely
 let actual_size_before_migration = mem.size();
 mem_mgr.reclaim_memory(mem_id_a);   // Free A's memory buckets for reuse
 
-let mut map_b: BTreeMap<u64, u8, _> = BTreeMap::init(mem_mgr.get(mem_id_b));
+// Reusing free memory buckets works best on newly created structures
+let mut map_b: BTreeMap<u64, u8, _> = BTreeMap::new(mem_mgr.get(mem_id_b));
 map_b.insert(1, data.unwrap());     // B reuses A's reclaimed memory buckets
 let actual_size_after_migration = mem.size();
                                     // Result: 1x memory usage
@@ -155,8 +162,6 @@ let actual_size_after_migration = mem.size();
 assert!(actual_size_before_migration == actual_size_after_migration);
 ```
 
-**Important**: Always drop the original structure before calling `reclaim_memory`.
-
 ## Example Canister
 
 Here's a fully working canister example that ties everything together.

From 2ef15ffe92a8af3c862b5ffdfa4efbb931181ba4 Mon Sep 17 00:00:00 2001
From: Maksym Arutyunyan <maksym.arutyunyan@dfinity.org>
Date: Thu, 21 Aug 2025 09:27:09 +0200
Subject: [PATCH 09/18] free_buckets btreeset

---
 benchmarks/memory_manager/canbench_results.yml | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/benchmarks/memory_manager/canbench_results.yml b/benchmarks/memory_manager/canbench_results.yml
index bf482bf0..f8b1942a 100644
--- a/benchmarks/memory_manager/canbench_results.yml
+++ b/benchmarks/memory_manager/canbench_results.yml
@@ -9,14 +9,14 @@ benches:
   memory_manager_grow:
     total:
       calls: 1
-      instructions: 347433966
+      instructions: 348682004
       heap_increase: 2
       stable_memory_increase: 32000
     scopes: {}
   memory_manager_overhead:
     total:
       calls: 1
-      instructions: 1181977502
+      instructions: 1181984577
       heap_increase: 0
       stable_memory_increase: 8320
     scopes: {}

From 1c20a788a4760ba40e801b35c96525948a5b1a77 Mon Sep 17 00:00:00 2001
From: Maksym Arutyunyan <maksym.arutyunyan@dfinity.org>
Date: Thu, 21 Aug 2025 09:27:32 +0200
Subject: [PATCH 10/18] free_buckets btreeset

---
 src/memory_manager.rs | 58 +++++++++++++++----------------------------
 1 file changed, 20 insertions(+), 38 deletions(-)

diff --git a/src/memory_manager.rs b/src/memory_manager.rs
index 0b074853..b82ab842 100644
--- a/src/memory_manager.rs
+++ b/src/memory_manager.rs
@@ -46,7 +46,8 @@ use crate::{
     write, write_struct, Memory, WASM_PAGE_SIZE,
 };
 use std::cell::{Cell, RefCell};
-use std::collections::VecDeque;
+use std::collections::BTreeSet;
+use std::ops::Bound::{Excluded, Unbounded};
 use std::rc::Rc;
 
 const MAGIC: &[u8; 3] = b"MGR";
@@ -270,7 +271,7 @@ struct MemoryManagerInner<M: Memory> {
     /// A pool of free buckets that can be reused by any memory.
     /// Maintained in increasing order (lowest bucket IDs at the front)
     /// to ensure optimal memory locality when reusing buckets.
-    free_buckets: VecDeque<BucketId>,
+    free_buckets: BTreeSet<BucketId>,
 }
 
 impl<M: Memory> MemoryManagerInner<M> {
@@ -299,7 +300,7 @@ impl<M: Memory> MemoryManagerInner<M> {
             memory_sizes_in_pages: [0; MAX_NUM_MEMORIES as usize],
             memory_buckets: vec![vec![]; MAX_NUM_MEMORIES as usize],
             bucket_size_in_pages,
-            free_buckets: VecDeque::new(),
+            free_buckets: BTreeSet::new(),
         };
 
         mem_mgr.save_header();
@@ -329,7 +330,7 @@ impl<M: Memory> MemoryManagerInner<M> {
         );
 
         let mut memory_buckets = vec![vec![]; MAX_NUM_MEMORIES as usize];
-        let mut free_buckets_vec = Vec::new();
+        let mut free_buckets = BTreeSet::new();
         for (bucket_idx, memory_id) in buckets.into_iter().enumerate() {
             let bucket_id = BucketId(bucket_idx as u16);
             if memory_id != UNALLOCATED_BUCKET_MARKER {
@@ -337,14 +338,10 @@ impl<M: Memory> MemoryManagerInner<M> {
             } else if bucket_id.0 < header.num_allocated_buckets {
                 // Only buckets with indices less than `num_allocated_buckets` were ever allocated.
                 // If such a bucket is now marked as unallocated, it is free to reuse.
-                free_buckets_vec.push(bucket_id);
+                free_buckets.insert(bucket_id);
             }
         }
 
-        // Sort free buckets in increasing order for optimal memory locality
-        free_buckets_vec.sort_by_key(|bucket| bucket.0);
-        let free_buckets = VecDeque::from(free_buckets_vec);
-
         Self {
             memory,
             allocated_buckets: header.num_allocated_buckets,
@@ -392,38 +389,30 @@ impl<M: Memory> MemoryManagerInner<M> {
         }
 
         let memory_bucket = &mut self.memory_buckets[id.0 as usize];
-        // Allocate new buckets as needed.
         memory_bucket.reserve(new_buckets_needed as usize);
+
         for _ in 0..new_buckets_needed {
             let new_bucket_id = if memory_bucket.is_empty() {
                 // Fresh memory can reuse any free bucket (prefer lowest ID for locality)
-                if let Some(free_bucket) = self.free_buckets.pop_front() {
+                if let Some(free_bucket) = self.free_buckets.pop_first() {
                     free_bucket
                 } else {
-                    // No free buckets available, allocate a new one
                     let bucket = BucketId(self.allocated_buckets);
                     self.allocated_buckets += 1;
                     bucket
                 }
             } else {
-                // Find the first free bucket with ID higher than current max for this memory
-                // This ensures bucket ordering is preserved across memory manager reloads
-                let current_max_bucket = memory_bucket.iter().map(|b| b.0).max().unwrap();
-
-                // Search through all free buckets to find the first one with ID > current_max_bucket
-                let mut reusable_bucket_index = None;
-                for (i, free_bucket) in self.free_buckets.iter().enumerate() {
-                    if free_bucket.0 > current_max_bucket {
-                        reusable_bucket_index = Some(i);
-                        break;
-                    }
-                }
-
-                if let Some(index) = reusable_bucket_index {
-                    // Remove and return the found bucket
-                    self.free_buckets.remove(index).unwrap()
+                // Buckets are sorted — last element is the max
+                let current_max_bucket = memory_bucket.last().unwrap().0;
+
+                // Find the smallest free bucket with ID > current_max_bucket
+                if let Some(&candidate) = self
+                    .free_buckets
+                    .range((Excluded(BucketId(current_max_bucket)), Unbounded))
+                    .next()
+                {
+                    self.free_buckets.take(&candidate).unwrap()
                 } else {
-                    // No suitable free bucket found, allocate a new one
                     let bucket = BucketId(self.allocated_buckets);
                     self.allocated_buckets += 1;
                     bucket
@@ -475,22 +464,15 @@ impl<M: Memory> MemoryManagerInner<M> {
         }
 
         // Mark all buckets as unallocated in stable storage and collect them
-        let mut reclaimed_buckets = Vec::new();
         for &bucket_id in memory_buckets.iter() {
             write(
                 &self.memory,
                 bucket_allocations_address(bucket_id).get(),
                 &[UNALLOCATED_BUCKET_MARKER],
             );
-            reclaimed_buckets.push(bucket_id);
+            self.free_buckets.insert(bucket_id);
         }
 
-        // Add reclaimed buckets to free pool and sort to maintain increasing order
-        self.free_buckets.extend(reclaimed_buckets);
-        let mut sorted_buckets: Vec<_> = self.free_buckets.drain(..).collect();
-        sorted_buckets.sort_by_key(|bucket| bucket.0);
-        self.free_buckets = VecDeque::from(sorted_buckets);
-
         // Clear the memory's bucket list
         memory_buckets.clear();
 
@@ -715,7 +697,7 @@ impl MemoryId {
 }
 
 // Referring to a bucket.
-#[derive(Clone, Copy, Debug, PartialEq)]
+#[derive(Copy, Clone, Eq, PartialEq, Ord, PartialOrd, Debug)]
 struct BucketId(u16);
 
 fn bucket_allocations_address(id: BucketId) -> Address {

From 350eb9dac123897f0957c3042cecd050a9e90227 Mon Sep 17 00:00:00 2001
From: Maksym Arutyunyan <maksym.arutyunyan@dfinity.org>
Date: Thu, 21 Aug 2025 09:46:52 +0200
Subject: [PATCH 11/18] cleanup

---
 .../memory_manager/canbench_results.yml       |  4 +-
 src/memory_manager.rs                         | 47 ++++++++-----------
 2 files changed, 22 insertions(+), 29 deletions(-)

diff --git a/benchmarks/memory_manager/canbench_results.yml b/benchmarks/memory_manager/canbench_results.yml
index f8b1942a..2da1eb60 100644
--- a/benchmarks/memory_manager/canbench_results.yml
+++ b/benchmarks/memory_manager/canbench_results.yml
@@ -9,14 +9,14 @@ benches:
   memory_manager_grow:
     total:
       calls: 1
-      instructions: 348682004
+      instructions: 348618011
       heap_increase: 2
       stable_memory_increase: 32000
     scopes: {}
   memory_manager_overhead:
     total:
       calls: 1
-      instructions: 1181984577
+      instructions: 1181984482
       heap_increase: 0
       stable_memory_increase: 8320
     scopes: {}
diff --git a/src/memory_manager.rs b/src/memory_manager.rs
index b82ab842..05b7ea26 100644
--- a/src/memory_manager.rs
+++ b/src/memory_manager.rs
@@ -390,38 +390,31 @@ impl<M: Memory> MemoryManagerInner<M> {
 
         let memory_bucket = &mut self.memory_buckets[id.0 as usize];
         memory_bucket.reserve(new_buckets_needed as usize);
-
         for _ in 0..new_buckets_needed {
-            let new_bucket_id = if memory_bucket.is_empty() {
-                // Fresh memory can reuse any free bucket (prefer lowest ID for locality)
-                if let Some(free_bucket) = self.free_buckets.pop_first() {
-                    free_bucket
-                } else {
-                    let bucket = BucketId(self.allocated_buckets);
-                    self.allocated_buckets += 1;
-                    bucket
-                }
-            } else {
-                // Buckets are sorted — last element is the max
-                let current_max_bucket = memory_bucket.last().unwrap().0;
-
-                // Find the smallest free bucket with ID > current_max_bucket
-                if let Some(&candidate) = self
-                    .free_buckets
-                    .range((Excluded(BucketId(current_max_bucket)), Unbounded))
-                    .next()
-                {
-                    self.free_buckets.take(&candidate).unwrap()
-                } else {
-                    let bucket = BucketId(self.allocated_buckets);
-                    self.allocated_buckets += 1;
-                    bucket
+            // Try to reuse a free bucket.
+            let maybe_free_bucket: Option<BucketId> = match memory_bucket.last() {
+                // Fresh memory -- use smallest free bucket.
+                None => self.free_buckets.pop_first(),
+                // Existing memory -- use smallest free bucket greater than current max.
+                Some(max_b) => {
+                    let next_gt = {
+                        let mut it = self
+                            .free_buckets
+                            .range((Excluded(BucketId(max_b.0)), Unbounded));
+                        it.next().copied()
+                    };
+                    next_gt.and_then(|b| self.free_buckets.take(&b))
                 }
             };
 
-            memory_bucket.push(new_bucket_id);
+            // If there is no free bucket available, allocate a new one.
+            let new_bucket_id = maybe_free_bucket.unwrap_or_else(|| {
+                let bucket = BucketId(self.allocated_buckets);
+                self.allocated_buckets += 1;
+                bucket
+            });
 
-            // Write in stable store that this bucket belongs to the memory with the provided `id`.
+            memory_bucket.push(new_bucket_id);
             write(
                 &self.memory,
                 bucket_allocations_address(new_bucket_id).get(),

From 4da3b09f9c61312fb467dc1f5dee2fed2e79a4dd Mon Sep 17 00:00:00 2001
From: Maksym Arutyunyan <maksym.arutyunyan@dfinity.org>
Date: Thu, 21 Aug 2025 09:47:32 +0200
Subject: [PATCH 12/18] .

---
 src/memory_manager.rs | 1 +
 1 file changed, 1 insertion(+)

diff --git a/src/memory_manager.rs b/src/memory_manager.rs
index 05b7ea26..0ce73dd4 100644
--- a/src/memory_manager.rs
+++ b/src/memory_manager.rs
@@ -389,6 +389,7 @@ impl<M: Memory> MemoryManagerInner<M> {
         }
 
         let memory_bucket = &mut self.memory_buckets[id.0 as usize];
+        // Allocate new buckets as needed.
         memory_bucket.reserve(new_buckets_needed as usize);
         for _ in 0..new_buckets_needed {
             // Try to reuse a free bucket.

From fd01ab3664e957fcc2a5b3b1fff460c7ff6f23ba Mon Sep 17 00:00:00 2001
From: Maksym Arutyunyan <maksym.arutyunyan@dfinity.org>
Date: Thu, 21 Aug 2025 09:48:24 +0200
Subject: [PATCH 13/18] .

---
 src/memory_manager.rs | 1 +
 1 file changed, 1 insertion(+)

diff --git a/src/memory_manager.rs b/src/memory_manager.rs
index 0ce73dd4..2c57a792 100644
--- a/src/memory_manager.rs
+++ b/src/memory_manager.rs
@@ -396,6 +396,7 @@ impl<M: Memory> MemoryManagerInner<M> {
             let maybe_free_bucket: Option<BucketId> = match memory_bucket.last() {
                 // Fresh memory -- use smallest free bucket.
                 None => self.free_buckets.pop_first(),
+
                 // Existing memory -- use smallest free bucket greater than current max.
                 Some(max_b) => {
                     let next_gt = {

From 94984b66de0191a6d1d2b0d5ae6878456495b0f2 Mon Sep 17 00:00:00 2001
From: Maksym Arutyunyan <maksym.arutyunyan@dfinity.org>
Date: Thu, 21 Aug 2025 09:53:53 +0200
Subject: [PATCH 14/18] cleanup

---
 .../memory_manager/canbench_results.yml       |  4 ++--
 src/memory_manager.rs                         | 23 +++++++++++--------
 2 files changed, 15 insertions(+), 12 deletions(-)

diff --git a/benchmarks/memory_manager/canbench_results.yml b/benchmarks/memory_manager/canbench_results.yml
index 2da1eb60..bd8ad406 100644
--- a/benchmarks/memory_manager/canbench_results.yml
+++ b/benchmarks/memory_manager/canbench_results.yml
@@ -9,14 +9,14 @@ benches:
   memory_manager_grow:
     total:
       calls: 1
-      instructions: 348618011
+      instructions: 348618009
       heap_increase: 2
       stable_memory_increase: 32000
     scopes: {}
   memory_manager_overhead:
     total:
       calls: 1
-      instructions: 1181984482
+      instructions: 1181984472
       heap_increase: 0
       stable_memory_increase: 8320
     scopes: {}
diff --git a/src/memory_manager.rs b/src/memory_manager.rs
index 2c57a792..b5d684f2 100644
--- a/src/memory_manager.rs
+++ b/src/memory_manager.rs
@@ -393,19 +393,22 @@ impl<M: Memory> MemoryManagerInner<M> {
         memory_bucket.reserve(new_buckets_needed as usize);
         for _ in 0..new_buckets_needed {
             // Try to reuse a free bucket.
-            let maybe_free_bucket: Option<BucketId> = match memory_bucket.last() {
-                // Fresh memory -- use smallest free bucket.
+            let maybe_free_bucket = match memory_bucket.last() {
+                // Fresh memory — smallest free bucket.
                 None => self.free_buckets.pop_first(),
 
-                // Existing memory -- use smallest free bucket greater than current max.
+                // Existing memory — smallest free bucket > current max.
                 Some(max_b) => {
-                    let next_gt = {
-                        let mut it = self
-                            .free_buckets
-                            .range((Excluded(BucketId(max_b.0)), Unbounded));
-                        it.next().copied()
-                    };
-                    next_gt.and_then(|b| self.free_buckets.take(&b))
+                    if let Some(candidate) = self
+                        .free_buckets
+                        .range((Excluded(max_b), Unbounded))
+                        .next()
+                        .copied()
+                    {
+                        self.free_buckets.take(&candidate)
+                    } else {
+                        None
+                    }
                 }
             };
 

From 98744f60cce30f76344d7b804ef511312acabd77 Mon Sep 17 00:00:00 2001
From: Maksym Arutyunyan <maksym.arutyunyan@dfinity.org>
Date: Thu, 21 Aug 2025 09:58:07 +0200
Subject: [PATCH 15/18] add debug assert

---
 src/memory_manager.rs | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/src/memory_manager.rs b/src/memory_manager.rs
index b5d684f2..8d82d52e 100644
--- a/src/memory_manager.rs
+++ b/src/memory_manager.rs
@@ -419,6 +419,11 @@ impl<M: Memory> MemoryManagerInner<M> {
                 bucket
             });
 
+            // Assert ascending order (strictly increasing).
+            if let Some(prev_max) = memory_bucket.last().copied() {
+                debug_assert!(prev_max < new_bucket_id);
+            }
+
             memory_bucket.push(new_bucket_id);
             write(
                 &self.memory,

From 93c20907237fa9202b843b5517c75b9c6d365604 Mon Sep 17 00:00:00 2001
From: Maksym Arutyunyan <maksym.arutyunyan@dfinity.org>
Date: Thu, 21 Aug 2025 10:14:45 +0200
Subject: [PATCH 16/18] docs

---
 docs/src/concepts/memory-manager.md | 19 ++++++++++++++++---
 1 file changed, 16 insertions(+), 3 deletions(-)

diff --git a/docs/src/concepts/memory-manager.md b/docs/src/concepts/memory-manager.md
index bf4452b9..59690170 100644
--- a/docs/src/concepts/memory-manager.md
+++ b/docs/src/concepts/memory-manager.md
@@ -40,7 +40,19 @@ Each memory instance must be assigned to exactly one stable structure.
 
 ## Memory Reclamation
 
-During data migration scenarios, you often need to create a new data structure and populate it with data from an existing structure. Without memory reclamation, this process doubles memory usage even after the original structure is no longer needed.
+During migrations you often create a new structure (B) and copy data from an existing one (A). Without reclamation, this can double memory usage even after A is no longer needed.
+
+Bucket IDs are an internal implementation detail — hidden and not user-controllable — and each virtual memory must receive bucket IDs in strictly ascending order. Because of this, reuse of freed buckets is guaranteed when allocating into a newly created (empty) structure. For existing structures, reuse may or may not work: it succeeds only if there is a free bucket with an ID greater than the structure’s current maximum; otherwise a new bucket is allocated.
+
+Example: A = `[0, 4, 5]`, B = `[1, 2, 3]`. After releasing A, `free = [0, 4, 5]`. When B grows, it can’t take `0` (must be `> 3`) but can take `4` → `B = [1, 2, 3, 4]`, `free = [0, 5]`.
+
+**Recommendation:** for predictable reuse migrate into a newly created structure rather than relying on reuse with a populated one.
+
+```admonish info ""
+**⚠️ CRITICAL SAFETY REQUIREMENT:**
+- **MUST** drop the original structure object before calling `reclaim_memory`.
+- **NEVER** use the original structure after reclamation — doing so corrupts data.
+```
 
 The `MemoryManager` provides a `reclaim_memory` method to efficiently handle these scenarios:
 
@@ -64,7 +76,7 @@ map_a.clear_new();                  // A is now empty
 drop(map_a);                        // Memory stays allocated to mem_id_a
 let actual_size_before_migration = mem.size();
 
-let mut map_b: BTreeMap<u64, u8, _> = BTreeMap::init(mem_mgr.get(mem_id_b));
+let mut map_b: BTreeMap<u64, u8, _> = BTreeMap::new(mem_mgr.get(mem_id_b));
 map_b.insert(1, data.unwrap());     // B allocates NEW memory
 let actual_size_after_migration = mem.size();
                                     // Result: ~2x memory usage
@@ -82,7 +94,8 @@ drop(map_a);                        // Drop A completely
 let actual_size_before_migration = mem.size();
 mem_mgr.reclaim_memory(mem_id_a);   // Free A's memory buckets for reuse
 
-let mut map_b: BTreeMap<u64, u8, _> = BTreeMap::init(mem_mgr.get(mem_id_b));
+// Reusing free memory buckets works best on newly created structures
+let mut map_b: BTreeMap<u64, u8, _> = BTreeMap::new(mem_mgr.get(mem_id_b));
 map_b.insert(1, data.unwrap());     // B reuses A's reclaimed memory buckets
 let actual_size_after_migration = mem.size();
                                     // Result: 1x memory usage

From 88ffbb6d17504baf135ba8d601c8c88333ce58b9 Mon Sep 17 00:00:00 2001
From: Maksym Arutyunyan <maksym.arutyunyan@dfinity.org>
Date: Thu, 21 Aug 2025 10:15:29 +0200
Subject: [PATCH 17/18] docs

---
 docs/src/concepts/memory-manager.md | 5 -----
 1 file changed, 5 deletions(-)

diff --git a/docs/src/concepts/memory-manager.md b/docs/src/concepts/memory-manager.md
index 59690170..4fa27edd 100644
--- a/docs/src/concepts/memory-manager.md
+++ b/docs/src/concepts/memory-manager.md
@@ -102,8 +102,3 @@ let actual_size_after_migration = mem.size();
                                     // Memory allocation stayed the same (no waste)
 assert!(actual_size_before_migration == actual_size_after_migration);
 ```
-
-```admonish info ""
-**Important**: Always drop the original structure before calling `reclaim_memory`.
-The method returns the number of pages that were reclaimed and made available for reuse.
-```

From a52350336a970b5c39677dfe1e63b72ce9b19f94 Mon Sep 17 00:00:00 2001
From: Maksym Arutyunyan <maksym.arutyunyan@dfinity.org>
Date: Thu, 21 Aug 2025 10:16:20 +0200
Subject: [PATCH 18/18] docs

---
 README.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/README.md b/README.md
index f8654c2e..e14e461a 100644
--- a/README.md
+++ b/README.md
@@ -113,7 +113,7 @@ Example: A = `[0, 4, 5]`, B = `[1, 2, 3]`. After releasing A, `free = [0, 4, 5]`
 > - **MUST** drop the original structure object before calling `reclaim_memory`.
 > - **NEVER** use the original structure after reclamation — doing so corrupts data.
 
-Consider this migration scenario:
+The `MemoryManager` provides a `reclaim_memory` method to efficiently handle these scenarios:
 
 ```rust
 use ic_stable_structures::{