zombie_rs/

record.rs

//! Record nodes for tracking computation state during bind operations.
//!
//! The record system implements a stack-based state machine following C++ semantics:
//! - `Record::Root`: Top-level record, creates RootContext when suspended
//! - `Record::Head`: Active record for a bind computation
//! - `Record::Value`: Completed record holding a result tock
//!
//! # Time Measurement
//!
//! Currently uses PLANK_TIME (constant) for time cost estimation, matching the
//! C++ reference implementation's `use_measured_time = false` setting. This provides:
//! - Deterministic behavior across runs
//! - No timing overhead
//! - Simpler implementation
//!
//! For measured time support, see `ZombieConfig::use_measured_time` (currently
//! the config option exists but measured time is not yet implemented in the
//! record completion path).

use std::any::Any;
use std::rc::Rc;

use smallvec::SmallVec;

use crate::context::{EZombieNode, FullContext, RootContext, ContextNode, TockVec, ValueVec};
use crate::meter::{Space, Time, PLANK_TIME_NS};
use crate::tock::Tock;

/// Type-erased replay function.
///
/// Takes a vector of pointers to input values and executes the computation.
/// The function should call `finish_record` when complete.
pub type ReplayFn = Box<dyn Fn(&[*const dyn Any])>;

/// Inputs for replay: each entry is (tock, optional cached node).
/// Uses SmallVec for inline storage when input count is small (typical case).
pub type ReplayInputs = SmallVec<[(Tock, Option<Rc<dyn EZombieNode>>); 4]>;

/// Dependencies stored as SmallVec - most binds have only 1-4 dependencies.
/// This avoids HashMap/HashSet overhead which is significant for small collections.
/// Duplicates are allowed since we only need to iterate, not lookup.
pub type DependencyVec = SmallVec<[Tock; 4]>;

/// Stores information needed to replay a computation.
#[derive(Clone)]
pub struct Replayer {
    pub f: Rc<ReplayFn>,
    pub inputs: TockVec,
}

impl Replayer {
    pub fn new(f: ReplayFn, inputs: TockVec) -> Rc<Self> {
        Rc::new(Self {
            f: Rc::new(f),
            inputs,
        })
    }
}

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

/// Common record state shared by Root and Head records.
pub struct RecordState {
    pub start_tock: Tock,
    pub values: ValueVec,
    pub space_taken: Space,
    /// Dependencies stored as SmallVec for efficiency.
    /// Most binds have only 1-4 dependencies, so inline storage avoids allocation.
    /// We use a simple append-only list since we only need iteration, not lookup.
    pub dependencies: DependencyVec,
}

impl std::fmt::Debug for RecordState {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("RecordState")
            .field("start_tock", &self.start_tock)
            .field("values_count", &self.values.len())
            .field("space_taken", &self.space_taken)
            .field("dependencies_count", &self.dependencies.len())
            .finish()
    }
}

impl RecordState {
    pub fn new(start_tock: Tock) -> Self {
        Self {
            start_tock,
            values: ValueVec::new(),
            space_taken: Space::ZERO,
            dependencies: DependencyVec::new(),
        }
    }

    pub fn push_value(&mut self, node: Rc<dyn EZombieNode>) {
        self.space_taken += Space::from_bytes(node.size());
        self.values.push(Some(node));
    }

    /// Register a dependency on a tock.
    /// Only registers if tock is before this record's start (external dependency).
    /// Duplicates are allowed - we only iterate, never lookup.
    #[inline]
    pub fn register_dependency(&mut self, tock: Tock) {
        if tock < self.start_tock {
            self.dependencies.push(tock);
        }
    }
}

/// Record enum - stack entries for tracking computation state.
///
/// This enum replaces the previous `RecordNode` trait, providing:
/// - No dynamic dispatch overhead
/// - Exhaustive pattern matching
/// - Better memory layout (single allocation)
#[derive(Debug)]
pub enum Record {
    /// Root record - top level, not evictable.
    Root(RootRecord),
    /// Head record - active bind computation.
    Head(HeadRecord),
    /// Value record - completed computation.
    Value(ValueRecord),
}

impl Record {
    /// Create a new root record.
    pub fn new_root(start_tock: Tock) -> Self {
        Record::Root(RootRecord::new(start_tock))
    }

