[safety] Don't allow a root to be used with a frozen collector.

Also update `unfreeze` to properly wait until a collection is over.

This is incorrect in the face of compacting collectors,
since the collector might want to rewrite pointers
while the mutator is actively using them (via the "frozen" context).
Any parents using `safepoint_recurse` are fine.
Their roots will still be relocated by the collector.

I was actually tipped off by the fact by the fact the collector
wants an `&mut` reference to all the roots in the shadow stack.
Its impossible for the frozen collector
to have a mutable reference while the mutator is still using it!!!
The original &mut reference we were passing to `freeze_safepoint`
was only valid for the initial safepoint. The collector can't assume
that its valid to mutate (or even access) for any longer than that.

The old implementation was invalid even though it wasn't a moving collection.
However the bug only got triggered with optimization enabled :p

We'll have to implement GC Handles as a replacement: #9

Cargo.toml

@@ -21,3 +21,6 @@ ndarray = { version = "^0.10.12", optional = true }
 
 [workspace]
 members = ["libs/simple", "libs/derive"]
+
+[profile.dev]
+opt-level = 1

libs/simple/examples/binary_trees_parallel.rs

@@ -2,7 +2,7 @@
     arbitrary_self_types, // Unfortunately this is required for methods on Gc refs
 )]
 use zerogc::{
-    safepoint, safepoint_recurse, freeze_safepoint, unfreeze,
+    safepoint, safepoint_recurse, freeze_context, unfreeze_context,
     GcSimpleAlloc, GcCell, GcSafe
 };
 
