zombie_rs/

context.rs

//! Context nodes for storing computation state and enabling replay.
//!
//! Contexts are the fundamental building blocks of the tock tree (akasha).
//! They store zombie values and the metadata needed to recompute them.
//!
//! # Context Types
//!
//! - [`RootContext`]: Non-evictable values at the top level. Created when
//!   `Zombie::new` is called outside of any `bind` operation.
//!
//! - [`FullContext`]: Evictable values from `bind` operations. Contains a
//!   [`Replayer`] that can recreate all values in this context.
//!
//! # Paper Reference
//!
//! Section 3 (PVZ): "A context at tock t contain the following:
//! - A function pointer
//! - A list of tock representing the input of said function pointer
//! - A Timespan recording how long the execution take
//! - A vector of shared ptr, storing ZombieNode"
//!
//! # Memory Layout
//!
//! Each context owns an array of `Option<Rc<dyn EZombieNode>>`. When a value
//! is evicted, its slot becomes `None` but the context remains to preserve
//! the replayer for future recomputation.

use std::any::Any;
use std::cell::{Cell, RefCell};
use std::num::NonZeroUsize;
use std::rc::{Rc, Weak};

use smallvec::SmallVec;

use crate::config::Cost;
use crate::meter::{Space, Time, ZombieOps};
use crate::record::Replayer;
use crate::tock::Tock;
use crate::union_find::UF;

/// SmallVec type alias for storing Tock dependencies.
/// Most bind operations have 1-4 dependencies, so inline storage for 4 elements
/// avoids heap allocation in the common case. sizeof(Tock) = 8, so 4*8 = 32 bytes inline.
pub type TockVec = SmallVec<[Tock; 4]>;

/// SmallVec type alias for storing zombie node values.
/// Most contexts contain 1-2 values, so inline storage for 2 elements
/// avoids heap allocation. sizeof(Option<Rc<...>>) = 16, so 2*16 = 32 bytes inline.
pub type ValueVec = SmallVec<[Option<Rc<dyn EZombieNode>>; 2]>;

/// Reference-counted context node (type-erased).
pub type ContextRef = Rc<dyn ContextNode>;

/// Weak reference to context node (type-erased).
pub type ContextWeakRef = Weak<dyn ContextNode>;

/// Type-erased zombie node trait.
pub trait EZombieNode: Any {
    fn created_time(&self) -> Tock;
    fn size(&self) -> usize;
    fn as_any(&self) -> &dyn Any;
    fn get_value_ptr(&self) -> *const dyn Any;
    fn set_context(&self, ctx: ContextWeakRef);
    fn get_context(&self) -> Option<ContextRef>;
    /// Get a clone of the weak context reference without upgrading.
    /// This is more efficient than get_context() when you only need the weak reference.
    fn clone_context_weak(&self) -> Option<ContextWeakRef>;
    fn notify_accessed(&self);
}

/// Concrete zombie node holding a value of type T.
///
/// Memory layout optimized to match C++ implementation:
/// - created_time: 8 bytes (Tock = u64)
/// - value: sizeof(T)
/// - context: 16 bytes (Weak pointer, but dyn makes it fat pointer)
///
/// Note: Unlike C++ which stores context as a simple weak_ptr (8 bytes),
/// Rust requires a fat pointer for dyn ContextNode (16 bytes).
/// This is an unavoidable overhead in Rust's type system.
pub struct ZombieNode<T: ZombieOps> {
    pub created_time: Tock,
    pub value: T,
    // Use Cell for smaller size (RefCell adds 8 bytes for borrow flag)
    context: std::cell::Cell<Option<ContextWeakRef>>,
}

impl<T: 'static + Clone + ZombieOps> ZombieNode<T> {
    #[inline]
    pub fn new(created_time: Tock, value: T) -> Rc<Self> {
        Rc::new(Self {
            created_time,
            value,
            context: std::cell::Cell::new(None),
        })
    }

    #[inline]
    pub fn get(&self) -> &T {
        self.notify_accessed();
        &self.value
    }
}

impl<T: 'static + ZombieOps> EZombieNode for ZombieNode<T> {
    fn created_time(&self) -> Tock {
        self.created_time
    }

    fn size(&self) -> usize {
        // Compute size dynamically like C++ get_size() virtual function
        self.value.get_size()
    }

    fn as_any(&self) -> &dyn Any {
        self
    }

    fn get_value_ptr(&self) -> *const dyn Any {
        &self.value as &dyn Any as *const dyn Any
    }

    fn set_context(&self, ctx: ContextWeakRef) {
        self.context.set(Some(ctx));
    }

    fn get_context(&self) -> Option<ContextRef> {
        // Take the weak pointer out, try to upgrade, then put it back
        let weak = self.context.take();
        let result = weak.as_ref().and_then(|w| w.upgrade());
        self.context.set(weak);
        result
    }

    fn clone_context_weak(&self) -> Option<ContextWeakRef> {
        // Take the weak pointer out, clone it, then put it back
        let weak = self.context.take();
        let result = weak.clone();
        self.context.set(weak);
        result
    }

    fn notify_accessed(&self) {
        if let Some(ctx) = self.get_context() {
            ctx.accessed();
        }
    }
}