    /// Create a new head record.
    pub fn new_head(replayer: Rc<Replayer>, start_tock: Tock, start_time: Time) -> Self {
        Record::Head(HeadRecord::new(replayer, start_tock, start_time))
    }

    /// Create a new value record.
    pub fn new_value(result_tock: Tock, context_start_tock: Tock) -> Self {
        Record::Value(ValueRecord::new(result_tock, context_start_tock))
    }

    /// Get start tock.
    pub fn start_tock(&self) -> Tock {
        match self {
            Record::Root(r) => r.state.start_tock,
            Record::Head(r) => r.state.start_tock,
            Record::Value(r) => r.result_tock,
        }
    }

    /// Get values in this record.
    pub fn values(&self) -> &[Option<Rc<dyn EZombieNode>>] {
        match self {
            Record::Root(r) => &r.state.values,
            Record::Head(r) => &r.state.values,
            Record::Value(_) => &[],
        }
    }

    /// Push a value to this record.
    pub fn push_value(&mut self, node: Rc<dyn EZombieNode>) {
        match self {
            Record::Root(r) => r.state.push_value(node),
            Record::Head(r) => r.state.push_value(node),
            Record::Value(_) => panic!("ValueRecord cannot push values"),
        }
    }

    /// Register a dependency on the current record.
    pub fn register_dependency(&mut self, tock: Tock) {
        match self {
            Record::Root(r) => r.state.register_dependency(tock),
            Record::Head(r) => r.state.register_dependency(tock),
            Record::Value(_) => {} // No-op
        }
    }

    /// Get dependencies as a slice.
    pub fn dependencies(&self) -> &[Tock] {
        match self {
            Record::Root(r) => &r.state.dependencies,
            Record::Head(r) => &r.state.dependencies,
            Record::Value(_) => &[],
        }
    }

    /// Get space taken.
    pub fn space_taken(&self) -> Space {
        match self {
            Record::Root(r) => r.state.space_taken,
            Record::Head(r) => r.state.space_taken,
            Record::Value(_) => Space::ZERO,
        }
    }

    /// Check if this is a head record.
    pub fn is_head(&self) -> bool {
        matches!(self, Record::Head(_))
    }

    /// Check if this is a value record.
    pub fn is_value(&self) -> bool {
        matches!(self, Record::Value(_))
    }

    /// Check if this is a tailcall record (HeadRecord).
    pub fn is_tailcall(&self) -> bool {
        matches!(self, Record::Head(_))
    }

    /// Get result tock (for ValueRecord).
    pub fn get_result_tock(&self) -> Option<Tock> {
        match self {
            Record::Value(r) => Some(r.result_tock),
            _ => None,
        }
    }

    /// Get the start_tock of the context containing this record's result (for ValueRecord).
    pub fn get_context_start_tock(&self) -> Option<Tock> {
        match self {
            Record::Value(r) => Some(r.context_start_tock),
            _ => None,
        }
    }

    /// Get replayer (for HeadRecord).
    pub fn get_replayer(&self) -> Option<Rc<Replayer>> {
        match self {
            Record::Head(r) => Some(r.replayer.clone()),
            _ => None,
        }
    }

    /// Handle suspension (creates context).
    ///
    /// For Root: creates RootContext
    /// For Head: creates FullContext (tailcall semantics)
    ///
    /// # Returns
    /// For HeadRecord: returns (dependencies, forward_uf) for backedge registration.
    /// For RootRecord: returns None (RootContext doesn't need backedges).
    pub fn suspend(
        &mut self,
        replayer: Rc<Replayer>,
        current_tock: Tock,
        akasha_insert: &mut dyn FnMut(Tock, Rc<dyn ContextNode>),
        heap_push: Option<&mut dyn FnMut(Rc<FullContext>)>,
    ) -> Option<(TockVec, crate::union_find::UF)> {
        match self {
            Record::Root(r) => {
                // Create RootContext with accumulated values
                let ctx = RootContext::new_with_values(
                    r.state.start_tock,
                    std::mem::take(&mut r.state.values),
                    r.state.space_taken,
                    Some(replayer),
                );
                akasha_insert(r.state.start_tock, ctx);
                None
            }
            Record::Head(r) => {
                // HeadRecord suspend uses tailcall semantics:
                // Complete the current computation and save state to FullContext,
                // then the caller will push a new HeadRecord for the inner bind.
                let deps: TockVec = std::mem::take(&mut r.state.dependencies);
                let ctx = FullContext::new(
                    r.state.start_tock,
                    current_tock,
                    std::mem::take(&mut r.state.values),
                    r.state.space_taken,
                    Time::from_duration(std::time::Duration::from_nanos(PLANK_TIME_NS)),
                    r.replayer.clone(),
                    deps.clone(),
                );
                // Clone forward_uf before moving ctx
                let forward_uf = *ctx.forward_uf();
                akasha_insert(r.state.start_tock, ctx.clone());
                if let Some(push_fn) = heap_push {
                    push_fn(ctx);
                }
                Some((deps, forward_uf))
            }
            Record::Value(_) => panic!("ValueRecord cannot be suspended"),
        }
    }

