cu29_runtime/

pool.rs

use arrayvec::ArrayString;
use bincode::de::Decoder;
use bincode::enc::Encoder;
use bincode::error::{DecodeError, EncodeError};
use bincode::{Decode, Encode};
use cu29_traits::CuResult;
use hashbrown::HashMap;
use object_pool::{Pool, ReusableOwned};
use serde::de::{self, MapAccess, SeqAccess, Visitor};
use serde::{Deserialize, Deserializer, Serialize, Serializer};
use smallvec::SmallVec;
use std::alloc::{Layout, alloc, dealloc};
use std::cell::Cell;
#[cfg(feature = "remote-debug")]
use std::cell::RefCell;
use std::cell::UnsafeCell;
use std::fmt::Debug;
use std::fs::OpenOptions;
use std::marker::PhantomData;
use std::mem::{align_of, size_of};
use std::ops::{Deref, DerefMut};
use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicBool, AtomicU8, Ordering};
use std::sync::{Arc, Mutex, MutexGuard, OnceLock};

use memmap2::{MmapMut, MmapOptions};
use tempfile::NamedTempFile;

type PoolID = ArrayString<64>;

/// Trait for a Pool to exposed to be monitored by the monitoring API.
pub trait PoolMonitor: Send + Sync {
    /// A unique and descriptive identifier for the pool.
    fn id(&self) -> PoolID;

    /// Number of buffer slots left in the pool.
    fn space_left(&self) -> usize;

    /// Total size of the pool in number of buffers.
    fn total_size(&self) -> usize;

    /// Size of one buffer
    fn buffer_size(&self) -> usize;
}

static POOL_REGISTRY: OnceLock<Mutex<HashMap<String, Arc<dyn PoolMonitor>>>> = OnceLock::new();
const MAX_POOLS: usize = 16;

fn lock_unpoison<T>(mutex: &Mutex<T>) -> MutexGuard<'_, T> {
    match mutex.lock() {
        Ok(guard) => guard,
        Err(poison) => poison.into_inner(),
    }
}

// Register a pool to the global registry.
fn register_pool(pool: Arc<dyn PoolMonitor>) {
    POOL_REGISTRY
        .get_or_init(|| Mutex::new(HashMap::new()))
        .lock()
        .unwrap_or_else(|poison| poison.into_inner())
        .insert(pool.id().to_string(), pool);
}

type PoolStats = (PoolID, usize, usize, usize);

/// Get the list of pools and their statistics.
/// We use SmallVec here to avoid heap allocations while the stack is running.
pub fn pools_statistics() -> SmallVec<[PoolStats; MAX_POOLS]> {
    // Safely get the registry, returning empty stats if not initialized.
    let registry_lock = match POOL_REGISTRY.get() {
        Some(lock) => lock_unpoison(lock),
        None => return SmallVec::new(), // Return empty if registry is not initialized
    };
    let mut result = SmallVec::with_capacity(MAX_POOLS);
    for pool in registry_lock.values() {
        result.push((
            pool.id(),
            pool.space_left(),
            pool.total_size(),
            pool.buffer_size(),
        ));
    }
    result
}

/// Basic Type that can be used in a buffer in a CuPool.
pub trait ElementType: Default + Sized + Copy + Debug + Unpin + Send + Sync {
    fn encode<E: Encoder>(&self, encoder: &mut E) -> Result<(), EncodeError>;
    fn decode<D: Decoder<Context = ()>>(decoder: &mut D) -> Result<Self, DecodeError>;
}

/// Blanket implementation for all types that are Sized, Copy, Encode, Decode and Debug.
impl<T> ElementType for T
where
    T: Default + Sized + Copy + Debug + Unpin + Send + Sync,
    T: Encode,
    T: Decode<()>,
{
    fn encode<E: Encoder>(&self, encoder: &mut E) -> Result<(), EncodeError> {
        self.encode(encoder)
    }

    fn decode<D: Decoder<Context = ()>>(decoder: &mut D) -> Result<Self, DecodeError> {
        Self::decode(decoder)
    }
}

pub trait ArrayLike: Deref<Target = [Self::Element]> + DerefMut + Debug + Sync + Send {
    type Element: ElementType;
}

thread_local! {
    static SHARED_HANDLE_SERIALIZATION_ENABLED: Cell<bool> = const { Cell::new(false) };
}

pub struct SharedHandleSerializationGuard {
    previous: bool,
}

impl Drop for SharedHandleSerializationGuard {
    fn drop(&mut self) {
        SHARED_HANDLE_SERIALIZATION_ENABLED.with(|enabled| enabled.set(self.previous));
    }
}

pub fn enable_shared_handle_serialization() -> SharedHandleSerializationGuard {
    let previous = SHARED_HANDLE_SERIALIZATION_ENABLED.with(|enabled| {
        let previous = enabled.get();
        enabled.set(true);
        previous
    });
    SharedHandleSerializationGuard { previous }
}

fn shared_handle_serialization_enabled() -> bool {
    SHARED_HANDLE_SERIALIZATION_ENABLED.with(Cell::get)
}

#[cfg(feature = "remote-debug")]
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum DebugHandleEncoding {
    RawLittleEndian,
    Cbor,
}

#[cfg(feature = "remote-debug")]
#[derive(Clone, Debug)]
pub(crate) struct CollectedDebugHandleAttachment {
    pub id: u32,
    pub encoding: DebugHandleEncoding,
    pub element_type: Option<CuSharedMemoryElementType>,
    pub len_elements: Option<usize>,
    pub data: Vec<u8>,
}

#[cfg(feature = "remote-debug")]
#[derive(Serialize)]
struct DebugHandleDescriptor {
    #[serde(rename = "__cu_handle__")]
    marker: bool,
    attachment_id: Option<u32>,
    encoding: DebugHandleEncoding,
    element_type: Option<CuSharedMemoryElementType>,
    len_elements: Option<usize>,
    byte_len: Option<usize>,
}

#[cfg(feature = "remote-debug")]
#[derive(Default)]
struct DebugHandleSerializationState {
    include_contents: bool,
    attachments: Vec<CollectedDebugHandleAttachment>,
}

#[cfg(feature = "remote-debug")]
thread_local! {
    static DEBUG_HANDLE_SERIALIZATION_STATE: RefCell<Option<DebugHandleSerializationState>> =
        const { RefCell::new(None) };
}

/// Runs debugger payload serialization with handle contents deferred by default.
///
/// This function and all state it uses are compiled only for `remote-debug`. It is
/// deliberately independent from Copper's `Encode` path so unified logging and task
/// execution never consult debugger policy.
#[cfg(feature = "remote-debug")]
pub(crate) fn collect_debug_handle_attachments<R>(
    f: impl FnOnce() -> R,
) -> (R, Vec<CollectedDebugHandleAttachment>) {
    let previous = DEBUG_HANDLE_SERIALIZATION_STATE
        .with(|state| state.replace(Some(DebugHandleSerializationState::default())));
    let result = f();
    let collected = DEBUG_HANDLE_SERIALIZATION_STATE.with(|state| {
        let current = state.replace(previous);
        current
            .map(|current| current.attachments)
            .unwrap_or_default()
    });
    (result, collected)
}

#[cfg(feature = "remote-debug")]
pub(crate) fn with_debug_handle_contents<R>(include_contents: bool, f: impl FnOnce() -> R) -> R {
    let previous = DEBUG_HANDLE_SERIALIZATION_STATE.with(|state| {
        let mut state = state.borrow_mut();
        state.as_mut().map(|state| {
            let previous = state.include_contents;
            state.include_contents = include_contents;
            previous
        })
    });
    let result = f();
    if let Some(previous) = previous {
        DEBUG_HANDLE_SERIALIZATION_STATE.with(|state| {
            if let Some(state) = state.borrow_mut().as_mut() {
                state.include_contents = previous;
            }
        });
    }
    result
}