/// Trait for context nodes (tock tree entries).
///
/// Note: This trait must be dyn-compatible, so no generic methods allowed.
/// Type-specific operations are done via downcasting through as_any().
pub trait ContextNode {
    fn start_tock(&self) -> Tock;
    fn end_tock(&self) -> Tock;
    fn evictable(&self) -> bool;
    fn evict(&self);
    fn evict_at(&self, idx: usize);
    fn accessed(&self);
    fn get_replayer(&self) -> Option<Rc<Replayer>>;
    fn get_value_at(&self, idx: usize) -> Option<Rc<dyn EZombieNode>>;
    /// For downcasting to concrete context types.
    fn as_any(&self) -> &dyn Any;
    /// Get forward_uf for backedge registration. Returns None for RootContext.
    fn get_forward_uf(&self) -> Option<UF> {
        None
    }
}

impl std::fmt::Debug for dyn ContextNode {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "ContextNode(start={}, end={})", self.start_tock(), self.end_tock())
    }
}

/// Root context - top level, not evictable.
pub struct RootContext {
    start: Tock,
    end: RefCell<Tock>,
    values: RefCell<ValueVec>,
    // Using Cell<Space> instead of RefCell<Space> saves 8 bytes (Space is Copy)
    space_taken: Cell<Space>,
    replayer: Option<Rc<Replayer>>,
}

impl std::fmt::Debug for RootContext {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("RootContext")
            .field("start", &self.start)
            .field("end", &*self.end.borrow())
            .field("values_count", &self.values.borrow().len())
            .finish()
    }
}

impl RootContext {
    pub fn new(start: Tock) -> Rc<Self> {
        Rc::new(Self {
            start,
            end: RefCell::new(start),
            values: RefCell::new(ValueVec::new()),
            space_taken: Cell::new(Space::ZERO),
            replayer: None,
        })
    }

    /// Create with pre-recorded values.
    pub fn new_with_values(
        start: Tock,
        values: ValueVec,
        space: Space,
        replayer: Option<Rc<Replayer>>,
    ) -> Rc<Self> {
        let end = Tock(start.0 + values.len() as i64 + 1);
        let ctx = Rc::new(Self {
            start,
            end: RefCell::new(end),
            values: RefCell::new(values),
            space_taken: Cell::new(space),
            replayer,
        });

        // Set context reference for all values
        let weak_ctx: ContextWeakRef = Rc::downgrade(&ctx) as ContextWeakRef;
        for value in ctx.values.borrow().iter().flatten() {
            value.set_context(weak_ctx.clone());
        }

        ctx
    }

    /// Push a new value to this context.
    pub fn push(&self, node: Rc<dyn EZombieNode>) {
        self.space_taken.set(self.space_taken.get() + Space::from_bytes(node.size()));
        self.values.borrow_mut().push(Some(node));
        *self.end.borrow_mut() = self.start + self.values.borrow().len() + 1;
    }
}

impl ContextNode for RootContext {
    fn start_tock(&self) -> Tock {
        self.start
    }

    fn end_tock(&self) -> Tock {
        *self.end.borrow()
    }

    fn evictable(&self) -> bool {
        false
    }

    fn evict(&self) {
        // Root context cannot be evicted
    }

    fn evict_at(&self, _idx: usize) {
        // Root context values cannot be evicted
    }

    fn accessed(&self) {
        // No-op for root context
    }

    fn get_replayer(&self) -> Option<Rc<Replayer>> {
        self.replayer.clone()
    }

    fn get_value_at(&self, idx: usize) -> Option<Rc<dyn EZombieNode>> {
        self.values.borrow().get(idx).cloned().flatten()
    }

    fn as_any(&self) -> &dyn Any {
        self
    }
}

/// Full context - evictable, created by bind operations.
pub struct FullContext {
    start: Tock,
    end: Tock,
    values: RefCell<ValueVec>,
    // Using Cell<Space> instead of RefCell<Space> saves 8 bytes (Space is Copy)
    space_taken: Cell<Space>,
    time_taken: Time,
    /// Last accessed time in nanoseconds since UNIX epoch.
    /// Using Cell<u64> instead of RefCell<Duration> saves 16 bytes.
    last_accessed_ns: Cell<u64>,
    replayer: Rc<Replayer>,
    dependencies: TockVec,