@@ -74,7 +74,7 @@ fn main() {
     safepoint!(gc, ());
 
     let long_lived_tree = bottom_up_tree(&gc, max_depth);
-    let (frozen, long_lived_tree) = freeze_safepoint!(gc, long_lived_tree);
+    let frozen = freeze_context!(gc);
 
     (min_depth / 2..max_depth / 2 + 1).into_par_iter().for_each(|half_depth| {
         let depth = half_depth * 2;
@@ -83,7 +83,7 @@ fn main() {
         let message = inner(&collector, depth, iterations);
         println!("{}", message);
     });
-    unfreeze!(frozen => new_context, long_lived_tree);
+    let new_context = unfreeze_context!(frozen);
 
     println!("long lived tree of depth {}\t check: {}", max_depth, item_check(&long_lived_tree));
 }

libs/simple/src/context.rs

@@ -4,7 +4,7 @@ use std::sync::atomic::{Ordering};
 
 use slog::{Logger, FnValue, trace, Drain, o};
 
-use zerogc::{GcContext, Trace, FrozenContext};
+use zerogc::{GcContext, Trace};
 use crate::{SimpleCollector, RawSimpleCollector, DynTrace};
 use crate::utils::ThreadId;
 use std::mem::ManuallyDrop;
@@ -46,18 +46,12 @@ enum ContextState {
     /// Because no allocation or mutation can happen,
     /// its shadow_stack stack is guarenteed to
     /// accurately reflect the roots of the context.
-    Frozen {
-        /// The pointer to the last element in the
-        /// frozen shadow stack.
-        ///
-        /// Used for internal sanity checking.
-        last_ptr: *mut dyn DynTrace
-    }
+    Frozen,
 }
 impl ContextState {
     fn is_frozen(&self) -> bool {
         match *self {
-            ContextState::Frozen { .. } => true,
+            ContextState::Frozen => true,
             _ => false
         }
     }
@@ -309,46 +303,25 @@ unsafe impl GcContext for SimpleCollectorContext {
         debug_assert_eq!(popped, dyn_ptr);
     }
 
-    unsafe fn frozen_safepoint<T: Trace>(&mut self, value: &mut &mut T)
-                                         -> FrozenContext<'_, Self> {
+    unsafe fn freeze(&mut self) {
         assert_eq!(self.raw.state.get(), ContextState::Active);
-        let dyn_ptr = (*self.raw.shadow_stack.get())
-            .push(value);
-        // NOTE: Implicitly trigger safepoint before we freeze ourselves
-        if self.raw.collector.should_collect() {
-            /*
-             * NOTE: This should implicitly notify threads
-             * in case someone was waiting on us.
-             */
-            self.raw.trigger_safepoint();
-        }
-        assert_eq!(self.raw.state.get(), ContextState::Active);
-        self.raw.state.set(ContextState::Frozen {
-            last_ptr: dyn_ptr
-        });
+        self.raw.state.set(ContextState::Frozen);
         // We may need to notify others that we are frozen
         if self.raw.collector.collecting.load(Ordering::Acquire) {
             self.raw.collector.notify_waiting_safepoints();
         }
-        FrozenContext::new(self)
     }
 
-    unsafe fn unfreeze(frozen: FrozenContext<'_, Self>) -> &'_ mut Self {
-        let ctx = FrozenContext::into_inner(frozen);
+    unsafe fn unfreeze(&mut self) {
         /*
-         * TODO: Ensure there is no collection in progess.
          * A pending collection might be relying in the validity of this
          * context's shadow stack, so unfreezing it while in progress
          * could trigger undefined behavior!!!!!
          */
-        let dyn_ptr = match ctx.raw.state.get() {
-            ContextState::Frozen { last_ptr } => last_ptr,
-            state => unreachable!("{:?}", state)
-        };
-        let actual_ptr = (*ctx.raw.shadow_stack.get()).pop();
-        debug_assert_eq!(actual_ptr, Some(dyn_ptr));
-        ctx.raw.state.set(ContextState::Active);
-        ctx
+        self.collector().prevent_collection(|_| {
+            assert_eq!(self.raw.state.get(), ContextState::Frozen);
+            self.raw.state.set(ContextState::Active);
+        })
     }
 
     #[inline]
@@ -408,10 +381,6 @@ impl ShadowStack {
         long_ptr
     }
     #[inline]
-    pub(crate) fn pop(&mut self) -> Option<*mut dyn DynTrace> {
-        self.elements.pop()
-    }
-    #[inline]
     pub(crate) unsafe fn pop_unchecked(&mut self) -> *mut dyn DynTrace {
         match self.elements.pop() {
             Some(value) => value,

src/lib.rs

@@ -128,55 +128,49 @@ macro_rules! __recurse_context {
 /// This macro is completely safe, although it expands to unsafe code internally.
 #[macro_export]
 macro_rules! safepoint {
-    ($collector:ident, $value:expr) => {unsafe {
+    ($context:ident, $value:expr) => {unsafe {
         use $crate::{GcContext};
         // TODO: What happens if we panic during a collection
         /*
          * Some collectors support multiple running instances
          * with different ids, handing out different GC pointers.
          * TODO: Should we be checking somehow that the ids match?
          */
-        let mut erased = $collector.rebrand_static($value);
-        $collector.basic_safepoint(&mut &mut erased);
-        $collector.rebrand_self(erased)
+        let mut erased = $context.rebrand_static($value);
+        $context.basic_safepoint(&mut &mut erased);
+        $context.rebrand_self(erased)
     }};
 }
 
 /// Indicate its safe to begin a garbage collection (like [safepoint!])
 /// and then "freeze" the specified context.
 ///
 /// Until it's unfrozen, the context can't be used for allocation.
-/// However its roots are still considered valid and it allows other
-/// threads to perform collections *without blocking this thread*.
+/// Its roots are marked invalid, since the collector could be relocating them.
+/// However, the roots of any parent contexts are still considered valid.
+///
+/// This allows other threads to perform collections *without blocking this thread*.
 #[macro_export]
-macro_rules! freeze_safepoint {
-    ($collector:ident, $value:expr) => {unsafe {
-        use $crate::{GcContext};
-        // See also `safepoint!` macro
-        let mut erased = $collector.rebrand_static($value);
-        let frozen = $collector.frozen_safepoint(&mut &mut erased);
-        // NOTE: This is branded to `frozen` now
-        let new_root = frozen.rebrand_self(erased);
-        // TODO: Undefined behavior if we don't unfreeze -_-
-        (frozen, new_root)
+macro_rules! freeze_context {
+    ($context:ident) => {unsafe {
+        use $crate::{GcContext, FrozenContext};
+        let mut context = $context;
+        $context.freeze();
+        FrozenContext::new(context)
     }};
 }
 
 /// Unfreeze the context, allowing it to be used again
 ///
-/// This macro is a little weird, since it automatically
-/// declares a `let` binding to the new context's name.
+/// Returns a [zerogc::GcContext] struct.
 #[macro_export]
-macro_rules! unfreeze {
-    ($frozen:ident => $new_context:ident, $value:expr) => {
-        let $new_context;
-        unsafe {
-            use $crate::{GcContext, FrozenContext};
-            let erased = $frozen.rebrand_static($value);
-            $new_context = FrozenContext::unfreeze($frozen);
-            $new_context.rebrand_self(erased)
-        }
-    };
+macro_rules! unfreeze_context {
+    ($frozen:ident) => {unsafe {
+        use $crate::{FrozenContext, GcContext};
+        let mut context = FrozenContext::into_context($frozen);
+        context.unfreeze();
+        context
+    }};
 }
 
 /// A garbage collector implementation.
@@ -213,20 +207,27 @@ pub unsafe trait GcContext: Sized {
     /// This allows other threds to perform collections while this
     /// thread does other work (without using the GC).
     ///
+    /// The current contexts roots are considered invalid
+    /// for the duration of the collection,
+    /// since the collector could potentially relocate them.
+    ///
+    /// Any parent contexts are fine and their roots will be
+    /// preserved by collections.
+    ///
     /// ## Safety
-    /// Unsafe for the same reasons as `basic_safepoint`.
+    /// Assumes this context is valid and not already frozen.
     ///
-    /// It's assumed that this context's roots are valid for the entire
-    /// duration of the freeze.
-    unsafe fn frozen_safepoint<T: Trace>(&mut self, value: &mut &mut T) -> FrozenContext<'_, Self>;
+    /// Don't invoke this directly
+    unsafe fn freeze(&mut self);
 
-    /// Unfreeze this context
+    /// Unfreeze this context, allowing it to be used again.
     ///
     /// ## Safety
     /// Must be a valid context!
+    /// Must be currently frozen!
     ///
     /// Don't invoke this directly
-    unsafe fn unfreeze(frozen: FrozenContext<'_, Self>) -> &'_ mut Self;
+    unsafe fn unfreeze(&mut self);
 
     #[inline(always)]
     #[doc(hidden)]
@@ -289,64 +290,20 @@ pub unsafe trait GcSimpleAlloc<'gc, T: GcSafe + 'gc>: 'gc {
 /// Don't use this directly!!!
 #[doc(hidden)]
 #[must_use]
-pub struct FrozenContext<'a, C: GcContext>(&'a mut C);
-impl<'a, C: GcContext> FrozenContext<'a, C> {
-    #[inline]
-    pub unsafe fn new(ctx: &'a mut C) -> Self {
-        FrozenContext(ctx)
-    }
-    /// Get the underlying reference to the context
-    ///
-    /// This is an associated function because I don't want users
-    /// invoking it directly.
-    ///
-    /// ## Safety
-    /// It's unsafe to actually allocate or mutate a frozen collector.
-    ///
-    /// This should only be used internally
-    #[inline]
-    pub unsafe fn into_inner(ctx: Self) -> &'a mut C {
-        /*
-         * This is needed because we have a dummy destructor
-         * TODO: Remove once we resolve issues with use after free
-         */
-        let taken = std::ptr::read(&ctx.0);
-        mem::forget(ctx); // Don't panic (42)
-        taken
-    }
-    /// Unfreezes the context, dispatching to `GcContext::unfreeze`.
-    ///
-    /// This is an associated function because I don't want users
-    /// invoking it directly.
-    ///
-    /// ## Safety
-    /// Don't invoke this directly!
-    ///
-    /// This is an implementation detail
-    #[inline]
-    pub unsafe fn unfreeze(ctx: Self) -> &'a mut C {
-        C::unfreeze(ctx)
-    }
-    #[inline(always)]
+pub struct FrozenContext<C: GcContext> {
+    /// The frozen context
+    context: C,
+}
+impl<C: GcContext> FrozenContext<C> {
     #[doc(hidden)]
-    /// Copied from GcContext::rebrand_self
-    pub unsafe fn rebrand_self<T: GcBrand<'a, C::System>>(&self, value: T) -> T::Branded {
-        let branded = mem::transmute_copy(&value);
-        mem::forget(value);
-        branded
+    #[inline]
+    pub unsafe fn new(context: C) -> Self {
+        FrozenContext { context }
     }
-    /// Likewise copied from GcContext::rebrand_static
-    #[inline(always)]
     #[doc(hidden)]
-    pub unsafe fn rebrand_static<T: GcBrand<'static, C::System>>(&self, value: T) -> T::Branded {
-        let branded = mem::transmute_copy(&value);
-        mem::forget(value);
-        branded
-    }
-}
-impl<'a, C: GcContext> Drop for FrozenContext<'a, C> {
-    fn drop(&mut self) {
-        todo!("Undefined behavior to forget to unfreeze")
+    #[inline]
+    pub unsafe fn into_context(self) -> C {
+        self.context
     }
 }