feat: replace HashMap with HashTable for WeakMap (#49)

Author: shruti2522

oscars/src/collectors/mark_sweep/pointers/weak_map.rs

@@ -1,14 +1,22 @@
-use rustc_hash::FxHashMap;
+use hashbrown::HashTable;
+use rustc_hash::FxHasher;
 
 use crate::{
     alloc::mempool3::PoolPointer,
     collectors::collector::Collector,
     collectors::mark_sweep::{Finalize, TraceColor, internals::Ephemeron, trace::Trace},
 };
-use core::ptr::NonNull;
+use core::{hash::Hasher, ptr::NonNull};
 
 use super::Gc;
 
+#[inline]
+fn hash_addr(addr: usize) -> u64 {
+    let mut h = FxHasher::default();
+    h.write_usize(addr);
+    h.finish()
+}
+
 // type erased trait so the collector can prune any WeakMap without knowing K/V
 #[doc(hidden)]
 pub trait ErasedWeakMap {
@@ -17,23 +25,27 @@ pub trait ErasedWeakMap {
 }
 
 // the actual weak map store, managed by the collector
-//
-// TODO: a HashTable might be a better approach here
 struct WeakMapInner<K: Trace + 'static, V: Trace + 'static> {
-    entries: FxHashMap<usize, PoolPointer<'static, Ephemeron<K, V>>>,
+    // keyed by the raw pointer address of the GC object; stored inline as
+    // `(addr, ptr)` so HashTable needs no separate key allocation
+    entries: HashTable<(usize, PoolPointer<'static, Ephemeron<K, V>>)>,
     is_alive: core::cell::Cell<bool>,
 }
 
 impl<K: Trace, V: Trace> WeakMapInner<K, V> {
     fn new() -> Self {
         Self {
-            entries: FxHashMap::default(),
+            entries: HashTable::new(),
             is_alive: core::cell::Cell::new(true),
         }
     }
 
     fn remove_and_invalidate(&mut self, key_addr: usize) {
-        if let Some(old_ephemeron) = self.entries.remove(&key_addr) {
+        if let Ok(entry) = self
+            .entries
+            .find_entry(hash_addr(key_addr), |e| e.0 == key_addr)
+        {
+            let ((_, old_ephemeron), _) = entry.remove();
             old_ephemeron.as_inner_ref().invalidate();
         }
     }
@@ -43,40 +55,48 @@ impl<K: Trace, V: Trace> WeakMapInner<K, V> {
         key_addr: usize,
         ephemeron_ptr: PoolPointer<'static, Ephemeron<K, V>>,
     ) {
-        self.entries.insert(key_addr, ephemeron_ptr);
+        // caller guarantees no duplicate exists since remove_and_invalidate was called first
+        self.entries
+            .insert_unique(hash_addr(key_addr), (key_addr, ephemeron_ptr), |e| {
+                hash_addr(e.0)
+            });
     }
 
     fn get(&self, key: &Gc<K>) -> Option<&V> {
         let key_addr = key.inner_ptr.as_non_null().as_ptr() as usize;
         self.entries
-            .get(&key_addr)
-            .map(|p| p.as_inner_ref().value())
+            .find(hash_addr(key_addr), |e| e.0 == key_addr)
+            .map(|(_, p)| p.as_inner_ref().value())
     }
 
     fn is_key_alive(&self, key: &Gc<K>) -> bool {
         let key_addr = key.inner_ptr.as_non_null().as_ptr() as usize;
-        self.entries.contains_key(&key_addr)
+        self.entries
+            .find(hash_addr(key_addr), |e| e.0 == key_addr)
+            .is_some()
     }
 
     fn remove(&mut self, key: &Gc<K>) -> bool {
         let key_addr = key.inner_ptr.as_non_null().as_ptr() as usize;
         // the backing ephemeron stays in the collector queue and gets swept
         // when the key is collected
-        self.entries
-            .remove(&key_addr)
-            .map(|p| {
-                p.as_inner_ref().invalidate();
-            })
-            .is_some()
+        if let Ok(entry) = self
+            .entries
+            .find_entry(hash_addr(key_addr), |e| e.0 == key_addr)
+        {
+            let ((_, ptr), _) = entry.remove();
+            ptr.as_inner_ref().invalidate();
+            true
+        } else {
+            false
+        }
     }
 }
 
 impl<K: Trace, V: Trace> ErasedWeakMap for WeakMapInner<K, V> {
     fn prune_dead_entries(&mut self, color: TraceColor) {
-        self.entries.retain(|_, ephemeron_ptr| {
-            let ephemeron = ephemeron_ptr.as_inner_ref();
-            ephemeron.is_reachable(color)
-        });
+        self.entries
+            .retain(|(_, ephemeron_ptr)| ephemeron_ptr.as_inner_ref().is_reachable(color));
     }
 
     fn is_alive(&self) -> bool {

oscars/src/collectors/mark_sweep/tests.rs

@@ -422,6 +422,118 @@ fn alive_wm() {
     );
 }
 
+#[test]
+fn two_distinct_keys_wm() {
+    // Ensure two distinct Gc keys produce unique addresses and don't collide
+    // in the HashTable or corrupt each other's entries
+    let collector = &mut MarkSweepGarbageCollector::default()
+        .with_page_size(256)
+        .with_heap_threshold(512);
+
+    let mut map = WeakMap::new(collector);
+    let key1 = Gc::new_in(1u64, collector);
+    let key2 = Gc::new_in(2u64, collector);
+
+    map.insert(&key1, 10u64, collector);
+    map.insert(&key2, 20u64, collector);
+
+    // Verify both entries are independent and accessible
+    assert_eq!(
+        map.get(&key1),
+        Some(&10u64),
+        "key1 lookup returned wrong value"
+    );
+    assert_eq!(
+        map.get(&key2),
+        Some(&20u64),
+        "key2 lookup returned wrong value"
+    );
+    assert!(map.is_key_alive(&key1), "key1 should be alive");
+    assert!(map.is_key_alive(&key2), "key2 should be alive");
+
+    // Dropping one key must prune only its entry, leaving the other intact
+    drop(key1);
+    collector.collect();
+
+    assert_eq!(
+        map.get(&key2),
+        Some(&20u64),
+        "key2 incorrectly pruned after key1 was collected",
+    );
+    assert!(
+        map.is_key_alive(&key2),
+        "key2 reported dead while still rooted",
+    );
+
+    // Both keys dead and collected means no leaks
+    drop(key2);
+    collector.collect();
+
+    assert_eq!(
+        collector.allocator.borrow().pools_len(),
+        0,
+        "ephemerons leaked after both keys collected",
+    );
+}
+
+#[test]
+fn two_maps_same_key_wm() {
+    // Same Gc key registered in two independent WeakMaps, each must hold its
+    // own value and neither map's operations must bleed into the other
+    let collector = &mut MarkSweepGarbageCollector::default()
+        .with_page_size(256)
+        .with_heap_threshold(512);
+
+    let key = Gc::new_in(1u64, collector);
+    let mut map1 = WeakMap::new(collector);
+    let mut map2 = WeakMap::new(collector);
+
+    map1.insert(&key, 10u64, collector);
+    map2.insert(&key, 20u64, collector);
+
+    assert_eq!(map1.get(&key), Some(&10u64));
+    assert_eq!(map2.get(&key), Some(&20u64));
+
+    drop(key);
+    collector.collect();
+
+    assert_eq!(
+        collector.allocator.borrow().pools_len(),
+        0,
+        "ephemerons leaked after key collected"
+    );
+}
+
+#[test]
+fn drop_map_with_live_key_wm() {
+    // Dropping a WeakMap while its key is still alive must not corrupt the
+    // collector, it skips dead maps during prune_dead_entries
+    let collector = &mut MarkSweepGarbageCollector::default()
+        .with_page_size(256)
+        .with_heap_threshold(512);
+
+    let key = Gc::new_in(42u64, collector);
+
+    {
+        let mut map = WeakMap::new(collector);
+        map.insert(&key, 100u64, collector);
+        // Map dropped here, WeakMap::drop sets is_alive = false
+    }
+
+    // Collector must handle the dead map entry without panic
+    collector.collect();
+
+    // Key still live, drop it then collect to verify no leak
+    drop(key);
+    collector.collect();
+
+    assert_eq!(
+        collector.allocator.borrow().pools_len(),
+        0,
+        "ephemerons leaked after map dropped with live key"
+    );
+}
+
 /// Edge-case stability tests for the mark-sweep garbage collector.
 ///
 /// These tests exercise corner cases that could cause crashes, stack overflows,

oscars/src/collectors/mark_sweep_arena2/pointers/weak_map.rs

@@ -1,13 +1,21 @@
-use rustc_hash::FxHashMap;
+use hashbrown::HashTable;
+use rustc_hash::FxHasher;
 
 use crate::{
     alloc::arena2::ArenaPointer,
     collectors::mark_sweep_arena2::{Finalize, TraceColor, internals::Ephemeron, trace::Trace},
 };
-use core::ptr::NonNull;
+use core::{hash::Hasher, ptr::NonNull};
 
 use super::Gc;
 
+#[inline]
+fn hash_addr(addr: usize) -> u64 {
+    let mut h = FxHasher::default();
+    h.write_usize(addr);
+    h.finish()
+}
+
 // type erased trait so the collector can prune any WeakMap without knowing K/V
 #[doc(hidden)]
 pub trait ErasedWeakMap {
@@ -16,23 +24,27 @@ pub trait ErasedWeakMap {
 }
 
 // the actual weak map store, managed by the collector
-//
-// TODO: a HashTable might be a better approach here
 struct WeakMapInner<K: Trace + 'static, V: Trace + 'static> {
-    entries: FxHashMap<usize, ArenaPointer<'static, Ephemeron<K, V>>>,
+    // keyed by the raw pointer address of the GC object, stored inline as
+    // `(addr, ptr)` so HashTable needs no separate key allocation
+    entries: HashTable<(usize, ArenaPointer<'static, Ephemeron<K, V>>)>,
     is_alive: core::cell::Cell<bool>,
 }
 
 impl<K: Trace, V: Trace> WeakMapInner<K, V> {
     fn new() -> Self {
         Self {
-            entries: FxHashMap::default(),
+            entries: HashTable::new(),
             is_alive: core::cell::Cell::new(true),
         }
     }
 
     fn remove_and_invalidate(&mut self, key_addr: usize) {
-        if let Some(old_ephemeron) = self.entries.remove(&key_addr) {
+        if let Ok(entry) = self
+            .entries
+            .find_entry(hash_addr(key_addr), |e| e.0 == key_addr)
+        {
+            let ((_, old_ephemeron), _) = entry.remove();
             old_ephemeron.as_inner_ref().invalidate();
         }
     }
@@ -42,40 +54,48 @@ impl<K: Trace, V: Trace> WeakMapInner<K, V> {
         key_addr: usize,
         ephemeron_ptr: ArenaPointer<'static, Ephemeron<K, V>>,
     ) {
-        self.entries.insert(key_addr, ephemeron_ptr);
+        // caller guarantees no duplicate exists since remove_and_invalidate was called first
+        self.entries
+            .insert_unique(hash_addr(key_addr), (key_addr, ephemeron_ptr), |e| {
+                hash_addr(e.0)
+            });
     }
 
     fn get(&self, key: &Gc<K>) -> Option<&V> {
         let key_addr = key.inner_ptr.as_non_null().as_ptr() as usize;
         self.entries
-            .get(&key_addr)
-            .map(|p| p.as_inner_ref().value())
+            .find(hash_addr(key_addr), |e| e.0 == key_addr)
+            .map(|(_, p)| p.as_inner_ref().value())
     }
 
     fn is_key_alive(&self, key: &Gc<K>) -> bool {
         let key_addr = key.inner_ptr.as_non_null().as_ptr() as usize;
-        self.entries.contains_key(&key_addr)
+        self.entries
+            .find(hash_addr(key_addr), |e| e.0 == key_addr)
+            .is_some()
     }
 
     fn remove(&mut self, key: &Gc<K>) -> bool {
         let key_addr = key.inner_ptr.as_non_null().as_ptr() as usize;
         // the backing ephemeron stays in the collector queue and gets swept
         // when the key is collected
-        self.entries
-            .remove(&key_addr)
-            .map(|p| {
-                p.as_inner_ref().invalidate();
-            })
-            .is_some()
+        if let Ok(entry) = self
+            .entries
+            .find_entry(hash_addr(key_addr), |e| e.0 == key_addr)
+        {
+            let ((_, ptr), _) = entry.remove();
+            ptr.as_inner_ref().invalidate();
+            true
+        } else {
+            false
+        }
     }
 }
 
 impl<K: Trace, V: Trace> ErasedWeakMap for WeakMapInner<K, V> {
     fn prune_dead_entries(&mut self, color: TraceColor) {
-        self.entries.retain(|_, ephemeron_ptr| {
-            let ephemeron = ephemeron_ptr.as_inner_ref();
-            ephemeron.is_reachable(color)
-        });
+        self.entries
+            .retain(|(_, ephemeron_ptr)| ephemeron_ptr.as_inner_ref().is_reachable(color));
     }
 
     fn is_alive(&self) -> bool {