FrameRef for all types of frames

ostd/src/mm/frame/frame_ref.rs

@@ -0,0 +1,81 @@
+// SPDX-License-Identifier: MPL-2.0
+
+use core::{marker::PhantomData, mem::ManuallyDrop, ops::Deref, ptr::NonNull};
+
+use super::{
+    meta::{AnyFrameMeta, MetaSlot},
+    Frame,
+};
+use crate::{mm::Paddr, sync::non_null::NonNullPtr};
+
+/// A struct that can work as `&'a Frame<M>`.
+pub struct FrameRef<'a, M: AnyFrameMeta + ?Sized> {
+    inner: ManuallyDrop<Frame<M>>,
+    _marker: PhantomData<&'a Frame<M>>,
+}
+
+impl<M: AnyFrameMeta + ?Sized> FrameRef<'_, M> {
+    /// Borrows the [`Frame`] at the physical address as a [`FrameRef`].
+    ///
+    /// # Safety
+    ///
+    /// The caller must ensure that:
+    ///  - the frame outlives the created reference, so that the reference can
+    ///    be seen as borrowed from that frame.
+    ///  - the type of the [`FrameRef`] (`M`) matches the borrowed frame.
+    pub(in crate::mm) unsafe fn borrow_paddr(raw: Paddr) -> Self {
+        Self {
+            // SAFETY: The caller ensures the safety.
+            inner: ManuallyDrop::new(unsafe { Frame::from_raw(raw) }),
+            _marker: PhantomData,
+        }
+    }
+}
+
+impl<M: AnyFrameMeta + ?Sized> Deref for FrameRef<'_, M> {
+    type Target = Frame<M>;
+
+    fn deref(&self) -> &Self::Target {
+        &self.inner
+    }
+}
+
+// SAFETY: `Frame` is essentially a `*const MetaSlot` that could be used as a non-null
+// `*const` pointer.
+unsafe impl<M: AnyFrameMeta + ?Sized> NonNullPtr for Frame<M> {
+    type Target = PhantomData<Self>;
+
+    type Ref<'a>
+        = FrameRef<'a, M>
+    where
+        Self: 'a;
+
+    const ALIGN_BITS: u32 = core::mem::align_of::<MetaSlot>().trailing_zeros();
+
+    fn into_raw(self) -> NonNull<Self::Target> {
+        let ptr = NonNull::new(self.ptr.cast_mut()).unwrap();
+        let _ = ManuallyDrop::new(self);
+        ptr.cast()
+    }
+
+    unsafe fn from_raw(raw: NonNull<Self::Target>) -> Self {
+        Self {
+            ptr: raw.as_ptr().cast_const().cast(),
+            _marker: PhantomData,
+        }
+    }
+
+    unsafe fn raw_as_ref<'a>(raw: NonNull<Self::Target>) -> Self::Ref<'a> {
+        Self::Ref {
+            inner: ManuallyDrop::new(Frame {
+                ptr: raw.as_ptr().cast_const().cast(),
+                _marker: PhantomData,
+            }),
+            _marker: PhantomData,
+        }
+    }
+
+    fn ref_as_raw(ptr_ref: Self::Ref<'_>) -> core::ptr::NonNull<Self::Target> {
+        NonNull::new(ptr_ref.inner.ptr.cast_mut()).unwrap().cast()
+    }
+}

ostd/src/mm/frame/mod.rs

@@ -38,6 +38,9 @@ pub mod segment;
 pub mod unique;
 pub mod untyped;
 
+mod frame_ref;
+pub use frame_ref::FrameRef;
+
 #[cfg(ktest)]
 mod test;
 
@@ -171,6 +174,12 @@ impl<M: AnyFrameMeta + ?Sized> Frame<M> {
         refcnt
     }
 
+    /// Borrows a reference from the given frame.
+    pub fn borrow(&self) -> FrameRef<'_, M> {
+        // SAFETY: Both the lifetime and the type matches `self`.
+        unsafe { FrameRef::borrow_paddr(self.start_paddr()) }
+    }
+
     /// Forgets the handle to the frame.
     ///
     /// This will result in the frame being leaked without calling the custom dropper.

ostd/src/mm/frame/test.rs

