celox 0.11.0

use std::{any::Any, cell::UnsafeCell, sync::Arc};

/// Fixed-capacity storage for one host simulation memory image.
///
/// The allocation is reference counted separately from the backend so an FFI
/// view can keep the bytes alive without retaining compiled code, scheduler
/// state, or background compilation workers.
struct MemoryAllocation {
    words: Box<[UnsafeCell<u64>]>,
    start_word: usize,
    capacity_words: usize,
}

// SAFETY: `MemoryAllocation` only owns storage. Safe access to its contents is
// mediated by `MemoryImage`, which requires `&self` for reads and `&mut self`
// for writes. Cloned `Arc`s are opaque lifetime leases and expose no accessors.
// FFI callers that receive a raw pointer remain responsible for synchronizing
// external access with simulator execution, as they were before this owner was
// introduced.
unsafe impl Sync for MemoryAllocation {}

impl MemoryAllocation {
    fn zeroed(capacity_words: usize) -> Self {
        let alignment = celox_state_layout::LANE_SEGMENT_ALIGN;
        let extra_words = alignment / std::mem::size_of::<u64>();
        let allocation_words = capacity_words
            .checked_add(extra_words)
            .expect("simulation memory allocation size overflow");
        let words = std::iter::repeat_with(|| UnsafeCell::new(0u64))
            .take(allocation_words)
            .collect::<Vec<_>>()
            .into_boxed_slice();
        // Offset alignment in MemoryLayout only isolates cache lines if the
        // exposed allocation base has the same alignment. Keep the original
        // Box as owner so allocation and deallocation layouts still match.
        let start_word = words.as_ptr().align_offset(alignment);
        assert!(start_word < extra_words);
        Self {
            words,
            start_word,
            capacity_words,
        }
    }

    fn as_ptr(&self) -> *const u64 {
        // SAFETY: the initialized padding covers start_word and the entire
        // requested capacity; even an empty image has an in-allocation base.
        unsafe { self.words.as_ptr().add(self.start_word).cast::<u64>() }
    }

    fn as_mut_ptr(&self) -> *mut u64 {
        // SAFETY: start_word lies within the allocated initialized padding.
        unsafe { UnsafeCell::raw_get(self.words.as_ptr().add(self.start_word)) }
    }

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

/// A logical memory image backed by a fixed-capacity allocation.
///
/// Growing the logical image is permitted only within its allocation. Tiered
/// execution reserves its maximum required capacity before exposing a pointer,
/// ensuring promotion never invalidates external zero-copy views.
pub(crate) struct MemoryImage {
    allocation: Arc<MemoryAllocation>,
    logical_words: usize,
    // Keep this outside the exposed allocation: restoring or clearing the
    // whole state image must not forget an escaped writable pointer.
    raw_view_exposed: bool,
}

impl MemoryImage {
    pub(crate) fn zeroed(logical_words: usize) -> Self {
        Self {
            allocation: Arc::new(MemoryAllocation::zeroed(logical_words)),
            logical_words,
            raw_view_exposed: false,
        }
    }

    pub(crate) fn as_ptr(&self) -> *const u64 {
        self.allocation.as_ptr()
    }

    pub(crate) fn as_mut_ptr(&mut self) -> *mut u64 {
        self.allocation.as_mut_ptr()
    }

    pub(crate) fn expose_mut_ptr(&mut self) -> *mut u64 {
        self.raw_view_exposed = true;
        self.as_mut_ptr()
    }

    /// Overwrite bytes of the image without exposing a writable view.
    pub(crate) fn write_bytes(&mut self, offset: usize, bytes: &[u8]) {
        assert!(
            offset + bytes.len() <= self.logical_words * 8,
            "write past the memory image"
        );
        // Safety: the range lies inside the logical image, which this value
        // owns mutably, and `bytes` is a separate allocation.
        unsafe {
            std::ptr::copy_nonoverlapping(
                bytes.as_ptr(),
                self.as_mut_ptr().cast::<u8>().add(offset),
                bytes.len(),
            );
        }
    }

    /// Store a native-endian word in a state-header slot, such as the
    /// address of a host buffer that generated code reads.
    pub(crate) fn write_header_word(&mut self, offset: usize, value: u64) {
        self.write_bytes(offset, &value.to_ne_bytes());
    }

    pub(crate) fn vcd_tracking_enabled(&self) -> bool {
        !self.raw_view_exposed
    }

    #[cfg(test)]
    fn as_slice(&self) -> &[u64] {
        // SAFETY: the allocation contains at least `logical_words` initialized
        // `u64`s and remains alive for the returned borrow.
        unsafe { std::slice::from_raw_parts(self.as_ptr(), self.logical_words) }
    }

