feat: support weak map (#13)

Author: shruti2522

notes/2026-02-26.md

@@ -0,0 +1,91 @@
+# WeakMap Support
+
+Note author: shruti2522
+
+This document summarizes all the changes made to add weak map support to the mark sweep collector.
+
+This adds `WeakMap<K, V>` to the mark sweep collector backed by ephemerons.
+Each entry's key is held weakly and when the key is collected, the entry is pruned from the map automatically.
+
+## changes made in PR #13
+
+### `WeakGcBox` refactor
+
+Previously `WeakGcBox<T>` allocated its own `GcBox<T>` and used an `IS_WEAK`
+bit in the header flags to mark it as weak. this approach made sense as a first pass but a weak reference to an existing `Gc<K>` should point at the same allocation, not create a second one.
+The GC would never see the key as dead if `WeakGcBox` held its own copy.
+
+So I refactored `WeakGcBox` to hold an `ErasedArenaPointer` into an existing
+allocation, also removed `IS_WEAK` flag, `weak_white`/`weak_black` header constructors
+and the related `GcBox` code from `gc_header.rs` and `gc_box.rs`
+
+### `Ephemeron::new_in` now uses `WeakGcBox<K>`
+
+`Ephemeron::new_in` used to accept a value `K` and create a `WeakGcBox`
+internally. Now it takes a `WeakGcBox<K>` directly. The caller creates the weak
+reference from an existing `Gc<K>`. This keeps things organized by
+separating how we create the weak link from how the ephemeron uses it
+
+### new cleanup functions added to `EphemeronVTable`
+
+Added `is_reachable_fn` and `finalize_fn` to `EphemeronVTable`. This
+lets the collector's sweep phase check if a key is still alive and run any
+finalizers without needing to know the specific types involved.
+
+### `WeakMap<K, V>`
+
+It's a `HashMap<usize, ArenaPointer<Ephemeron<K, V>>>`,
+keyed by the raw pointer address of `Gc<K>`. This gives O(1) average time
+for insert, lookup and remove
+
+- `insert` removes any existing entry for the key before allocating a new ephemeron.
+This prevents the old one from leaking into the collector queue when a value is
+updated
+
+- `remove` takes the entry out of the map but the backing ephemeron stays in the
+collector queue, it gets swept when the key is collected.
+
+`prune_dead_entries` only visits entries still in `self.entries` so there is no
+risk of reading freed memory.
+
+- `Trace` is handled, the ephemerons are in the collector's own queue so `WeakMap` itself doesn't need to do anything extra
+
+## how the collector owns the map
+
+`WeakMap::new(collector)` boxes the `WeakMapInner<K, V>`, grabs a raw pointer
+from the box before erasing it to `Box<dyn ErasedWeakMap>`, then pushes that
+erased box into `collector.weak_maps`. The `WeakMap` handle just holds that raw
+pointer and is valid for as long as the collector lives
+
+`ErasedWeakMap` is an internal helper with one method: `prune_dead_entries`.
+during `collect()` the collector calls it on every map in `weak_maps` after the
+sweep phase but before the dead arenas are freed. this ensures we can still read
+the dropped flag on ephemerons before their memory is gone.
+
+In my fisrt attempt, I made users manually register their maps with the collector, this
+was awkward because it required `unsafe` code. It was also hard for the collector to safely
+keep track of maps that lived outside its own memory. Worse, if a user forgot to
+unregister a map before deleting it, the collector was left holding a bad pointer
+which could cause the whole program to crash
+
+By allocating `WeakMapInner` on the collector's heap and giving the collector ownership,
+we eliminated the manual registration step completely. the user gets a handle that
+is valid for the collector's lifetime, the aliasing
+concerns are resolved because the collector owns the memory it prunes
+
+## potential improvements
+
+**`weak_map.rs`**
+
+- explore using `HashTable` instead of `HashMap` to save memory.
+- consider whether `insert` should take an `Ephemeron` directly instead of a key/value pair.
+
+## conclusion
+
+This approach gives us a map that cleans itself up automatically. It lives exactly
+as long as the collector does, users don't have to fiddle with manual registration
+or write `unsafe` code, and it plugs right into the collector's existing trace
+and sweep phases.
+
+In the future, the best improvement would be switching from `HashMap` to `HashTable`
+to save memory. Until then, this first version works well and gives us the weak map behavior we need

oscars/src/alloc/arena2/alloc.rs

@@ -353,6 +353,19 @@ impl<'arena> Arena<'arena> {
         }
         true
     }