#[cfg(feature = "remote-debug")]
fn primitive_slice_bytes<U: ElementType + 'static>(
    values: &[U],
    element_type: CuSharedMemoryElementType,
) -> Vec<u8> {
    macro_rules! cast_slice {
        ($ty:ty) => {{
            // SAFETY: `element_type` was obtained from `TypeId::of::<U>()`, so
            // `U` and the selected primitive are the same type.
            unsafe { core::slice::from_raw_parts(values.as_ptr().cast::<$ty>(), values.len()) }
        }};
    }

    match element_type {
        CuSharedMemoryElementType::U8 => cast_slice!(u8).to_vec(),
        CuSharedMemoryElementType::I8 => cast_slice!(i8).iter().map(|value| *value as u8).collect(),
        CuSharedMemoryElementType::U16 => cast_slice!(u16)
            .iter()
            .flat_map(|value| value.to_le_bytes())
            .collect(),
        CuSharedMemoryElementType::U32 => cast_slice!(u32)
            .iter()
            .flat_map(|value| value.to_le_bytes())
            .collect(),
        CuSharedMemoryElementType::U64 => cast_slice!(u64)
            .iter()
            .flat_map(|value| value.to_le_bytes())
            .collect(),
        CuSharedMemoryElementType::I16 => cast_slice!(i16)
            .iter()
            .flat_map(|value| value.to_le_bytes())
            .collect(),
        CuSharedMemoryElementType::I32 => cast_slice!(i32)
            .iter()
            .flat_map(|value| value.to_le_bytes())
            .collect(),
        CuSharedMemoryElementType::I64 => cast_slice!(i64)
            .iter()
            .flat_map(|value| value.to_le_bytes())
            .collect(),
        CuSharedMemoryElementType::F32 => cast_slice!(f32)
            .iter()
            .flat_map(|value| value.to_le_bytes())
            .collect(),
        CuSharedMemoryElementType::F64 => cast_slice!(f64)
            .iter()
            .flat_map(|value| value.to_le_bytes())
            .collect(),
    }
}

#[cfg(feature = "remote-debug")]
fn debug_handle_descriptor<U>(values: &[U]) -> Result<Option<DebugHandleDescriptor>, String>
where
    U: ElementType + Serialize + 'static,
{
    DEBUG_HANDLE_SERIALIZATION_STATE.with(|state| {
        let mut state = state.borrow_mut();
        let Some(state) = state.as_mut() else {
            return Ok(None);
        };

        let element_type = CuSharedMemoryElementType::of::<U>();
        let encoding = if element_type.is_some() {
            DebugHandleEncoding::RawLittleEndian
        } else {
            DebugHandleEncoding::Cbor
        };
        let mut byte_len = None;
        let attachment_id = if state.include_contents {
            let data = match element_type {
                Some(element_type) => primitive_slice_bytes(values, element_type),
                None => minicbor_serde::to_vec(values)
                    .map_err(|err| format!("failed to encode debug handle contents: {err}"))?,
            };
            byte_len = Some(data.len());
            let id = state.attachments.len() as u32;
            state.attachments.push(CollectedDebugHandleAttachment {
                id,
                encoding,
                element_type,
                len_elements: Some(values.len()),
                data,
            });
            Some(id)
        } else {
            None
        };

        Ok(Some(DebugHandleDescriptor {
            marker: true,
            attachment_id,
            encoding,
            element_type,
            len_elements: Some(values.len()),
            byte_len,
        }))
    })
}

#[cfg(feature = "remote-debug")]
fn debug_handle_value_descriptor<T>(value: &T) -> Result<Option<DebugHandleDescriptor>, String>
where
    T: Serialize,
{
    DEBUG_HANDLE_SERIALIZATION_STATE.with(|state| {
        let mut state = state.borrow_mut();
        let Some(state) = state.as_mut() else {
            return Ok(None);
        };

        let mut byte_len = None;
        let attachment_id = if state.include_contents {
            let data = minicbor_serde::to_vec(value)
                .map_err(|err| format!("failed to encode debug handle contents: {err}"))?;
            byte_len = Some(data.len());
            let id = state.attachments.len() as u32;
            state.attachments.push(CollectedDebugHandleAttachment {
                id,
                encoding: DebugHandleEncoding::Cbor,
                element_type: None,
                len_elements: None,
                data,
            });
            Some(id)
        } else {
            None
        };

        Ok(Some(DebugHandleDescriptor {
            marker: true,
            attachment_id,
            encoding: DebugHandleEncoding::Cbor,
            element_type: None,
            len_elements: None,
            byte_len,
        }))
    })
}

#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum CuSharedMemoryElementType {
    U8,
    U16,
    U32,
    U64,
    I8,
    I16,
    I32,
    I64,
    F32,
    F64,
}

impl CuSharedMemoryElementType {
    pub fn of<E: ElementType + 'static>() -> Option<Self> {
        let type_id = core::any::TypeId::of::<E>();
        if type_id == core::any::TypeId::of::<u8>() {
            Some(Self::U8)
        } else if type_id == core::any::TypeId::of::<u16>() {
            Some(Self::U16)
        } else if type_id == core::any::TypeId::of::<u32>() {
            Some(Self::U32)
        } else if type_id == core::any::TypeId::of::<u64>() {
            Some(Self::U64)
        } else if type_id == core::any::TypeId::of::<i8>() {
            Some(Self::I8)
        } else if type_id == core::any::TypeId::of::<i16>() {
            Some(Self::I16)
        } else if type_id == core::any::TypeId::of::<i32>() {
            Some(Self::I32)
        } else if type_id == core::any::TypeId::of::<i64>() {
            Some(Self::I64)
        } else if type_id == core::any::TypeId::of::<f32>() {
            Some(Self::F32)
        } else if type_id == core::any::TypeId::of::<f64>() {
            Some(Self::F64)
        } else {
            None
        }
    }
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct CuSharedMemoryHandleDescriptor {
    #[serde(rename = "__cu_shm_handle__")]
    pub marker: bool,
    pub path: String,
    pub offset_bytes: usize,
    pub len_elements: usize,
    pub element_type: CuSharedMemoryElementType,
}

impl CuSharedMemoryHandleDescriptor {
    fn new(
        path: String,
        offset_bytes: usize,
        len_elements: usize,
        element_type: CuSharedMemoryElementType,
    ) -> Self {
        Self {
            marker: true,
            path,
            offset_bytes,
            len_elements,
            element_type,
        }
    }
}

struct CuSharedMemoryRegion {
    path: PathBuf,
    mmap: UnsafeCell<MmapMut>,
    _backing_file: Option<NamedTempFile>,
}

impl CuSharedMemoryRegion {
    fn create(byte_len: usize) -> CuResult<Arc<Self>> {
        let file = NamedTempFile::new()
            .map_err(|e| cu29_traits::CuError::new_with_cause("create shared memory file", e))?;
        file.as_file()
            .set_len(byte_len as u64)
            .map_err(|e| cu29_traits::CuError::new_with_cause("size shared memory file", e))?;
        let mmap = unsafe {
            MmapOptions::new()
                .len(byte_len)
                .map_mut(file.as_file())
                .map_err(|e| cu29_traits::CuError::new_with_cause("map shared memory file", e))?
        };
        let region = Arc::new(Self {
            path: file.path().to_path_buf(),
            mmap: UnsafeCell::new(mmap),
            _backing_file: Some(file),
        });
        cache_shared_region(region.clone());
        Ok(region)
    }

    fn open(path: &Path) -> CuResult<Arc<Self>> {
        if let Some(region) = cached_shared_region(path) {
            return Ok(region);
        }

        let file = OpenOptions::new()
            .read(true)
            .write(true)
            .open(path)
            .map_err(|e| cu29_traits::CuError::new_with_cause("open shared memory file", e))?;
        let len = file
            .metadata()
            .map_err(|e| cu29_traits::CuError::new_with_cause("stat shared memory file", e))?
            .len() as usize;
        let mmap = unsafe {
            MmapOptions::new()
                .len(len)
                .map_mut(&file)
                .map_err(|e| cu29_traits::CuError::new_with_cause("map shared memory file", e))?
        };
        let region = Arc::new(Self {
            path: path.to_path_buf(),
            mmap: UnsafeCell::new(mmap),
            _backing_file: None,
        });
        cache_shared_region(region.clone());
        Ok(region)
    }
}

impl Debug for CuSharedMemoryRegion {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("CuSharedMemoryRegion")
            .field("path", &self.path)
            .finish_non_exhaustive()
    }
}

