wrap Payloads that are dropped eraly in ManuallyDrop

src/arena.rs

@@ -310,6 +310,13 @@ impl<T: ?Sized + ArenaAllocated> TypedArenaPtr<T> {
     }
 }
 
+impl<P, T: ?Sized + ArenaAllocated<Payload = ManuallyDrop<P>>> TypedArenaPtr<T> {
+    pub fn drop_payload(&mut self) {
+        self.set_tag(ArenaHeaderTag::Dropped);
+        unsafe { ManuallyDrop::drop(&mut *self.as_ptr()) }
+    }
+}
+
 impl<T: ?Sized + ArenaAllocated> TypedArenaPtr<T>
 where
     T::Payload: Sized,
@@ -582,16 +589,34 @@ impl ArenaAllocated for Rational {
     }
 }
 
+impl AllocateInArena<LiveLoadState> for LiveLoadState {
+    fn arena_allocate(self, arena: &mut Arena) -> TypedArenaPtr<LiveLoadState> {
+        LiveLoadState::alloc(arena, ManuallyDrop::new(self))
+    }
+}
+
 impl ArenaAllocated for LiveLoadState {
-    type Payload = Self;
+    type Payload = ManuallyDrop<Self>;
     #[inline]
     fn tag() -> ArenaHeaderTag {
         ArenaHeaderTag::LiveLoadState
     }
+
+    unsafe fn dealloc(ptr: NonNull<TypedAllocSlab<Self>>) {
+        let mut slab = unsafe { Box::from_raw(ptr.as_ptr()) };
+        unsafe { ManuallyDrop::drop(&mut slab.payload) };
+        drop(slab);
+    }
+}
+
+impl AllocateInArena<TcpListener> for TcpListener {
+    fn arena_allocate(self, arena: &mut Arena) -> TypedArenaPtr<TcpListener> {
+        TcpListener::alloc(arena, ManuallyDrop::new(self))
+    }
 }
 
 impl ArenaAllocated for TcpListener {
-    type Payload = Self;
+    type Payload = ManuallyDrop<Self>;
     #[inline]
     fn tag() -> ArenaHeaderTag {
         ArenaHeaderTag::TcpListener
@@ -701,6 +726,10 @@ pub struct TypedAllocSlab<T: ?Sized + ArenaAllocated> {
 }
 
 impl<T: ?Sized + ArenaAllocated> TypedAllocSlab<T> {
+    pub fn payload(&mut self) -> &mut T::Payload {
+        &mut self.payload
+    }
+
     /// # Safety
     /// - ptr points to a valid allocation of Self
     #[inline]

src/machine/loader.rs

@@ -323,7 +323,7 @@ impl<'a> LoadState<'a> for LiveLoadAndMachineState<'a> {
     #[inline(always)]
     fn reset_machine(loader: &mut Loader<'a, Self>) {
         if loader.payload.load_state.get_tag() != ArenaHeaderTag::Dropped {
-            loader.payload.load_state.set_tag(ArenaHeaderTag::Dropped);
+            loader.payload.load_state.drop_payload();
             loader.reset_machine();
         }
     }
@@ -1788,11 +1788,8 @@ impl Machine {
             (HeapCellValueTag::Cons, cons_ptr) => {
                 match_untyped_arena_ptr!(cons_ptr,
                     (ArenaHeaderTag::LiveLoadState, payload) => {
-                        unsafe {
+                        let mut payload = payload;
-                            std::ptr::drop_in_place(
+                        payload.drop_payload()
-                                payload.as_ptr() as *mut LiveLoadState,
-                            );
-                        }
                     }
                     _ => {}
                 );

src/machine/streams.rs

@@ -454,13 +454,25 @@ impl<T> DerefMut for StreamLayout<T> {
 
 macro_rules! arena_allocated_impl_for_stream {
     ($stream_type:ty, $stream_tag:ident) => {
+        impl $crate::arena::AllocateInArena<$stream_tag> for StreamLayout<$stream_type> {
+            fn arena_allocate(self, arena: &mut Arena) -> TypedArenaPtr<$stream_tag> {
+                $stream_tag::alloc(arena, core::mem::ManuallyDrop::new(self))
+            }
+        }
+
         impl ArenaAllocated for $stream_tag {
-            type Payload = StreamLayout<$stream_type>;
+            type Payload = core::mem::ManuallyDrop<StreamLayout<$stream_type>>;
 
             #[inline]
             fn tag() -> ArenaHeaderTag {
                 ArenaHeaderTag::$stream_tag
             }
+
+            unsafe fn dealloc(ptr: std::ptr::NonNull<TypedAllocSlab<Self>>) {
+                let mut slab = unsafe { Box::from_raw(ptr.as_ptr()) };
+                unsafe { std::mem::ManuallyDrop::drop(slab.payload()) };
+                drop(slab);
+            }
         }
     };
 }
@@ -994,7 +1006,7 @@ impl Stream {
                 past_end_of_stream,
                 stream,
                 ..
-            } = &mut **stream_layout;
+            } = &mut ***stream_layout;
 
             stream
                 .get_mut()
@@ -1068,7 +1080,7 @@ impl Stream {
                     past_end_of_stream,
                     stream,
                     ..
-                } = &mut **stream_layout;
+                } = &mut ***stream_layout;
 
                 let cursor_len = stream.get_ref().0.get_ref().len() as u64;
                 cursor_position(past_end_of_stream, &stream.get_ref().0, cursor_len)
@@ -1078,7 +1090,7 @@ impl Stream {
                     past_end_of_stream,
                     stream,
                     ..
-                } = &mut **stream_layout;
+                } = &mut ***stream_layout;
 
                 let cursor_len = stream.stream.get_ref().len() as u64;
                 cursor_position(past_end_of_stream, &stream.stream, cursor_len)
@@ -1090,7 +1102,7 @@ impl Stream {
                     past_end_of_stream,
                     stream,
                     ..
-                } = &mut **stream_layout;
+                } = &mut ***stream_layout;
 
                 match stream.get_ref().file.metadata() {
                     Ok(metadata) => {
@@ -1270,38 +1282,25 @@ impl Stream {
             Stream::NamedTls(ref mut tls_stream) => tls_stream.inner_mut().tls_stream.shutdown(),
             #[cfg(feature = "http")]
             Stream::HttpRead(ref mut http_stream) => {
-                unsafe {
+                http_stream.drop_payload();
-                    http_stream.set_tag(ArenaHeaderTag::Dropped);
-                    std::ptr::drop_in_place(&mut http_stream.inner_mut().body_reader as *mut _);
-                }
 
                 Ok(())
             }
             #[cfg(feature = "http")]
-            Stream::HttpWrite(ref mut http_stream) => {
+            Stream::HttpWrite(mut http_stream) => {
-                http_stream.inner_mut().drop();
+                http_stream.drop_payload();
-                unsafe {
-                    http_stream.set_tag(ArenaHeaderTag::Dropped);
-                    std::ptr::drop_in_place(&mut http_stream.inner_mut().buffer as *mut _);
-                }
 
                 Ok(())
             }
             Stream::InputFile(mut file_stream) => {
                 // close the stream by dropping the inner File.
-                unsafe {
+                file_stream.drop_payload();
-                    file_stream.set_tag(ArenaHeaderTag::Dropped);
-                    std::ptr::drop_in_place(&mut file_stream.inner_mut().file as *mut _);
-                }
 
                 Ok(())
             }
             Stream::OutputFile(mut file_stream) => {
                 // close the stream by dropping the inner File.
-                unsafe {
+                file_stream.drop_payload();
-                    file_stream.set_tag(ArenaHeaderTag::Dropped);
-                    std::ptr::drop_in_place(&mut file_stream.file as *mut _);
-                }
 
                 Ok(())
             }

src/machine/system_calls.rs

@@ -6737,12 +6737,8 @@ impl Machine {
             (HeapCellValueTag::Cons, cons_ptr) => {
                 match_untyped_arena_ptr!(cons_ptr,
                     (ArenaHeaderTag::TcpListener, tcp_listener) => {
-                        unsafe {
+                        tcp_listener.drop_payload();
-                            // dropping closes the instance.
-                            std::ptr::drop_in_place(&mut tcp_listener as *mut _);
-                        }
 
-                        tcp_listener.set_tag(ArenaHeaderTag::Dropped);
                         return Ok(());
                     }
                     _ => {