+
+    // checks dropped items in this arena
+    #[cfg(test)]
+    pub fn item_drop_states(&self) -> rust_alloc::vec::Vec<bool> {
+        let mut result = rust_alloc::vec::Vec::new();
+        let mut unchecked_ptr = self.last_allocation.get();
+        while let Some(node) = NonNull::new(unchecked_ptr) {
+            let item = unsafe { node.as_ref() };
+            result.push(item.is_dropped());
+            unchecked_ptr = item.next.as_ptr() as *mut ErasedHeapItem
+        }
+        result
+    }
 }
 
 impl<'arena> Drop for Arena<'arena> {

oscars/src/alloc/arena2/mod.rs

@@ -6,7 +6,7 @@ use rust_alloc::collections::LinkedList;
 mod alloc;
 
 use alloc::Arena;
-pub use alloc::{ArenaAllocationData, ArenaHeapItem, ArenaPointer, ErasedArenaPointer};
+pub use alloc::{ArenaAllocationData, ArenaHeapItem, ArenaPointer, ErasedArenaPointer, ErasedHeapItem};
 
 #[cfg(test)]
 mod tests;
@@ -142,4 +142,13 @@ impl<'alloc> ArenaAllocator<'alloc> {
             drop(dead_arenas)
         }
     }
+
+    // checks dropped items across all arenas
+    #[cfg(test)]
+    pub fn arena_drop_states(&self) -> rust_alloc::vec::Vec<rust_alloc::vec::Vec<bool>> {
+        self.arenas
+            .iter()
+            .map(|arena| arena.item_drop_states())
+            .collect()
+    }
 }

oscars/src/collectors/mark_sweep/README.md

@@ -4,8 +4,8 @@ This is a basic mark-sweep collector using an underlying arena allocator.
 
 ## TODO list
 
-- [ ] Support weak maps
-- [ ] Add Tests
+- [x] Support weak maps
+- [x] Add Tests
 
 
 ## Areas of improvement

oscars/src/collectors/mark_sweep/internals/ephemeron.rs

@@ -7,33 +7,29 @@ use crate::{
     collectors::mark_sweep::{
         CollectionState, ErasedEphemeron, MarkSweepGarbageCollector, TraceColor,
         internals::{GcBox, WeakGcBox, gc_header::HeaderColor},
+        pointers::Gc,
         trace::Trace,
     },
 };
 
 use crate::collectors::mark_sweep::Finalize;
 
-// TODO: key's GcBox should be notably a weak box
 pub struct Ephemeron<K: Trace + ?Sized + 'static, V: Trace + 'static> {
     pub(crate) value: GcBox<V>,
     vtable: &'static EphemeronVTable,
     pub(crate) key: WeakGcBox<K>,
 }
 