// SAFETY:
// Access to the mapped bytes is mediated through Copper handles and pool slot
// leasing, so cross-thread aliasing follows the same external synchronization as
// other mutable payload buffers.
unsafe impl Send for CuSharedMemoryRegion {}
// SAFETY:
// See `Send` rationale above.
unsafe impl Sync for CuSharedMemoryRegion {}

fn shared_region_cache() -> &'static Mutex<HashMap<PathBuf, std::sync::Weak<CuSharedMemoryRegion>>>
{
    static CACHE: OnceLock<Mutex<HashMap<PathBuf, std::sync::Weak<CuSharedMemoryRegion>>>> =
        OnceLock::new();
    CACHE.get_or_init(|| Mutex::new(HashMap::new()))
}

fn cache_shared_region(region: Arc<CuSharedMemoryRegion>) {
    lock_unpoison(shared_region_cache()).insert(region.path.clone(), Arc::downgrade(&region));
}

fn cached_shared_region(path: &Path) -> Option<Arc<CuSharedMemoryRegion>> {
    lock_unpoison(shared_region_cache())
        .get(path)
        .and_then(std::sync::Weak::upgrade)
}

fn shared_slot_stride<E: ElementType>(len_elements: usize) -> usize {
    let raw_bytes = len_elements
        .checked_mul(size_of::<E>())
        .expect("shared memory slot size overflow");
    let alignment = align_of::<E>().max(1);
    raw_bytes.div_ceil(alignment) * alignment
}

#[derive(Debug)]
pub struct CuSharedMemoryBuffer<E: ElementType> {
    region: Arc<CuSharedMemoryRegion>,
    offset_bytes: usize,
    len_elements: usize,
    _marker: PhantomData<E>,
}

impl<E: ElementType + 'static> CuSharedMemoryBuffer<E> {
    fn from_region(
        region: Arc<CuSharedMemoryRegion>,
        offset_bytes: usize,
        len_elements: usize,
    ) -> Self {
        Self {
            region,
            offset_bytes,
            len_elements,
            _marker: PhantomData,
        }
    }

    pub fn from_vec_detached(data: Vec<E>) -> CuResult<Self> {
        let len_elements = data.len();
        let slot_stride = shared_slot_stride::<E>(len_elements.max(1));
        let region = CuSharedMemoryRegion::create(slot_stride)?;
        let mut buffer = Self::from_region(region, 0, len_elements);
        if !data.is_empty() {
            buffer.copy_from_slice(&data);
        }
        Ok(buffer)
    }

    pub fn from_descriptor(descriptor: &CuSharedMemoryHandleDescriptor) -> CuResult<Self> {
        let expected = CuSharedMemoryElementType::of::<E>()
            .ok_or_else(|| cu29_traits::CuError::from("unsupported shared memory element type"))?;
        if descriptor.element_type != expected {
            return Err(cu29_traits::CuError::from(
                "shared memory descriptor element type mismatch",
            ));
        }
        let region = CuSharedMemoryRegion::open(Path::new(&descriptor.path))?;
        Ok(Self::from_region(
            region,
            descriptor.offset_bytes,
            descriptor.len_elements,
        ))
    }

    pub fn descriptor(&self) -> Option<CuSharedMemoryHandleDescriptor>
    where
        E: 'static,
    {
        CuSharedMemoryElementType::of::<E>().map(|element_type| {
            CuSharedMemoryHandleDescriptor::new(
                self.region.path.display().to_string(),
                self.offset_bytes,
                self.len_elements,
                element_type,
            )
        })
    }
}

impl<E: ElementType> Deref for CuSharedMemoryBuffer<E> {
    type Target = [E];

    fn deref(&self) -> &Self::Target {
        let ptr = unsafe { (*self.region.mmap.get()).as_ptr().add(self.offset_bytes) as *const E };
        unsafe { std::slice::from_raw_parts(ptr, self.len_elements) }
    }
}

impl<E: ElementType> DerefMut for CuSharedMemoryBuffer<E> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        let ptr = unsafe {
            (*self.region.mmap.get())
                .as_mut_ptr()
                .add(self.offset_bytes) as *mut E
        };
        unsafe { std::slice::from_raw_parts_mut(ptr, self.len_elements) }
    }
}

impl<E: ElementType> ArrayLike for CuSharedMemoryBuffer<E> {
    type Element = E;
}

impl<E: ElementType> Encode for CuSharedMemoryBuffer<E> {
    fn encode<Enc: Encoder>(&self, encoder: &mut Enc) -> Result<(), EncodeError> {
        let len = self.len_elements as u64;
        Encode::encode(&len, encoder)?;
        for value in self.deref() {
            value.encode(encoder)?;
        }
        Ok(())
    }
}

impl<E: ElementType + 'static> Decode<()> for CuSharedMemoryBuffer<E> {
    fn decode<D: Decoder<Context = ()>>(decoder: &mut D) -> Result<Self, DecodeError> {
        let len = <u64 as Decode<()>>::decode(decoder)? as usize;
        let mut vec = Vec::with_capacity(len);
        for _ in 0..len {
            vec.push(E::decode(decoder)?);
        }
        Self::from_vec_detached(vec).map_err(|e| DecodeError::OtherString(e.to_string()))
    }
}

/// A handle to a pooled or detached object.
///
/// For onboard usages, large payloads should typically be pooled. The detached form exists for
/// offline/deserialization flows and for payloads that are intentionally heap-backed instead of
/// pool-backed.
pub enum CuHandleInner<T: Debug + Send + Sync> {
    Pooled(ReusableOwned<Box<T>>),
    Detached(Box<T>), // Should only be used in offline cases (e.g. deserialization)
}

impl<T> Debug for CuHandleInner<T>
where
    T: Debug + Send + Sync,
{
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            CuHandleInner::Pooled(r) => {
                write!(f, "Pooled: {:?}", r.deref().deref())
            }
            CuHandleInner::Detached(r) => write!(f, "Detached: {r:?}"),
        }
    }
}

impl<T> CuHandleInner<T>
where
    T: Debug + Send + Sync,
{
    fn inner_ref(&self) -> &T {
        match self {
            CuHandleInner::Pooled(pooled) => pooled.deref().as_ref(),
            CuHandleInner::Detached(detached) => detached.deref(),
        }
    }

    fn inner_mut(&mut self) -> &mut T {
        match self {
            CuHandleInner::Pooled(pooled) => pooled.deref_mut().as_mut(),
            CuHandleInner::Detached(detached) => detached.deref_mut(),
        }
    }
}

impl<T> AsRef<T> for CuHandleInner<T>
where
    T: Debug + Send + Sync,
{
    fn as_ref(&self) -> &T {
        self.inner_ref()
    }
}

impl<T> AsMut<T> for CuHandleInner<T>
where
    T: Debug + Send + Sync,
{
    fn as_mut(&mut self) -> &mut T {
        self.inner_mut()
    }
}

impl<T: ArrayLike> Deref for CuHandleInner<T> {
    type Target = [T::Element];

    fn deref(&self) -> &Self::Target {
        self.inner_ref().deref()
    }
}

impl<T: ArrayLike> DerefMut for CuHandleInner<T> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        self.inner_mut().deref_mut()
    }
}

// Re-export the configuration policy here for pool consumers that don't otherwise
// reach for the config module.
pub use crate::config::HandleContent;

/// Backing storage for a [`CuHandle`]: the payload mutex plus the per-handle touched flag
/// and logging mode. Shared across handle clones via [`Arc`].
///
/// `mode` is stored as an [`AtomicU8`] so the runtime can override it once after the
/// source's `process()` returns (when the configured `handle_content` policy is known
/// but the source itself didn't construct the handle with that policy in mind).
#[derive(Debug)]
struct CuHandleCell<T: Debug + Send + Sync> {
    touched: AtomicBool,
    mode: AtomicU8,
    inner: Mutex<CuHandleInner<T>>,
}

/// A shareable handle to a pooled or detached object.
///
/// When `T: ArrayLike`, the handle also participates in Copper's buffer pool APIs.
#[derive(Debug)]
#[cfg_attr(feature = "reflect", derive(bevy_reflect::Reflect))]
#[cfg_attr(
    feature = "reflect",
    reflect(opaque, from_reflect = false, type_path = false)
)]
pub struct CuHandle<T: Debug + Send + Sync>(Arc<CuHandleCell<T>>);

