Implement ephemeron-based weak map (#3052)

* Implement ephemeron-based weak map

* Document code and improve ephemerons

* Fix docs

* cargo clippy

boa_engine/src/builtins/weak_set/mod.rs

@@ -151,7 +151,7 @@ impl WeakSet {
 
         // 3. If Type(value) is not Object, throw a TypeError exception.
         let value = args.get_or_undefined(0);
-        let Some(value) = args.get_or_undefined(0).as_object() else {
+        let Some(value) = value.as_object() else {
             return Err(JsNativeError::typ()
                 .with_message(format!(
                     "WeakSet.add: expected target argument of type `object`, got target of type `{}`",

boa_gc/Cargo.toml

@@ -18,4 +18,5 @@ thinvec = ["thin-vec"]
 boa_profiler.workspace = true
 boa_macros.workspace = true
 
-thin-vec = { version = "0.2.12", optional = true }
+thin-vec = { version = "0.2.12", optional = true }
+hashbrown = { version = "0.14.0", features = ["raw"] }

boa_gc/src/internals/ephemeron_box.rs

@@ -1,6 +1,6 @@
 use crate::{trace::Trace, Gc, GcBox};
 use std::{
-    cell::Cell,
+    cell::{Cell, UnsafeCell},
     ptr::{self, NonNull},
 };
 
@@ -89,51 +89,60 @@ impl core::fmt::Debug for EphemeronBoxHeader {
 }
 
 /// The inner allocation of an [`Ephemeron`][crate::Ephemeron] pointer.
-pub(crate) struct EphemeronBox<K: Trace + ?Sized + 'static, V: Trace + 'static> {
+pub(crate) struct EphemeronBox<K: Trace + 'static, V: Trace + 'static> {
     pub(crate) header: EphemeronBoxHeader,
-    data: Cell<Option<NonNull<Data<K, V>>>>,
+    data: UnsafeCell<Option<Data<K, V>>>,
 }
 
-impl<K: Trace + ?Sized + 'static, V: Trace + 'static> Drop for EphemeronBox<K, V> {
-    fn drop(&mut self) {
-        if let Some(data) = self.data.take() {
-            // SAFETY: `data` comes from an `into_raw` call, so this pointer is safe to pass to
-            // `from_raw`.
-            drop(unsafe { Box::from_raw(data.as_ptr()) });
-        }
-    }
-}
-
-struct Data<K: Trace + ?Sized + 'static, V: Trace + 'static> {
+struct Data<K: Trace + 'static, V: Trace + 'static> {
     key: NonNull<GcBox<K>>,
     value: V,
 }
 
-impl<K: Trace + ?Sized, V: Trace> EphemeronBox<K, V> {
+impl<K: Trace, V: Trace> EphemeronBox<K, V> {
     pub(crate) fn new(key: &Gc<K>, value: V) -> Self {
-        let data = Box::into_raw(Box::new(Data {
-            key: key.inner_ptr(),
-            value,
-        }));
-        // SAFETY: `Box::into_raw` must always return a non-null pointer.
-        let data = unsafe { NonNull::new_unchecked(data) };
         Self {
             header: EphemeronBoxHeader::new(),
-            data: Cell::new(Some(data)),
+            data: UnsafeCell::new(Some(Data {
+                key: key.inner_ptr(),
+                value,
+            })),
         }
     }
 
-    /// Returns `true` if the two references refer to the same `GcBox`.
+    /// Returns `true` if the two references refer to the same `EphemeronBox`.
     pub(crate) fn ptr_eq(this: &Self, other: &Self) -> bool {
         // Use .header to ignore fat pointer vtables, to work around
         // https://github.com/rust-lang/rust/issues/46139
         ptr::eq(&this.header, &other.header)
     }
 
     /// Returns a reference to the ephemeron's value or None.
-    pub(crate) fn value(&self) -> Option<&V> {
-        // SAFETY: the garbage collector ensures `ptr` is valid as long as `data` is `Some`.
-        unsafe { self.data.get().map(|ptr| &ptr.as_ref().value) }
+    ///
+    /// # Safety
+    ///
+    /// The garbage collector must not run between the call to this function and the eventual
+    /// drop of the returned reference, since that could free the inner value.
+    pub(crate) unsafe fn value(&self) -> Option<&V> {
+        // SAFETY: the garbage collector ensures the ephemeron doesn't mutate until
+        // finalization.
+        let data = unsafe { &*self.data.get() };
+        data.as_ref().map(|data| &data.value)
+    }
+
+    /// Returns a reference to the ephemeron's key or None.
+    ///
+    /// # Safety
+    ///
+    /// The garbage collector must not run between the call to this function and the eventual
+    /// drop of the returned reference, since that could free the inner value.
+    pub(crate) unsafe fn key(&self) -> Option<&GcBox<K>> {
+        // SAFETY: the garbage collector ensures the ephemeron doesn't mutate until
+        // finalization.
+        unsafe {
+            let data = &*self.data.get();
+            data.as_ref().map(|data| data.key.as_ref())
+        }
     }
 
     /// Marks this `EphemeronBox` as live.
@@ -177,7 +186,7 @@ pub(crate) trait ErasedEphemeronBox {
     fn finalize_and_clear(&self);
 }
 
-impl<K: Trace + ?Sized, V: Trace> ErasedEphemeronBox for EphemeronBox<K, V> {
+impl<K: Trace, V: Trace> ErasedEphemeronBox for EphemeronBox<K, V> {
     fn header(&self) -> &EphemeronBoxHeader {
         &self.header
     }
@@ -187,12 +196,13 @@ impl<K: Trace + ?Sized, V: Trace> ErasedEphemeronBox for EphemeronBox<K, V> {
             return false;
         }
 
-        let Some(data) = self.data.get() else {
+        // SAFETY: the garbage collector ensures the ephemeron doesn't mutate until
+        // finalization.
+        let data = unsafe { &*self.data.get() };
+        let Some(data) = data.as_ref() else {
             return true;
         };
 
-        // SAFETY: `data` comes from a `Box`, so it is safe to dereference.
-        let data = unsafe { data.as_ref() };
         // SAFETY: `key` comes from a `Gc`, and the garbage collector only invalidates
         // `key` when it is unreachable, making `key` always valid.
         let key = unsafe { data.key.as_ref() };
@@ -209,20 +219,18 @@ impl<K: Trace + ?Sized, V: Trace> ErasedEphemeronBox for EphemeronBox<K, V> {
     }
 
     fn trace_non_roots(&self) {
-        let Some(data) = self.data.get() else {
-            return;
-        };
-        // SAFETY: `data` comes from a `Box`, so it is safe to dereference.
+        // SAFETY: Tracing always executes before collecting, meaning this cannot cause
+        // use after free.
         unsafe {
-            data.as_ref().value.trace_non_roots();
-        };
+            if let Some(value) = self.value() {
+                value.trace_non_roots();
+            }
+        }
     }
 
     fn finalize_and_clear(&self) {
-        if let Some(data) = self.data.take() {
-            // SAFETY: `data` comes from an `into_raw` call, so this pointer is safe to pass to
-            // `from_raw`.
-            let _contents = unsafe { Box::from_raw(data.as_ptr()) };
-        }
+        // SAFETY: the invariants of the garbage collector ensures this is only executed when
+        // there are no remaining references to the inner data.
+        unsafe { (*self.data.get()).take() };
     }
 }

boa_gc/src/internals/weak_map_box.rs

@@ -1,9 +1,9 @@
-use crate::{GcRefCell, Trace, WeakGc};
-use std::{cell::Cell, collections::HashMap, ptr::NonNull};
+use crate::{pointers::RawWeakMap, GcRefCell, Trace, WeakGc};
+use std::{cell::Cell, ptr::NonNull};
 
 /// A box that is used to track [`WeakMap`][`crate::WeakMap`]s.
 pub(crate) struct WeakMapBox<K: Trace + Sized + 'static, V: Trace + Sized + 'static> {
-    pub(crate) map: WeakGc<GcRefCell<HashMap<WeakGc<K>, V>>>,
+    pub(crate) map: WeakGc<GcRefCell<RawWeakMap<K, V>>>,
     pub(crate) next: Cell<Option<NonNull<dyn ErasedWeakMapBox>>>,
 }
 
@@ -18,15 +18,16 @@ pub(crate) trait ErasedWeakMapBox {
     /// Returns `true` if the [`WeakMapBox`] is live.
     fn is_live(&self) -> bool;
 
-    /// Traces the weak reference inside of the [`WeakMapBox`] it the weak map is live.
+    /// Traces the weak reference inside of the [`WeakMapBox`] if the weak map is live.
     unsafe fn trace(&self);
 }
 
-impl<K: Trace, V: Trace> ErasedWeakMapBox for WeakMapBox<K, V> {
+impl<K: Trace, V: Trace + Clone> ErasedWeakMapBox for WeakMapBox<K, V> {
     fn clear_dead_entries(&self) {
         if let Some(map) = self.map.upgrade() {
-            let mut map = map.borrow_mut();
-            map.retain(|k, _| k.upgrade().is_some());
+            if let Ok(mut map) = map.try_borrow_mut() {
+                map.clear_expired();
+            }
         }
     }
 

boa_gc/src/lib.rs

@@ -78,9 +78,9 @@ pub(crate) mod internals;
 
 use boa_profiler::Profiler;
 use internals::{EphemeronBox, ErasedEphemeronBox, ErasedWeakMapBox, WeakMapBox};
+use pointers::RawWeakMap;
 use std::{
     cell::{Cell, RefCell},
-    collections::HashMap,
     mem,
     ptr::NonNull,
 };
@@ -198,7 +198,7 @@ impl Allocator {
         })
     }
 
-    fn alloc_ephemeron<K: Trace + ?Sized, V: Trace>(
+    fn alloc_ephemeron<K: Trace, V: Trace>(
         value: EphemeronBox<K, V>,
     ) -> NonNull<EphemeronBox<K, V>> {
         let _timer = Profiler::global().start_event("New EphemeronBox", "BoaAlloc");
@@ -219,11 +219,11 @@ impl Allocator {
         })
     }
 
-    fn alloc_weak_map<K: Trace, V: Trace>() -> WeakMap<K, V> {
+    fn alloc_weak_map<K: Trace, V: Trace + Clone>() -> WeakMap<K, V> {
         let _timer = Profiler::global().start_event("New WeakMap", "BoaAlloc");
 
         let weak_map = WeakMap {
-            inner: Gc::new(GcRefCell::new(HashMap::new())),
+            inner: Gc::new(GcRefCell::new(RawWeakMap::new())),
         };
         let weak = WeakGc::new(&weak_map.inner);
 

boa_gc/src/pointers/ephemeron.rs

@@ -8,18 +8,20 @@ use std::ptr::NonNull;
 
 /// A key-value pair where the value becomes unaccesible when the key is garbage collected.
 ///
-/// See Racket's explanation on [**ephemerons**][eph] for a brief overview or read Barry Hayes'
-/// [_Ephemerons_: a new finalization mechanism][acm].
+/// You can read more about ephemerons on:
+/// - Racket's page about [**ephemerons**][eph], which gives a brief overview.
+/// - Barry Hayes' paper ["_Ephemerons_: a new finalization mechanism"][acm] which explains the topic
+/// in full detail.
 ///
 ///
 /// [eph]: https://docs.racket-lang.org/reference/ephemerons.html
 /// [acm]: https://dl.acm.org/doi/10.1145/263700.263733
 #[derive(Debug)]
-pub struct Ephemeron<K: Trace + ?Sized + 'static, V: Trace + 'static> {
+pub struct Ephemeron<K: Trace + 'static, V: Trace + 'static> {
     inner_ptr: NonNull<EphemeronBox<K, V>>,
 }
 
-impl<K: Trace + ?Sized, V: Trace + Clone> Ephemeron<K, V> {
+impl<K: Trace, V: Trace + Clone> Ephemeron<K, V> {
     /// Gets the stored value of this `Ephemeron`, or `None` if the key was already garbage collected.
     ///
     /// This needs to return a clone of the value because holding a reference to it between
@@ -40,7 +42,7 @@ impl<K: Trace + ?Sized, V: Trace + Clone> Ephemeron<K, V> {
     }
 }
 
-impl<K: Trace + ?Sized, V: Trace> Ephemeron<K, V> {
+impl<K: Trace, V: Trace> Ephemeron<K, V> {
     /// Creates a new `Ephemeron`.
     pub fn new(key: &Gc<K>, value: V) -> Self {
         let inner_ptr = Allocator::alloc_ephemeron(EphemeronBox::new(key, value));
@@ -58,7 +60,7 @@ impl<K: Trace + ?Sized, V: Trace> Ephemeron<K, V> {
         self.inner_ptr
     }
 
-    fn inner(&self) -> &EphemeronBox<K, V> {
+    pub(crate) fn inner(&self) -> &EphemeronBox<K, V> {
         // SAFETY: Please see Gc::inner_ptr()
         unsafe { self.inner_ptr().as_ref() }
     }
@@ -75,7 +77,7 @@ impl<K: Trace + ?Sized, V: Trace> Ephemeron<K, V> {
     }
 }
 
-impl<K: Trace + ?Sized, V: Trace> Finalize for Ephemeron<K, V> {
+impl<K: Trace, V: Trace> Finalize for Ephemeron<K, V> {
     fn finalize(&self) {
         // SAFETY: inner_ptr should be alive when calling finalize.
         // We don't call inner_ptr() to avoid overhead of calling finalizer_safe().
@@ -87,7 +89,7 @@ impl<K: Trace + ?Sized, V: Trace> Finalize for Ephemeron<K, V> {
 
 // SAFETY: `Ephemeron`s trace implementation only marks its inner box because we want to stop
 // tracing through weakly held pointers.
-unsafe impl<K: Trace + ?Sized, V: Trace> Trace for Ephemeron<K, V> {
+unsafe impl<K: Trace, V: Trace> Trace for Ephemeron<K, V> {
     unsafe fn trace(&self) {
         // SAFETY: We need to mark the inner box of the `Ephemeron` since it is reachable
         // from a root and this means it cannot be dropped.
@@ -105,7 +107,7 @@ unsafe impl<K: Trace + ?Sized, V: Trace> Trace for Ephemeron<K, V> {
     }
 }
 
-impl<K: Trace + ?Sized, V: Trace> Clone for Ephemeron<K, V> {
+impl<K: Trace, V: Trace> Clone for Ephemeron<K, V> {
     fn clone(&self) -> Self {
         let ptr = self.inner_ptr();
         self.inner().inc_ref_count();
@@ -114,7 +116,7 @@ impl<K: Trace + ?Sized, V: Trace> Clone for Ephemeron<K, V> {
     }
 }
 
-impl<K: Trace + ?Sized, V: Trace> Drop for Ephemeron<K, V> {
+impl<K: Trace, V: Trace> Drop for Ephemeron<K, V> {
     fn drop(&mut self) {
         if finalizer_safe() {
             Finalize::finalize(self);

boa_gc/src/pointers/mod.rs

@@ -9,3 +9,5 @@ pub use ephemeron::Ephemeron;
 pub use gc::Gc;
 pub use weak::WeakGc;
 pub use weak_map::WeakMap;
+
+pub(crate) use weak_map::RawWeakMap;

boa_gc/src/pointers/weak.rs

@@ -7,7 +7,7 @@ use std::hash::{Hash, Hasher};
 /// garbage collections. However, this also means [`WeakGc::upgrade`] could return `None` at any moment.
 #[derive(Debug, Trace, Finalize)]
 #[repr(transparent)]
-pub struct WeakGc<T: Trace + ?Sized + 'static> {
+pub struct WeakGc<T: Trace + 'static> {
     inner: Ephemeron<T, Gc<T>>,
 }