zombie_rs/

runtime.rs

//! Runtime - Global state manager for zombie computations.
//!
//! The global runtime manages:
//! - Current tock (logical time)
//! - Timeline (context tree / tock tree)
//! - Eviction heap (GD algorithm)
//! - Record stack
//! - Replay stack

// Allow Box<HashMap> because Option<Box<HashMap>> is 8 bytes when None,
// while Option<HashMap> is 48 bytes (HashMap has no niche optimization).
#![allow(clippy::box_collection)]

use std::cell::RefCell;
use std::collections::HashMap;
use std::rc::{Rc, Weak};

use crate::config::ZombieConfig;
use crate::context::{ContextNode, EZombieNode, FullContext, ZombieNode};
use crate::heap::{GDHeap, HeapIndexNotify};
use crate::meter::ZombieMeter;
use crate::record::{Record, Replayer};
use crate::splay_list::SplayList;
use crate::tock::Tock;
use crate::union_find::UFSet;

/// Result type for `find_context_and_value`: (context_start_tock, context_ref, typed_node).
/// This combines the context lookup result with the downcast zombie node for efficient caching.
pub type ContextLookupResult<'a, T> = (Tock, &'a Rc<dyn ContextNode>, Rc<ZombieNode<T>>);

/// Replay state for tracking recomputation target.
pub struct ReplayState {
    pub target: Tock,
    pub result: Rc<RefCell<Option<Rc<dyn EZombieNode>>>>,
}

impl ReplayState {
    pub fn new(target: Tock, result: Rc<RefCell<Option<Rc<dyn EZombieNode>>>>) -> Self {
        Self { target, result }
    }

    pub fn has_result(&self) -> bool {
        self.result.borrow().is_some()
    }

    pub fn set_result(&self, node: Rc<dyn EZombieNode>) {
        *self.result.borrow_mut() = Some(node);
    }

    pub fn take_result(&self) -> Option<Rc<dyn EZombieNode>> {
        self.result.borrow_mut().take()
    }
}

/// Handle for evictable context in the heap.
///
/// Implements `HeapIndexNotify` to track its position in the heap,
/// enabling O(log n) touch operations via `pool_index`.
struct EvictHandle {
    context: Rc<FullContext>,
}

impl HeapIndexNotify for EvictHandle {
    fn notify_index_changed(&self, new_index: usize) {
        self.context.set_pool_index(new_index);
    }

    fn notify_removed(&self) {
        self.context.clear_pool_index();
    }
}

impl PartialEq for EvictHandle {
    fn eq(&self, other: &Self) -> bool {
        Rc::ptr_eq(&self.context, &other.context)
    }
}

impl Eq for EvictHandle {}

impl PartialOrd for EvictHandle {
    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
        Some(self.cmp(other))
    }
}

impl Ord for EvictHandle {
    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
        self.context.start_tock().cmp(&other.context.start_tock())
    }
}

/// Global state manager for the Zombie runtime.
pub struct Runtime {
    config: ZombieConfig,

    /// Current logical time.
    current_tock: Tock,

    /// Timeline tree mapping tocks to contexts.
    timeline: SplayList<Tock, Rc<dyn ContextNode>>,

    /// Eviction heap using GD algorithm.
    heap: GDHeap<EvictHandle>,

    /// Record stack.
    ///
    /// # Invariant
    ///
    /// This stack is NEVER empty after initialization. It always contains at least
    /// the `RootRecord` at the bottom. This invariant is established in `new()` and
    /// maintained by all operations. Methods like `current_record()` rely on this
    /// invariant and use `unwrap()` knowing it cannot fail in correctly initialized state.
    records: Vec<Record>,

    /// Replay stack.
    replays: Vec<ReplayState>,

    /// Time meter.
    meter: ZombieMeter,

    /// Total space used by evictable contexts (O(1) tracking instead of O(N) iteration).
    total_space_used: usize,

    /// Backedges map: stores backedges for each context (keyed by start_tock).
    /// Only allocated when memory_limit is set (eviction enabled).
    /// Using Option<Box<...>> ensures None variant is only 8 bytes (pointer size).
    backedges: Option<Box<HashMap<Tock, UFSet>>>,
}

thread_local! {
    static RUNTIME: RefCell<Option<Runtime>> = const { RefCell::new(None) };
}

impl Runtime {
    /// Create a new Runtime with default config.
    pub fn new() -> Self {
        Self::with_config(ZombieConfig::default())
    }

    /// Create with custom config.
    ///
    /// Note: This also initializes UFArena if not already done.
    /// This ensures Runtime::new() can be used directly in tests.
    pub fn with_config(config: ZombieConfig) -> Self {
        // Ensure UFArena is initialized (idempotent)
        if !crate::union_find::UFArena::is_initialized() {
            crate::union_find::UFArena::init();
        }

        // Initialize with root record at tock 0
        let root_record = Record::new_root(Tock(0));

        // Default replay state with MAX target (no active replay)
        let default_replay = ReplayState::new(
            Tock::MAX,
            Rc::new(RefCell::new(None)),
        );

        // Only create backedges map if eviction is enabled
        let backedges = if config.memory_limit.is_some() {
            Some(Box::new(HashMap::new()))
        } else {
            None
        };

        // Create GDHeap with config's approx_factor and staleness
        let heap = GDHeap::with_config(config.approx_factor, config.staleness);

        Self {
            config,
            current_tock: Tock(1), // Start at 1 like C++
            timeline: SplayList::new(),
            heap,
            records: vec![root_record],
            replays: vec![default_replay],
            meter: ZombieMeter::new(),
            total_space_used: 0,
            backedges,
        }
    }