#[cfg(feature = "reflect")]
impl<T: Debug + Send + Sync + 'static> bevy_reflect::TypePath for CuHandle<T> {
    fn type_path() -> &'static str {
        core::any::type_name::<Self>()
    }
    fn short_type_path() -> &'static str {
        Self::type_path()
    }
}

impl<T: Debug + Send + Sync + bincode::ValueDecode> bincode::ValueDecode for CuHandle<T> {
    const DECODE: &'static bincode::ValueDecodeSpec = T::DECODE;
}

impl<T: Debug + Send + Sync> Clone for CuHandle<T> {
    fn clone(&self) -> Self {
        Self(self.0.clone())
    }
}

impl<T: Debug + Send + Sync> Deref for CuHandle<T> {
    type Target = Mutex<CuHandleInner<T>>;

    fn deref(&self) -> &Self::Target {
        &self.0.inner
    }
}

impl<T: Debug + Send + Sync> CuHandle<T> {
    /// Wrap a raw [`CuHandleInner`] into a fresh handle with the given logging mode.
    fn from_inner(inner: CuHandleInner<T>, mode: HandleContent) -> Self {
        CuHandle(Arc::new(CuHandleCell {
            touched: AtomicBool::new(false),
            mode: AtomicU8::new(mode as u8),
            inner: Mutex::new(inner),
        }))
    }

    /// Create a new CuHandle not part of a Pool (not for onboard usages, use pools instead)
    pub fn new_detached(inner: T) -> Self {
        Self::new_detached_box(Box::new(inner))
    }

    /// Create a detached handle from an already heap-allocated object.
    pub fn new_detached_box(inner: Box<T>) -> Self {
        Self::from_inner(CuHandleInner::Detached(inner), HandleContent::default())
    }

    /// Create a detached handle with a non-default logging mode.
    ///
    /// Mostly useful for tests and for sources that want to mint detached handles
    /// (instead of pool-acquired ones) under a specific [`HandleContent`] policy.
    pub fn new_detached_with_mode(inner: T, mode: HandleContent) -> Self {
        Self::from_inner(CuHandleInner::Detached(Box::new(inner)), mode)
    }

    /// Safely access the inner value, applying a closure to it.
    pub fn with_inner<R>(&self, f: impl FnOnce(&CuHandleInner<T>) -> R) -> R {
        let lock = lock_unpoison(&self.0.inner);
        f(&*lock)
    }

    /// Mutably access the inner value, applying a closure to it.
    pub fn with_inner_mut<R>(&self, f: impl FnOnce(&mut CuHandleInner<T>) -> R) -> R {
        let mut lock = lock_unpoison(&self.0.inner);
        f(&mut *lock)
    }

    /// Serialize a handle-backed value while allowing the remote debugger to replace
    /// the value with a deferred descriptor.
    ///
    /// Outside a `remote-debug` build this compiles down to normal inner-value
    /// serialization. Handle-backed payload wrappers should use this instead of
    /// locking and serializing the inner value themselves.
    pub fn serialize_value<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        T: Serialize,
        S: Serializer,
    {
        let inner = lock_unpoison(&self.0.inner);
        let value = inner.inner_ref();

        #[cfg(feature = "remote-debug")]
        if let Some(descriptor) =
            debug_handle_value_descriptor(value).map_err(<S::Error as serde::ser::Error>::custom)?
        {
            return descriptor.serialize(serializer);
        }

        value.serialize(serializer)
    }

    /// Returns the number of handles sharing this payload storage.
    pub fn strong_count(&self) -> usize {
        Arc::strong_count(&self.0)
    }

    /// Returns true when this is the only handle to this payload storage.
    pub fn is_unique(&self) -> bool {
        self.strong_count() == 1
    }

    /// Returns a process-local identifier for the shared storage behind this handle.
    ///
    /// Clones of the same handle return the same identifier. This is intended for
    /// diagnostics such as proving that a task graph forwarded a pooled allocation
    /// without replacing it with a copy; it is not a device or host memory address.
    pub fn storage_id(&self) -> usize {
        Arc::as_ptr(&self.0) as usize
    }

    /// Mark this handle as read by a downstream consumer.
    ///
    /// When the source is configured with [`HandleContent::TouchedOnly`], the unified-log
    /// encoder will only write the payload bytes if at least one consumer called this.
    /// Cheap (one relaxed atomic store); safe to call multiple times.
    pub fn mark_touched(&self) {
        self.0.touched.store(true, Ordering::Relaxed);
    }

    /// Returns true if [`mark_touched`](Self::mark_touched) has ever been called on this
    /// handle (or any clone of it).
    pub fn was_touched(&self) -> bool {
        self.0.touched.load(Ordering::Relaxed)
    }

    /// Logging mode currently in effect for this handle.
    pub fn logging_mode(&self) -> HandleContent {
        HandleContent::from_u8(self.0.mode.load(Ordering::Relaxed))
    }

    /// Convenience: [`with_inner`](Self::with_inner) plus [`mark_touched`](Self::mark_touched)
    /// in one call, for the common consumer-side access pattern.
    pub fn with_touched_inner<R>(&self, f: impl FnOnce(&CuHandleInner<T>) -> R) -> R {
        self.mark_touched();
        self.with_inner(f)
    }

    /// Decides whether the unified-log encoder should write this handle's payload bytes
    /// for the current frame.
    ///
    /// This is the inherent "specific" arm of the autoref-specialization pattern; the
    /// encoder resolves to this method for handle payloads (or composite payloads that
    /// forward to it) and to [`PayloadDefaultLoggingPolicy::payload_should_log`] for
    /// every other payload type.
    pub fn payload_should_log(&self) -> bool {
        match self.logging_mode() {
            HandleContent::All => true,
            HandleContent::None => false,
            HandleContent::TouchedOnly => self.was_touched(),
        }
    }

    /// Apply a source's configured [`HandleContent`] policy to this handle. Inherent
    /// "specific" arm of the autoref-specialization pattern; visible to every clone.
    /// `Relaxed` is sufficient — synchronization piggy-backs on the copperlist handoff.
    pub fn apply_handle_content_policy(&self, mode: HandleContent) {
        self.0.mode.store(mode as u8, Ordering::Relaxed);
    }
}

/// Opt-in marker required when `NodeLogging.handle_content` is non-default. Codegen
/// emits a compile-time bound check against this trait, so an unmarked payload
/// surfaces as a clear compile error instead of a silent runtime no-op.
///
/// ```ignore
/// impl cu29::pool::HandleContentAware for MyPayload {}
/// ```
///
/// The impl is the gate; for the policy to actually fire, the payload must also
/// forward [`apply_handle_content_policy`](CuHandle::apply_handle_content_policy)
/// and [`payload_should_log`](CuHandle::payload_should_log) to an inner [`CuHandle`]
/// (see `CuImage`). [`CuHandle`] itself satisfies both halves.
pub trait HandleContentAware {}

impl<T: Debug + Send + Sync> HandleContentAware for CuHandle<T> {}

/// Default arm of the autoref-specialization pattern used by the unified-log encoder.
///
/// Blanket-impl'd for every type, so any payload that doesn't define its own inherent
/// `payload_should_log` method inherits this default-true implementation. Types like
/// [`CuHandle`] (and composite payloads wrapping one) provide an inherent method with
/// the same name, which wins method resolution because inherent methods are tried
/// before trait methods at the same candidate self-type.
pub trait PayloadDefaultLoggingPolicy {
    fn payload_should_log(&self) -> bool {
        true
    }
}

impl<T: ?Sized> PayloadDefaultLoggingPolicy for T {}

/// Default arm of the autoref-specialization pattern used by the runtime to push a
/// source's configured [`HandleContent`] policy into the [`CuHandle`]s that live inside
/// a payload. Default is a no-op; [`CuHandle`] and composite payloads provide inherent
/// overrides that propagate the mode.
pub trait PayloadDefaultHandlePolicyApply {
    fn apply_handle_content_policy(&self, _mode: HandleContent) {}
}

impl<T: ?Sized> PayloadDefaultHandlePolicyApply for T {}

impl<U> Serialize for CuHandle<Vec<U>>
where
    U: ElementType + Serialize + 'static,
{
    fn serialize<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
        let inner = lock_unpoison(&self.0.inner);
        let values = inner.inner_ref();

        #[cfg(feature = "remote-debug")]
        if let Some(descriptor) = debug_handle_descriptor(values.as_slice())
            .map_err(<S::Error as serde::ser::Error>::custom)?
        {
            return descriptor.serialize(serializer);
        }

        values.serialize(serializer)
    }
}

