From 6aba270a9b2f0ee2752506c36811d24a7f61c34d Mon Sep 17 00:00:00 2001
From: Chen Chengjun <chenchengjun.ccj@antgroup.com>
Date: Thu, 10 Apr 2025 10:41:33 +0800
Subject: [PATCH] Enable RCU to support Weak

---
 ostd/src/sync/rcu/non_null.rs | 53 +++++++++++++++++++++++++++++++++++
 1 file changed, 53 insertions(+)

diff --git a/ostd/src/sync/rcu/non_null.rs b/ostd/src/sync/rcu/non_null.rs
index f33c5765d..22a2c04b5 100644
--- a/ostd/src/sync/rcu/non_null.rs
+++ b/ostd/src/sync/rcu/non_null.rs
@@ -3,6 +3,7 @@
 //! This module provides a trait and some auxiliary types to help abstract and
 //! work with non-null pointers.
 
+use alloc::sync::Weak;
 use core::{marker::PhantomData, mem::ManuallyDrop, ops::Deref, ptr::NonNull};
 
 use crate::prelude::*;
@@ -188,3 +189,55 @@ unsafe impl<T: Send + Sync + 'static> NonNullPtr for Arc<T> {
         NonNullPtr::into_raw(ManuallyDrop::into_inner(ptr_ref.inner))
     }
 }
+
+/// A type that represents `&'a Weak<T>`.
+#[derive(Debug)]
+pub struct WeakRef<'a, T> {
+    inner: ManuallyDrop<Weak<T>>,
+    _marker: PhantomData<&'a Weak<T>>,
+}
+
+impl<T> Deref for WeakRef<'_, T> {
+    type Target = Weak<T>;
+
+    fn deref(&self) -> &Self::Target {
+        &self.inner
+    }
+}
+
+unsafe impl<T: Send + Sync + 'static> NonNullPtr for Weak<T> {
+    type Target = T;
+
+    type Ref<'a>
+        = WeakRef<'a, T>
+    where
+        Self: 'a;
+
+    fn into_raw(self) -> NonNull<Self::Target> {
+        let ptr = Weak::into_raw(self).cast_mut();
+
+        // SAFETY: The pointer representing an `Weak` can never be NULL.
+        unsafe { NonNull::new_unchecked(ptr) }
+    }
+
+    unsafe fn from_raw(ptr: NonNull<Self::Target>) -> Self {
+        let ptr = ptr.as_ptr().cast_const();
+
+        // SAFETY: The safety is upheld by the caller.
+        unsafe { Weak::from_raw(ptr) }
+    }
+
+    unsafe fn raw_as_ref<'a>(raw: NonNull<Self::Target>) -> Self::Ref<'a> {
+        // SAFETY: The safety is upheld by the caller.
+        unsafe {
+            WeakRef {
+                inner: ManuallyDrop::new(Weak::from_raw(raw.as_ptr())),
+                _marker: PhantomData,
+            }
+        }
+    }
+
+    fn ref_as_raw(ptr_ref: Self::Ref<'_>) -> NonNull<Self::Target> {
+        NonNullPtr::into_raw(ManuallyDrop::into_inner(ptr_ref.inner))
+    }
+}