    // For eviction cost calculation (Union-Find based)
    // Using Cell<Option<NonZeroUsize>> instead of RefCell<Option<usize>> saves 16 bytes:
    // - RefCell<Option<usize>> = 24 bytes (borrow flag + padding + option + value)
    // - Cell<Option<NonZeroUsize>> = 8 bytes (niche optimization: None = 0)
    // The index is stored as (actual_index + 1) to allow 0 to represent None.
    pool_index: Cell<Option<NonZeroUsize>>,
    forward_uf: UF,
    backward_uf: UF,
    // NOTE: Backedges are stored in Trailokya's backedges HashMap, not here.
    // This achieves zero memory overhead when eviction is disabled.
}

impl std::fmt::Debug for FullContext {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("FullContext")
            .field("start", &self.start)
            .field("end", &self.end)
            .field("values_count", &self.values.borrow().len())
            .finish()
    }
}

impl FullContext {
    /// Create a new FullContext.
    ///
    /// Note: Backedge registration should be done by the caller after creation
    /// but before inserting into akasha. This matches the C++ design where
    /// backedge registration happens in the constructor (before the context
    /// is visible in akasha).
    pub fn new(
        start: Tock,
        end: Tock,
        values: ValueVec,
        space: Space,
        time_taken: Time,
        replayer: Rc<Replayer>,
        dependencies: TockVec,
    ) -> Rc<Self> {
        let ctx = Rc::new(Self {
            start,
            end,
            values: RefCell::new(values),
            space_taken: Cell::new(space),
            time_taken,
            last_accessed_ns: Cell::new(0),
            replayer,
            dependencies,
            pool_index: Cell::new(None),
            forward_uf: UF::new(Time::ZERO),
            backward_uf: UF::new(Time::ZERO),
        });

        // Set context reference for all values
        let weak_ctx: ContextWeakRef = Rc::downgrade(&ctx) as ContextWeakRef;
        for value in ctx.values.borrow().iter().flatten() {
            value.set_context(weak_ctx.clone());
        }

        ctx
    }

    pub fn set_pool_index(&self, idx: usize) {
        // Store as idx + 1 to use NonZeroUsize (allows niche optimization)
        self.pool_index.set(NonZeroUsize::new(idx + 1));
    }

    pub fn clear_pool_index(&self) {
        self.pool_index.set(None);
    }

    pub fn pool_index(&self) -> Option<usize> {
        // Convert back: stored value - 1 = actual index
        self.pool_index.get().map(|n| n.get() - 1)
    }

    /// Compute time cost for eviction decision.
    pub fn time_cost(&self) -> Time {
        self.time_taken + self.forward_uf.value() + self.backward_uf.value()
    }

    /// Compute eviction cost using metric function.
    pub fn cost(&self, metric: fn(Time, Time, Space) -> Cost) -> Cost {
        metric(self.time_taken, self.time_cost(), self.space_taken.get())
    }

    pub fn space(&self) -> Space {
        self.space_taken.get()
    }

    pub fn dependencies(&self) -> &[Tock] {
        &self.dependencies
    }

    /// Get the time taken for this computation (used for cost calculation).
    pub fn time_taken(&self) -> Time {
        self.time_taken
    }

    /// Get the backward UF for this context.
    pub fn backward_uf(&self) -> &UF {
        &self.backward_uf
    }

    /// Get the forward UF for this context.
    pub fn forward_uf(&self) -> &UF {
        &self.forward_uf
    }
}

impl ContextNode for FullContext {
    fn start_tock(&self) -> Tock {
        self.start
    }

    fn end_tock(&self) -> Tock {
        self.end
    }

    fn evictable(&self) -> bool {
        true
    }

    fn evict(&self) {
        // Clear all values to free memory
        for v in self.values.borrow_mut().iter_mut() {
            *v = None;
        }
        self.space_taken.set(Space::ZERO);

        // Merge time cost into this context's own UF structures
        // This is the basic cost propagation within the context itself.
        //
        // NOTE: Full cost propagation to other contexts (dependencies, parent,
        // dependent contexts) is handled in Trailokya::evict_one() since it
        // requires access to the akasha (tock tree) to find related contexts.
        let cost = UF::new(self.time_taken);
        self.forward_uf.merge(&cost);
        self.backward_uf.merge(&cost);
    }

    fn evict_at(&self, idx: usize) {
        let mut values = self.values.borrow_mut();
        if let Some(node) = values.get_mut(idx).and_then(|slot| slot.take()) {
            self.space_taken.set(self.space_taken.get() - Space::from_bytes(node.size()));
        }
    }