impl<'de, U> Deserialize<'de> for CuHandle<Vec<U>>
where
    U: ElementType + Deserialize<'de> + 'static,
{
    fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
        Vec::<U>::deserialize(deserializer).map(CuHandle::new_detached)
    }
}

impl<U> Serialize for CuHandle<CuSharedMemoryBuffer<U>>
where
    U: ElementType + Serialize + 'static,
{
    fn serialize<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
        let inner = lock_unpoison(&self.0.inner);
        let buffer = inner.inner_ref();

        #[cfg(feature = "remote-debug")]
        if let Some(descriptor) = debug_handle_descriptor(buffer.deref())
            .map_err(<S::Error as serde::ser::Error>::custom)?
        {
            return descriptor.serialize(serializer);
        }

        if shared_handle_serialization_enabled()
            && let Some(descriptor) = buffer.descriptor()
        {
            return descriptor.serialize(serializer);
        }

        buffer.deref().serialize(serializer)
    }
}

impl<'de, U> Deserialize<'de> for CuHandle<CuSharedMemoryBuffer<U>>
where
    U: ElementType + Deserialize<'de> + 'static,
{
    fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
        enum Repr<U> {
            Descriptor(CuSharedMemoryHandleDescriptor),
            Data(Vec<U>),
        }

        impl<'de, U> Deserialize<'de> for Repr<U>
        where
            U: ElementType + Deserialize<'de>,
        {
            fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
                struct ReprVisitor<U>(PhantomData<U>);

                impl<'de, U> Visitor<'de> for ReprVisitor<U>
                where
                    U: ElementType + Deserialize<'de>,
                {
                    type Value = Repr<U>;

                    fn expecting(
                        &self,
                        formatter: &mut std::fmt::Formatter<'_>,
                    ) -> std::fmt::Result {
                        formatter
                            .write_str("a shared-memory handle descriptor or an element sequence")
                    }

                    fn visit_seq<A: SeqAccess<'de>>(self, seq: A) -> Result<Self::Value, A::Error> {
                        let data =
                            Vec::<U>::deserialize(de::value::SeqAccessDeserializer::new(seq))?;
                        Ok(Repr::Data(data))
                    }

                    fn visit_map<A: MapAccess<'de>>(self, map: A) -> Result<Self::Value, A::Error> {
                        let descriptor = CuSharedMemoryHandleDescriptor::deserialize(
                            de::value::MapAccessDeserializer::new(map),
                        )?;
                        Ok(Repr::Descriptor(descriptor))
                    }
                }

                deserializer.deserialize_any(ReprVisitor(PhantomData))
            }
        }

        match Repr::<U>::deserialize(deserializer)? {
            Repr::Descriptor(descriptor) => CuSharedMemoryBuffer::from_descriptor(&descriptor)
                .map(CuHandle::new_detached)
                .map_err(de::Error::custom),
            Repr::Data(data) => CuSharedMemoryBuffer::from_vec_detached(data)
                .map(CuHandle::new_detached)
                .map_err(de::Error::custom),
        }
    }
}

impl<T: ArrayLike + Encode> Encode for CuHandle<T>
where
    <T as ArrayLike>::Element: 'static,
{
    fn encode<E: Encoder>(&self, encoder: &mut E) -> Result<(), EncodeError> {
        let inner = lock_unpoison(&self.0.inner);
        crate::monitoring::record_payload_handle_bytes(
            inner.inner_ref().len() * size_of::<T::Element>(),
        );
        inner.inner_ref().encode(encoder)
    }
}

impl<T: Debug + Send + Sync> Default for CuHandle<T> {
    fn default() -> Self {
        panic!("Cannot create a default CuHandle")
    }
}

impl<U: ElementType + Decode<()> + 'static> Decode<()> for CuHandle<Vec<U>> {
    fn decode<D: Decoder<Context = ()>>(decoder: &mut D) -> Result<Self, DecodeError> {
        let vec: Vec<U> = Vec::decode(decoder)?;
        Ok(CuHandle::new_detached(vec))
    }
}

impl<U: ElementType + Decode<()> + 'static> Decode<()> for CuHandle<CuSharedMemoryBuffer<U>> {
    fn decode<D: Decoder<Context = ()>>(decoder: &mut D) -> Result<Self, DecodeError> {
        let buffer = CuSharedMemoryBuffer::<U>::decode(decoder)?;
        Ok(CuHandle::new_detached(buffer))
    }
}

/// A CuPool is a pool of buffers that can be shared between different parts of the code.
/// Handles can be stored locally in the tasks and shared between them.
pub trait CuPool<T: ArrayLike>: PoolMonitor {
    /// Acquire a buffer from the pool.
    fn acquire(&self) -> Option<CuHandle<T>>;

    /// Copy data from a handle to a new handle from the pool.
    fn copy_from<O>(&self, from: &mut CuHandle<O>) -> CuHandle<T>
    where
        O: ArrayLike<Element = T::Element>;
}

/// A device memory pool can copy data from a device to a host memory pool on top.
pub trait DeviceCuPool<T: ArrayLike>: CuPool<T> {
    /// Takes a handle to a device buffer and copies it into a host buffer pool.
    /// It returns a new handle from the host pool with the data from the device handle given.
    fn copy_to_host_pool<O>(
        &self,
        from_device_handle: &CuHandle<T>,
        to_host_handle: &mut CuHandle<O>,
    ) -> CuResult<()>
    where
        O: ArrayLike<Element = T::Element>;
}

/// A pool of host memory buffers.
pub struct CuHostMemoryPool<T> {
    /// Underlying pool of host buffers.
    // Being an Arc is a requirement of try_pull_owned() so buffers can refer back to the pool.
    id: PoolID,
    pool: Arc<Pool<Box<T>>>,
    size: usize,
    buffer_size: usize,
}

impl<T: ArrayLike + 'static> CuHostMemoryPool<T> {
    pub fn new<F>(id: &str, size: usize, buffer_initializer: F) -> CuResult<Arc<Self>>
    where
        F: Fn() -> T,
    {
        let pool = Arc::new(Pool::new(size, move || Box::new(buffer_initializer())));
        let buffer_size = pool.try_pull().unwrap().len() * size_of::<T::Element>();

        let og = Self {
            id: PoolID::from(id).map_err(|_| "Failed to create PoolID")?,
            pool,
            size,
            buffer_size,
        };
        let og = Arc::new(og);
        register_pool(og.clone());
        Ok(og)
    }
}

impl<T: ArrayLike> PoolMonitor for CuHostMemoryPool<T> {
    fn id(&self) -> PoolID {
        self.id
    }

    fn space_left(&self) -> usize {
        self.pool.len()
    }

    fn total_size(&self) -> usize {
        self.size
    }

    fn buffer_size(&self) -> usize {
        self.buffer_size
    }
}

impl<T: ArrayLike> CuPool<T> for CuHostMemoryPool<T> {
    fn acquire(&self) -> Option<CuHandle<T>> {
        let owned_object = self.pool.try_pull_owned(); // Use the owned version

        owned_object.map(|reusable| {
            CuHandle::from_inner(CuHandleInner::Pooled(reusable), HandleContent::default())
        })
    }

    fn copy_from<O: ArrayLike<Element = T::Element>>(&self, from: &mut CuHandle<O>) -> CuHandle<T> {
        let to_handle = self.acquire().expect("No available buffers in the pool");
        {
            let from_lock = lock_unpoison(&from.0.inner);
            let mut to_lock = lock_unpoison(&to_handle.0.inner);
            to_lock.inner_mut().copy_from_slice(from_lock.inner_ref());
        }
        to_handle
    }
}

/// A pool of fixed-size shared-memory buffers that can be leased to a child
/// process without copying the underlying bytes.
pub struct CuSharedMemoryPool<E: ElementType> {
    id: PoolID,
    pool: Arc<Pool<Box<CuSharedMemoryBuffer<E>>>>,
    size: usize,
    buffer_size: usize,
}

