ksync/

kcell.rs

// Copyright 2026 The Fuchsia Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

use crate::lock_token::{LockTokenRead, LockTokenWrite};
use core::marker::PhantomData;

/// A container cell that safe-guards data of type `T` using type-level lock class tracking.
///
/// `KCell` holds data that can only be safely accessed (read or written) by proving that the
/// corresponding mutual exclusion lock of class `Class` is currently held by the current thread
/// via a `LockToken`.
#[repr(transparent)]
pub struct KCell<T, Class> {
    value: core::cell::UnsafeCell<T>,
    _marker: PhantomData<Class>,
}

// SAFETY: `KCell` synchronizes shared access via `LockToken`, so it is `Sync` when `T: Send`.
unsafe impl<T: Send, Class> Sync for KCell<T, Class> {}
// SAFETY: `KCell` owns `T`, so transferring it across threads is safe when `T: Send`.
unsafe impl<T: Send, Class> Send for KCell<T, Class> {}

impl<T, Class> KCell<T, Class> {
    /// Creates a new `KCell` containing the specified `value`.
    #[inline]
    pub const fn new(value: T) -> Self {
        Self { value: core::cell::UnsafeCell::new(value), _marker: PhantomData }
    }

    /// Access the guarded value immutably using a shared lock token.
    ///
    /// # Safety
    ///
    /// The caller must guarantee that the provided `LockToken` belongs to the specific lock
    /// instance that guards this `KCell` (rather than a different lock of the same lock class
    /// `Class`).
    #[inline]
    pub unsafe fn get<'b, Token, Idx>(&self, _token: &'b Token) -> &'b T
    where
        Token: LockTokenRead<Class, Idx>,
    {
        // SAFETY: The caller guarantees that the provided LockToken belongs to the specific lock
        // instance that guards this cell.
        unsafe { &*self.value.get() }
    }

    /// Access the guarded value mutably using a mutable lock token.
    ///
    /// # Safety
    ///
    /// The caller must guarantee that the provided `LockToken` belongs to the specific lock
    /// instance that guards this `KCell` (rather than a different lock of the same lock class
    /// `Class`).
    #[inline]
    pub unsafe fn get_mut<'b, Token, Idx>(&self, _token: &'b mut Token) -> &'b mut T
    where
        Token: LockTokenWrite<Class, Idx>,
    {
        // SAFETY: The caller guarantees that the provided LockToken belongs to the specific lock
        // instance that guards this cell, and the exclusive mutable borrow of the `LockToken`
        // ensures that no other active borrows of the same cell can co-exist.
        unsafe { &mut *self.value.get() }
    }

    /// Access the guarded value as a pinned mutable reference using a mutable lock token.
    ///
    /// # Safety
    ///
    /// The caller must guarantee that the provided `LockToken` belongs to the specific lock
    /// instance that guards this `KCell` (rather than a different lock of the same lock class
    /// `Class`), and that this `KCell` is pinned in memory so its inner value is never moved.
    #[inline]
    pub unsafe fn get_pinned_mut_locked<'b, Token, Idx>(
        &self,
        token: &'b mut Token,
    ) -> core::pin::Pin<&'b mut T>
    where
        Token: LockTokenWrite<Class, Idx>,
    {
        // SAFETY: The caller guarantees that the provided LockToken belongs to the specific lock
        // instance that guards this cell, and that the cell is pinned in memory.
        unsafe { core::pin::Pin::new_unchecked(self.get_mut(token)) }
    }

    /// Returns a const raw pointer to the guarded value.
    ///
    /// # Safety
    ///
    /// The caller must guarantee that a lock instance guarding this `KCell` is held continuously
    /// while dereferencing or accessing the returned raw pointer.
    #[inline]
    pub unsafe fn as_ptr<Token, Idx>(&self, _token: &Token) -> *const T
    where
        Token: LockTokenRead<Class, Idx>,
    {
        self.value.get()
    }

    /// Returns a mutable raw pointer to the guarded value.
    ///
    /// # Safety
    ///
    /// The caller must guarantee that the lock instance(s) guarding this `KCell` are held
    /// continuously while dereferencing or accessing the returned raw pointer.
    #[inline]
    pub unsafe fn as_mut_ptr<Token, Idx>(&self, _token: &mut Token) -> *mut T
    where
        Token: LockTokenWrite<Class, Idx>,
    {
        self.value.get()
    }

    /// Accesses the inner value mutably by bypassing the locking requirements using unique borrow
    /// ownership.
    #[inline]
    pub fn get_inner_mut(&mut self) -> &mut T {
        self.value.get_mut()
    }

    /// Projects a pinned reference to the cell to a pinned reference to the inner value.
    #[inline]
    pub fn get_pinned_mut(self: core::pin::Pin<&mut Self>) -> core::pin::Pin<&mut T> {
        // SAFETY: `KCell` is `repr(transparent)` and contains `T` at the same address.
        // Pinning is preserved because `KCell` does not move `T`.
        unsafe { self.map_unchecked_mut(|s| s.get_inner_mut()) }
    }

    /// Initializes a `KCell` in place with an `Init` initializer for its inner value.
    pub fn init<I, E>(init: I) -> impl pin_init::Init<Self, E>
    where
        I: pin_init::Init<T, E>,
    {
        kcell_init(init)
    }

    /// Initializes a `KCell` with a `PinInit` initializer for its inner value.
    pub fn pin_init<I, E>(init: I) -> impl pin_init::PinInit<Self, E>
    where
        I: pin_init::PinInit<T, E>,
    {
        kcell_init(init)
    }

    /// Unwraps the cell, returning the inner value.
    #[inline]
    pub fn into_inner(self) -> T {
        self.value.into_inner()
    }
}

