ruda-runtime 0.1.13

use crate::runtime::memory_management::MemoryHandle;
use alloc::sync::Arc;
use core::sync::atomic::AtomicUsize;
#[cfg(target_has_atomic = "64")]
use core::sync::atomic::{AtomicU64, Ordering};
#[cfg(not(target_has_atomic = "64"))]
use spin::Mutex;

/// Managed Memory handle
#[derive(Debug)]
pub struct ManagedMemoryHandle {
    descriptor: Arc<ManagedMemoryDescriptor>,
    // Holds only the reference counts of the handle.
    handle_count: Arc<()>,
}

/// Binding of a memory handle
#[derive(Debug)]
pub struct ManagedMemoryBinding {
    descriptor: Arc<ManagedMemoryDescriptor>,
}

impl Clone for ManagedMemoryHandle {
    fn clone(&self) -> Self {
        Self {
            descriptor: self.descriptor.clone(),
            handle_count: self.handle_count.clone(),
        }
    }
}

/// Managed memory descriptor.
///
/// Host-side diagnostics can read the location while the device thread updates
/// it. Protect the whole location so all readers observe a consistent snapshot.
/// Send/Sync are derived from the fields; no unchecked Cell sharing is needed.
pub(crate) struct ManagedMemoryDescriptor {
    pub(crate) id: ManagedMemoryId,
    pins: AtomicUsize,
    #[cfg(target_has_atomic = "64")]
    location: AtomicU64,
    #[cfg(not(target_has_atomic = "64"))]
    location: Mutex<MemoryLocation>,
}

impl core::fmt::Debug for ManagedMemoryDescriptor {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        f.debug_struct("ManagedMemoryDescriptor")
            .field("id", &self.id)
            .field("location", &self.location())
            .finish()
    }
}

#[derive(Debug, PartialEq, Eq, Clone, Copy, Hash)]
/// Managed memory unique identifier.
pub struct ManagedMemoryId {
    pub(crate) value: usize,
}

impl PartialEq for ManagedMemoryDescriptor {
    fn eq(&self, other: &Self) -> bool {
        self.id == other.id
    }
}

impl Eq for ManagedMemoryDescriptor {}

#[derive(Clone, Copy, Debug)]
/// Defines where the [`ManagedMemoryId`] is located.
pub(crate) struct MemoryLocation {
    /// The memory pool index in the global memory management.
    pub pool: u8,
    /// The memory page index in a memory pool.
    pub page: u16,
    /// The memory slice index in a memory page.
    pub slice: u32,
    /// Whether the memory location is known/initialized.
    pub init: u8,
}

impl ManagedMemoryDescriptor {
    /// Update the memory location for the given [`ManagedMemoryId`].
    pub(crate) fn update_location(&self, location: MemoryLocation) {
        #[cfg(target_has_atomic = "64")]
        {
            self.location.store(location.to_bits(), Ordering::Release);
        }
        #[cfg(not(target_has_atomic = "64"))]
        {
            *self.location.lock() = location;
        }
    }

    /// Update only the slice position for the given [`ManagedMemoryId`].
    pub(crate) fn update_slice(&self, slice: u32) {
        self.modify(|location| MemoryLocation { slice, ..location });
    }

    /// Update only the memory page position for the given [`ManagedMemoryId`].
    pub fn update_page(&self, page: u16) {
        self.modify(|location| MemoryLocation { page, ..location });
    }

    /// Retrieves the current location.
    pub(crate) fn location(&self) -> MemoryLocation {
        #[cfg(target_has_atomic = "64")]
        {
            MemoryLocation::from_bits(self.location.load(Ordering::Acquire))
        }
        #[cfg(not(target_has_atomic = "64"))]
        {
            *self.location.lock()
        }
    }

    pub(crate) fn slice(&self) -> usize {
        self.location().slice as usize
    }

    pub(crate) fn page(&self) -> usize {
        self.location().page as usize
    }