impl<E: ElementType + 'static> CuSharedMemoryPool<E> {
    pub fn new(id: &str, size: usize, elements_per_buffer: usize) -> CuResult<Arc<Self>> {
        let slot_stride = shared_slot_stride::<E>(elements_per_buffer.max(1));
        let region = CuSharedMemoryRegion::create(
            slot_stride
                .checked_mul(size)
                .ok_or_else(|| cu29_traits::CuError::from("shared memory pool size overflow"))?,
        )?;
        let next_slot = Arc::new(std::sync::atomic::AtomicUsize::new(0));
        let initializer_region = region.clone();
        let initializer_next_slot = next_slot.clone();
        let pool = Arc::new(Pool::new(size, move || {
            let slot = initializer_next_slot.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
            assert!(slot < size, "shared memory pool slot index overflow");
            Box::new(CuSharedMemoryBuffer::from_region(
                initializer_region.clone(),
                slot * slot_stride,
                elements_per_buffer,
            ))
        }));

        let pool = Arc::new(Self {
            id: PoolID::from(id).map_err(|_| "Failed to create PoolID")?,
            pool,
            size,
            buffer_size: elements_per_buffer * size_of::<E>(),
        });
        register_pool(pool.clone());
        Ok(pool)
    }
}

impl<E: ElementType> PoolMonitor for CuSharedMemoryPool<E> {
    fn id(&self) -> PoolID {
        self.id
    }

    fn space_left(&self) -> usize {
        self.pool.len()
    }

    fn total_size(&self) -> usize {
        self.size
    }

    fn buffer_size(&self) -> usize {
        self.buffer_size
    }
}

impl<E: ElementType> CuPool<CuSharedMemoryBuffer<E>> for CuSharedMemoryPool<E> {
    fn acquire(&self) -> Option<CuHandle<CuSharedMemoryBuffer<E>>> {
        self.pool.try_pull_owned().map(|reusable| {
            CuHandle::from_inner(CuHandleInner::Pooled(reusable), HandleContent::default())
        })
    }

    fn copy_from<O>(&self, from: &mut CuHandle<O>) -> CuHandle<CuSharedMemoryBuffer<E>>
    where
        O: ArrayLike<Element = E>,
    {
        let to_handle = self.acquire().expect("No available buffers in the pool");
        {
            let from_lock = lock_unpoison(&from.0.inner);
            let mut to_lock = lock_unpoison(&to_handle.0.inner);
            to_lock.inner_mut().copy_from_slice(from_lock.inner_ref());
        }
        to_handle
    }
}

impl<E: ElementType + 'static> ArrayLike for Vec<E> {
    type Element = E;
}

#[cfg(all(feature = "cuda", not(target_os = "macos")))]
pub mod cuda {
    use super::*;
    use cu29_traits::CuError;
    use cudarc::driver::{
        CudaContext, CudaSlice, CudaStream, DeviceRepr, HostSlice, SyncOnDrop, ValidAsZeroBits,
    };
    use std::sync::Arc;

    #[derive(Debug)]
    pub struct CudaSliceWrapper<E>(CudaSlice<E>);

    impl<E> Deref for CudaSliceWrapper<E>
    where
        E: ElementType,
    {
        type Target = [E];

        fn deref(&self) -> &Self::Target {
            // Implement logic to return a slice
            panic!("You need to copy data to host memory pool before accessing it.");
        }
    }

    impl<E> DerefMut for CudaSliceWrapper<E>
    where
        E: ElementType,
    {
        fn deref_mut(&mut self) -> &mut Self::Target {
            panic!("You need to copy data to host memory pool before accessing it.");
        }
    }

    impl<E: ElementType> ArrayLike for CudaSliceWrapper<E> {
        type Element = E;
    }

    impl<E> CudaSliceWrapper<E> {
        /// Number of elements in the device allocation.
        pub fn len(&self) -> usize {
            self.0.len()
        }

        /// Whether the device allocation is empty.
        pub fn is_empty(&self) -> bool {
            self.0.is_empty()
        }

        /// Size of the device allocation in bytes.
        pub fn num_bytes(&self) -> usize {
            self.0.num_bytes()
        }

        pub fn as_cuda_slice(&self) -> &CudaSlice<E> {
            &self.0
        }

        pub fn as_cuda_slice_mut(&mut self) -> &mut CudaSlice<E> {
            &mut self.0
        }
    }

    // Create a wrapper type to bridge between ArrayLike and HostSlice
    pub struct HostSliceWrapper<'a, T: ArrayLike> {
        inner: &'a T,
    }

    impl<T: ArrayLike> HostSlice<T::Element> for HostSliceWrapper<'_, T> {
        fn len(&self) -> usize {
            self.inner.len()
        }

        // SAFETY: HostSlice requires the returned slice to remain valid for 'b.
        unsafe fn stream_synced_slice<'b>(
            &'b self,
            stream: &'b CudaStream,
        ) -> (&'b [T::Element], SyncOnDrop<'b>) {
            (self.inner.deref(), SyncOnDrop::sync_stream(stream))
        }

        // SAFETY: This wrapper cannot provide mutable access; callers must not rely on this.
        unsafe fn stream_synced_mut_slice<'b>(
            &'b mut self,
            _stream: &'b CudaStream,
        ) -> (&'b mut [T::Element], SyncOnDrop<'b>) {
            panic!("Cannot get mutable reference from immutable wrapper")
        }
    }

    // Mutable wrapper
    pub struct HostSliceMutWrapper<'a, T: ArrayLike> {
        inner: &'a mut T,
    }

    impl<T: ArrayLike> HostSlice<T::Element> for HostSliceMutWrapper<'_, T> {
        fn len(&self) -> usize {
            self.inner.len()
        }

        // SAFETY: HostSlice requires the returned slice to remain valid for 'b.
        unsafe fn stream_synced_slice<'b>(
            &'b self,
            stream: &'b CudaStream,
        ) -> (&'b [T::Element], SyncOnDrop<'b>) {
            (self.inner.deref(), SyncOnDrop::sync_stream(stream))
        }

        // SAFETY: HostSlice requires the returned slice to remain valid for 'b.
        unsafe fn stream_synced_mut_slice<'b>(
            &'b mut self,
            stream: &'b CudaStream,
        ) -> (&'b mut [T::Element], SyncOnDrop<'b>) {
            (self.inner.deref_mut(), SyncOnDrop::sync_stream(stream))
        }
    }

    // Add helper methods to the CuCudaPool implementation
    impl<E: ElementType + ValidAsZeroBits + DeviceRepr> CuCudaPool<E> {
        // Helper method to get a HostSliceWrapper from a CuHandleInner
        fn get_host_slice_wrapper<O: ArrayLike<Element = E>>(
            handle_inner: &CuHandleInner<O>,
        ) -> HostSliceWrapper<'_, O> {
            HostSliceWrapper {
                inner: handle_inner.inner_ref(),
            }
        }

        // Helper method to get a HostSliceMutWrapper from a CuHandleInner
        fn get_host_slice_mut_wrapper<O: ArrayLike<Element = E>>(
            handle_inner: &mut CuHandleInner<O>,
        ) -> HostSliceMutWrapper<'_, O> {
            HostSliceMutWrapper {
                inner: handle_inner.inner_mut(),
            }
        }
    }
    /// A pool of CUDA memory buffers.
    pub struct CuCudaPool<E>
    where
        E: ElementType + ValidAsZeroBits + DeviceRepr + Unpin,
    {
        id: PoolID,
        stream: Arc<CudaStream>,
        pool: Arc<Pool<Box<CudaSliceWrapper<E>>>>,
        nb_buffers: usize,
        nb_element_per_buffer: usize,
    }

    impl<E: ElementType + ValidAsZeroBits + DeviceRepr + 'static> CuCudaPool<E> {
        pub fn new(
            id: &str,
            ctx: Arc<CudaContext>,
            nb_buffers: usize,
            nb_element_per_buffer: usize,
        ) -> CuResult<Arc<Self>> {
            let stream = ctx.default_stream();
            let pool = (0..nb_buffers)
                .map(|_| {
                    stream
                        .alloc_zeros(nb_element_per_buffer)
                        .map(CudaSliceWrapper)
                        .map(Box::new)
                        .map_err(|_| "Failed to allocate device memory")
                })
                .collect::<Result<Vec<_>, _>>()?;

            let pool = Arc::new(Self {
                id: PoolID::from(id).map_err(|_| "Failed to create PoolID")?,
                stream,
                pool: Arc::new(Pool::from_vec(pool)),
                nb_buffers,
                nb_element_per_buffer,
            });
            register_pool(pool.clone());
            Ok(pool)
        }
    }

    impl<E> PoolMonitor for CuCudaPool<E>
    where
        E: DeviceRepr + ElementType + ValidAsZeroBits,
    {
        fn id(&self) -> PoolID {
            self.id
        }

        fn space_left(&self) -> usize {
            self.pool.len()
        }

        fn total_size(&self) -> usize {
            self.nb_buffers
        }

        fn buffer_size(&self) -> usize {
            self.nb_element_per_buffer * size_of::<E>()
        }
    }

    impl<E> CuPool<CudaSliceWrapper<E>> for CuCudaPool<E>
    where
        E: DeviceRepr + ElementType + ValidAsZeroBits,
    {
        fn acquire(&self) -> Option<CuHandle<CudaSliceWrapper<E>>> {
            self.pool
                .try_pull_owned()
                .map(|x| CuHandle::from_inner(CuHandleInner::Pooled(x), HandleContent::default()))
        }

        fn copy_from<O>(&self, from_handle: &mut CuHandle<O>) -> CuHandle<CudaSliceWrapper<E>>
        where
            O: ArrayLike<Element = E>,
        {
            let to_handle = self.acquire().expect("No available buffers in the pool");

            {
                let from_lock = lock_unpoison(&from_handle.0.inner);
                let mut to_lock = lock_unpoison(&to_handle.0.inner);

                match &mut *to_lock {
                    CuHandleInner::Detached(to) => {
                        let wrapper = Self::get_host_slice_wrapper(&*from_lock);
                        self.stream
                            .memcpy_htod(&wrapper, to.deref_mut().as_cuda_slice_mut())
                            .expect("Failed to copy data to device");
                    }
                    CuHandleInner::Pooled(to) => {
                        let wrapper = Self::get_host_slice_wrapper(&*from_lock);
                        self.stream
                            .memcpy_htod(&wrapper, to.deref_mut().as_mut().as_cuda_slice_mut())
                            .expect("Failed to copy data to device");
                    }
                }
            } // locks are dropped here
            to_handle // now we can safely return to_handle
        }
    }

    impl<E> DeviceCuPool<CudaSliceWrapper<E>> for CuCudaPool<E>
    where
        E: ElementType + ValidAsZeroBits + DeviceRepr,
    {
        /// Copy from device to host
        fn copy_to_host_pool<O>(
            &self,
            device_handle: &CuHandle<CudaSliceWrapper<E>>,
            host_handle: &mut CuHandle<O>,
        ) -> Result<(), CuError>
        where
            O: ArrayLike<Element = E>,
        {
            let device_lock = device_handle.lock().map_err(|e| {
                CuError::from("Device handle mutex poisoned").add_cause(&e.to_string())
            })?;
            let mut host_lock = host_handle.lock().map_err(|e| {
                CuError::from("Host handle mutex poisoned").add_cause(&e.to_string())
            })?;
            let src = match &*device_lock {
                CuHandleInner::Pooled(source) => source.deref().as_ref().as_cuda_slice(),
                CuHandleInner::Detached(source) => source.deref().as_cuda_slice(),
            };
            let mut wrapper = Self::get_host_slice_mut_wrapper(&mut *host_lock);
            self.stream.memcpy_dtoh(src, &mut wrapper).map_err(|e| {
                CuError::from("Failed to copy data from device to host").add_cause(&e.to_string())
            })?;
            Ok(())
        }
    }
}