    /// Initialize thread-local Runtime.
    ///
    /// # Panics
    ///
    /// This function does not panic. If already initialized, the old state
    /// is replaced (useful for test isolation).
    pub fn init() {
        // Initialize UFArena for GhostCell-based Union-Find
        if !crate::union_find::UFArena::is_initialized() {
            crate::union_find::UFArena::init();
        }
        RUNTIME.with(|t| {
            *t.borrow_mut() = Some(Runtime::new());
        });
    }

    /// Initialize with custom config.
    ///
    /// # Panics
    ///
    /// This function does not panic. If already initialized, the old state
    /// is replaced (useful for test isolation).
    pub fn init_with_config(config: ZombieConfig) {
        // Initialize UFArena for GhostCell-based Union-Find
        if !crate::union_find::UFArena::is_initialized() {
            crate::union_find::UFArena::init();
        }
        RUNTIME.with(|t| {
            *t.borrow_mut() = Some(Runtime::with_config(config));
        });
    }

    /// Check if Runtime is initialized on the current thread.
    #[inline]
    pub fn is_initialized() -> bool {
        RUNTIME.with(|t| t.borrow().is_some())
    }

    /// Access thread-local Runtime (immutable).
    ///
    /// # Panics
    ///
    /// Panics if `Runtime::init()` has not been called on this thread.
    #[inline]
    pub fn with<R, F: FnOnce(&Runtime) -> R>(f: F) -> R {
        RUNTIME.with(|t| {
            let guard = t.borrow();
            let trailokya = guard.as_ref().expect(
                "Runtime not initialized. Call Runtime::init() before using Zombie."
            );
            f(trailokya)
        })
    }

    /// Access thread-local Runtime (mutable).
    ///
    /// # Panics
    ///
    /// Panics if `Runtime::init()` has not been called on this thread.
    #[inline]
    pub fn with_mut<R, F: FnOnce(&mut Runtime) -> R>(f: F) -> R {
        RUNTIME.with(|t| {
            let mut guard = t.borrow_mut();
            let trailokya = guard.as_mut().expect(
                "Runtime not initialized. Call Runtime::init() before using Zombie."
            );
            f(trailokya)
        })
    }

    /// Try to access thread-local Runtime without panicking.
    ///
    /// Returns `None` if not initialized.
    #[inline]
    pub fn try_with<R, F: FnOnce(&Runtime) -> R>(f: F) -> Option<R> {
        RUNTIME.with(|t| {
            let guard = t.borrow();
            guard.as_ref().map(f)
        })
    }

    /// Try to access thread-local Runtime mutably without panicking.
    ///
    /// Returns `None` if not initialized.
    #[inline]
    pub fn try_with_mut<R, F: FnOnce(&mut Runtime) -> R>(f: F) -> Option<R> {
        RUNTIME.with(|t| {
            let mut guard = t.borrow_mut();
            guard.as_mut().map(f)
        })
    }

    // ========== Config Queries ==========

    /// Check if memory_limit is set (eviction enabled).
    ///
    /// This is a simple field access, used to skip expensive operations
    /// when eviction is disabled.
    #[inline]
    pub fn has_memory_limit(&self) -> bool {
        self.config.memory_limit.is_some()
    }

    // ========== Tock Management ==========

    /// Allocate a tock and return the old value (post-increment).
    #[inline]
    pub fn tick(&mut self) -> Tock {
        let old = self.current_tock;
        self.current_tock = Tock(self.current_tock.0 + 1);
        old
    }

    /// Get current tock without advancing.
    #[inline]
    pub fn current_tock(&self) -> Tock {
        self.current_tock
    }

    /// Set current tock (for replay).
    pub fn set_tock(&mut self, tock: Tock) {
        self.current_tock = tock;
    }

    // ========== Timeline (Context Tree) ==========

    /// Find context containing a tock (with its start key).
    pub fn find_context_with_key(&mut self, tock: Tock) -> Option<(Tock, &Rc<dyn ContextNode>)> {
        self.timeline.find_le_with_key(&tock).map(|(k, v)| (*k, v))
    }