    fn accessed(&self) {
        // Update last accessed time for recency tracking.
        // The GD heap uses this during eviction to prioritize recently accessed items.
        use std::time::{SystemTime, UNIX_EPOCH};
        let now_ns = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .map(|d| d.as_nanos() as u64)
            .unwrap_or(0);
        self.last_accessed_ns.set(now_ns);
    }

    fn get_replayer(&self) -> Option<Rc<Replayer>> {
        Some(self.replayer.clone())
    }

    fn get_value_at(&self, idx: usize) -> Option<Rc<dyn EZombieNode>> {
        self.values.borrow().get(idx).cloned().flatten()
    }

    fn as_any(&self) -> &dyn Any {
        self
    }

    fn get_forward_uf(&self) -> Option<UF> {
        Some(self.forward_uf)
    }
}

/// Downcast an `Rc<dyn EZombieNode>` to `Rc<ZombieNode<T>>`.
///
/// This is a critical hot-path function used during `get()` and `bind()` operations.
///
/// # Type Safety
///
/// Uses Rust's `TypeId` mechanism (via `Any::is::<T>()`) to verify the concrete type
/// before downcasting. This is the standard, safe way to perform runtime type checking.
///
/// # Memory Safety (unsafe block justification)
///
/// The unsafe block is necessary because Rust doesn't provide a safe way to convert
/// `Rc<dyn Trait>` to `Rc<ConcreteType>` even after verifying the type. Here's why
/// the conversion is safe:
///
/// 1. **Type Verification**: We've confirmed via `is::<ZombieNode<T>>()` that the
///    concrete type is exactly `ZombieNode<T>`.
///
/// 2. **Rc Layout**: `Rc<dyn Trait>` stores a fat pointer (data + vtable) to an `RcBox`.
///    `Rc<ConcreteType>` stores a thin pointer to the same `RcBox`. Both point to the
///    same allocation; only the pointer representation differs.
///
/// 3. **Clone Before Convert**: We clone the `Rc` first, incrementing the strong count.
///    `Rc::into_raw` gives us the data pointer without decrementing count. `Rc::from_raw`
///    reconstructs the `Rc` from that data pointer. The net effect: count stays correct.
///
/// 4. **Fat-to-Thin Conversion**: `raw as *const ZombieNode<T>` extracts just the data
///    pointer from the fat pointer. This is well-defined and portable.
///
/// This pattern is widely used in the Rust ecosystem for trait object downcasting.
/// See also: `downcast-rs`, `intertrait`, and similar crates.
///
/// # Example
///
/// ```ignore
/// let node: Rc<dyn EZombieNode> = ZombieNode::new(Tock(1), 42i32);
/// let typed: Option<Rc<ZombieNode<i32>>> = downcast_zombie(&node);
/// assert!(typed.is_some());
/// ```
#[inline]
pub fn downcast_zombie<T: 'static + ZombieOps>(node: &Rc<dyn EZombieNode>) -> Option<Rc<ZombieNode<T>>> {
    // Type check via Any - this is the standard Rust way to do runtime type checking
    let any_ref = node.as_any();
    if !any_ref.is::<ZombieNode<T>>() {
        return None;
    }

    // SAFETY: See function documentation above for detailed justification.
    //
    // Summary:
    // - Type verified via TypeId comparison
    // - Clone ensures reference count is incremented before into_raw
    // - Fat-to-thin pointer cast preserves the data pointer
    // - from_raw reconstructs Rc with correct RcBox address
    let cloned = node.clone();
    let raw: *const dyn EZombieNode = Rc::into_raw(cloned);
    let data_ptr = raw as *const ZombieNode<T>;
    Some(unsafe { Rc::from_raw(data_ptr) })
}

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

    #[test]
    fn test_root_context() {
        let ctx = RootContext::new(Tock(0));
        assert!(!ctx.evictable());
        assert_eq!(ctx.start_tock(), Tock(0));
    }

    #[test]
    fn test_zombie_node_context_link() {
        let ctx = RootContext::new(Tock(0));
        let node = ZombieNode::new(Tock(1), 42);

        // Set context
        let weak_ctx: ContextWeakRef = Rc::downgrade(&ctx) as ContextWeakRef;
        node.set_context(weak_ctx);

        // Verify link
        assert!(node.get_context().is_some());
    }

    #[test]
    fn test_downcast_zombie() {
        let node: Rc<dyn EZombieNode> = ZombieNode::new(Tock(1), 42i32);

        // Correct type
        let typed = downcast_zombie::<i32>(&node);
        assert!(typed.is_some());
        assert_eq!(typed.unwrap().value, 42);

        // Wrong type
        let wrong = downcast_zombie::<String>(&node);
        assert!(wrong.is_none());
    }
}