feat(jstz_engine): implement a write-barriered GC pointer abstraction

crates/jstz_engine/src/gc/mod.rs

@@ -0,0 +1 @@
+mod ptr;

crates/jstz_engine/src/gc/ptr.rs

@@ -0,0 +1,271 @@
+//! A garbage-collected heap pointer used to refer to on-heap objects.
+//! All garbage-collected pointers should be wrapped in a `GcPtr`
+//! for safety purposes.  
+
+use std::{cell::UnsafeCell, marker::PhantomPinned, mem, pin::Pin, ptr};
+
+use mozjs::{
+    jsapi::{
+        jsid, HeapBigIntWriteBarriers, HeapObjectWriteBarriers, HeapScriptWriteBarriers,
+        HeapStringWriteBarriers, HeapValueWriteBarriers, JSFunction, JSObject, JSScript,
+        JSString, JS::BigInt as JSBigInt, JS::Symbol as JSSymbol,
+    },
+    jsid::VoidId,
+    jsval::{JSVal, UndefinedValue},
+};
+
+/// A GC barrier is a mechanism used to ensure that the garbage collector maintains
+/// a valid set of reachable objects.
+///
+/// A write barrier is a mechanism used to ensure that the garbage collector is notified
+/// when a reference to an object is changed. In general, a write barrier should be invoked
+/// whenever a write can cause the set of things traced by the GC to change.
+///
+/// Every barriered write should have the following form:
+/// ```notrust
+///   field = new_value;
+///   <write-barrier>
+/// ```
+pub unsafe trait WriteBarrieredPtr: Copy {
+    /// Creates a uninitialized value
+    unsafe fn uninit() -> Self;
+
+    /// Perform a write barrier on the given GC value
+    unsafe fn write_barrier(v: *mut Self, prev: Self, next: Self);
+}
+
+/// A garbage-collected pointer used to refer to on-heap objects
+///
+/// # Safety
+///
+/// `GcPtr<T>` should only be used by values on the heap. Garbage collected pointers
+/// on the stack should be rooted.
+pub struct GcPtr<T: WriteBarrieredPtr> {
+    // # Safety
+    //
+    // For garbage collection to work correctly, when modifying
+    // the wrapped value that points to a GC cell, the write barrier
+    // must be invoked.
+    //
+    // This means after calling the `set` method, the `GcPtr` *must not*
+    // be moved in memory. Doing so would invalidate the local reference.
+    // For safety, we use `Box::pin` to pin the `UnsafeCell` to the Rust heap.
+    inner_ptr: Pin<Box<UnsafeCell<T>>>,
+    _unused: PhantomPinned,
+}
+
+impl<T: WriteBarrieredPtr> GcPtr<T> {
+    /// Creates a new [`GcPtr`] from an existing pointer.
+    ///
+    /// # Safety
+    ///
+    /// The raw pointer `ptr` must point to an object that extends a `js::gc::Cell`.
+    pub fn new(ptr: T) -> Self {
+        let inner_ptr = Box::pin(UnsafeCell::new(ptr));
+        unsafe { T::write_barrier(inner_ptr.get(), T::uninit(), ptr) };
+
+        Self {
+            inner_ptr,
+            _unused: PhantomPinned,
+        }
+    }
+
+    /// Creates a new [`GcPtr`] from an existing pointer without performing a write barrier.
+    ///
+    /// # Safety
+    ///
+    /// It isn't safe to create a new [`GcPtr`] from an existing pointer without performing a write barrier.
+    /// This function should only be used when the pointer is known have a barrier.
+    pub fn new_unbarriered(ptr: T) -> Self {
+        let inner_ptr = Box::pin(UnsafeCell::new(ptr));
+
+        Self {
+            inner_ptr,
+            _unused: PhantomPinned,
+        }
+    }
+
+    /// Compares two pointers for equality
+    #[allow(dead_code)]
+    fn ptr_eq(&self, other: &Self) -> bool {
+        self.inner_ptr.get() == other.inner_ptr.get()
+    }
+
+    /// Returns the raw pointer
+    pub fn get(&self) -> T {
+        unsafe { *self.inner_ptr.get() }
+    }
+
+    /// Sets the pointer to a new value
+    pub fn set(&self, next: T) {
+        let self_ptr = self.inner_ptr.get();
+        unsafe {
+            let prev = *self_ptr;
+
+            *self_ptr = next;
+            T::write_barrier(self_ptr, prev, next)
+        }
+    }
+}
+
+impl<T: WriteBarrieredPtr> Clone for GcPtr<T> {
+    fn clone(&self) -> Self {
+        Self::new(self.get())
+    }
+}
+
+impl<T: WriteBarrieredPtr> Drop for GcPtr<T> {
+    fn drop(&mut self) {
+        unsafe {
+            let inner_ptr = self.inner_ptr.get();
+            T::write_barrier(inner_ptr, *inner_ptr, T::uninit())
+        }
+    }
+}
+
+unsafe impl WriteBarrieredPtr for *mut JSObject {
+    unsafe fn uninit() -> Self {
+        ptr::null_mut()
+    }
+
+    unsafe fn write_barrier(v: *mut Self, prev: Self, next: Self) {
+        HeapObjectWriteBarriers(v, prev, next)
+    }
+}
+
+unsafe impl WriteBarrieredPtr for *mut JSString {
+    unsafe fn uninit() -> Self {
+        ptr::null_mut()
+    }
+
+    unsafe fn write_barrier(v: *mut Self, prev: Self, next: Self) {
+        HeapStringWriteBarriers(v, prev, next)
+    }
+}
+
+unsafe impl WriteBarrieredPtr for *mut JSFunction {
+    unsafe fn uninit() -> Self {
+        ptr::null_mut()
+    }
+
+    unsafe fn write_barrier(v: *mut Self, prev: Self, next: Self) {
+        HeapObjectWriteBarriers(
+            mem::transmute(v),
+            mem::transmute(prev),
+            mem::transmute(next),
+        )
+    }
+}
+
+unsafe impl WriteBarrieredPtr for *mut JSSymbol {
+    unsafe fn uninit() -> Self {
+        ptr::null_mut()
+    }
+
+    unsafe fn write_barrier(_v: *mut Self, _prev: Self, _next: Self) {
+        // No write barrier needed for JSSymbol
+    }
+}
+
+unsafe impl WriteBarrieredPtr for *mut JSBigInt {
+    unsafe fn uninit() -> Self {
+        ptr::null_mut()
+    }
+
+    unsafe fn write_barrier(v: *mut Self, prev: Self, next: Self) {
+        HeapBigIntWriteBarriers(v, prev, next)
+    }
+}
+
+unsafe impl WriteBarrieredPtr for *mut JSScript {
+    unsafe fn uninit() -> Self {
+        ptr::null_mut()
+    }
+
+    unsafe fn write_barrier(v: *mut Self, prev: Self, next: Self) {
+        HeapScriptWriteBarriers(v, prev, next)
+    }
+}
+
+unsafe impl WriteBarrieredPtr for jsid {
+    unsafe fn uninit() -> Self {
+        VoidId()
+    }
+
+    unsafe fn write_barrier(_v: *mut Self, _prev: Self, _next: Self) {
+        // No write barrier needed for jsid
+    }
+}
+
+unsafe impl WriteBarrieredPtr for JSVal {
+    unsafe fn uninit() -> Self {
+        UndefinedValue()
+    }
+
+    unsafe fn write_barrier(v: *mut Self, prev: Self, next: Self) {
+        HeapValueWriteBarriers(v, &prev, &next)
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use std::sync::Mutex;
+
+    use crate::gc::ptr::{GcPtr, WriteBarrieredPtr};
+
+    #[derive(Debug, PartialEq, Eq, Copy, Clone)]
+    pub struct TestPtr {
+        value: i32,
+    }
+
+    const TEST_PTR_UNINIT: TestPtr = TestPtr { value: 0 };
+
+    static WRITE_BARRIER_LOG: Mutex<Vec<(TestPtr, TestPtr)>> = Mutex::new(Vec::new());
+
+    unsafe impl WriteBarrieredPtr for TestPtr {
+        unsafe fn uninit() -> Self {
+            TEST_PTR_UNINIT
+        }
+
+        unsafe fn write_barrier(_v: *mut Self, _prev: Self, _next: Self) {
+            // No write barrier needed for TestPtr
+
+            WRITE_BARRIER_LOG.lock().unwrap().push((_prev, _next));
+        }
+    }
+
+    #[test]
+    fn test_new_triggers_barrier() {
+        WRITE_BARRIER_LOG.lock().unwrap().clear();
+
+        let _ptr = GcPtr::new(TestPtr { value: 42 });
+
+        let write_barrier_log = WRITE_BARRIER_LOG.lock().unwrap();
+        assert_eq!(write_barrier_log.len(), 1);
+        assert_eq!(
+            write_barrier_log[0],
+            (TEST_PTR_UNINIT, TestPtr { value: 42 })
+        );
+    }
+
+    #[test]
+    fn test_set_calls_write_barrier() {
+        WRITE_BARRIER_LOG.lock().unwrap().clear();
+
+        let ptr = GcPtr::new(TestPtr { value: 42 });
+        let new_ptr = TestPtr { value: 43 };
+
+        ptr.set(new_ptr);
+
+        let write_barrier_log = WRITE_BARRIER_LOG.lock().unwrap();
+        assert_eq!(write_barrier_log.len(), 2);
+        assert_eq!(
+            write_barrier_log[0],
+            (TEST_PTR_UNINIT, TestPtr { value: 42 })
+        );
+        assert_eq!(
+            write_barrier_log[1],
+            (TestPtr { value: 42 }, TestPtr { value: 43 })
+        );
+    }
+}

crates/jstz_engine/src/lib.rs

@@ -1,5 +1,6 @@
 mod compartment;
 mod context;
+mod gc;
 mod realm;
 mod script;