    pub(crate) fn as_mut_slice(&mut self) -> &mut [u64] {
        // SAFETY: `&mut self` provides the backend-side exclusive borrow. Raw
        // FFI users must separately obey the shared-memory synchronization
        // contract documented on `MemoryAllocation`.
        unsafe { std::slice::from_raw_parts_mut(self.as_mut_ptr(), self.logical_words) }
    }

    pub(crate) fn len_words(&self) -> usize {
        self.logical_words
    }

    pub(crate) fn capacity_words(&self) -> usize {
        self.allocation.capacity_words()
    }

    /// Reserve a larger fixed allocation before any external owner is issued.
    pub(crate) fn reserve_total(&mut self, total_words: usize) {
        if total_words <= self.capacity_words() {
            return;
        }
        assert_eq!(
            Arc::strong_count(&self.allocation),
            1,
            "cannot move a simulation memory image after exposing an owner"
        );

        let replacement = Arc::new(MemoryAllocation::zeroed(total_words));
        // SAFETY: both allocations are valid for `logical_words` elements,
        // they cannot overlap, and no external lease exists by the assertion.
        unsafe {
            std::ptr::copy_nonoverlapping(
                self.allocation.as_ptr(),
                replacement.as_mut_ptr(),
                self.logical_words,
            );
        }
        self.allocation = replacement;
    }

    /// Change the logical length without moving the allocation.
    pub(crate) fn resize_zeroed_within_capacity(&mut self, new_words: usize) {
        assert!(
            new_words <= self.capacity_words(),
            "simulation memory image exceeded its reserved capacity"
        );
        if new_words > self.logical_words {
            // SAFETY: the asserted capacity covers the extension and the
            // destination lies within this allocation.
            unsafe {
                std::ptr::write_bytes(
                    self.allocation.as_mut_ptr().add(self.logical_words),
                    0,
                    new_words - self.logical_words,
                );
            }
        }
        self.logical_words = new_words;
    }

    pub(crate) fn owner(&self) -> Arc<dyn Any + Send + Sync> {
        self.allocation.clone()
    }
}

impl Default for MemoryImage {
    fn default() -> Self {
        Self::zeroed(0)
    }
}

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

    #[test]
    fn allocation_and_reserved_growth_preserve_coherence_alignment() {
        for size in [0, 1, 2, 31, 32, 33, 129, 1025] {
            let mut image = MemoryImage::zeroed(size);
            assert_eq!(
                image.as_ptr() as usize % celox_state_layout::LANE_SEGMENT_ALIGN,
                0
            );
            assert_eq!(image.capacity_words(), size);
            image.as_mut_slice().fill(17);
            image.reserve_total(size + 33);
            assert_eq!(
                image.as_ptr() as usize % celox_state_layout::LANE_SEGMENT_ALIGN,
                0
            );
            assert_eq!(image.as_slice(), vec![17; size]);
            let pointer = image.as_ptr();
            let owner = image.owner();
            image.resize_zeroed_within_capacity(size + 33);
            assert_eq!(image.as_ptr(), pointer);
            assert!(image.as_slice()[size..].iter().all(|word| *word == 0));
            drop(image);
            // An escaped view remains valid for the lifetime of its owner.
            assert_eq!(unsafe { *pointer.add(size) }, 0);
            drop(owner);
        }
    }

    #[test]
    fn owner_keeps_only_the_allocation_alive() {
        let image = MemoryImage::zeroed(2);
        let allocation = Arc::downgrade(&image.allocation);
        let first = image.owner();
        let second = image.owner();

        drop(image);
        assert!(allocation.upgrade().is_some());
        drop(first);
        assert!(allocation.upgrade().is_some());
        drop(second);
        assert!(allocation.upgrade().is_none());
    }

    #[test]
    fn reserved_growth_preserves_pointer_and_zeroes_extension() {
        let mut image = MemoryImage::zeroed(2);
        image.as_mut_slice().copy_from_slice(&[11, 22]);
        image.reserve_total(4);
        let ptr = image.as_ptr();

        image.resize_zeroed_within_capacity(4);

        assert_eq!(image.as_ptr(), ptr);
        assert_eq!(image.as_slice(), &[11, 22, 0, 0]);
    }

    #[test]
    #[should_panic(expected = "after exposing an owner")]
    fn reserve_rejects_moving_an_exposed_allocation() {
        let mut image = MemoryImage::zeroed(1);
        let _owner = image.owner();
        image.reserve_total(2);
    }

    #[test]
    fn empty_image_has_no_logical_words() {
        let mut image = MemoryImage::default();
        assert!(image.as_slice().is_empty());
        assert!(image.as_mut_slice().is_empty());
    }
}