impl<T: Default, Class> Default for KCell<T, Class> {
    #[inline]
    fn default() -> Self {
        Self::new(T::default())
    }
}

impl<T, Class> From<T> for KCell<T, Class> {
    #[inline]
    fn from(value: T) -> Self {
        Self::new(value)
    }
}

impl<T, Class> AsMut<T> for KCell<T, Class> {
    #[inline]
    fn as_mut(&mut self) -> &mut T {
        self.get_inner_mut()
    }
}

impl<T, Class> core::fmt::Debug for KCell<T, Class> {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        f.debug_struct("KCell")
            .field("value", &"<locked>")
            .field("class", &core::any::type_name::<Class>())
            .finish()
    }
}

/// Creates an initializer wrapper for initializing a `KCell` with an inner `Init` or `PinInit`
/// initializer.
#[inline]
pub fn kcell_init<I, Class>(init: I) -> KCellInit<I, Class> {
    KCellInit(init, PhantomData)
}

/// Initializer for `KCell`.
pub struct KCellInit<I, Class>(I, PhantomData<Class>);

// SAFETY: KCell is repr(transparent) and contains T at the same address.
// Pinning is preserved because KCell doesn't move T.
unsafe impl<T, Class, I, E> pin_init::PinInit<KCell<T, Class>, E> for KCellInit<I, Class>
where
    I: pin_init::PinInit<T, E>,
{
    unsafe fn __pinned_init(self, slot: *mut KCell<T, Class>) -> Result<(), E> {
        // SAFETY: The caller guarantees slot is valid. KCell is repr(transparent)
        // so slot has the same address and layout as T.
        unsafe { self.0.__pinned_init(slot as *mut T) }
    }
}

// SAFETY: KCell is repr(transparent) and contains T at the same address and layout.
unsafe impl<T, Class, I, E> pin_init::Init<KCell<T, Class>, E> for KCellInit<I, Class>
where
    I: pin_init::Init<T, E>,
{
    unsafe fn __init(self, slot: *mut KCell<T, Class>) -> Result<(), E> {
        // SAFETY: The caller guarantees slot is valid. KCell is repr(transparent)
        // so slot has the same address and layout as T.
        unsafe { self.0.__init(slot as *mut T) }
    }
}

#[cfg(not(feature = "kernel"))]
#[cfg(test)]
mod tests {
    use super::*;
    use crate::lock_token::LockToken;
    use lockdep::LockClass;

    struct MyClass;
    impl LockClass for MyClass {
        const ID: *mut core::ffi::c_void = core::ptr::null_mut();
    }

    #[test]
    fn test_kcell_default() {
        let cell: KCell<u32, MyClass> = KCell::default();
        // SAFETY: Single-threaded test with exclusive access to `cell`.
        unsafe {
            let token = LockToken::new();
            assert_eq!(*cell.get(&token), 0);
        }
    }

    #[test]
    fn test_kcell_exclusive_access() {
        let mut cell: KCell<u32, MyClass> = KCell::new(10);
        *cell.get_inner_mut() = 20;
        let reference: &mut u32 = cell.as_mut();
        *reference = 30;
        let value = cell.into_inner();
        assert_eq!(value, 30);
    }

    #[test]
    fn test_kcell_as_mut_ptr() {
        let cell: KCell<u32, MyClass> = KCell::new(100u32);
        // SAFETY: Single-threaded test with exclusive access to `cell`.
        unsafe {
            let mut token = LockToken::new();
            let ptr = cell.as_mut_ptr(&mut token);
            assert_eq!(*ptr, 100);
        }
    }

