From 3f6078945cb5c8986089d92d9f19c0dacdddd834 Mon Sep 17 00:00:00 2001
From: Johannes Hostert <jhostert@ethz.ch>
Date: Mon, 26 Aug 2024 16:19:54 +0200
Subject: [PATCH] fix approach 2: use unsafecell

---
 core/arena/unique_arena.rs              | 28 +++++++------------------
 core/runtime/op_driver/erased_future.rs |  9 +++++---
 core/runtime/op_driver/future_arena.rs  | 14 +++++++------
 3 files changed, 22 insertions(+), 29 deletions(-)

diff --git a/core/arena/unique_arena.rs b/core/arena/unique_arena.rs
index 8d8e16ca5..a07a0d939 100644
--- a/core/arena/unique_arena.rs
+++ b/core/arena/unique_arena.rs
@@ -59,8 +59,8 @@ impl<T: 'static> ArenaBox<T> {
   /// This function returns a raw pointer without managing the memory, potentially leading to
   /// memory leaks if the pointer is not properly handled or deallocated.
   #[inline(always)]
-  pub fn into_raw(mut alloc: ArenaBox<T>) -> NonNull<T> {
-    let ptr = NonNull::from(alloc.data_mut());
+  pub fn into_raw(alloc: ArenaBox<T>) -> NonNull<T> {
+    let ptr: NonNull<ArenaBoxData<T>> = alloc.ptr;
     std::mem::forget(alloc);
     unsafe { Self::ptr_from_data(ptr) }
   }
@@ -91,8 +91,10 @@ impl<T> ArenaBox<T> {
   }
 
   #[inline(always)]
-  fn data_mut(&mut self) -> &mut ArenaBoxData<T> {
-    unsafe { self.ptr.as_mut() }
+  pub(crate) fn deref_data(&self) -> NonNull<T> {
+    unsafe {
+      NonNull::new_unchecked(std::ptr::addr_of_mut!((*self.ptr.as_ptr()).data))
+    }
   }
 }
 
@@ -120,20 +122,6 @@ impl<T> std::convert::AsRef<T> for ArenaBox<T> {
   }
 }
 
-impl<T> std::ops::DerefMut for ArenaBox<T> {
-  #[inline(always)]
-  fn deref_mut(&mut self) -> &mut Self::Target {
-    &mut self.data_mut().data
-  }
-}
-
-impl<T> std::convert::AsMut<T> for ArenaBox<T> {
-  #[inline(always)]
-  fn as_mut(&mut self) -> &mut T {
-    &mut self.data_mut().data
-  }
-}
-
 impl<F, R> std::future::Future for ArenaBox<F>
 where
   F: Future<Output = R>,
@@ -142,10 +130,10 @@ where
 
   #[inline(always)]
   fn poll(
-    mut self: Pin<&mut Self>,
+    self: Pin<&mut Self>,
     cx: &mut std::task::Context<'_>,
   ) -> std::task::Poll<Self::Output> {
-    unsafe { F::poll(Pin::new_unchecked(&mut self.data_mut().data), cx) }
+    unsafe { F::poll(Pin::new_unchecked(self.deref_data().as_mut()), cx) }
   }
 }
 
diff --git a/core/runtime/op_driver/erased_future.rs b/core/runtime/op_driver/erased_future.rs
index 2399992b9..08bc6e2fd 100644
--- a/core/runtime/op_driver/erased_future.rs
+++ b/core/runtime/op_driver/erased_future.rs
@@ -38,8 +38,11 @@ impl<const MAX_SIZE: usize> TypeErased<MAX_SIZE> {
   }
 
   #[inline(always)]
-  pub fn raw_ptr<R>(&mut self) -> NonNull<R> {
-    unsafe { NonNull::new_unchecked(self.memory.as_mut_ptr() as *mut _) }
+  /// Safety: This uses a place projection to `this.memory`, which must be in-bounds.
+  pub unsafe fn raw_ptr<R>(this: *mut Self) -> NonNull<R> {
+    unsafe {
+      NonNull::new_unchecked(std::ptr::addr_of_mut!((*this).memory) as *mut _)
+    }
   }
 
   #[inline(always)]
@@ -70,7 +73,7 @@ impl<const MAX_SIZE: usize> TypeErased<MAX_SIZE> {
 
 impl<const MAX_SIZE: usize> Drop for TypeErased<MAX_SIZE> {
   fn drop(&mut self) {
-    (self.drop)(self.raw_ptr::<()>())
+    (self.drop)(unsafe { Self::raw_ptr(self) })
   }
 }
 
diff --git a/core/runtime/op_driver/future_arena.rs b/core/runtime/op_driver/future_arena.rs
index 26c170102..339dff36e 100644
--- a/core/runtime/op_driver/future_arena.rs
+++ b/core/runtime/op_driver/future_arena.rs
@@ -3,6 +3,7 @@ use super::erased_future::TypeErased;
 use crate::arena::ArenaBox;
 use crate::arena::ArenaUnique;
 use pin_project::pin_project;
+use std::cell::UnsafeCell;
 use std::future::Future;
 use std::marker::PhantomData;
 use std::mem::MaybeUninit;
@@ -25,7 +26,7 @@ pub trait FutureContextMapper<T, C, R> {
 
 struct DynFutureInfoErased<T, C> {
   ptr: MaybeUninit<NonNull<dyn ContextFuture<T, C>>>,
-  data: TypeErased<MAX_ARENA_FUTURE_SIZE>,
+  data: UnsafeCell<TypeErased<MAX_ARENA_FUTURE_SIZE>>,
 }
 
 pub trait ContextFuture<T, C>: Future<Output = T> {
@@ -190,19 +191,20 @@ impl<T, C: Clone> FutureArena<T, C> {
     {
       unsafe {
         if let Some(reservation) = self.arena.reserve_space() {
-          let mut alloc = self.arena.complete_reservation(
+          let alloc = self.arena.complete_reservation(
             reservation,
             DynFutureInfoErased {
               ptr: MaybeUninit::uninit(),
-              data: TypeErased::new(DynFutureInfo {
+              data: UnsafeCell::new(TypeErased::new(DynFutureInfo {
                 context,
                 future,
                 _phantom: PhantomData,
-              }),
+              })),
             },
           );
-          let ptr = alloc.data.raw_ptr::<DynFutureInfo<T, C, M, F>>();
-          alloc.ptr.write(ptr);
+          let ptr =
+            TypeErased::raw_ptr::<DynFutureInfo<T, C, M, F>>(alloc.data.get());
+          (*alloc.deref_data().as_ptr()).ptr.write(ptr);
           return TypedFutureAllocation {
             inner: FutureAllocation::Arena(alloc),
             ptr,