From 864150231cd860b10a645e2717c5a0c7d197c2f7 Mon Sep 17 00:00:00 2001
From: Wedson Almeida Filho <wedsonaf@google.com>
Date: Tue, 6 Jul 2021 03:53:21 +0100
Subject: [PATCH] rust: switch away from `Box` in `Ref` implementation.

This is needed because we will introduce `TryFrom<Vec<T>>` for
`Ref<[T]>`, which cannot be implemented with `Box`. For the drop
implementation to be compatible with both constructors, we need them to
use the same allocator.

This is in preparation for deprecating `Arc<T>`.

Signed-off-by: Wedson Almeida Filho <wedsonaf@google.com>
---
 rust/kernel/sync/arc.rs | 37 ++++++++++++++++++++++++++-----------
 1 file changed, 26 insertions(+), 11 deletions(-)

diff --git a/rust/kernel/sync/arc.rs b/rust/kernel/sync/arc.rs
index 1e0650165d987d..d0a6b4be234759 100644
--- a/rust/kernel/sync/arc.rs
+++ b/rust/kernel/sync/arc.rs
@@ -15,8 +15,8 @@
 //!
 //! [`Arc`]: https://doc.rust-lang.org/std/sync/struct.Arc.html
 
-use crate::{bindings, Result};
-use alloc::boxed::Box;
+use crate::{bindings, Error, Result};
+use alloc::alloc::{alloc, dealloc};
 use core::{
     alloc::Layout,
     cell::UnsafeCell,
@@ -79,23 +79,32 @@ impl<T> Ref<T> {
     ///
     /// This is useful because it provides a mutable reference to `T` at its final location.
     pub fn try_new_and_init<U: FnOnce(Pin<&mut T>)>(contents: T, init: U) -> Result<Self> {
+        let layout = Layout::new::<RefInner<T>>();
+        // SAFETY: The layout size is guaranteed to be non-zero because `RefInner` contains the
+        // reference count.
+        let mut inner = NonNull::new(unsafe { alloc(layout) })
+            .ok_or(Error::ENOMEM)?
+            .cast::<RefInner<T>>();
+
         // INVARIANT: The refcount is initialised to a non-zero value.
-        let mut inner = Box::try_new(RefInner {
+        let value = RefInner {
             // SAFETY: Just an FFI call that returns a `refcount_t` initialised to 1.
             refcount: UnsafeCell::new(unsafe { bindings::REFCOUNT_INIT(1) }),
             data: contents,
-        })?;
+        };
+        // SAFETY: `inner` is writable and properly aligned.
+        unsafe { inner.as_ptr().write(value) };
 
         // SAFETY: By the invariant, `RefInner` is pinned and `T` is also pinned.
-        let pinned = unsafe { Pin::new_unchecked(&mut inner.data) };
+        let pinned = unsafe { Pin::new_unchecked(&mut inner.as_mut().data) };
 
         // INVARIANT: The only places where `&mut T` is available are here, which is explicitly
         // pinned, and in `drop`. Both are compatible with the pin requirements.
         init(pinned);
 
-        // SAFETY: We just created `inner` with a reference count of 1 and we're leaking it. So the
-        // new `Ref` object owns the reference.
-        Ok(unsafe { Self::from_inner(NonNull::from(Box::leak(inner))) })
+        // SAFETY: We just created `inner` with a reference count of 1, which is owned by the new
+        // `Ref` object.
+        Ok(unsafe { Self::from_inner(inner) })
     }
 
     /// Deconstructs a [`Ref`] object into a `usize`.
@@ -243,9 +252,15 @@ impl<T: ?Sized> Drop for Ref<T> {
         let is_zero = unsafe { bindings::refcount_dec_and_test(refcount) };
         if is_zero {
             // The count reached zero, we must free the memory.
-            //
-            // SAFETY: The pointer was initialised from the result of `Box::leak`.
-            unsafe { Box::from_raw(self.ptr.as_ptr()) };
+
+            // SAFETY: This thread holds the only remaining reference to `self`, so it is safe to
+            // get a mutable reference to it.
+            let inner = unsafe { self.ptr.as_mut() };
+            let layout = Layout::for_value(inner);
+            // SAFETY: The value stored in inner is valid.
+            unsafe { core::ptr::drop_in_place(inner) };
+            // SAFETY: The pointer was initialised from the result of a call to `alloc`.
+            unsafe { dealloc(self.ptr.cast().as_ptr(), layout) };
         }
     }
 }