    #[test]
    fn test_kcell_debug() {
        extern crate std;
        let cell: KCell<u32, MyClass> = KCell::new(5);
        let debug_str = std::format!("{:?}", cell);
        assert!(debug_str.contains("KCell"));
        assert!(debug_str.contains("<locked>"));
        assert!(debug_str.contains("MyClass"));
    }

    #[test]
    fn test_kcell_init() {
        // SAFETY: Closure initializes `slot` in place before returning `Ok(())`.
        let init = unsafe {
            pin_init::pin_init_from_closure(|slot: *mut u32| {
                slot.write(42);
                Ok::<(), core::convert::Infallible>(())
            })
        };
        let cell_init = kcell_init::<_, MyClass>(init);

        pin_init::stack_pin_init!(let cell = cell_init);
        let cell: core::pin::Pin<&mut KCell<u32, MyClass>> = cell;

        // SAFETY: Single-threaded test with exclusive access to `cell`.
        unsafe {
            let token = LockToken::<MyClass>::new();
            let val = cell.as_ref().get_ref().get(&token);
            assert_eq!(*val, 42);
        }

        #[pin_init::pin_data]
        struct UnpinnedHolder {
            cell: KCell<[usize; 4], MyClass>,
        }

        pin_init::stack_pin_init!(let holder = pin_init::pin_init!(UnpinnedHolder {
            cell <- KCell::<[usize; 4], MyClass>::init(pin_init::init_zeroed::<[usize; 4]>()),
        }));
        // SAFETY: Single-threaded test with exclusive access to `holder.cell`.
        unsafe {
            let token = LockToken::<MyClass>::new();
            assert_eq!(*holder.as_ref().get_ref().cell.get(&token), [0usize; 4]);
        }
    }

    #[test]
    fn test_kcell_get_pinned_mut_locked() {
        let cell = core::pin::pin!(KCell::<u32, MyClass>::new(7));
        let cell: &KCell<u32, MyClass> = cell.as_ref().get_ref();

        // SAFETY: Single-threaded test with exclusive access to pinned `cell`.
        unsafe {
            let mut token = LockToken::<MyClass>::new();
            let mut pinned: core::pin::Pin<&mut u32> = cell.get_pinned_mut_locked(&mut token);
            assert_eq!(*pinned, 7);
            pinned.set(8);
            *pinned.as_mut().get_mut() += 1;
            assert_eq!(*cell.get(&token), 9);
        }
    }

    struct SecondClass;
    impl LockClass for SecondClass {
        const ID: *mut core::ffi::c_void = core::ptr::null_mut();
    }

    struct ThirdClass;
    impl LockClass for ThirdClass {
        const ID: *mut core::ffi::c_void = core::ptr::null_mut();
    }

    #[test]
    fn test_kcell_multi_lock() {
        let cell = core::pin::pin!(KCell::<u32, (MyClass, SecondClass)>::new(10));
        let cell: &KCell<u32, (MyClass, SecondClass)> = cell.as_ref().get_ref();
        // SAFETY: Single-threaded test with exclusive access to `cell`.
        unsafe {
            let mut token1 = LockToken::<MyClass>::new();
            let mut token2 = LockToken::<SecondClass>::new();

            // Read with either single token via `get` and `as_ptr`.
            assert_eq!(*cell.get(&token1), 10);
            assert_eq!(*cell.get(&token2), 10);
            assert_eq!(*cell.as_ptr(&token1), 10);
            assert_eq!(*cell.as_ptr(&token2), 10);

            // Mutate with joined tokens in either order.
            *cell.get_mut(token1.join(&mut token2)) = 20;
            assert_eq!(*cell.get(&token1), 20);

            *cell.get_mut(token2.join(&mut token1)) = 30;
            assert_eq!(*cell.get(&token2), 30);

            // Test `as_mut_ptr` and `get_pinned_mut_locked` with joined tokens.
            *cell.as_mut_ptr(token1.join(&mut token2)) = 40;
            assert_eq!(*cell.get(&token1), 40);

            let mut pinned = cell.get_pinned_mut_locked(token2.join(&mut token1));
            pinned.set(50);
            assert_eq!(*cell.get(&token2), 50);
        }

        let cell3: KCell<u32, (MyClass, SecondClass, ThirdClass)> = KCell::new(100);
        // SAFETY: Single-threaded test with exclusive access to `cell3`.
        unsafe {
            let mut t1 = LockToken::<MyClass>::new();
            let mut t2 = LockToken::<SecondClass>::new();
            let mut t3 = LockToken::<ThirdClass>::new();

            assert_eq!(*cell3.get(&t1), 100);
            assert_eq!(*cell3.get(&t2), 100);
            assert_eq!(*cell3.get(&t3), 100);

            *cell3.get_mut(t1.join(&mut t2).join(&mut t3)) = 200;
            assert_eq!(*cell3.get(&t3), 200);
        }
    }
}