    fn modify(&self, update: impl Fn(MemoryLocation) -> MemoryLocation) {
        #[cfg(target_has_atomic = "64")]
        {
            let _ = self.location.fetch_update(Ordering::AcqRel, Ordering::Acquire, |bits| {
                Some(update(MemoryLocation::from_bits(bits)).to_bits())
            });
        }
        #[cfg(not(target_has_atomic = "64"))]
        {
            let mut location = self.location.lock();
            *location = update(*location);
        }
    }
}

impl MemoryLocation {
    #[cfg(target_has_atomic = "64")]
    fn to_bits(self) -> u64 {
        self.pool as u64
            | (self.page as u64) << 8
            | (self.slice as u64) << 24
            | (self.init as u64) << 56
    }

    #[cfg(target_has_atomic = "64")]
    fn from_bits(bits: u64) -> Self {
        Self {
            pool: bits as u8,
            page: (bits >> 8) as u16,
            slice: (bits >> 24) as u32,
            init: (bits >> 56) as u8,
        }
    }

    /// Creates a new memory location.
    pub(crate) fn new(pool: u8, page: u16, slice: u32) -> Self {
        Self {
            pool,
            page,
            slice,
            init: 1,
        }
    }

    /// Creates a new uninitialized memory location.
    pub(crate) fn uninit() -> Self {
        Self {
            pool: 0,
            page: 0,
            slice: 0,
            init: 0,
        }
    }
}

impl ManagedMemoryHandle {
    /// Creates a new managed memory handle.
    pub fn new() -> Self {
        let value = Self::gen_id();

        Self {
            descriptor: Arc::new(ManagedMemoryDescriptor {
                id: ManagedMemoryId { value },
                pins: AtomicUsize::new(0),
                #[cfg(target_has_atomic = "64")]
                location: AtomicU64::new(MemoryLocation::uninit().to_bits()),
                #[cfg(not(target_has_atomic = "64"))]
                location: Mutex::new(MemoryLocation::uninit()),
            }),
            handle_count: Arc::new(()),
        }
    }

    /// Retrieves the descriptor for the current handle.
    pub(crate) fn descriptor(&self) -> &ManagedMemoryDescriptor {
        &self.descriptor
    }

    /// Return whether the current handle can be modified in-place.
    pub fn can_mut(&self) -> bool {
        Arc::strong_count(&self.handle_count) <= 2
    }

    /// Return whether the current handle is free.
    pub fn is_free(&self) -> bool {
        Arc::strong_count(&self.descriptor) <= 1
    }

    pub(crate) fn is_pinned(&self) -> bool {
        self.descriptor.pins.load(core::sync::atomic::Ordering::Acquire) != 0
    }

    /// Returns the binding for the current handle.
    pub fn binding(self) -> ManagedMemoryBinding {
        ManagedMemoryBinding {
            descriptor: self.descriptor.clone(),
        }
    }

    fn gen_id() -> usize {
        static COUNTER: core::sync::atomic::AtomicUsize = core::sync::atomic::AtomicUsize::new(0);
        let value = COUNTER.fetch_add(1, core::sync::atomic::Ordering::Relaxed);
        if value == usize::MAX {
            core::panic!("Memory ID overflowed");
        }
        value
    }
}

impl ManagedMemoryBinding {
    /// Keep this allocation's address fixed until the returned pin is dropped.
    pub fn pin(&self) -> MemoryResourcePin {
        self.descriptor.pins.fetch_add(1, core::sync::atomic::Ordering::AcqRel);
        MemoryResourcePin { binding: self.clone() }
    }
    /// Stable allocation identity, independent of the device address or view.
    /// This is an identity token, not a pointer or proof that a resource is ready.
    pub fn id(&self) -> ManagedMemoryId { self.descriptor.id }

    /// Retrieves the descriptor for the current binding.
    pub(crate) fn descriptor(&self) -> &ManagedMemoryDescriptor {
        &self.descriptor
    }
}