    /// Called when child record returns.
    pub fn resume(&mut self, new_tock: Tock) {
        match self {
            Record::Root(r) => {
                r.state = RecordState::new(new_tock);
            }
            Record::Head(r) => {
                r.state = RecordState::new(new_tock);
            }
            Record::Value(_) => panic!("ValueRecord cannot be resumed"),
        }
    }

    /// Handle completion (creates FullContext).
    ///
    /// # Returns
    /// Returns (dependencies, forward_uf) for backedge registration, or None if not a HeadRecord.
    pub fn complete(
        &mut self,
        current_tock: Tock,
        replayer: Option<Rc<Replayer>>,
        akasha_insert: &mut dyn FnMut(Tock, Rc<dyn ContextNode>),
        heap_push: &mut dyn FnMut(Rc<FullContext>),
    ) -> Option<(TockVec, crate::union_find::UF)> {
        match self {
            Record::Root(_) => panic!("RootRecord cannot be completed"),
            Record::Head(r) => {
                let deps: TockVec = std::mem::take(&mut r.state.dependencies);
                let ctx = FullContext::new(
                    r.state.start_tock,
                    current_tock,
                    std::mem::take(&mut r.state.values),
                    r.state.space_taken,
                    Time::from_duration(std::time::Duration::from_nanos(PLANK_TIME_NS)),
                    replayer.unwrap_or_else(|| r.replayer.clone()),
                    deps.clone(),
                );
                // Clone forward_uf before moving ctx
                let forward_uf = *ctx.forward_uf();
                akasha_insert(r.state.start_tock, ctx.clone());
                heap_push(ctx);
                Some((deps, forward_uf))
            }
            Record::Value(_) => {
                // ValueRecord is already complete - no-op
                None
            }
        }
    }

    /// Prepare replay by collecting inputs.
    /// Returns the replayer and inputs (some may be missing).
    pub fn prepare_replay(
        &mut self,
        get_input: &mut dyn FnMut(Tock) -> Option<Rc<dyn EZombieNode>>,
    ) -> (Rc<Replayer>, ReplayInputs) {
        match self {
            Record::Root(_) => panic!("RootRecord cannot replay"),
            Record::Head(r) => {
                if r.played {
                    panic!("HeadRecord already played");
                }
                let mut inputs = ReplayInputs::new();
                for &input_tock in r.replayer.inputs.iter() {
                    let node = get_input(input_tock);
                    inputs.push((input_tock, node));
                }
                r.played = true;
                (r.replayer.clone(), inputs)
            }
            Record::Value(_) => panic!("ValueRecord cannot replay"),
        }
    }
}

// Keep the struct definitions for internal use

/// Root record - top level, not evictable.
#[derive(Debug)]
pub struct RootRecord {
    state: RecordState,
}

impl RootRecord {
    pub fn new(start_tock: Tock) -> Self {
        Self {
            state: RecordState::new(start_tock),
        }
    }
}

/// Head record - active bind computation.
#[derive(Debug)]
pub struct HeadRecord {
    state: RecordState,
    replayer: Rc<Replayer>,
    #[allow(dead_code)]
    start_time: Time,
    played: bool,
}