@@ -435,7 +435,6 @@ mod segment {
 // Untyped frame/segment tests
 mod untyped {
     use super::*;
-    use crate::mm::frame::untyped::FrameRef;
 
     #[ktest]
     fn untyped_frame_reader_writer() {
@@ -484,11 +483,33 @@ mod untyped {
         reader.read(&mut buffer.as_mut_slice().into());
         assert_eq!(buffer, data);
     }
+}
+
+mod frame_ref {
+    use super::*;
+    use crate::sync::non_null::NonNullPtr;
+
+    #[ktest]
+    fn frame_ref_preserves_refcnt() {
+        let init_val = 42;
+        let frame = FrameAllocOptions::new()
+            .alloc_frame_with(MockUFrameMeta { value: init_val })
+            .expect("Failed to allocate frame");
+
+        assert_eq!(frame.reference_count(), 1);
+
+        {
+            let frame_ref = frame.borrow();
+            assert_eq!(frame_ref.meta().value, init_val);
+            assert_eq!(frame_ref.reference_count(), 1);
+            assert_eq!(frame.reference_count(), 1);
+        }
+
+        assert_eq!(frame.reference_count(), 1);
+    }
 
     #[ktest]
     fn frame_impls_non_null_ptr() {
-        use crate::sync::non_null::NonNullPtr;
-
         let init_val = 42;
         let frame = FrameAllocOptions::new()
             .alloc_frame_with(MockUFrameMeta { value: init_val })

ostd/src/mm/frame/untyped.rs

@@ -7,16 +7,12 @@
 //! the declaration of untyped frames and segments, and the implementation of
 //! extra functionalities (such as [`VmIo`]) for them.
 
-use super::{
-    meta::{AnyFrameMeta, MetaSlot},
-    Frame, Segment,
-};
+use super::{meta::AnyFrameMeta, Frame, Segment};
 use crate::{
     mm::{
         io::{FallibleVmRead, FallibleVmWrite, VmIo, VmReader, VmWriter},
         paddr_to_vaddr, Infallible,
     },
-    sync::non_null::NonNullPtr,
     Error, Result,
 };
 
@@ -135,63 +131,3 @@ macro_rules! impl_untyped_for {
 
 impl_untyped_for!(Frame);
 impl_untyped_for!(Segment);
-
-// Here are implementations for `crate::sync::rcu`.
-
-use core::{marker::PhantomData, mem::ManuallyDrop, ops::Deref, ptr::NonNull};
-
-/// `FrameRef` is a struct that can work as `&'a Frame<m>`.
-///
-/// This is useful for [`crate::sync::rcu`].
-pub struct FrameRef<'a, M: AnyUFrameMeta + ?Sized> {
-    inner: ManuallyDrop<Frame<M>>,
-    _marker: PhantomData<&'a Frame<M>>,
-}
-
-impl<M: AnyUFrameMeta + ?Sized> Deref for FrameRef<'_, M> {
-    type Target = Frame<M>;
-
-    fn deref(&self) -> &Self::Target {
-        &self.inner
-    }
-}
-
-// SAFETY: `Frame` is essentially an `*const MetaSlot` that could be used as a non-null
-// `*const` pointer.
-unsafe impl<M: AnyUFrameMeta + ?Sized> NonNullPtr for Frame<M> {
-    type Target = PhantomData<Self>;
-
-    type Ref<'a>
-        = FrameRef<'a, M>
-    where
-        Self: 'a;
-
-    const ALIGN_BITS: u32 = core::mem::align_of::<MetaSlot>().trailing_zeros();
-
-    fn into_raw(self) -> NonNull<Self::Target> {
-        let ptr = NonNull::new(self.ptr.cast_mut()).unwrap();
-        let _ = ManuallyDrop::new(self);
-        ptr.cast()
-    }
-
-    unsafe fn from_raw(raw: NonNull<Self::Target>) -> Self {
-        Self {
-            ptr: raw.as_ptr().cast_const().cast(),
-            _marker: PhantomData,
-        }
-    }
-
-    unsafe fn raw_as_ref<'a>(raw: NonNull<Self::Target>) -> Self::Ref<'a> {
-        Self::Ref {
-            inner: ManuallyDrop::new(Frame {
-                ptr: raw.as_ptr().cast_const().cast(),
-                _marker: PhantomData,
-            }),
-            _marker: PhantomData,
-        }
-    }
-
-    fn ref_as_raw(ptr_ref: Self::Ref<'_>) -> core::ptr::NonNull<Self::Target> {
-        NonNull::new(ptr_ref.inner.ptr.cast_mut()).unwrap().cast()
-    }
-}