-// NOTE: There is going to be an issue here in that we initialize the GC
-// box to the wrong state.
-//
-// So we either need the color to be global that is provided to the allocation
-// or we need state access
 impl<K: Trace, V: Trace> Ephemeron<K, V> {
-    pub fn new_in(key: K, value: V, collector: &mut MarkSweepGarbageCollector) -> Self
-    where
-        K: Sized,
+    // Creates a new [`Ephemeron`] with given key and value
+    //
+    // The [`WeakGcBox`] for the key is created internally from the provided [`Gc`] pointer
+    pub fn new_in(key: &Gc<K>, value: V, collector: &mut MarkSweepGarbageCollector) -> Self
     {
-        let key = WeakGcBox::new_in(key, &collector.state);
-        let value = GcBox::new_in(value, &collector.state);
+        let weak_key = WeakGcBox::new(key.inner_ptr);
+        let value = GcBox::new(value, &collector.state);
         let vtable = vtable_of::<K, V>();
-        Self { key, value, vtable }
+        Self { key: weak_key, value, vtable }
     }
 
     pub fn key(&self) -> &K {
@@ -61,6 +57,14 @@ impl<K: Trace, V: Trace> Ephemeron<K, V> {
     pub(crate) fn drop_fn(&self) -> EphemeronDropFn {
         self.vtable.drop_fn
     }
+
+    pub(crate) fn is_reachable_fn(&self) -> EphemeronIsReachableFn {
+        self.vtable.is_reachable_fn
+    }
+	
+    pub(crate) fn finalize_fn(&self) -> EphemeronFinalizeFn {
+        self.vtable.finalize_fn
+    }
 }
 
 impl<K: Trace, V: Trace> Finalize for Ephemeron<K, V> {}
@@ -120,12 +124,42 @@ pub(crate) const fn vtable_of<K: Trace + 'static, V: Trace + 'static>() -> &'sta
             // SAFETY: The caller must ensure the erased pointer is not dropped or deallocated.
             unsafe { this.as_mut().mark_dropped() };
         }
+
+        // SAFETY: Cast back to concrete types to check reachability
+        unsafe fn is_reachable_fn<K: Trace + 'static, V: Trace + 'static>(
+            this: ErasedEphemeron,
+            color: TraceColor,
+        ) -> bool {
+            // SAFETY: The caller must ensure that the passed erased pointer is
+            // `ArenaHeapItem<Ephemeron<K, V>>`
+            let ephemeron = unsafe {
+                this.cast::<ArenaHeapItem<Ephemeron<K, V>>>()
+                    .as_ref()
+                    .value()
+            };
+            ephemeron.is_reachable(color)
+        }
+
+        // SAFETY: Cast back to concrete types to run finalizers
+        unsafe fn finalize_fn<K: Trace + 'static, V: Trace + 'static>(this: ErasedEphemeron) {
+            // SAFETY: The caller must ensure that the passed erased pointer is
+            // `ArenaHeapItem<Ephemeron<K, V>>`
+            let ephemeron = unsafe {
+                this.cast::<ArenaHeapItem<Ephemeron<K, V>>>()
+                    .as_ref()
+                    .value()
+            };
+            Finalize::finalize(ephemeron.key());
+            Finalize::finalize(ephemeron.value());
+        }
     }
 
     impl<K: Trace + 'static, V: Trace + 'static> HasVTable for EphemeronMarker<K, V> {
         const VTABLE: &'static EphemeronVTable = &EphemeronVTable {
             trace_fn: EphemeronMarker::<K, V>::trace_fn::<K, V>,
             drop_fn: EphemeronMarker::<K, V>::drop_fn::<K, V>,
+            is_reachable_fn: EphemeronMarker::<K, V>::is_reachable_fn::<K, V>,
+            finalize_fn: EphemeronMarker::<K, V>::finalize_fn::<K, V>,
             _key_type_id: TypeId::of::<K>(),
             _key_size: size_of::<WeakGcBox<K>>(),
             _value_type_id: TypeId::of::<V>(),
@@ -138,10 +172,14 @@ pub(crate) const fn vtable_of<K: Trace + 'static, V: Trace + 'static>() -> &'sta
 
 type EphemeronTraceFn = unsafe fn(this: ErasedEphemeron, color: TraceColor);
 type EphemeronDropFn = unsafe fn(this: ErasedEphemeron);
+type EphemeronIsReachableFn = unsafe fn(this: ErasedEphemeron, color: TraceColor) -> bool;
+type EphemeronFinalizeFn = unsafe fn(this: ErasedEphemeron);
 
 pub struct EphemeronVTable {
     trace_fn: EphemeronTraceFn,
     drop_fn: EphemeronDropFn,
+    is_reachable_fn: EphemeronIsReachableFn,
+    finalize_fn: EphemeronFinalizeFn,
     _key_type_id: TypeId,
     _key_size: usize,
     _value_type_id: TypeId,