#[derive(Debug)]
/// A buffer that is aligned to a specific size with the Element of type E.
pub struct AlignedBuffer<E: ElementType> {
    ptr: *mut E,
    size: usize,
    layout: Layout,
}

impl<E: ElementType> AlignedBuffer<E> {
    pub fn new(num_elements: usize, alignment: usize) -> Self {
        assert!(
            num_elements > 0 && size_of::<E>() > 0,
            "AlignedBuffer requires a non-zero element count and non-zero-sized element type"
        );
        let alignment = alignment.max(align_of::<E>());
        let alloc_size = num_elements
            .checked_mul(size_of::<E>())
            .expect("AlignedBuffer allocation size overflow");
        let layout = Layout::from_size_align(alloc_size, alignment).unwrap();
        // SAFETY: layout describes a valid, non-zero allocation request.
        let ptr = unsafe { alloc(layout) as *mut E };
        if ptr.is_null() {
            panic!("Failed to allocate memory");
        }
        // SAFETY: ptr is valid for writes of `num_elements` elements.
        unsafe {
            for i in 0..num_elements {
                std::ptr::write(ptr.add(i), E::default());
            }
        }
        Self {
            ptr,
            size: num_elements,
            layout,
        }
    }
}

impl<E: ElementType> Deref for AlignedBuffer<E> {
    type Target = [E];

    fn deref(&self) -> &Self::Target {
        // SAFETY: `new` initializes all elements and keeps the pointer aligned.
        unsafe { std::slice::from_raw_parts(self.ptr, self.size) }
    }
}

impl<E: ElementType> DerefMut for AlignedBuffer<E> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        // SAFETY: `new` initializes all elements and keeps the pointer aligned.
        unsafe { std::slice::from_raw_parts_mut(self.ptr, self.size) }
    }
}