    /// Find context and typed value for a given tock.
    /// Returns (context_start_tock, context_ref, typed_node) if found.
    /// This is an optimized version that returns the context for caching during bind setup.
    pub fn find_context_and_value<T: 'static + Clone + crate::meter::ZombieOps>(
        &mut self,
        tock: Tock,
    ) -> Option<ContextLookupResult<'_, T>> {
        let (ctx_start, ctx) = self.timeline.find_le_with_key(&tock)?;
        let ctx_start = *ctx_start;

        // Compute index within context
        let idx = crate::tock::try_tock_to_index(tock, ctx_start)?;

        // Verify tock is within this context's range
        if tock >= ctx.end_tock() {
            return None;
        }

        // Get the value at this index
        let node = ctx.get_value_at(idx)?;

        // Downcast from type-erased to typed node
        let typed_rc = crate::context::downcast_zombie::<T>(&node)?;

        Some((ctx_start, ctx, typed_rc))
    }

    /// Finds the context to replay from for a given target tock.
    ///
    /// Returns the context that can replay to recreate the target value.
    /// There are two cases:
    /// 1. Target is in a context that still exists: use that context's start_tock
    /// 2. Target was in a context that was deleted (evicted): use predecessor's end_tock
    ///
    /// Returns (replay_from_tock, context) for replay initialization.
    pub fn find_replay_context(&mut self, target_tock: Tock) -> Option<(Tock, &Rc<dyn ContextNode>)> {
        // Find the context containing or preceding the target tock
        let ctx = self.timeline.find_le(&target_tock)?;

        // Case 1: Context still contains target [start, end) where end > target
        if ctx.end_tock() > target_tock {
            // Context exists and contains target - replay from its start
            return Some((ctx.start_tock(), ctx));
        }

        // Case 2: Context ends before or at target (end <= target)
        // This means target was in a context that was deleted (evicted from timeline)
        // We need to replay from this context's end_tock
        // The replayer must exist for replay to work
        if ctx.get_replayer().is_some() {
            return Some((ctx.end_tock(), ctx));
        }

        // Fallback: no replayer available (e.g., RootContext)
        None
    }

    /// Insert a context.
    pub fn insert_context(&mut self, start: Tock, ctx: Rc<dyn ContextNode>) {
        self.timeline.insert(start, ctx);
    }

    /// Remove a context.
    pub fn remove_context(&mut self, start: Tock) {
        self.timeline.remove(&start);
    }

    /// Look up a zombie node at a specific tock.
    ///
    /// This is a common pattern used in multiple places:
    /// 1. Find context containing the tock
    /// 2. Compute index within context
    /// 3. Verify tock is within range
    /// 4. Get the value
    ///
    /// Returns None if tock is not in any context or context doesn't contain value.
    #[inline]
    fn lookup_node_at(&mut self, tock: Tock) -> Option<Rc<dyn EZombieNode>> {
        let (ctx_start, ctx) = self.timeline.find_le_with_key(&tock)?;
        let idx = crate::tock::try_tock_to_index(tock, *ctx_start)?;
        if tock < ctx.end_tock() {
            ctx.get_value_at(idx)
        } else {
            None
        }
    }

    // ========== Record Stack ==========

    /// Get current record (top of record stack).
    ///
    /// # Safety Invariant
    ///
    /// Uses `unwrap()` because the record stack is guaranteed non-empty after
    /// initialization. See the documentation on the `records` field.
    #[inline]
    pub fn current_record(&self) -> &Record {
        self.records.last().unwrap()
    }

    /// Get current record mutably.
    ///
    /// # Safety Invariant
    ///
    /// Uses `unwrap()` because the record stack is guaranteed non-empty after
    /// initialization. See the documentation on the `records` field.
    #[inline]
    pub fn current_record_mut(&mut self) -> &mut Record {
        self.records.last_mut().unwrap()
    }

    /// Push a HeadRecord for a new computation.
    pub fn push_head_record(&mut self, replayer: Rc<Replayer>, start_tock: Tock) {
        let time = self.meter.time();
        self.records.push(Record::new_head(replayer, start_tock, time));
    }

    /// Check if current record is a HeadRecord (tailcall context)
    pub fn is_current_record_tailcall(&self) -> bool {
        self.records.last().is_some_and(|r| r.is_tailcall())
    }

    /// Suspend current record (creates context).
    /// Returns true if this was a tailcall (HeadRecord suspend).
    pub fn suspend_record(&mut self, replayer: Rc<Replayer>) -> bool {
        let current_tock = self.current_tock;
        let is_tailcall = self.is_current_record_tailcall();
        let record = self.records.last_mut().unwrap();

        // Create heap push closure for HeadRecord suspend (adds FullContext to heap)
        let heap = &mut self.heap;
        let metric = self.config.metric;
        let total_space_used = &mut self.total_space_used;

        let mut heap_push_fn = |ctx: Rc<FullContext>| {
            *total_space_used += ctx.space().as_bytes();
            let cost = ctx.cost(metric);
            heap.push_notify(EvictHandle { context: ctx }, cost);
        };

        // Create timeline insert closure
        let timeline = &mut self.timeline;
        let mut insert_fn = |start: Tock, ctx: Rc<dyn ContextNode>| {
            timeline.insert(start, ctx);
        };

        // Suspend record - creates RootContext or FullContext
        // HeadRecord suspend saves state to FullContext for potential eviction
        let backedge_info = record.suspend(replayer, current_tock, &mut insert_fn, Some(&mut heap_push_fn));

        // Register backedges if memory_limit is set (eviction enabled).
        // Uses timeline lookup with splay tree locality - recently accessed contexts
        // are near the root, making lookups nearly O(1).
        // Backedges are stored in Runtime's HashMap, not in FullContext.
        if let Some(ref mut backedges_map) = self.backedges
            && let Some((deps, forward_uf)) = backedge_info
        {
            for dep_tock in deps.iter() {
                if let Some((dep_start, _dep_ctx)) = self.timeline.find_le_with_key(dep_tock) {
                    // Store backedge using dependency's start_tock as key
                    backedges_map
                        .entry(*dep_start)
                        .or_insert_with(UFSet::new)
                        .insert(forward_uf);
                }
            }
        }

        is_tailcall
    }

    /// Replace current record with a new HeadRecord (for tailcall semantics).
    /// This is used when an inner bind happens inside a HeadRecord.
    pub fn replace_head_record(&mut self, replayer: Rc<Replayer>, start_tock: Tock) {
        let time = self.meter.time();
        // Pop the old HeadRecord (which was already suspended/completed)
        self.records.pop();
        // Push the new HeadRecord
        self.records.push(Record::new_head(replayer, start_tock, time));
    }

    /// Finish record with result (creates ValueRecord).
    /// Returns true if actually finished, false if record was already replaced by tailcall.
    pub fn finish_record(&mut self, result_tock: Tock) -> bool {
        // Check if current record is a HeadRecord.
        // If not, our HeadRecord was replaced by an inner bind's tailcall.
        if !self.records.last().is_some_and(|r| r.is_head()) {
            return false;
        }

        // Complete current record
        let current_tock = self.current_tock;
        let metric = self.config.metric;

        let record = self.records.last_mut().unwrap();

        let forward_uf_result = {
            let timeline = &mut self.timeline;
            let heap = &mut self.heap;
            let total_space_used = &mut self.total_space_used;

            let mut insert_fn = |start: Tock, ctx: Rc<dyn ContextNode>| {
                timeline.insert(start, ctx);
            };

            let mut heap_push_fn = |ctx: Rc<FullContext>| {
                // Track space for O(1) total_space queries
                *total_space_used += ctx.space().as_bytes();
                let cost = ctx.cost(metric);
                heap.push_notify(EvictHandle { context: ctx }, cost);
            };

            record.complete(current_tock, None, &mut insert_fn, &mut heap_push_fn)
        };

        // Only register backedges if memory_limit is set (eviction is enabled).
        // Uses timeline lookup with splay tree locality for O(1) amortized performance.
        // Backedges are stored in Runtime's HashMap, not in FullContext.
        if let Some(ref mut backedges_map) = self.backedges
            && let Some((deps, new_ctx_forward_uf)) = forward_uf_result
        {
            for dep_tock in deps.iter() {
                if let Some((dep_start, _dep_ctx)) = self.timeline.find_le_with_key(dep_tock) {
                    // Store backedge using dependency's start_tock as key
                    backedges_map
                        .entry(*dep_start)
                        .or_insert_with(UFSet::new)
                        .insert(new_ctx_forward_uf);
                }
            }
        }

        // Capture the start_tock of the just-created context before replacing record
        let context_start_tock = record.start_tock();

        // Replace with ValueRecord (includes context_start_tock for eviction protection)
        self.records.pop();
        self.records.push(Record::new_value(result_tock, context_start_tock));

        // During replay: explicitly set the result node for the replay target.
        // This is necessary because nested binds create zombies with new tocks during replay
        // that don't match the original target.
        if self.replays.len() > 1
            && result_tock != Tock::MAX
            && let Some(node) = self.lookup_node_at(result_tock)
        {
            self.set_replay_result(node);
        }

        // Check memory limit and evict if necessary.
        self.maybe_evict_with_protection(Some(context_start_tock));

        true
    }

    /// Get a weak reference to a zombie node by its tock.
    pub fn get_weak(&mut self, tock: Tock) -> Option<Weak<dyn EZombieNode>> {
        self.lookup_node_at(tock).map(|rc| Rc::downgrade(&rc))
    }

    /// Pop value record and resume parent, return result tock.
    pub fn pop_value_record(&mut self) -> Tock {
        let value_record = self.records.pop().unwrap();
        let result_tock = value_record.get_result_tock().expect("Expected ValueRecord");
        let context_start_tock = value_record.get_context_start_tock();

        // Resume parent record
        let new_tock = self.tick();
        if let Some(parent) = self.records.last_mut() {
            parent.resume(new_tock);
        }

        // Check memory limit and evict if necessary.
        // The context containing the result is protected from eviction because
        // the Zombie hasn't been fully returned yet - its ptr_cache (a Weak reference)
        // would become invalid if the context is evicted during nested bind operations.
        self.maybe_evict_with_protection(context_start_tock);

        result_tock
    }

    /// Prepare replay for current record.
    pub fn prepare_replay_current_record(
        &mut self,
    ) -> (Rc<Replayer>, crate::record::ReplayInputs) {
        let record = self.records.last_mut().unwrap();

        // Create input getter
        let timeline = &mut self.timeline;
        let mut get_input = |tock: Tock| -> Option<Rc<dyn EZombieNode>> {
            let (ctx_start, ctx) = timeline.find_le_with_key(&tock)?;
            // Use try_tock_to_index which safely handles boundary cases
            let idx = crate::tock::try_tock_to_index(tock, *ctx_start)?;
            if tock < ctx.end_tock() {
                ctx.get_value_at(idx)
            } else {
                None
            }
        };

        record.prepare_replay(&mut get_input)
    }

    /// Get replay inputs for a replayer without marking as played.
    /// Used when resuming from a nested replay (WaitingForInput state).
    pub fn get_replay_inputs_for_replayer(
        &mut self,
        replayer: &Rc<Replayer>,
    ) -> crate::record::ReplayInputs {
        let mut inputs = crate::record::ReplayInputs::new();
        for &input_tock in replayer.inputs.iter() {
            let node = self.lookup_node_at(input_tock);
            inputs.push((input_tock, node));
        }
        inputs
    }

    /// Pop record after replay completes (should be empty HeadRecord).
    pub fn pop_record_after_replay(&mut self) {
        let record = self.records.pop().unwrap();
        debug_assert!(record.values().is_empty(), "Record should be empty after replay");
    }

    /// Register a dependency on the current record.
    pub fn register_dependency(&mut self, tock: Tock) {
        self.current_record_mut().register_dependency(tock);
    }

    // ========== Replay Stack ==========

    /// Push replay state.
    pub fn push_replay(&mut self, target: Tock, result: Rc<RefCell<Option<Rc<dyn EZombieNode>>>>) {
        self.replays.push(ReplayState::new(target, result));
    }

    /// Pop replay state and return result.
    pub fn pop_replay(&mut self) -> Option<Rc<dyn EZombieNode>> {
        self.replays.pop().and_then(|state| state.take_result())
    }

    /// Check if we've reached current replay target.
    pub fn at_replay_target(&self, tock: Tock) -> bool {
        self.replays.last().is_some_and(|r| r.target == tock)
    }

    /// Set replay result.
    pub fn set_replay_result(&self, node: Rc<dyn EZombieNode>) {
        if let Some(replay) = self.replays.last() {
            replay.set_result(node);
        }
    }

    /// Check if current replay has result.
    pub fn replay_has_result(&self) -> bool {
        self.replays.last().is_none_or(|r| r.has_result())
    }

    /// Get current replay target.
    pub fn current_replay_target(&self) -> Tock {
        self.replays.last().map_or(Tock::MAX, |r| r.target)
    }

    /// Replay depth (1 = not replaying).
    pub fn replay_depth(&self) -> usize {
        self.replays.len()
    }

    // ========== Unroll Optimization ==========

    /// Check if we should unroll (skip creating context).
    ///
    /// C++ implementation:
    /// ```cpp
    /// if (t.records.back()->is_tailcall() && t.current_tock - t.records.back()->t < unroll_factor) {
    ///     // unroll
    /// }
    /// ```
    ///
    /// We only unroll when:
    /// 1. Current record is a tailcall (HeadRecord)
    /// 2. The tock delta is within unroll_factor
    #[inline]
    pub fn should_unroll(&self) -> bool {
        let record = self.current_record();
        // Only unroll in tailcall context (HeadRecord)
        if !record.is_tailcall() {
            return false;
        }
        let tock_delta = (self.current_tock.0 - record.start_tock().0).max(0) as usize;
        tock_delta < self.config.unroll_factor
    }

    /// Get unroll factor from config.
    pub fn unroll_factor(&self) -> usize {
        self.config.unroll_factor
    }

    // ========== GD Touch ==========

    /// Touch a context in the GD heap to update its priority - O(log n).
    ///
    /// This implements the "access recency" part of the Greedy-Dual algorithm.
    /// When a context is accessed (its values read), we touch it to delay eviction.
    ///
    /// Uses `pool_index` stored in FullContext for O(1) lookup + O(log n) rebalance,
    /// matching the C++ implementation's performance characteristics.
    ///
    /// # Performance
    ///
    /// - O(log n) using pool_index (like C++ implementation)
    /// - Only called when `memory_limit` is set (eviction enabled)
    /// - Only called from slow path (cache miss)
    #[inline]
    pub fn touch_context(&mut self, pool_index: usize) {
        // Update last_accessed timestamp (C++ does this in accessed())
        if let Some(handle) = self.heap.get_at(pool_index) {
            handle.context.accessed();
        }

        // C++ touch() only updates L_at_insert, doesn't recalculate cost
        // Matching C++ behavior exactly - cost is only recalculated during eviction (pop)
        self.heap.touch_at(pool_index);
    }

    // ========== Full Cost Calculation ==========

    /// Compute the complete time cost for a context, matching C++ implementation.
    ///
    /// C++ `FullContextNode::time_cost()` includes:
    /// 1. time_taken (base computation cost)
    /// 2. forward_uf.value() (propagated forward cost)
    /// 3. backward_uf.value() (propagated backward cost)
    /// 4. backedges.sum(counted) (cost from contexts depending on this one)
    /// 5. Each dependency's backward_uf (cost from dependencies)
    /// 6. Parent context's backward_uf
    ///
    /// All UF values are deduplicated using a HashSet to avoid double-counting
    /// merged sets.
    ///
    /// Note: Requires &mut self because SplayList::find_le_with_key performs splay operations.
    fn compute_full_time_cost(&mut self, ctx: &FullContext) -> crate::meter::Time {
        use std::collections::HashSet;
        use crate::union_find::UF;

        let mut counted: HashSet<UF> = HashSet::new();
        let mut cost = ctx.time_taken();

        // 1. Add forward_uf value (if not already counted)
        let forward_uf = *ctx.forward_uf();
        if counted.insert(forward_uf) {
            cost = cost + forward_uf.value();
        }

        // 2. Add backward_uf value (if not already counted)
        let backward_uf = *ctx.backward_uf();
        if counted.insert(backward_uf) {
            cost = cost + backward_uf.value();
        }

        // 3. Add backedges sum (contexts that depend on this one)
        if let Some(ref mut backedges_map) = self.backedges
            && let Some(backedges) = backedges_map.get_mut(&ctx.start_tock())
        {
            cost = cost + backedges.sum_excluding(&mut counted);
        }

        // 4. Add each dependency's backward_uf
        // Clone dependencies to avoid borrow issues with self.timeline
        let dependencies: Vec<_> = ctx.dependencies().to_vec();
        for dep_tock in dependencies {
            if let Some((_, dep_ctx)) = self.timeline.find_le_with_key(&dep_tock)
                && let Some(full_ctx) = dep_ctx.as_any().downcast_ref::<FullContext>()
            {
                let dep_backward_uf = *full_ctx.backward_uf();
                if counted.insert(dep_backward_uf) {
                    cost = cost + dep_backward_uf.value();
                }
            }
        }

        // 5. Add parent context's backward_uf
        let parent_tock = crate::tock::Tock(ctx.start_tock().0 - 1);
        if parent_tock >= crate::tock::Tock(0)
            && let Some((_, parent_ctx)) = self.timeline.find_le_with_key(&parent_tock)
            && let Some(full_ctx) = parent_ctx.as_any().downcast_ref::<FullContext>()
        {
            let parent_backward_uf = *full_ctx.backward_uf();
            if counted.insert(parent_backward_uf) {
                cost = cost + parent_backward_uf.value();
            }
        }

        cost
    }

    /// Compute full eviction cost using the metric function.
    ///
    /// This is the complete cost calculation matching C++ behavior.
    /// Note: Requires &mut self because it calls compute_full_time_cost.
    #[allow(dead_code)]  // Available for testing/debugging
    fn compute_full_cost(&mut self, ctx: &FullContext) -> crate::config::Cost {
        let time_cost = self.compute_full_time_cost(ctx);
        (self.config.metric)(ctx.time_taken(), time_cost, ctx.space())
    }

    // ========== Eviction ==========

    /// Evict one context from the heap.
    /// `protected_tock` specifies a context start_tock that should be skipped.
    pub fn evict_one_with_protection(&mut self, protected_tock: Option<Tock>) -> bool {
        let metric = self.config.metric;

        // Pop contexts from heap, skipping protected one
        // Track if we've already skipped the protected context once
        // to prevent infinite loop when it's the only context in the heap
        let mut skipped_protected = false;

        let result = loop {
            let handle = match self.heap.pop_with_cost_check_notify(|h| h.context.cost(metric)) {
                Some(h) => h,
                None => break None,
            };

            // Check if this context is protected
            if let Some(pt) = protected_tock
                && handle.context.start_tock() == pt
            {
                // If we've already skipped this context once, it means it's the only
                // context in the heap. Return None to avoid infinite loop.
                if skipped_protected {
                    // Put it back and return None - can't evict anything
                    let cost = handle.context.cost(metric);
                    self.heap.push_notify(handle, cost);
                    break None;
                }
                // First time seeing this protected context - skip it
                skipped_protected = true;
                let cost = handle.context.cost(metric);
                self.heap.push_notify(handle, cost);
                continue;
            }

            break Some(handle);
        };

        if let Some(handle) = result {
            // Decrease total space used (before evict clears the space)
            let space_freed = handle.context.space().as_bytes();
            self.total_space_used = self.total_space_used.saturating_sub(space_freed);

            // Create UF with this context's time cost for propagation
            let cost_uf = crate::union_find::UF::new(handle.context.time_taken());

            // Propagate cost to dependencies' backward_uf
            // When evicting A, all contexts that A depends on get increased backward cost
            // (because replaying A will need to access them)
            for &dep_tock in handle.context.dependencies() {
                if let Some((_, dep_ctx)) = self.timeline.find_le_with_key(&dep_tock)
                    && let Some(full_ctx) = dep_ctx.as_any().downcast_ref::<FullContext>()
                {
                    full_ctx.backward_uf().merge(&cost_uf);
                }
            }

            // Propagate cost to backedges' forward_uf (O(1) merge operation)
            // When evicting A, all contexts that depend on A get increased forward cost.
            // Backedges are stored in Runtime's HashMap.
            let ctx_start = handle.context.start_tock();
            if let Some(ref backedges_map) = self.backedges
                && let Some(backedges) = backedges_map.get(&ctx_start)
            {
                backedges.merge_all(&cost_uf);
            }

            // Propagate cost to parent context's backward_uf
            let parent_tock = crate::tock::Tock(handle.context.start_tock().0 - 1);
            if parent_tock >= crate::tock::Tock(0)
                && let Some((_, parent_ctx)) = self.timeline.find_le_with_key(&parent_tock)
                && let Some(full_ctx) = parent_ctx.as_any().downcast_ref::<FullContext>()
            {
                full_ctx.backward_uf().merge(&cost_uf);
            }

            // Evict the context itself (clears values and merges internal UFs)
            handle.context.evict();

            // Remove backedges for the evicted context
            if let Some(ref mut backedges_map) = self.backedges {
                backedges_map.remove(&ctx_start);
            }

            // Remove from timeline to avoid metadata leaks (O(N) memory).
            // Replay will find the parent context via find_le_with_key.
            self.remove_context(handle.context.start_tock());
            true
        } else {
            false
        }
    }

    /// Get total space used by all contexts in the heap (O(1)).
    pub fn total_space(&self) -> usize {
        self.total_space_used
    }

    /// Check memory limit and evict if necessary.
    /// This is called after creating new zombie values.
    #[inline]
    pub fn maybe_evict(&mut self) {
        self.maybe_evict_with_protection(None);
    }

    /// Check memory limit and evict if necessary, with a protected context.
    /// The context with start_tock = protected_tock will not be evicted.
    pub fn maybe_evict_with_protection(&mut self, protected_tock: Option<Tock>) {

        if let Some(limit) = self.config.memory_limit {
            while self.total_space_used > limit && self.evict_one_with_protection(protected_tock) {
                // Keep evicting until under limit or no more to evict
            }
        }
    }

    /// Evict one context (convenience wrapper).
    #[inline]
    pub fn evict_one(&mut self) -> bool {
        self.evict_one_with_protection(None)
    }

    // ========== Meter ==========

    /// Get meter reference.
    pub fn meter(&self) -> &ZombieMeter {
        &self.meter
    }

    /// Get meter mutable reference.
    pub fn meter_mut(&mut self) -> &mut ZombieMeter {
        &mut self.meter
    }

    // ========== Stats ==========

    /// Number of contexts in tock tree.
    pub fn context_count(&self) -> usize {
        self.timeline.len()
    }

    /// Estimated metadata memory usage in bytes.
    ///
    /// This includes:
    /// - Contexts (FullContext, RootContext) with their SmallVecs
    /// - GD heap entries
    /// - Tock tree (SplayList) nodes
    ///
    /// Does NOT include user data (the actual Zombie values).
    pub fn metadata_bytes(&self) -> usize {
        use std::mem::size_of;

        // Use actual struct sizes - SmallVec vs Vec difference is reflected here
        let context_size = size_of::<FullContext>(); // Includes SmallVec inline storage
        let context_overhead = self.timeline.len() * context_size;

        // Heap entry: (Cost, Weak<FullContext>) ≈ 24 bytes
        let heap_overhead = self.heap.len() * 24;

        // SplayList node overhead (ptr + key + value ref)
        let splay_overhead = self.timeline.len() * 48;

        // Record: RecordNode enum with Replayer reference
        let record_overhead = self.records.len() * 64;

        context_overhead + heap_overhead + splay_overhead + record_overhead
    }

    /// Number of entries in the GD heap.
    pub fn heap_len(&self) -> usize {
        self.heap.len()
    }
}

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

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

    #[test]
    fn test_tick() {
        let mut t = Runtime::new();
        assert_eq!(t.tick(), Tock(1));
        assert_eq!(t.tick(), Tock(2));
        assert_eq!(t.current_tock(), Tock(3));
    }

    #[test]
    fn test_thread_local() {
        Runtime::init();
        Runtime::with_mut(|t| {
            assert_eq!(t.tick(), Tock(1));
        });
        Runtime::with(|t| {
            assert_eq!(t.current_tock(), Tock(2));
        });
    }

    #[test]
    fn test_replay_depth() {
        let t = Runtime::new();
        assert_eq!(t.replay_depth(), 1); // Default replay state
    }

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

    /// Test that empty timeline operations don't panic
    #[test]
    fn test_empty_timeline_operations() {
        let mut runtime = Runtime::new();

        // lookup_node_at on empty timeline
        assert!(runtime.lookup_node_at(Tock(0)).is_none());
        assert!(runtime.lookup_node_at(Tock(100)).is_none());

        // find_context_with_key on empty timeline
        assert!(runtime.find_context_with_key(Tock(0)).is_none());

        // find_replay_context on empty timeline
        assert!(runtime.find_replay_context(Tock(0)).is_none());

        // get_weak on empty timeline
        assert!(runtime.get_weak(Tock(0)).is_none());

        // remove_context on empty timeline doesn't panic
        runtime.remove_context(Tock(0));

        // Stats
        assert_eq!(runtime.context_count(), 0);
        assert_eq!(runtime.total_space(), 0);
    }

    /// Test tock boundary values
    #[test]
    fn test_tock_boundary_values() {
        let mut runtime = Runtime::new();

        // Tock(0) as lookup key
        assert!(runtime.lookup_node_at(Tock(0)).is_none());

        // Tock::MAX as lookup key
        assert!(runtime.lookup_node_at(Tock::MAX).is_none());

        // set_tock and current_tock consistency
        runtime.set_tock(Tock(100));
        assert_eq!(runtime.current_tock(), Tock(100));

        // tick increments tock
        assert_eq!(runtime.tick(), Tock(100));
        assert_eq!(runtime.current_tock(), Tock(101));
    }

    /// Test record stack invariant
    #[test]
    fn test_record_stack_invariant() {
        let runtime = Runtime::new();

        // Record stack is non-empty after init
        let record = runtime.current_record();
        // RootRecord starts at Tock(0)
        assert_eq!(record.start_tock(), Tock(0));

        // Initial record is not HeadRecord
        assert!(!record.is_head());
        assert!(!record.is_value());
    }

    /// Test memory limit configuration
    #[test]
    fn test_memory_limit_config() {
        use crate::config::ZombieConfig;

        // No memory limit
        let config = ZombieConfig::default();
        let runtime = Runtime::with_config(config);
        assert!(!runtime.has_memory_limit());

        // With memory limit
        let config = ZombieConfig::default().with_memory_limit(1024 * 1024);
        let runtime = Runtime::with_config(config);
        assert!(runtime.has_memory_limit());
    }

    /// Test thread-local initialization idempotency
    #[test]
    fn test_init_idempotent() {
        // Multiple init calls don't panic
        Runtime::init();
        Runtime::init();
        Runtime::init();

        // Should still work normally
        Runtime::with(|t| {
            assert!(t.current_tock() >= Tock(1));
        });
    }

    /// Test metadata_bytes calculation
    #[test]
    fn test_metadata_bytes() {
        let runtime = Runtime::new();

        // Empty runtime should have minimal metadata
        let bytes = runtime.metadata_bytes();
        assert!(bytes > 0); // At least record stack overhead
    }

    /// Test that try_with returns None when Runtime is uninitialized
    #[test]
    fn test_try_with_uninitialized() {
        // Create a new thread where Runtime hasn't been initialized
        let handle = std::thread::spawn(|| {
            // try_with should return None before init
            let result = Runtime::try_with(|_| 42);
            // Note: After the fix, Runtime::new() now calls UFArena::init(),
            // but try_with still returns None because the thread-local Runtime
            // is not initialized until Runtime::init() is called.
            result.is_none()
        });

        let was_none = handle.join().expect("Thread should not panic");
        assert!(was_none, "try_with should return None in uninitialized thread");
    }

    /// Test replay stack operations
    #[test]
    fn test_replay_stack_operations() {
        let mut runtime = Runtime::new();

        // Initially replay stack has one entry
        assert!(runtime.replays.len() == 1);
        assert_eq!(runtime.replays[0].target, Tock::MAX);

        // Push a new replay state
        let new_replay = ReplayState::new(Tock(100), Rc::new(RefCell::new(None)));
        runtime.replays.push(new_replay);
        assert_eq!(runtime.replays.len(), 2);

        // Pop replay state
        runtime.replays.pop();
        assert_eq!(runtime.replays.len(), 1);
    }

    /// Test should_unroll with various staleness thresholds
    #[test]
    fn test_should_unroll() {
        use crate::config::ZombieConfig;

        // Default config with high staleness threshold (unroll_factor = 32)
        let config = ZombieConfig::default();
        let runtime = Runtime::with_config(config);

        // RootRecord is NOT a tailcall, so should_unroll returns false regardless of tock delta.
        // This matches C++ behavior: unrolling only happens inside TailCall contexts.
        assert!(!runtime.should_unroll());

        // Config with low staleness threshold shouldn't affect unrolling logic directly
        // but let's just check the unroll factor config
        let config = ZombieConfig::default().with_unroll_factor(0);
        let runtime = Runtime::with_config(config);
        // With unroll_factor=0 and non-tailcall record, still false
        assert!(!runtime.should_unroll());
    }

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

    /// Test timeline find_context_with_key boundary cases
    #[test]
    fn test_timeline_find_boundary() {
        use crate::context::RootContext;

        let mut runtime = Runtime::new();

        // Insert a context at Tock(10) with end at Tock(15)
        let ctx = RootContext::new(Tock(10));
        // Simulate end tock by pushing some values
        runtime.insert_context(Tock(10), ctx.clone());

        // Exact match should find
        let result = runtime.find_context_with_key(Tock(10));
        assert!(result.is_some());
        assert_eq!(result.unwrap().0, Tock(10));

        // Key after context start should still find it (find_le)
        let result = runtime.find_context_with_key(Tock(12));
        assert!(result.is_some());
        assert_eq!(result.unwrap().0, Tock(10));

        // Key before any context should return None
        let result = runtime.find_context_with_key(Tock(5));
        assert!(result.is_none());
    }

    /// Test lookup_node_at returns None for out-of-range tock
    #[test]
    fn test_lookup_node_out_of_range() {
        let mut runtime = Runtime::new();

        // Lookup on empty timeline
        assert!(runtime.lookup_node_at(Tock(0)).is_none());
        assert!(runtime.lookup_node_at(Tock(100)).is_none());
        assert!(runtime.lookup_node_at(Tock::MAX).is_none());
    }

    /// Test at_replay_target correctly identifies target
    #[test]
    fn test_at_replay_target() {
        let mut runtime = Runtime::new();

        // Initial state: target is MAX
        assert!(runtime.at_replay_target(Tock::MAX));
        assert!(!runtime.at_replay_target(Tock(100)));

        // Push new replay with target 100
        runtime.push_replay(Tock(100), Rc::new(RefCell::new(None)));
        assert!(runtime.at_replay_target(Tock(100)));
        assert!(!runtime.at_replay_target(Tock(50)));
        assert!(!runtime.at_replay_target(Tock::MAX));

        // Pop replay - back to MAX
        runtime.pop_replay();
        assert!(runtime.at_replay_target(Tock::MAX));
    }

    /// Test total_space tracking
    #[test]
    fn test_total_space_tracking() {
        use crate::config::ZombieConfig;

        // With memory limit enabled
        let config = ZombieConfig::default().with_memory_limit(1024 * 1024);
        let runtime = Runtime::with_config(config);

        // Initially zero
        assert_eq!(runtime.total_space(), 0);
    }
}