impl HeadRecord {
    /// Create a new HeadRecord.
    /// Dependencies are registered from replayer inputs.
    pub fn new(replayer: Rc<Replayer>, start_tock: Tock, start_time: Time) -> Self {
        let mut state = RecordState::new(start_tock);

        // Register dependencies from replayer inputs
        for &input_tock in &replayer.inputs {
            state.register_dependency(input_tock);
        }

        Self {
            state,
            replayer,
            start_time,
            played: false,
        }
    }
}

/// Value record - completed computation.
#[derive(Debug)]
pub struct ValueRecord {
    result_tock: Tock,
    context_start_tock: Tock,
}

impl ValueRecord {
    pub fn new(result_tock: Tock, context_start_tock: Tock) -> Self {
        Self { result_tock, context_start_tock }
    }

    pub fn context_start_tock(&self) -> Tock {
        self.context_start_tock
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::context::ZombieNode;
    use crate::meter::Time;

    #[test]
    fn test_root_record() {
        let record = Record::new_root(Tock(0));
        assert_eq!(record.start_tock(), Tock(0));
        assert!(record.values().is_empty());
        assert!(!record.is_head());
        assert!(!record.is_value());
    }

    #[test]
    fn test_value_record() {
        let record = Record::new_value(Tock(42), Tock(40));
        assert!(record.is_value());
        assert_eq!(record.get_result_tock(), Some(Tock(42)));
        assert_eq!(record.get_context_start_tock(), Some(Tock(40)));
    }

    // =========================================================================
    // Edge Case Tests
    // =========================================================================

    /// Test that only external dependencies are registered
    #[test]
    fn test_record_state_dependency_filtering() {
        let mut state = RecordState::new(Tock(10));

        // Dependencies before start_tock should be registered
        state.register_dependency(Tock(5));
        state.register_dependency(Tock(0));

        // Dependencies at or after start_tock should NOT be registered
        state.register_dependency(Tock(10)); // Same as start
        state.register_dependency(Tock(15)); // After start

        // Only 2 dependencies should be registered
        assert_eq!(state.dependencies.len(), 2);
        assert!(state.dependencies.contains(&Tock(5)));
        assert!(state.dependencies.contains(&Tock(0)));
    }

    /// Test space tracking accumulates correctly
    #[test]
    fn test_record_space_tracking() {
        let mut record = Record::new_root(Tock(0));

        assert_eq!(record.space_taken().as_bytes(), 0);

        // Push nodes with known sizes
        let node1 = ZombieNode::new(Tock(1), vec![0u8; 100]);
        let node2 = ZombieNode::new(Tock(2), vec![0u8; 200]);

        record.push_value(node1 as Rc<dyn crate::context::EZombieNode>);
        assert!(record.space_taken().as_bytes() >= 100);

        let space_after_first = record.space_taken().as_bytes();
        record.push_value(node2 as Rc<dyn crate::context::EZombieNode>);
        assert!(record.space_taken().as_bytes() >= space_after_first + 200);
    }

    /// Test HeadRecord initialization and replayer access
    #[test]
    fn test_head_record_replayer() {
        let inputs = {
            let mut v = TockVec::new();
            v.push(Tock(1));
            v.push(Tock(5));
            v
        };

        let replayer = Replayer::new(
            Box::new(|_| {}),
            inputs.clone(),
        );

        let record = Record::new_head(
            replayer.clone(),
            Tock(10),
            Time::ZERO,
        );

        assert!(record.is_head());
        assert!(!record.is_value());
        assert_eq!(record.start_tock(), Tock(10));

        // Replayer should be accessible
        let got_replayer = record.get_replayer();
        assert!(got_replayer.is_some());
        assert_eq!(got_replayer.unwrap().inputs.len(), 2);

        // Dependencies from replayer inputs should be registered
        let deps = record.dependencies();
        assert_eq!(deps.len(), 2);
        assert!(deps.contains(&Tock(1)));
        assert!(deps.contains(&Tock(5)));
    }

    /// Test ValueRecord has no values or dependencies
    #[test]
    fn test_value_record_is_terminal() {
        let record = Record::new_value(Tock(42), Tock(40));

        // ValueRecord should have no values
        assert!(record.values().is_empty());

        // ValueRecord should have no dependencies
        assert!(record.dependencies().is_empty());

        // ValueRecord space is zero
        assert_eq!(record.space_taken().as_bytes(), 0);

        // Replayer should be None
        assert!(record.get_replayer().is_none());
    }
}