impl<E: ElementType> Drop for AlignedBuffer<E> {
    fn drop(&mut self) {
        // SAFETY: `ptr` was allocated with `layout` in `new`.
        unsafe { dealloc(self.ptr as *mut u8, self.layout) }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[cfg(feature = "self-describing-logs")]
    #[test]
    fn test_catalog_decodes_native_handle_buffer() {
        use cu29_value::Value;
        use cu29_value::decode::{ValueDecodeDescription, ValueDecodeLimits};

        let handle = CuHandle::new_detached(vec![1u16, 300, 65535]);
        let config = bincode::config::standard();
        let bytes = bincode::encode_to_vec(&handle, config).unwrap();
        let description = ValueDecodeDescription::from_type::<CuHandle<Vec<u16>>>().unwrap();
        let (decoded, used) = description
            .decode(&bytes, config, ValueDecodeLimits::default())
            .unwrap();
        assert_eq!(used, bytes.len());
        assert_eq!(
            decoded,
            Value::Seq(vec![Value::U16(1), Value::U16(300), Value::U16(65535)])
        );
        assert!(!handle.was_touched());
    }

    #[test]
    fn test_handle_touched_flag_defaults_false() {
        let h: CuHandle<Vec<u8>> = CuHandle::new_detached(vec![1, 2, 3]);
        assert!(!h.was_touched());
        assert_eq!(h.logging_mode(), HandleContent::All);
    }

    #[test]
    fn test_handle_mark_touched_propagates_across_clones() {
        let h: CuHandle<Vec<u8>> = CuHandle::new_detached(vec![1, 2, 3]);
        let clone = h.clone();
        assert!(!h.was_touched());
        assert!(!clone.was_touched());

        clone.mark_touched();
        // Shared Arc<CuHandleCell>: any clone sees the flag flip.
        assert!(h.was_touched());
        assert!(clone.was_touched());
    }

    #[test]
    fn test_handle_strong_count_tracks_clones() {
        let h: CuHandle<Vec<u8>> = CuHandle::new_detached(vec![1, 2, 3]);
        assert_eq!(h.strong_count(), 1);
        assert!(h.is_unique());

        let clone = h.clone();
        assert_eq!(h.strong_count(), 2);
        assert_eq!(clone.strong_count(), 2);
        assert!(!h.is_unique());
        assert!(!clone.is_unique());

        drop(clone);
        assert_eq!(h.strong_count(), 1);
        assert!(h.is_unique());
    }

    #[test]
    fn test_handle_storage_id_tracks_shared_allocation() {
        let first: CuHandle<Vec<u8>> = CuHandle::new_detached(vec![1]);
        let clone = first.clone();
        let second: CuHandle<Vec<u8>> = CuHandle::new_detached(vec![1]);

        assert_eq!(first.storage_id(), clone.storage_id());
        assert_ne!(first.storage_id(), second.storage_id());
    }

    #[test]
    fn test_with_touched_inner_marks_and_reads() {
        let h: CuHandle<Vec<u8>> = CuHandle::new_detached(vec![10, 20]);
        let first = h.with_touched_inner(|inner| inner.as_ref()[0]);
        assert_eq!(first, 10);
        assert!(h.was_touched());
    }

    #[test]
    fn test_with_inner_does_not_mark_touched() {
        let h: CuHandle<Vec<u8>> = CuHandle::new_detached(vec![10, 20]);
        let _ = h.with_inner(|inner| inner.as_ref()[0]);
        assert!(
            !h.was_touched(),
            "with_inner must not flip the touched flag"
        );
    }

    #[test]
    fn test_payload_should_log_mode_all() {
        let h: CuHandle<Vec<u8>> = CuHandle::new_detached_with_mode(vec![1], HandleContent::All);
        assert!(h.payload_should_log());
        h.mark_touched();
        assert!(h.payload_should_log());
    }

    #[test]
    fn test_payload_should_log_mode_none() {
        let h: CuHandle<Vec<u8>> = CuHandle::new_detached_with_mode(vec![1], HandleContent::None);
        assert!(!h.payload_should_log());
        h.mark_touched();
        assert!(
            !h.payload_should_log(),
            "HandleContent::None must never log payload, even when touched"
        );
    }

    #[test]
    fn test_payload_should_log_mode_touched_only() {
        let h: CuHandle<Vec<u8>> =
            CuHandle::new_detached_with_mode(vec![1], HandleContent::TouchedOnly);
        assert!(!h.payload_should_log(), "untouched + TouchedOnly => skip");
        h.mark_touched();
        assert!(h.payload_should_log(), "touched + TouchedOnly => log");
    }

    #[test]
    fn test_default_policy_returns_true_for_non_handle_types() {
        use crate::pool::PayloadDefaultLoggingPolicy as _;
        // The autoref-specialization fallback applies to any type that doesn't define
        // its own inherent payload_should_log. A plain integer is a fine stand-in.
        let v: u64 = 42;
        assert!(v.payload_should_log());
    }

    #[test]
    fn test_apply_handle_content_policy_overrides_mode() {
        // Handle minted with the default All policy; runtime hook flips it to
        // TouchedOnly before downstream consumers see it. Encoder decision tracks.
        let h: CuHandle<Vec<u8>> = CuHandle::new_detached(vec![1, 2, 3]);
        assert_eq!(h.logging_mode(), HandleContent::All);
        assert!(h.payload_should_log());

        h.apply_handle_content_policy(HandleContent::TouchedOnly);
        assert_eq!(h.logging_mode(), HandleContent::TouchedOnly);
        assert!(
            !h.payload_should_log(),
            "after switch to TouchedOnly an untouched handle must skip"
        );

        h.mark_touched();
        assert!(h.payload_should_log(), "touched + TouchedOnly => log");
    }

    #[test]
    fn test_default_apply_policy_is_noop_for_non_handle_types() {
        use crate::pool::PayloadDefaultHandlePolicyApply as _;
        // No-op for any payload that doesn't wrap a handle — compile-time check that
        // codegen can call this on every slot without knowing the payload shape.
        let v: u64 = 42;
        v.apply_handle_content_policy(HandleContent::TouchedOnly);
        // No observable change; the test just proves the call compiles and returns.
    }

    #[cfg(feature = "remote-debug")]
    #[test]
    fn debug_handle_contents_are_deferred_until_explicitly_included() {
        let handle = CuHandle::new_detached(vec![1.25_f32, -2.5_f32]);

        let (deferred, attachments) = collect_debug_handle_attachments(|| {
            with_debug_handle_contents(false, || {
                minicbor_serde::to_vec(&handle).expect("serialize deferred handle")
            })
        });
        let deferred: serde_json::Value =
            minicbor_serde::from_slice(&deferred).expect("decode deferred descriptor");
        assert_eq!(
            deferred.get("__cu_handle__"),
            Some(&serde_json::json!(true))
        );
        assert_eq!(
            deferred.get("attachment_id"),
            Some(&serde_json::Value::Null)
        );
        assert!(attachments.is_empty());

        let (included, attachments) = collect_debug_handle_attachments(|| {
            with_debug_handle_contents(true, || {
                minicbor_serde::to_vec(&handle).expect("serialize included handle")
            })
        });
        let included: serde_json::Value =
            minicbor_serde::from_slice(&included).expect("decode included descriptor");
        assert_eq!(included.get("attachment_id"), Some(&serde_json::json!(0)));
        assert_eq!(attachments.len(), 1);
        assert_eq!(
            attachments[0].encoding,
            DebugHandleEncoding::RawLittleEndian
        );
        assert_eq!(
            attachments[0].element_type,
            Some(CuSharedMemoryElementType::F32)
        );
        assert_eq!(
            attachments[0].data,
            [1.25_f32.to_le_bytes(), (-2.5_f32).to_le_bytes()].concat()
        );
    }

    #[test]
    fn test_pool() {
        use std::cell::RefCell;
        let objs = RefCell::new(vec![vec![1], vec![2], vec![3]]);
        let holding = objs.borrow().clone();
        let objs_as_slices = holding.iter().map(|x| x.as_slice()).collect::<Vec<_>>();
        let pool = CuHostMemoryPool::new("mytestcudapool", 3, || objs.borrow_mut().pop().unwrap())
            .unwrap();

        let obj1 = pool.acquire().unwrap();
        {
            let obj2 = pool.acquire().unwrap();
            assert!(objs_as_slices.contains(&obj1.lock().unwrap().deref().deref()));
            assert!(objs_as_slices.contains(&obj2.lock().unwrap().deref().deref()));
            assert_eq!(pool.space_left(), 1);
        }
        assert_eq!(pool.space_left(), 2);

        let obj3 = pool.acquire().unwrap();
        assert!(objs_as_slices.contains(&obj3.lock().unwrap().deref().deref()));

        assert_eq!(pool.space_left(), 1);

        let _obj4 = pool.acquire().unwrap();
        assert_eq!(pool.space_left(), 0);

        let obj5 = pool.acquire();
        assert!(obj5.is_none());
    }

    #[cfg(all(feature = "cuda", has_nvidia_gpu))]
    #[test]
    fn test_cuda_pool() {
        use crate::pool::cuda::CuCudaPool;
        use cudarc::driver::CudaContext;
        let ctx = CudaContext::new(0).unwrap();
        let pool = CuCudaPool::<f32>::new("mytestcudapool", ctx, 3, 1).unwrap();

        let _obj1 = pool.acquire().unwrap();

        {
            let _obj2 = pool.acquire().unwrap();
            assert_eq!(pool.space_left(), 1);
        }
        assert_eq!(pool.space_left(), 2);

        let _obj3 = pool.acquire().unwrap();

        assert_eq!(pool.space_left(), 1);

        let _obj4 = pool.acquire().unwrap();
        assert_eq!(pool.space_left(), 0);

        let obj5 = pool.acquire();
        assert!(obj5.is_none());
    }

    #[cfg(all(feature = "cuda", has_nvidia_gpu))]
    #[test]
    fn test_copy_roundtrip() {
        use crate::pool::cuda::CuCudaPool;
        use cudarc::driver::CudaContext;
        let ctx = CudaContext::new(0).unwrap();
        let host_pool = CuHostMemoryPool::new("mytesthostpool", 3, || vec![0.0; 1]).unwrap();
        let cuda_pool = CuCudaPool::<f32>::new("mytestcudapool", ctx, 3, 1).unwrap();

        let cuda_handle = {
            let mut initial_handle = host_pool.acquire().unwrap();
            {
                let mut inner_initial_handle = initial_handle.lock().unwrap();
                if let CuHandleInner::Pooled(ref mut pooled) = *inner_initial_handle {
                    pooled[0] = 42.0;
                } else {
                    panic!();
                }
            }

            // send that to the GPU
            cuda_pool.copy_from(&mut initial_handle)
        };

        // get it back to the host
        let mut final_handle = host_pool.acquire().unwrap();
        cuda_pool
            .copy_to_host_pool(&cuda_handle, &mut final_handle)
            .unwrap();

        let value = final_handle.lock().unwrap().deref().deref()[0];
        assert_eq!(value, 42.0);
    }
}