/// A resolved resource or native graph's fixed-address allocation lease.
#[derive(Debug)]
pub struct MemoryResourcePin {
    binding: ManagedMemoryBinding,
}

impl Clone for MemoryResourcePin {
    fn clone(&self) -> Self { self.binding.pin() }
}

impl Drop for MemoryResourcePin {
    fn drop(&mut self) {
        self.binding.descriptor.pins.fetch_sub(1, core::sync::atomic::Ordering::AcqRel);
    }
}

impl Default for ManagedMemoryHandle {
    fn default() -> Self {
        Self::new()
    }
}

impl Clone for ManagedMemoryBinding {
    fn clone(&self) -> Self {
        Self {
            descriptor: self.descriptor.clone(),
        }
    }
}

impl MemoryHandle<ManagedMemoryBinding> for ManagedMemoryHandle {
    fn can_mut(&self) -> bool {
        self.can_mut()
    }

    fn binding(self) -> ManagedMemoryBinding {
        self.binding()
    }
}

/// Calculates a best-effort heuristic for the alignment of row-aligned tensors.
/// Prefers contiguous alignments for unit dimensions, 16-byte minimum alignment for non-unit,
/// scaling with input size up to `buffer_align`.
pub fn optimal_align(shape: usize, elem_size: usize, buffer_align: usize) -> usize {
    if shape == 1 {
        elem_size
    } else {
        (shape * elem_size)
            .next_power_of_two()
            .clamp(16, buffer_align)
    }
}

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

    #[test]
    fn test_memory_id_mutability() {
        let handle1 = ManagedMemoryHandle::new();
        handle1.descriptor().update_slice(4);
        assert_eq!(handle1.descriptor().slice(), 4);

        let handle2 = ManagedMemoryHandle::new();
        handle2
            .clone()
            .descriptor()
            .update_location(handle1.descriptor().location());
        assert_eq!(handle2.descriptor().slice(), 4);
    }

    #[test]
    fn test_location_visible_through_shared_arc() {
        let handle = ManagedMemoryHandle::new();
        let handle2 = handle.clone();

        let location = MemoryLocation::new(1, 2, 3);
        handle.descriptor().update_location(location);

        assert_eq!(handle2.descriptor().location().pool, 1);
        assert_eq!(handle2.descriptor().location().page, 2);
        assert_eq!(handle2.descriptor().location().slice, 3);
        assert_eq!(handle2.descriptor().location().init, 1);

        handle.descriptor().update_slice(42);
        assert_eq!(handle2.descriptor().slice(), 42);
    }

    #[test]
    #[cfg(feature = "runtime-std")]
    fn concurrent_debug_reads_consistent_location_snapshots() {
        let handle = ManagedMemoryHandle::new();
        let writer = handle.clone();
        let task = std::thread::spawn(move || {
            for i in 1..=128_u32 {
                writer.descriptor().update_location(MemoryLocation::new(i as u8, i as u16, i));
                std::thread::yield_now();
            }
        });
        for _ in 0..128 {
            let _ = alloc::format!("{handle:?}");
            let location = handle.descriptor().location();
            assert_eq!(location.page as u32, location.slice);
            assert_eq!(location.pool as u32, location.slice);
            std::thread::yield_now();
        }
        task.join().unwrap();
        assert_eq!(handle.descriptor().slice(), 128);
    }

    #[test]
    #[cfg(feature = "runtime-std")]
    fn concurrent_field_updates_do_not_overwrite_each_other() {
        let handle = ManagedMemoryHandle::new();
        let first = handle.clone();
        let second = handle.clone();
        let page_writer = std::thread::spawn(move || {
            for page in 1..=128 { first.descriptor().update_page(page); }
        });
        let slice_writer = std::thread::spawn(move || {
            for slice in 1..=128 { second.descriptor().update_slice(slice); }
        });
        page_writer.join().unwrap();
        slice_writer.join().unwrap();
        assert_eq!(handle.descriptor().page(), 128);
        assert_eq!(handle.descriptor().slice(), 128);
    }

}