infotheory/backends/

ctw.rs

//! Context Tree Weighting (CTW) and Factorized Action-Conditional CTW (FAC-CTW).
//!
//! This implementation stores maximal non-root unary runs as chains while keeping
//! the root and true branching points explicit. Updates and reverts rebuild only
//! the touched context path from exact per-node state, which preserves predictive
//! semantics while avoiding the `O(depth)` explicit-node blow-up for singleton
//! paths.

use std::cell::{Cell, RefCell};
use std::f64;
use std::mem::size_of;
use std::sync::OnceLock;

type Symbol = bool;
const HISTORY_WORD_BITS: usize = u64::BITS as usize;

#[inline(always)]
fn history_word_len(bits: usize) -> usize {
    bits.div_ceil(HISTORY_WORD_BITS)
}

trait HistoryAccess {
    fn len(&self) -> usize;
    fn bit(&self, index: usize) -> Symbol;

    #[inline(always)]
    fn recent_bit(&self, _depth: usize) -> Option<Symbol> {
        None
    }

    #[inline(always)]
    fn recent_path_bits(&self, _depth: usize, _len: usize) -> Option<u64> {
        None
    }
}

impl HistoryAccess for [Symbol] {
    #[inline(always)]
    fn len(&self) -> usize {
        <[Symbol]>::len(self)
    }

    #[inline(always)]
    fn bit(&self, index: usize) -> Symbol {
        debug_assert!(index < self.len());
        unsafe { *self.get_unchecked(index) }
    }
}

impl HistoryAccess for Vec<Symbol> {
    #[inline(always)]
    fn len(&self) -> usize {
        self.as_slice().len()
    }

    #[inline(always)]
    fn bit(&self, index: usize) -> Symbol {
        self.as_slice().bit(index)
    }
}

#[derive(Clone, Debug, Default, PartialEq, Eq)]
struct BitHistory {
    words: Vec<u64>,
    len: usize,
    recent: u64,
}

impl BitHistory {
    #[inline(always)]
    fn len(&self) -> usize {
        self.len
    }

    #[inline(always)]
    fn is_empty(&self) -> bool {
        self.len == 0
    }

    #[inline(always)]
    fn memory_usage(&self) -> usize {
        self.words.capacity() * size_of::<u64>()
    }

    #[inline]
    fn rebuild_recent(&mut self) {
        self.recent = 0;
        let tail_len = self.len.min(HISTORY_WORD_BITS);
        for depth in 0..tail_len {
            let idx = self.len - depth - 1;
            if self.bit(idx) {
                self.recent |= 1u64 << depth;
            }
        }
    }

    #[inline]
    fn reserve_exact(&mut self, additional_bits: usize) {
        let required_bits = self.len.saturating_add(additional_bits);
        let required_words = history_word_len(required_bits);
        if required_words > self.words.capacity() {
            self.words
                .reserve_exact(required_words.saturating_sub(self.words.len()));
        }
    }

    #[inline]
    fn push(&mut self, bit: Symbol) {
        let word_idx = self.len / HISTORY_WORD_BITS;
        let bit_idx = self.len % HISTORY_WORD_BITS;
        if word_idx == self.words.len() {
            self.words.push(0);
        }
        let mask = 1u64 << bit_idx;
        if bit {
            self.words[word_idx] |= mask;
        } else {
            self.words[word_idx] &= !mask;
        }
        self.recent = (self.recent << 1) | (bit as u64);
        self.len += 1;
    }

    #[inline]
    fn pop(&mut self) -> Option<Symbol> {
        if self.len == 0 {
            return None;
        }
        let next_len = self.len - 1;
        let word_idx = next_len / HISTORY_WORD_BITS;
        let bit_idx = next_len % HISTORY_WORD_BITS;
        let mask = 1u64 << bit_idx;
        let bit = (self.words[word_idx] & mask) != 0;
        if bit {
            self.words[word_idx] &= !mask;
        }
        self.len = next_len;
        if bit_idx == 0 {
            self.words.truncate(word_idx);
        }
        self.recent >>= 1;
        if self.len >= HISTORY_WORD_BITS {
            let exposed_idx = self.len - HISTORY_WORD_BITS;
            if self.bit(exposed_idx) {
                self.recent |= 1u64 << (HISTORY_WORD_BITS - 1);
            }
        }
        Some(bit)
    }

    #[inline]
    fn extend_from_slice(&mut self, symbols: &[Symbol]) {
        self.reserve_exact(symbols.len());
        for &symbol in symbols {
            self.push(symbol);
        }
    }

    #[inline]
    fn truncate(&mut self, new_len: usize) {
        if new_len >= self.len {
            return;
        }
        self.len = new_len;
        self.words.truncate(history_word_len(new_len));
        let rem = new_len % HISTORY_WORD_BITS;
        if rem != 0 {
            let mask = (1u64 << rem) - 1;
            if let Some(last) = self.words.last_mut() {
                *last &= mask;
            }
        }
        self.rebuild_recent();
    }

    #[inline]
    fn clear(&mut self) {
        self.words.clear();
        self.len = 0;
        self.recent = 0;
    }

    #[cfg(test)]
    fn to_vec(&self) -> Vec<Symbol> {
        (0..self.len).map(|idx| self.bit(idx)).collect()
    }
}

impl HistoryAccess for BitHistory {
    #[inline(always)]
    fn len(&self) -> usize {
        self.len
    }

    #[inline(always)]
    fn bit(&self, index: usize) -> Symbol {
        debug_assert!(index < self.len);
        let word = unsafe { *self.words.get_unchecked(index / HISTORY_WORD_BITS) };
        ((word >> (index % HISTORY_WORD_BITS)) & 1) != 0
    }

    #[inline(always)]
    fn recent_bit(&self, depth: usize) -> Option<Symbol> {
        if depth < self.len.min(HISTORY_WORD_BITS) {
            Some(((self.recent >> depth) & 1) != 0)
        } else {
            None
        }
    }

    #[inline(always)]
    fn recent_path_bits(&self, depth: usize, len: usize) -> Option<u64> {
        let available = self.len.saturating_sub(depth).min(len);
        if available == 0 {
            return Some(0);
        }
        if depth + available <= self.len.min(HISTORY_WORD_BITS) {
            let mask = if available >= HISTORY_WORD_BITS {
                u64::MAX
            } else {
                (1u64 << available) - 1
            };
            Some((self.recent >> depth) & mask)
        } else {
            None
        }
    }
}

const CTW_LOG_CACHE_LIMIT: usize = 1 << 24;
const CTW_HOT_PREFIX_DEPTH_DEFAULT: usize = 12;
const CTW_LOG_OVERFLOW_CACHE_SLOTS: usize = 1 << 14;

#[cfg(not(test))]
#[inline(always)]
fn ctw_log_cache_limit() -> usize {
    CTW_LOG_CACHE_LIMIT
}

#[cfg(test)]
#[inline(always)]
fn ctw_log_cache_limit() -> usize {
    CTW_TEST_LOG_CACHE_LIMIT.with(|limit| limit.borrow().unwrap_or(CTW_LOG_CACHE_LIMIT))
}

fn ctw_hot_prefix_depth_limit() -> usize {
    static HOT_PREFIX_DEPTH: OnceLock<usize> = OnceLock::new();
    *HOT_PREFIX_DEPTH.get_or_init(|| {
        std::env::var("INFOTHEORY_CTW_HOT_PREFIX_DEPTH")
            .ok()
            .and_then(|raw| raw.parse::<usize>().ok())
            .unwrap_or(CTW_HOT_PREFIX_DEPTH_DEFAULT)
    })
}

#[cfg(not(test))]
#[inline(always)]
fn ctw_log_overflow_cache_slots() -> usize {
    CTW_LOG_OVERFLOW_CACHE_SLOTS
}

#[cfg(test)]
#[inline(always)]
fn ctw_log_overflow_cache_slots() -> usize {
    CTW_TEST_LOG_OVERFLOW_CACHE_SLOTS
        .with(|slots| slots.borrow().unwrap_or(CTW_LOG_OVERFLOW_CACHE_SLOTS))
}

#[cfg(test)]
#[inline(always)]
fn ensure_log_caches(log_int: &mut Vec<f64>, log_half: &mut Vec<f64>, upto: usize) {
    if upto < log_int.len() {
        return;
    }
    let required_len = upto + 1;
    log_int.reserve(required_len - log_int.len());
    log_half.reserve(required_len - log_half.len());
    push_log_cache_entries(log_int, log_half, upto);
}

#[inline(always)]
fn reserve_bounded_log_cache(cache: &mut Vec<f64>, required_len: usize, max_len: usize) {
    if cache.capacity() >= required_len {
        return;
    }
    let target_capacity = cache
        .capacity()
        .saturating_mul(2)
        .max(required_len)
        .min(max_len);
    cache.reserve_exact(target_capacity - cache.len());
}

#[inline(always)]
fn ensure_bounded_log_caches(
    log_int: &mut Vec<f64>,
    log_half: &mut Vec<f64>,
    upto: usize,
    limit: usize,
) {
    let target = upto.min(limit);
    if target < log_int.len() {
        return;
    }
    let required_len = target + 1;
    let max_len = limit + 1;
    reserve_bounded_log_cache(log_int, required_len, max_len);
    reserve_bounded_log_cache(log_half, required_len, max_len);
    push_log_cache_entries(log_int, log_half, target);
}

#[inline(always)]
fn push_log_cache_entries(log_int: &mut Vec<f64>, log_half: &mut Vec<f64>, upto: usize) {
    let start = log_int.len();
    for n in start..=upto {
        if n == 0 {
            log_int.push(f64::NEG_INFINITY);
        } else {
            log_int.push((n as f64).ln());
        }
        log_half.push((n as f64 + 0.5).ln());
    }
}

struct LogCacheSlot {
    key: Cell<usize>,
    value: Cell<f64>,
}

impl LogCacheSlot {
    #[inline(always)]
    fn empty() -> Self {
        Self {
            key: Cell::new(usize::MAX),
            value: Cell::new(0.0),
        }
    }
}

#[inline(always)]
fn ensure_overflow_log_cache(cache: &mut Vec<LogCacheSlot>, slots: usize) {
    if slots == 0 {
        cache.clear();
        cache.shrink_to(0);
        return;
    }
    if cache.len() == slots {
        return;
    }
    cache.clear();
    if cache.capacity() < slots {
        cache.reserve_exact(slots);
    }
    cache.resize_with(slots, LogCacheSlot::empty);
    if cache.capacity() > slots {
        cache.shrink_to(slots);
    }
}

trait CtLogAccess: Copy {
    fn log_int(self, n: usize) -> f64;
    fn log_half(self, n: usize) -> f64;
}

#[derive(Clone, Copy)]
struct CachedLogs<'a> {
    log_int: &'a [f64],
    log_half: &'a [f64],
}

impl<'a> CachedLogs<'a> {
    #[cfg(test)]
    #[inline(always)]
    fn new(log_int: &'a [f64], log_half: &'a [f64]) -> Self {
        Self { log_int, log_half }
    }
}

impl CtLogAccess for CachedLogs<'_> {
    #[inline(always)]
    fn log_int(self, n: usize) -> f64 {
        debug_assert!(n < self.log_int.len());
        // SAFETY: `CachedLogs` is used only when the shared prefix cache already
        // contains every value through the current visit bound.
        unsafe { *self.log_int.get_unchecked(n) }
    }

    #[inline(always)]
    fn log_half(self, n: usize) -> f64 {
        debug_assert!(n < self.log_half.len());
        // SAFETY: `CachedLogs` is used only when the shared prefix cache already
        // contains every value through the current visit bound.
        unsafe { *self.log_half.get_unchecked(n) }
    }
}

#[derive(Clone, Copy)]
struct BoundedLogs<'a> {
    log_int: &'a [f64],
    log_half: &'a [f64],
    overflow_log_int: &'a [LogCacheSlot],
    overflow_log_half: &'a [LogCacheSlot],
}

impl<'a> BoundedLogs<'a> {
    #[cfg(test)]
    #[inline(always)]
    fn new(log_int: &'a [f64], log_half: &'a [f64]) -> Self {
        Self {
            log_int,
            log_half,
            overflow_log_int: &[],
            overflow_log_half: &[],
        }
    }

    #[inline(always)]
    fn with_overflow(
        log_int: &'a [f64],
        log_half: &'a [f64],
        overflow_log_int: &'a [LogCacheSlot],
        overflow_log_half: &'a [LogCacheSlot],
    ) -> Self {
        Self {
            log_int,
            log_half,
            overflow_log_int,
            overflow_log_half,
        }
    }

    #[inline(always)]
    fn lookup_overflow(slots: &'a [LogCacheSlot], n: usize, compute: impl FnOnce() -> f64) -> f64 {
        if slots.is_empty() {
            return compute();
        }
        let len = slots.len();
        let slot_idx = if len.is_power_of_two() {
            n & (len - 1)
        } else {
            n % len
        };
        let slot = &slots[slot_idx];
        if slot.key.get() == n {
            slot.value.get()
        } else {
            let value = compute();
            slot.key.set(n);
            slot.value.set(value);
            value
        }
    }
}

impl CtLogAccess for BoundedLogs<'_> {
    #[inline(always)]
    fn log_int(self, n: usize) -> f64 {
        if n < self.log_int.len() {
            self.log_int[n]
        } else if n == 0 {
            f64::NEG_INFINITY
        } else {
            Self::lookup_overflow(self.overflow_log_int, n, || (n as f64).ln())
        }
    }

    #[inline(always)]
    fn log_half(self, n: usize) -> f64 {
        if n < self.log_half.len() {
            self.log_half[n]
        } else {
            Self::lookup_overflow(self.overflow_log_half, n, || (n as f64 + 0.5).ln())
        }
    }
}

#[derive(Default)]
struct SharedLogCache {
    log_int: Vec<f64>,
    log_half: Vec<f64>,
    overflow_log_int: Vec<LogCacheSlot>,
    overflow_log_half: Vec<LogCacheSlot>,
}

impl SharedLogCache {
    fn new() -> Self {
        Self {
            log_int: vec![f64::NEG_INFINITY],
            log_half: vec![(0.5f64).ln()],
            overflow_log_int: Vec::new(),
            overflow_log_half: Vec::new(),
        }
    }

    #[inline(always)]
    fn ensure(&mut self, upto: usize) {
        let limit = ctw_log_cache_limit();
        ensure_bounded_log_caches(&mut self.log_int, &mut self.log_half, upto, limit);
        if upto > limit {
            let slots = ctw_log_overflow_cache_slots();
            ensure_overflow_log_cache(&mut self.overflow_log_int, slots);
            ensure_overflow_log_cache(&mut self.overflow_log_half, slots);
        }
    }

    #[inline(always)]
    fn memory_usage(&self) -> usize {
        self.log_int.capacity() * size_of::<f64>()
            + self.log_half.capacity() * size_of::<f64>()
            + self.overflow_log_int.capacity() * size_of::<LogCacheSlot>()
            + self.overflow_log_half.capacity() * size_of::<LogCacheSlot>()
    }
}

thread_local! {
    static CTW_SHARED_LOG_CACHE: RefCell<SharedLogCache> =
        RefCell::new(SharedLogCache::new());
}

#[cfg(test)]
thread_local! {
    static CTW_TEST_LOG_CACHE_LIMIT: RefCell<Option<usize>> = const { RefCell::new(None) };
}

#[cfg(test)]
thread_local! {
    static CTW_TEST_LOG_OVERFLOW_CACHE_SLOTS: RefCell<Option<usize>> = const { RefCell::new(None) };
}

#[inline]
fn with_shared_cached_logs<R>(upto: usize, f: impl FnOnce(CachedLogs<'_>) -> R) -> R {
    debug_assert!(upto <= ctw_log_cache_limit());
    CTW_SHARED_LOG_CACHE.with(|cache_cell| {
        let mut cache = cache_cell.borrow_mut();
        cache.ensure(upto);
        f(CachedLogs {
            log_int: &cache.log_int,
            log_half: &cache.log_half,
        })
    })
}

#[inline]
fn with_shared_bounded_logs<R>(upto: usize, f: impl FnOnce(BoundedLogs<'_>) -> R) -> R {
    CTW_SHARED_LOG_CACHE.with(|cache_cell| {
        let mut cache = cache_cell.borrow_mut();
        cache.ensure(upto);
        f(BoundedLogs::with_overflow(
            &cache.log_int,
            &cache.log_half,
            &cache.overflow_log_int,
            &cache.overflow_log_half,
        ))
    })
}

#[inline]
fn shared_log_cache_memory_usage() -> usize {
    CTW_SHARED_LOG_CACHE.with(|cache_cell| cache_cell.borrow().memory_usage())
}

#[cfg(test)]
#[inline]
fn shared_log_cache_lens() -> (usize, usize) {
    CTW_SHARED_LOG_CACHE.with(|cache_cell| {
        let cache = cache_cell.borrow();
        (cache.log_int.len(), cache.log_half.len())
    })
}

#[cfg(test)]
#[inline]
fn shared_log_overflow_cache_lens() -> (usize, usize) {
    CTW_SHARED_LOG_CACHE.with(|cache_cell| {
        let cache = cache_cell.borrow();
        (cache.overflow_log_int.len(), cache.overflow_log_half.len())
    })
}

#[cfg(test)]
pub(crate) fn reset_shared_log_cache_for_test() {
    CTW_SHARED_LOG_CACHE.with(|cache_cell| {
        *cache_cell.borrow_mut() = SharedLogCache::new();
    });
}

#[inline(always)]
fn history_symbol<H: HistoryAccess + ?Sized>(history: &H, depth: usize) -> Symbol {
    if let Some(bit) = history.recent_bit(depth) {
        return bit;
    }
    let idx = history.len().wrapping_sub(depth + 1);
    if depth < history.len() {
        history.bit(idx)
    } else {
        false
    }
}

#[inline(always)]
fn history_at_or_zero<H: HistoryAccess + ?Sized>(
    history: &H,
    history_len: isize,
    idx: isize,
) -> Symbol {
    if idx >= 0 && idx < history_len {
        history.bit(idx as usize)
    } else {
        false
    }
}

const INDEX_BITS: u32 = 31;
const INDEX_LIMIT: usize = 1usize << INDEX_BITS;
const CHILD_SEGMENT_TAG: u32 = 1u32 << INDEX_BITS;
const CHILD_INDEX_MASK: u32 = CHILD_SEGMENT_TAG - 1;
const SEG_META_MODE_SHIFT: u32 = 30;
const SEG_META_MODE_MASK: u32 = 0b11 << SEG_META_MODE_SHIFT;
const SEG_LEN_MASK: u32 = !SEG_META_MODE_MASK;
const SEG_MODE_EXACT: u32 = 0 << SEG_META_MODE_SHIFT;
const SEG_MODE_HISTORY: u32 = 1 << SEG_META_MODE_SHIFT;
const SEG_MODE_HISTORY_INVERT: u32 = 2 << SEG_META_MODE_SHIFT;
const SEG_MODE_CONST: u32 = 3 << SEG_META_MODE_SHIFT;
const SEG_EXACT_MAX_LEN: u32 = 64;

/// Index into the explicit-node arena.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct NodeIndex(u32);

impl NodeIndex {
    #[cold]
    #[inline(never)]
    fn overflow() -> ! {
        panic!("ctw node index overflow");
    }

    #[inline(always)]
    fn from_usize(idx: usize) -> Self {
        if idx >= INDEX_LIMIT {
            Self::overflow();
        }
        Self(idx as u32)
    }

    #[inline(always)]
    fn get(self) -> usize {
        self.0 as usize
    }
}

/// Index into the unary-segment arena.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
struct SegmentIndex(u32);

impl SegmentIndex {
    #[cold]
    #[inline(never)]
    fn overflow() -> ! {
        panic!("ctw segment index overflow");
    }

    #[inline(always)]
    fn from_usize(idx: usize) -> Self {
        if idx >= INDEX_LIMIT {
            Self::overflow();
        }
        Self(idx as u32)
    }

    #[inline(always)]
    fn get(self) -> usize {
        self.0 as usize
    }
}

/// Tagged child reference: `None`, explicit node, or unary segment.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
struct ChildRef(u32);

impl ChildRef {
    const NONE: ChildRef = ChildRef(u32::MAX);

    #[inline(always)]
    fn from_node(idx: NodeIndex) -> Self {
        debug_assert!(idx.0 < CHILD_SEGMENT_TAG);
        Self(idx.0)
    }

    #[inline(always)]
    fn from_segment(idx: SegmentIndex) -> Self {
        debug_assert!(idx.0 < CHILD_SEGMENT_TAG);
        Self(CHILD_SEGMENT_TAG | idx.0)
    }

    #[inline(always)]
    fn is_none(self) -> bool {
        self.0 == u32::MAX
    }

    #[inline(always)]
    fn is_some(self) -> bool {
        self.0 != u32::MAX
    }

    #[inline(always)]
    fn as_node(self) -> Option<NodeIndex> {
        if self.is_none() || (self.0 & CHILD_SEGMENT_TAG) != 0 {
            None
        } else {
            Some(NodeIndex(self.0))
        }
    }

    #[inline(always)]
    fn as_segment(self) -> Option<SegmentIndex> {
        if self.is_none() || (self.0 & CHILD_SEGMENT_TAG) == 0 {
            None
        } else {
            Some(SegmentIndex(self.0 & CHILD_INDEX_MASK))
        }
    }
}

impl Default for ChildRef {
    fn default() -> Self {
        Self::NONE
    }
}

#[derive(Clone, Copy, Debug, Default)]
struct SegmentPayload {
    repr_lo: u32,
    repr_hi: u32,
    meta: u32,
}

impl SegmentPayload {
    #[inline(always)]
    fn exact(bits: u64, len: u32) -> Self {
        debug_assert!(len <= SEG_EXACT_MAX_LEN);
        debug_assert!(len <= SEG_LEN_MASK);
        Self {
            repr_lo: bits as u32,
            repr_hi: (bits >> 32) as u32,
            meta: SEG_MODE_EXACT | len,
        }
    }

    #[inline(always)]
    fn history(anchor: u32, len: u32, invert: bool) -> Self {
        debug_assert!(len <= SEG_LEN_MASK);
        Self {
            repr_lo: anchor,
            repr_hi: 0,
            meta: if invert {
                SEG_MODE_HISTORY_INVERT | len
            } else {
                SEG_MODE_HISTORY | len
            },
        }
    }

    #[inline(always)]
    fn constant(bit: bool, len: u32) -> Self {
        debug_assert!(len <= SEG_LEN_MASK);
        Self {
            repr_lo: bit as u32,
            repr_hi: 0,
            meta: SEG_MODE_CONST | len,
        }
    }

    #[inline(always)]
    fn len(self) -> u32 {
        self.meta & SEG_LEN_MASK
    }

    #[inline(always)]
    fn set_len(&mut self, len: u32) {
        debug_assert!(len <= SEG_LEN_MASK);
        self.meta = (self.meta & SEG_META_MODE_MASK) | len;
    }

    #[inline(always)]
    fn mode(self) -> u32 {
        self.meta & SEG_META_MODE_MASK
    }

    #[inline(always)]
    fn is_exact(self) -> bool {
        self.mode() == SEG_MODE_EXACT
    }

    #[inline(always)]
    fn exact_bits(self) -> u64 {
        (self.repr_lo as u64) | ((self.repr_hi as u64) << 32)
    }

    #[inline(always)]
    fn anchor_or_const(self) -> u32 {
        self.repr_lo
    }

    #[inline(always)]
    fn const_bit(self) -> bool {
        (self.repr_lo & 1) != 0
    }

    #[inline(always)]
    fn prepend_exact(self, edge: usize) -> Option<Self> {
        if !self.is_exact() || self.len() >= SEG_EXACT_MAX_LEN {
            return None;
        }
        let len = self.len() + 1;
        let bits = ((edge as u64) & 1) | (self.exact_bits() << 1);
        Some(Self::exact(bits, len))
    }

    #[inline(always)]
    fn prefix(self, len: u32) -> Self {
        debug_assert!(len <= self.len());
        match self.mode() {
            SEG_MODE_EXACT => Self::exact(self.exact_bits() & low_bits_mask_u64(len), len),
            SEG_MODE_HISTORY | SEG_MODE_HISTORY_INVERT => Self {
                meta: (self.meta & SEG_META_MODE_MASK) | len,
                ..self
            },
            SEG_MODE_CONST => Self::constant(self.const_bit(), len),
            _ => unreachable!("invalid ctw segment payload mode"),
        }
    }

    #[inline(always)]
    fn suffix_after(self, skip: u32) -> Self {
        debug_assert!(skip <= self.len());
        let new_len = self.len() - skip;
        match self.mode() {
            SEG_MODE_EXACT => Self::exact(self.exact_bits() >> skip, new_len),
            SEG_MODE_HISTORY | SEG_MODE_HISTORY_INVERT => Self {
                repr_lo: self
                    .anchor_or_const()
                    .checked_sub(skip)
                    .expect("ctw history segment anchor underflow"),
                meta: (self.meta & SEG_META_MODE_MASK) | new_len,
                ..self
            },
            SEG_MODE_CONST => Self::constant(self.const_bit(), new_len),
            _ => unreachable!("invalid ctw segment payload mode"),
        }
    }

    #[inline(always)]
    fn from_path<H: HistoryAccess + ?Sized>(history: &H, depth: usize, len: u32) -> Option<Self> {
        if len > SEG_EXACT_MAX_LEN {
            return None;
        }
        Some(Self::exact(
            path_bits_from_history(history, depth, len as usize),
            len,
        ))
    }
}

#[derive(Clone, Copy, Debug, Default)]
struct LevelState {
    symbol_count: [u32; 2],
    log_prob_kt: f64,
    sibling: ChildRef,
}

#[cfg(test)]
#[allow(dead_code)]
#[derive(Clone, Copy, Debug)]
struct PredictEntry {
    symbol_count: [u32; 2],
    log_prob_kt: f64,
    log_prob_weighted: f64,
    sibling_weight: f64,
    has_sibling: bool,
}

#[derive(Clone, Copy, Debug)]
enum Detach {
    NodeChild { node: NodeIndex, edge: usize },
    SegmentNext { segment: SegmentIndex, new_len: u32 },
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum ExistingSource {
    None,
    Node(NodeIndex),
    Segment(SegmentIndex, u32),
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum PreparedEnd {
    MaxDepth,
    MissingAtRoot,
    MissingAfterCurrent,
    MismatchAtCurrentSegment,
}

#[derive(Clone, Copy, Debug, PartialEq)]
struct PreparedStep {
    source: ExistingSource,
    span: u32,
    sibling_weight: f64,
    has_sibling: u8,
}

#[derive(Clone, Copy, Debug)]
/// Compact node payload used by [`CtArena`].
pub struct CtNode {
    children: [ChildRef; 2],
    log_prob_kt: f64,
    log_prob_weighted: f64,
    symbol_count: [u32; 2],
}

#[derive(Clone, Copy, Debug)]
struct CtSegment {
    tail: ChildRef,
    log_prob_kt: f64,
    head_log_prob_weighted: f64,
    symbol_count: [u32; 2],
    payload: SegmentPayload,
}

impl Default for CtSegment {
    fn default() -> Self {
        Self {
            tail: ChildRef::NONE,
            log_prob_kt: 0.0,
            head_log_prob_weighted: 0.0,
            symbol_count: [0, 0],
            payload: SegmentPayload::default(),
        }
    }
}

impl CtSegment {
    #[inline(always)]
    fn len(self) -> u32 {
        self.payload.len()
    }

    #[inline(always)]
    fn set_len(&mut self, len: u32) {
        self.payload.set_len(len);
    }
}

#[inline(always)]
fn low_bits_mask_u64(len: u32) -> u64 {
    if len >= 64 {
        u64::MAX
    } else {
        (1u64 << len) - 1
    }
}

#[inline(always)]
fn path_bits_from_history<H: HistoryAccess + ?Sized>(history: &H, depth: usize, len: usize) -> u64 {
    if let Some(bits) = history.recent_path_bits(depth, len) {
        return bits;
    }
    let history_len = history.len();
    let available = history_len.saturating_sub(depth).min(len);
    if available == 0 {
        return 0;
    }

    let mut bits = 0u64;
    let mut hist_idx = history_len - depth - 1;
    for offset in 0..available {
        bits |= (history.bit(hist_idx) as u64) << offset;
        if hist_idx == 0 {
            break;
        }
        hist_idx -= 1;
    }
    bits
}

#[inline(always)]
fn shift_path_bits(path_bits: u64, consumed: usize) -> u64 {
    if consumed >= 64 {
        0
    } else {
        path_bits >> consumed
    }
}

#[inline(always)]
fn first_exact_segment_mismatch(
    exact_bits: u64,
    path_bits: u64,
    comparable_len: usize,
) -> Option<(usize, bool, bool)> {
    if comparable_len == 0 {
        return None;
    }

    let diff = (exact_bits ^ path_bits) & low_bits_mask_u64(comparable_len as u32);
    if diff == 0 {
        None
    } else {
        let offset = diff.trailing_zeros() as usize;
        Some((
            offset,
            ((path_bits >> offset) & 1) != 0,
            ((exact_bits >> offset) & 1) != 0,
        ))
    }
}

#[inline(always)]
fn predict_ratio_kt(counts: [u32; 2], sym_idx: usize) -> f64 {
    let total = (counts[0] + counts[1]) as f64;
    let sym_count = counts[sym_idx] as f64;
    (sym_count + 0.5) / (total + 1.0)
}

#[inline(always)]
fn predict_ratio_kt_one(counts: [u32; 2]) -> f64 {
    let total = (counts[0] + counts[1]) as f64;
    let sym_count = counts[1] as f64;
    (sym_count + 0.5) / (total + 1.0)
}

#[inline(always)]
fn update_weighted_log_prob_non_leaf(kt_log_prob: f64, log_prob_w0: f64, log_prob_w1: f64) -> f64 {
    let child_log_prob = log_prob_w0 + log_prob_w1;
    if child_log_prob.to_bits() == kt_log_prob.to_bits() {
        return clamp_log_prob(kt_log_prob);
    }
    let delta = child_log_prob - kt_log_prob;
    let log_prob_weighted = if delta >= 0.0 {
        child_log_prob + (-delta).exp().ln_1p() - std::f64::consts::LN_2
    } else {
        kt_log_prob + delta.exp().ln_1p() - std::f64::consts::LN_2
    };
    clamp_log_prob(log_prob_weighted)
}

#[inline(always)]
fn update_weighted_log_prob(
    kt_log_prob: f64,
    log_prob_w0: f64,
    log_prob_w1: f64,
    is_leaf: bool,
) -> f64 {
    if is_leaf {
        clamp_log_prob(kt_log_prob)
    } else {
        update_weighted_log_prob_non_leaf(kt_log_prob, log_prob_w0, log_prob_w1)
    }
}

#[inline(always)]
fn clamp_log_prob(log_prob: f64) -> f64 {
    if log_prob > 1.0e-10 { 0.0 } else { log_prob }
}

#[inline(always)]
fn logsumexp_pair(lhs: f64, rhs: f64) -> f64 {
    if lhs == f64::NEG_INFINITY {
        return rhs;
    }
    if rhs == f64::NEG_INFINITY {
        return lhs;
    }
    let pivot = lhs.max(rhs);
    pivot + ((lhs - pivot).exp() + (rhs - pivot).exp()).ln()
}

#[inline(always)]
fn unary_chain_log_weight(kt_log_prob: f64, continuation_log_prob: f64, len: u32) -> f64 {
    debug_assert!(len > 0);
    if kt_log_prob.to_bits() == continuation_log_prob.to_bits() {
        return kt_log_prob;
    }
    let log_alpha = -(len as f64) * std::f64::consts::LN_2;
    let alpha = log_alpha.exp();
    let log_kt_mass = kt_log_prob + (-alpha).ln_1p();
    let log_cont_mass = continuation_log_prob + log_alpha;
    clamp_log_prob(logsumexp_pair(log_kt_mass, log_cont_mass))
}

#[inline(always)]
fn combined_weight_ratio_internal(
    kt_log_prob: f64,
    counts: [u32; 2],
    path_child_log_prob: f64,
    sibling_log_prob: f64,
    child_ratio: f64,
    sym_idx: usize,
) -> (f64, f64) {
    let kt_ratio = predict_ratio_kt(counts, sym_idx);
    let child_log_prob = path_child_log_prob + sibling_log_prob;
    if child_log_prob.to_bits() == kt_log_prob.to_bits() {
        return (clamp_log_prob(kt_log_prob), 0.5 * (kt_ratio + child_ratio));
    }
    let delta = child_log_prob - kt_log_prob;
    if delta >= 0.0 {
        let x = (-delta).exp();
        (
            clamp_log_prob(child_log_prob + x.ln_1p() - std::f64::consts::LN_2),
            (kt_ratio * x + child_ratio) / (1.0 + x),
        )
    } else {
        let x = delta.exp();
        (
            clamp_log_prob(kt_log_prob + x.ln_1p() - std::f64::consts::LN_2),
            (kt_ratio + x * child_ratio) / (1.0 + x),
        )
    }
}

#[inline(always)]
fn combined_weight_ratio_internal_one(
    kt_log_prob: f64,
    counts: [u32; 2],
    path_child_log_prob: f64,
    sibling_log_prob: f64,
    child_ratio: f64,
) -> (f64, f64) {
    let kt_ratio = predict_ratio_kt_one(counts);
    let child_log_prob = path_child_log_prob + sibling_log_prob;
    if child_log_prob.to_bits() == kt_log_prob.to_bits() {
        return (clamp_log_prob(kt_log_prob), 0.5 * (kt_ratio + child_ratio));
    }
    let delta = child_log_prob - kt_log_prob;
    if delta >= 0.0 {
        let x = (-delta).exp();
        (
            clamp_log_prob(child_log_prob + x.ln_1p() - std::f64::consts::LN_2),
            (kt_ratio * x + child_ratio) / (1.0 + x),
        )
    } else {
        let x = delta.exp();
        (
            clamp_log_prob(kt_log_prob + x.ln_1p() - std::f64::consts::LN_2),
            (kt_ratio + x * child_ratio) / (1.0 + x),
        )
    }
}

#[inline(always)]
fn unary_chain_log_weight_precomputed(
    kt_log_prob: f64,
    continuation_log_prob: f64,
    alpha: f64,
    log_alpha: f64,
    log_one_minus_alpha: f64,
) -> f64 {
    if kt_log_prob.to_bits() == continuation_log_prob.to_bits() {
        return clamp_log_prob(kt_log_prob);
    }

    let delta = continuation_log_prob - kt_log_prob;
    let log_prob_weighted = if delta >= 0.0 {
        let x = ((1.0 - alpha) / alpha) * (-delta).exp();
        continuation_log_prob + log_alpha + x.ln_1p()
    } else {
        let x = (alpha / (1.0 - alpha)) * delta.exp();
        kt_log_prob + log_one_minus_alpha + x.ln_1p()
    };
    clamp_log_prob(log_prob_weighted)
}

#[inline(always)]
// The CTW ratio transform is a hot scalar kernel; grouping these independent
// numeric inputs into a temporary struct would add ceremony without clarifying
// ownership, invariants, or call-site meaning.
#[allow(clippy::too_many_arguments)]
fn unary_chain_ratio_transform_precomputed(
    kt_log_prob: f64,
    counts: [u32; 2],
    continuation_log_prob: f64,
    continuation_ratio: f64,
    alpha: f64,
    log_alpha: f64,
    log_one_minus_alpha: f64,
    sym_idx: usize,
) -> (f64, f64) {
    let kt_ratio = predict_ratio_kt(counts, sym_idx);
    if kt_log_prob.to_bits() == continuation_log_prob.to_bits()
        && kt_ratio.to_bits() == continuation_ratio.to_bits()
    {
        return (clamp_log_prob(kt_log_prob), kt_ratio);
    }
    if kt_log_prob.to_bits() == continuation_log_prob.to_bits() {
        return (
            clamp_log_prob(kt_log_prob),
            (1.0 - alpha) * kt_ratio + alpha * continuation_ratio,
        );
    }

    let delta = continuation_log_prob - kt_log_prob;
    if delta >= 0.0 {
        let x = ((1.0 - alpha) / alpha) * (-delta).exp();
        (
            clamp_log_prob(continuation_log_prob + log_alpha + x.ln_1p()),
            (kt_ratio * x + continuation_ratio) / (1.0 + x),
        )
    } else {
        let x = (alpha / (1.0 - alpha)) * delta.exp();
        (
            clamp_log_prob(kt_log_prob + log_one_minus_alpha + x.ln_1p()),
            (kt_ratio + x * continuation_ratio) / (1.0 + x),
        )
    }
}

#[inline(always)]
fn unary_chain_ratio_transform_precomputed_one(
    kt_log_prob: f64,
    counts: [u32; 2],
    continuation_log_prob: f64,
    continuation_ratio: f64,
    alpha: f64,
    log_alpha: f64,
    log_one_minus_alpha: f64,
) -> (f64, f64) {
    let kt_ratio = predict_ratio_kt_one(counts);
    if kt_log_prob.to_bits() == continuation_log_prob.to_bits()
        && kt_ratio.to_bits() == continuation_ratio.to_bits()
    {
        return (clamp_log_prob(kt_log_prob), kt_ratio);
    }
    if kt_log_prob.to_bits() == continuation_log_prob.to_bits() {
        return (
            clamp_log_prob(kt_log_prob),
            (1.0 - alpha) * kt_ratio + alpha * continuation_ratio,
        );
    }

    let delta = continuation_log_prob - kt_log_prob;
    if delta >= 0.0 {
        let x = ((1.0 - alpha) / alpha) * (-delta).exp();
        (
            clamp_log_prob(continuation_log_prob + log_alpha + x.ln_1p()),
            (kt_ratio * x + continuation_ratio) / (1.0 + x),
        )
    } else {
        let x = (alpha / (1.0 - alpha)) * delta.exp();
        (
            clamp_log_prob(kt_log_prob + log_one_minus_alpha + x.ln_1p()),
            (kt_ratio + x * continuation_ratio) / (1.0 + x),
        )
    }
}

#[inline(always)]
fn predict_ratio_internal(
    kt_log_prob: f64,
    counts: [u32; 2],
    path_child_log_prob: f64,
    sibling_log_prob: f64,
    child_ratio: f64,
    sym_idx: usize,
) -> f64 {
    let kt_ratio = predict_ratio_kt(counts, sym_idx);
    let child_log_prob = path_child_log_prob + sibling_log_prob;
    if child_log_prob.to_bits() == kt_log_prob.to_bits() {
        return 0.5 * (kt_ratio + child_ratio);
    }
    let delta = child_log_prob - kt_log_prob;
    if delta >= 0.0 {
        let inv_rho = (-delta).exp();
        (kt_ratio * inv_rho + child_ratio) / (1.0 + inv_rho)
    } else {
        let rho = delta.exp();
        (kt_ratio + rho * child_ratio) / (1.0 + rho)
    }
}

#[inline(always)]
fn predict_ratio_internal_one(
    kt_log_prob: f64,
    counts: [u32; 2],
    path_child_log_prob: f64,
    sibling_log_prob: f64,
    child_ratio: f64,
) -> f64 {
    let kt_ratio = predict_ratio_kt_one(counts);
    let child_log_prob = path_child_log_prob + sibling_log_prob;
    if child_log_prob.to_bits() == kt_log_prob.to_bits() {
        return 0.5 * (kt_ratio + child_ratio);
    }
    let delta = child_log_prob - kt_log_prob;
    if delta >= 0.0 {
        let inv_rho = (-delta).exp();
        (kt_ratio * inv_rho + child_ratio) / (1.0 + inv_rho)
    } else {
        let rho = delta.exp();
        (kt_ratio + rho * child_ratio) / (1.0 + rho)
    }
}

#[inline(always)]
fn path_edge_at_depth<H: HistoryAccess + ?Sized>(
    history: &H,
    history_len: usize,
    depth: usize,
) -> bool {
    if depth < history_len {
        history.bit(history_len - depth - 1)
    } else {
        false
    }
}

#[inline(always)]
fn segment_edge_from_parts<H: HistoryAccess + ?Sized>(
    segment: CtSegment,
    offset: usize,
    history: &H,
    history_len: usize,
) -> bool {
    match segment.payload.mode() {
        SEG_MODE_EXACT => ((segment.payload.exact_bits() >> offset) & 1) != 0,
        SEG_MODE_HISTORY | SEG_MODE_HISTORY_INVERT => {
            if segment.payload.anchor_or_const() as usize >= offset {
                let hist_idx = segment.payload.anchor_or_const() as usize - offset;
                if hist_idx < history_len {
                    let raw = history.bit(hist_idx);
                    if segment.payload.mode() == SEG_MODE_HISTORY_INVERT {
                        !raw
                    } else {
                        raw
                    }
                } else {
                    segment.payload.mode() == SEG_MODE_HISTORY_INVERT
                }
            } else {
                segment.payload.mode() == SEG_MODE_HISTORY_INVERT
            }
        }
        SEG_MODE_CONST => segment.payload.const_bit(),
        _ => unreachable!("invalid ctw segment payload mode"),
    }
}

#[inline(always)]
fn first_segment_mismatch<H: HistoryAccess + ?Sized>(
    segment: CtSegment,
    depth: usize,
    history: &H,
    comparable_len: usize,
) -> Option<(usize, bool, bool)> {
    if comparable_len == 0 {
        return None;
    }

    match segment.payload.mode() {
        SEG_MODE_EXACT => first_exact_segment_mismatch(
            segment.payload.exact_bits(),
            path_bits_from_history(history, depth, comparable_len),
            comparable_len,
        ),
        SEG_MODE_HISTORY | SEG_MODE_HISTORY_INVERT => {
            let history_len = history.len() as isize;
            let mut path_hist_idx = history_len - depth as isize - 1;
            let mut seg_hist_idx = segment.payload.anchor_or_const() as isize;
            let invert = segment.payload.mode() == SEG_MODE_HISTORY_INVERT;
            for offset in 0..comparable_len {
                let path_edge = history_at_or_zero(history, history_len, path_hist_idx);
                let existing_raw = history_at_or_zero(history, history_len, seg_hist_idx);
                let existing_edge = if invert { !existing_raw } else { existing_raw };
                if existing_edge != path_edge {
                    return Some((offset, path_edge, existing_edge));
                }
                path_hist_idx -= 1;
                seg_hist_idx -= 1;
            }
            None
        }
        SEG_MODE_CONST => {
            let history_len = history.len() as isize;
            let mut path_hist_idx = history_len - depth as isize - 1;
            let existing_edge = segment.payload.const_bit();
            for offset in 0..comparable_len {
                let path_edge = history_at_or_zero(history, history_len, path_hist_idx);
                if existing_edge != path_edge {
                    return Some((offset, path_edge, existing_edge));
                }
                path_hist_idx -= 1;
            }
            None
        }
        _ => unreachable!("invalid ctw segment payload mode"),
    }
}

#[inline]
fn apply_update_to_state_raw<L: CtLogAccess>(
    logs: L,
    symbol_count: &mut [u32; 2],
    log_prob_kt: &mut f64,
    sym_idx: usize,
) {
    let total_before = (symbol_count[0] + symbol_count[1]) as usize;
    let sym_before = symbol_count[sym_idx] as usize;
    debug_assert!(sym_before <= total_before);
    let log_half_before = logs.log_half(sym_before);
    let log_total_after = logs.log_int(total_before + 1);
    *log_prob_kt += log_half_before - log_total_after;
    if *log_prob_kt > 1.0e-10 {
        *log_prob_kt = 0.0;
    }
    symbol_count[sym_idx] = symbol_count[sym_idx]
        .checked_add(1)
        .expect("ctw symbol count overflow");
}

#[inline]
fn apply_revert_to_state_raw<L: CtLogAccess>(
    logs: L,
    symbol_count: &mut [u32; 2],
    log_prob_kt: &mut f64,
    sym_idx: usize,
) {
    let total = (symbol_count[0] + symbol_count[1]) as usize;
    let sym_count = symbol_count[sym_idx] as usize;
    if sym_count > 0 && total > 0 {
        let log_half_before = logs.log_half(sym_count - 1);
        let log_total = logs.log_int(total);
        *log_prob_kt -= log_half_before - log_total;
        symbol_count[sym_idx] -= 1;
    }
    if *log_prob_kt > 1.0e-10 {
        *log_prob_kt = 0.0;
    }
}

#[derive(Clone, Debug)]
/// Arena allocator and storage for CTW nodes/segments.
///
/// This structure owns all backing memory for compressed CTW trees and provides
/// index-based access used by the update/revert engine.
pub struct CtArena {
    nodes: Vec<CtNode>,
    segments: Vec<CtSegment>,
    free_nodes: Vec<NodeIndex>,
    free_segments: Vec<SegmentIndex>,
}

impl CtArena {
    /// Create an empty arena with small default capacities.
    pub fn new() -> Self {
        Self {
            nodes: Vec::with_capacity(1024),
            segments: Vec::with_capacity(1024),
            free_nodes: Vec::new(),
            free_segments: Vec::new(),
        }
    }

    /// Create an empty arena with explicit node-oriented capacity hint.
    pub fn with_capacity(cap: usize) -> Self {
        Self {
            nodes: Vec::with_capacity(cap),
            segments: Vec::with_capacity(cap / 4 + 1),
            free_nodes: Vec::new(),
            free_segments: Vec::new(),
        }
    }

    #[inline]
    /// Reserve additional capacity for upcoming node/segment allocations.
    pub fn reserve_exact(&mut self, additional: usize) {
        self.nodes.reserve_exact(additional);
        self.segments.reserve_exact(additional / 4 + 1);
    }

    #[inline(always)]
    fn reset_node_slot(&mut self, idx: NodeIndex) {
        self.nodes[idx.get()] = CtNode {
            children: [ChildRef::NONE, ChildRef::NONE],
            log_prob_kt: 0.0,
            log_prob_weighted: 0.0,
            symbol_count: [0, 0],
        };
    }

    #[inline(always)]
    fn reset_segment_slot(&mut self, idx: SegmentIndex) {
        self.segments[idx.get()] = CtSegment::default();
    }

    #[inline(always)]
    fn alloc_node(&mut self) -> NodeIndex {
        if let Some(idx) = self.free_nodes.pop() {
            self.reset_node_slot(idx);
            idx
        } else {
            let idx = NodeIndex::from_usize(self.nodes.len());
            self.nodes.push(CtNode {
                children: [ChildRef::NONE, ChildRef::NONE],
                log_prob_kt: 0.0,
                log_prob_weighted: 0.0,
                symbol_count: [0, 0],
            });
            idx
        }
    }

    #[inline(always)]
    fn alloc_node_with_state(&mut self, symbol_count: [u32; 2], log_prob_kt: f64) -> NodeIndex {
        let idx = self.alloc_node();
        self.nodes[idx.get()].symbol_count = symbol_count;
        self.nodes[idx.get()].log_prob_kt = log_prob_kt;
        idx
    }

    #[inline(always)]
    fn free_node(&mut self, idx: NodeIndex) {
        self.free_nodes.push(idx);
    }

    #[inline(always)]
    fn alloc_segment(&mut self) -> SegmentIndex {
        if let Some(idx) = self.free_segments.pop() {
            self.reset_segment_slot(idx);
            idx
        } else {
            let idx = SegmentIndex::from_usize(self.segments.len());
            self.segments.push(CtSegment::default());
            idx
        }
    }

    #[inline(always)]
    fn free_segment(&mut self, idx: SegmentIndex) {
        self.reset_segment_slot(idx);
        self.free_segments.push(idx);
    }

    /// Drop all nodes/segments and free-list bookkeeping.
    pub fn clear(&mut self) {
        self.nodes.clear();
        self.segments.clear();
        self.free_nodes.clear();
        self.free_segments.clear();
    }

    #[inline(always)]
    fn child(&self, parent_idx: NodeIndex, child_idx: usize) -> ChildRef {
        debug_assert!(parent_idx.get() < self.nodes.len());
        debug_assert!(child_idx < 2);
        unsafe {
            *self
                .nodes
                .get_unchecked(parent_idx.get())
                .children
                .get_unchecked(child_idx)
        }
    }

    #[inline(always)]
    fn set_child(&mut self, parent_idx: NodeIndex, child_idx: usize, child: ChildRef) {
        debug_assert!(parent_idx.get() < self.nodes.len());
        debug_assert!(child_idx < 2);
        unsafe {
            *self
                .nodes
                .get_unchecked_mut(parent_idx.get())
                .children
                .get_unchecked_mut(child_idx) = child;
        }
    }

    #[inline(always)]
    fn set_segment_tail(&mut self, segment_idx: SegmentIndex, child: ChildRef) {
        self.segments[segment_idx.get()].tail = child;
    }

    #[inline(always)]
    fn counts(&self, idx: NodeIndex) -> [u32; 2] {
        self.nodes[idx.get()].symbol_count
    }

    #[inline(always)]
    fn visits(&self, idx: NodeIndex) -> u32 {
        let counts = self.nodes[idx.get()].symbol_count;
        counts[0] + counts[1]
    }

    #[inline(always)]
    fn segment_symbol_count(&self, segment_idx: SegmentIndex) -> [u32; 2] {
        self.segments[segment_idx.get()].symbol_count
    }

    #[inline(always)]
    fn segment_log_prob_kt(&self, segment_idx: SegmentIndex) -> f64 {
        self.segments[segment_idx.get()].log_prob_kt
    }

    #[inline(always)]
    fn segment_len(&self, segment_idx: SegmentIndex) -> u32 {
        self.segments[segment_idx.get()].len()
    }

    #[inline(always)]
    fn segment_has_child(&self, segment_idx: SegmentIndex, offset: u32) -> bool {
        let segment = self.segments[segment_idx.get()];
        offset + 1 < segment.len() || segment.tail.is_some()
    }

    #[inline(always)]
    fn log_prob_weighted(&self, idx: NodeIndex) -> f64 {
        self.nodes[idx.get()].log_prob_weighted
    }

    #[inline(always)]
    fn log_prob_kt(&self, idx: NodeIndex) -> f64 {
        self.nodes[idx.get()].log_prob_kt
    }

    #[inline(always)]
    unsafe fn child_ref_weighted_unchecked(&self, child: ChildRef) -> f64 {
        if child.is_none() {
            return 0.0;
        }

        let raw = child.0;
        if (raw & CHILD_SEGMENT_TAG) == 0 {
            debug_assert!((raw as usize) < self.nodes.len());
            self.nodes.get_unchecked(raw as usize).log_prob_weighted
        } else {
            let idx = (raw & CHILD_INDEX_MASK) as usize;
            debug_assert!(idx < self.segments.len());
            self.segments.get_unchecked(idx).head_log_prob_weighted
        }
    }

    #[inline(always)]
    fn child_ref_weighted(&self, child: ChildRef) -> f64 {
        unsafe { self.child_ref_weighted_unchecked(child) }
    }

    #[inline(always)]
    fn singleton_segment_payload(&self, edge: usize) -> SegmentPayload {
        SegmentPayload::exact((edge & 1) as u64, 1)
    }

    #[inline(always)]
    fn segment_edge(
        &self,
        segment_idx: SegmentIndex,
        offset: u32,
        history: &(impl HistoryAccess + ?Sized),
    ) -> usize {
        let segment = self.segments[segment_idx.get()];
        segment_edge_from_parts(segment, offset as usize, history, history.len()) as usize
    }

    fn segment_suffix_weight(&self, segment_idx: SegmentIndex, offset: u32) -> f64 {
        let segment = self.segments[segment_idx.get()];
        if offset >= segment.len() {
            return self.child_ref_weighted(segment.tail);
        }
        if segment.tail.is_none() {
            return segment.log_prob_kt;
        }
        let remaining = segment.len() - offset;
        unary_chain_log_weight(
            segment.log_prob_kt,
            self.child_ref_weighted(segment.tail),
            remaining,
        )
    }

    #[inline(always)]
    fn segment_continuation_weight(&self, segment_idx: SegmentIndex, offset: u32) -> f64 {
        let segment = self.segments[segment_idx.get()];
        if offset + 1 < segment.len() {
            self.segment_suffix_weight(segment_idx, offset + 1)
        } else {
            self.child_ref_weighted(segment.tail)
        }
    }

    fn recompute_segment_head(&mut self, segment_idx: SegmentIndex) {
        let segment = self.segments[segment_idx.get()];
        let head = if segment.tail.is_some() {
            unary_chain_log_weight(
                segment.log_prob_kt,
                self.child_ref_weighted(segment.tail),
                segment.len(),
            )
        } else {
            segment.log_prob_kt
        };
        self.segments[segment_idx.get()].head_log_prob_weighted = head;
    }

    fn recompute_node_weight(&mut self, idx: NodeIndex) {
        let slot = idx.get();
        debug_assert!(slot < self.nodes.len());
        let node = unsafe { *self.nodes.get_unchecked(slot) };
        let [left, right] = node.children;
        let weighted = if left.is_none() && right.is_none() {
            clamp_log_prob(node.log_prob_kt)
        } else {
            let w0 = unsafe { self.child_ref_weighted_unchecked(left) };
            let w1 = unsafe { self.child_ref_weighted_unchecked(right) };
            update_weighted_log_prob_non_leaf(node.log_prob_kt, w0, w1)
        };
        unsafe {
            self.nodes.get_unchecked_mut(slot).log_prob_weighted = weighted;
        }
    }

    fn alloc_segment_with_parts(
        &mut self,
        symbol_count: [u32; 2],
        log_prob_kt: f64,
        tail: ChildRef,
        payload: SegmentPayload,
    ) -> SegmentIndex {
        let segment_idx = self.alloc_segment();
        self.segments[segment_idx.get()] = CtSegment {
            tail,
            log_prob_kt,
            head_log_prob_weighted: 0.0,
            symbol_count,
            payload,
        };
        if payload.len() == 1 && tail.is_none() {
            self.segments[segment_idx.get()].head_log_prob_weighted = log_prob_kt;
        } else {
            self.recompute_segment_head(segment_idx);
        }
        segment_idx
    }

    fn detach_segment_continuation(
        &mut self,
        segment_idx: SegmentIndex,
        offset: u32,
        detaches: &mut Vec<Detach>,
    ) -> ChildRef {
        let segment = self.segments[segment_idx.get()];
        if offset + 1 < segment.len() {
            let suffix = self.alloc_segment_with_parts(
                segment.symbol_count,
                segment.log_prob_kt,
                segment.tail,
                segment.payload.suffix_after(offset + 1),
            );
            detaches.push(Detach::SegmentNext {
                segment: segment_idx,
                new_len: offset + 1,
            });
            ChildRef::from_segment(suffix)
        } else {
            let tail = segment.tail;
            if tail.is_some() {
                detaches.push(Detach::SegmentNext {
                    segment: segment_idx,
                    new_len: segment.len(),
                });
            }
            tail
        }
    }

    // These fields are the exact segment state plus insertion context. Keeping
    // them as scalar arguments avoids building a transient descriptor on this
    // path-compression hot path.
    #[allow(clippy::too_many_arguments)]
    fn prepend_or_alloc_segment(
        &mut self,
        history: &(impl HistoryAccess + ?Sized),
        depth: usize,
        symbol_count: [u32; 2],
        log_prob_kt: f64,
        child: ChildRef,
        edge: usize,
        allow_history_pattern: bool,
    ) -> ChildRef {
        let singleton_payload = self.singleton_segment_payload(edge);

        if let Some(segment_idx) = child.as_segment() {
            let segment = self.segments[segment_idx.get()];
            let same_state = segment.symbol_count == symbol_count
                && segment.log_prob_kt.to_bits() == log_prob_kt.to_bits();
            if same_state && segment.tail == child {
                let segment = &mut self.segments[segment_idx.get()];
                let extended_payload = if segment.payload.is_exact() {
                    segment.payload.prepend_exact(edge)
                } else if allow_history_pattern {
                    let path_payload =
                        SegmentPayload::from_path(history, depth, segment.len().saturating_add(1));
                    path_payload.filter(|payload| {
                        let mut matches = true;
                        for offset in 0..segment.len() as usize {
                            let seg_edge =
                                segment_edge_from_parts(*segment, offset, history, history.len());
                            let payload_edge = ((payload.exact_bits() >> (offset + 1)) & 1) != 0;
                            if seg_edge != payload_edge {
                                matches = false;
                                break;
                            }
                        }
                        matches
                    })
                } else {
                    None
                };
                if let Some(payload) = extended_payload {
                    let old_head = segment.head_log_prob_weighted;
                    segment.payload = payload;
                    segment.head_log_prob_weighted =
                        update_weighted_log_prob(log_prob_kt, old_head, 0.0, false);
                    return ChildRef::from_segment(segment_idx);
                }
            }
        }

        let segment_idx =
            self.alloc_segment_with_parts(symbol_count, log_prob_kt, child, singleton_payload);
        ChildRef::from_segment(segment_idx)
    }

    fn free_child_ref(&mut self, child: ChildRef) {
        let mut stack = Vec::with_capacity(16);
        if child.is_some() {
            stack.push(child);
        }
        while let Some(next) = stack.pop() {
            if let Some(node_idx) = next.as_node() {
                let children = self.nodes[node_idx.get()].children;
                if children[0].is_some() {
                    stack.push(children[0]);
                }
                if children[1].is_some() {
                    stack.push(children[1]);
                }
                self.free_node(node_idx);
            } else if let Some(segment_idx) = next.as_segment() {
                let tail = self.segments[segment_idx.get()].tail;
                if tail.is_some() {
                    stack.push(tail);
                }
                self.free_segment(segment_idx);
            }
        }
    }

    /// Approximate heap usage (bytes) for arena-owned storage.
    pub fn memory_usage(&self) -> usize {
        self.nodes.capacity() * size_of::<CtNode>()
            + self.segments.capacity() * size_of::<CtSegment>()
            + self.free_nodes.capacity() * size_of::<NodeIndex>()
            + self.free_segments.capacity() * size_of::<SegmentIndex>()
    }
}

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

#[derive(Clone)]
struct CtEngine {
    arena: CtArena,
    root: NodeIndex,
    max_depth: usize,
    segment_alpha: Vec<f64>,
    segment_log_alpha: Vec<f64>,
    segment_log_one_minus_alpha: Vec<f64>,
    levels: Vec<LevelState>,
    detaches: Vec<Detach>,
    prepared_steps: Vec<PreparedStep>,
    prepared_levels: usize,
    prepared_end: PreparedEnd,
}

impl CtEngine {
    const RESERVE_MIN_NODES: usize = 4 * 1024;
    const RESERVE_MAX_NODES: usize = 1 << 18;

    fn new(depth: usize) -> Self {
        let mut arena = CtArena::with_capacity(1024.min(1 << depth.min(16)));
        let root = arena.alloc_node();
        let mut segment_alpha = Vec::with_capacity(depth + 1);
        let mut segment_log_alpha = Vec::with_capacity(depth + 1);
        let mut segment_log_one_minus_alpha = Vec::with_capacity(depth + 1);
        segment_alpha.push(1.0);
        segment_log_alpha.push(0.0);
        segment_log_one_minus_alpha.push(f64::NEG_INFINITY);
        let mut alpha = 1.0f64;
        for len in 1..=depth {
            alpha *= 0.5;
            segment_alpha.push(alpha);
            segment_log_alpha.push(-(len as f64) * std::f64::consts::LN_2);
            segment_log_one_minus_alpha.push((-alpha).ln_1p());
        }
        Self {
            arena,
            root,
            max_depth: depth,
            segment_alpha,
            segment_log_alpha,
            segment_log_one_minus_alpha,
            levels: vec![LevelState::default(); depth],
            detaches: Vec::with_capacity(depth),
            prepared_steps: Vec::with_capacity(depth),
            prepared_levels: 0,
            prepared_end: PreparedEnd::MaxDepth,
        }
    }

    #[inline(always)]
    fn root_visits(&self) -> usize {
        self.arena.visits(self.root) as usize
    }

    #[inline(always)]
    fn hot_prefix_depth(&self) -> usize {
        self.max_depth.min(ctw_hot_prefix_depth_limit())
    }

    #[inline(always)]
    fn push_prepared_segment_step(
        &mut self,
        segment_idx: SegmentIndex,
        offset: usize,
        sibling_weight: f64,
        has_sibling: u8,
    ) {
        let span = (offset + 1) as u32;
        self.prepared_steps.push(PreparedStep {
            source: ExistingSource::Segment(segment_idx, offset as u32),
            span,
            sibling_weight,
            has_sibling,
        });
        self.prepared_levels += span as usize;
    }

    #[inline(always)]
    fn walk_prepared_exact_segment(
        &mut self,
        segment_idx: SegmentIndex,
        segment: CtSegment,
        depth: usize,
        path_bits: u64,
    ) -> Option<(usize, ExistingSource)> {
        let seg_len = segment.len() as usize;
        let terminal_offset = self.max_depth.saturating_sub(depth);
        let comparable_len = seg_len.min(terminal_offset);
        if let Some((offset, _, _)) =
            first_exact_segment_mismatch(segment.payload.exact_bits(), path_bits, comparable_len)
        {
            self.push_prepared_segment_step(
                segment_idx,
                offset,
                self.arena
                    .segment_continuation_weight(segment_idx, offset as u32),
                1,
            );
            self.prepared_end = PreparedEnd::MismatchAtCurrentSegment;
            return None;
        }

        if terminal_offset < seg_len {
            self.push_prepared_segment_step(segment_idx, terminal_offset, 0.0, 0);
            return None;
        }

        if segment.tail.is_none() {
            self.push_prepared_segment_step(segment_idx, seg_len - 1, 0.0, 0);
            self.prepared_end = PreparedEnd::MissingAfterCurrent;
            return None;
        }

        self.push_prepared_segment_step(segment_idx, seg_len - 1, 0.0, 0);
        let tail = segment.tail;
        Some((
            depth + seg_len,
            Self::child_to_existing_source(tail).unwrap_or(ExistingSource::None),
        ))
    }

    fn clear(&mut self) {
        self.arena.clear();
        self.root = self.arena.alloc_node();
        self.levels.fill(LevelState::default());
        self.detaches.clear();
        self.prepared_steps.clear();
        self.prepared_levels = 0;
        self.prepared_end = PreparedEnd::MaxDepth;
    }

    #[inline]
    fn reserve_for_symbols(&mut self, total_symbols: usize) {
        if total_symbols == 0 {
            return;
        }

        let depth_scale = self.max_depth.saturating_add(1);
        let reserve_nodes = total_symbols
            .saturating_div(depth_scale)
            .clamp(Self::RESERVE_MIN_NODES, Self::RESERVE_MAX_NODES);
        let free_nodes = self
            .arena
            .nodes
            .capacity()
            .saturating_sub(self.arena.nodes.len());
        if reserve_nodes > free_nodes {
            self.arena.reserve_exact(reserve_nodes - free_nodes);
        }
    }

    #[inline]
    fn get_log_block_probability(&self) -> f64 {
        self.arena.log_prob_weighted(self.root)
    }

    #[inline]
    fn with_cached_logs<R>(
        &mut self,
        upto: usize,
        f: impl FnOnce(&mut Self, CachedLogs<'_>) -> R,
    ) -> R {
        with_shared_cached_logs(upto, |logs| f(self, logs))
    }

    #[inline]
    fn with_bounded_logs<R>(
        &mut self,
        upto: usize,
        f: impl FnOnce(&mut Self, BoundedLogs<'_>) -> R,
    ) -> R {
        with_shared_bounded_logs(upto, |logs| f(self, logs))
    }

    #[inline]
    fn log_cache_memory_usage(&self) -> usize {
        shared_log_cache_memory_usage()
    }

    #[inline(always)]
    fn segment_constants(&self, len: u32) -> (f64, f64, f64) {
        let idx = len as usize;
        debug_assert!(idx < self.segment_alpha.len());
        debug_assert!(idx < self.segment_log_alpha.len());
        debug_assert!(idx < self.segment_log_one_minus_alpha.len());
        (
            unsafe { *self.segment_alpha.get_unchecked(idx) },
            unsafe { *self.segment_log_alpha.get_unchecked(idx) },
            unsafe { *self.segment_log_one_minus_alpha.get_unchecked(idx) },
        )
    }

    #[inline(always)]
    fn source_counts_and_kt_log_prob(&self, source: ExistingSource) -> ([u32; 2], f64) {
        match source {
            ExistingSource::Node(node_idx) => {
                let slot = node_idx.get();
                let node = unsafe { *self.arena.nodes.get_unchecked(slot) };
                (node.symbol_count, node.log_prob_kt)
            }
            ExistingSource::Segment(segment_idx, _) => {
                let slot = segment_idx.get();
                let segment = unsafe { *self.arena.segments.get_unchecked(slot) };
                (segment.symbol_count, segment.log_prob_kt)
            }
            ExistingSource::None => unreachable!("prepared step should never store None"),
        }
    }

    fn build_missing_segment_path(
        &mut self,
        depth: usize,
        history: &(impl HistoryAccess + ?Sized),
        sym_idx: usize,
        singleton_log_prob_kt: f64,
    ) -> ChildRef {
        if depth > self.max_depth {
            return ChildRef::NONE;
        }

        let mut counts = [0u32; 2];
        counts[sym_idx] = 1;
        let log_prob_kt = singleton_log_prob_kt;
        let total_len = self.max_depth - depth + 1;

        if let Some(payload) = SegmentPayload::from_path(history, depth, total_len as u32) {
            let segment =
                self.arena
                    .alloc_segment_with_parts(counts, log_prob_kt, ChildRef::NONE, payload);
            return ChildRef::from_segment(segment);
        }

        let history_nodes = if depth < history.len() {
            (self.max_depth.min(history.len() - 1) - depth) + 1
        } else {
            0
        };
        let const_nodes = total_len - history_nodes;

        let mut built = ChildRef::NONE;
        if const_nodes > 0 {
            let const_segment = self.arena.alloc_segment_with_parts(
                counts,
                log_prob_kt,
                ChildRef::NONE,
                SegmentPayload::constant(false, const_nodes as u32),
            );
            built = ChildRef::from_segment(const_segment);
        }
        if history_nodes > 0 {
            let history_segment = self.arena.alloc_segment_with_parts(
                counts,
                log_prob_kt,
                built,
                SegmentPayload::history(
                    (history.len() - depth - 1) as u32,
                    history_nodes as u32,
                    false,
                ),
            );
            built = ChildRef::from_segment(history_segment);
        }
        built
    }

    fn build_missing_path(
        &mut self,
        depth: usize,
        history: &(impl HistoryAccess + ?Sized),
        sym_idx: usize,
        singleton_log_prob_kt: f64,
    ) -> ChildRef {
        if depth > self.max_depth {
            return ChildRef::NONE;
        }

        let hot_prefix_depth = self.hot_prefix_depth();
        if depth > hot_prefix_depth {
            return self.build_missing_segment_path(depth, history, sym_idx, singleton_log_prob_kt);
        }

        let mut counts = [0u32; 2];
        counts[sym_idx] = 1;
        let mut built = if hot_prefix_depth < self.max_depth {
            self.build_missing_segment_path(
                hot_prefix_depth + 1,
                history,
                sym_idx,
                singleton_log_prob_kt,
            )
        } else {
            ChildRef::NONE
        };

        for node_depth in (depth..=hot_prefix_depth).rev() {
            let node = self
                .arena
                .alloc_node_with_state(counts, singleton_log_prob_kt);
            if node_depth < self.max_depth {
                let edge = history_symbol(history, node_depth) as usize;
                self.arena.set_child(node, edge, built);
            }
            self.arena.recompute_node_weight(node);
            built = ChildRef::from_node(node);
        }
        built
    }

    #[inline(always)]
    fn build_missing_segment_path_exact_bits(
        &mut self,
        depth: usize,
        path_bits: u64,
        sym_idx: usize,
        singleton_log_prob_kt: f64,
    ) -> ChildRef {
        debug_assert!(self.max_depth <= SEG_EXACT_MAX_LEN as usize);
        if depth > self.max_depth {
            return ChildRef::NONE;
        }

        let mut counts = [0u32; 2];
        counts[sym_idx] = 1;
        let total_len = self.max_depth - depth + 1;
        let payload = SegmentPayload::exact(
            path_bits & low_bits_mask_u64(total_len as u32),
            total_len as u32,
        );
        let segment = self.arena.alloc_segment_with_parts(
            counts,
            singleton_log_prob_kt,
            ChildRef::NONE,
            payload,
        );
        ChildRef::from_segment(segment)
    }

    #[inline(always)]
    fn build_missing_path_exact_bits(
        &mut self,
        depth: usize,
        path_bits: u64,
        sym_idx: usize,
        singleton_log_prob_kt: f64,
    ) -> ChildRef {
        debug_assert!(self.max_depth <= SEG_EXACT_MAX_LEN as usize);
        if depth > self.max_depth {
            return ChildRef::NONE;
        }

        let hot_prefix_depth = self.hot_prefix_depth();
        if depth > hot_prefix_depth {
            return self.build_missing_segment_path_exact_bits(
                depth,
                path_bits,
                sym_idx,
                singleton_log_prob_kt,
            );
        }

        let mut counts = [0u32; 2];
        counts[sym_idx] = 1;
        let mut built = if hot_prefix_depth < self.max_depth {
            self.build_missing_segment_path_exact_bits(
                hot_prefix_depth + 1,
                shift_path_bits(path_bits, hot_prefix_depth + 1 - depth),
                sym_idx,
                singleton_log_prob_kt,
            )
        } else {
            ChildRef::NONE
        };

        for node_depth in (depth..=hot_prefix_depth).rev() {
            let node = self
                .arena
                .alloc_node_with_state(counts, singleton_log_prob_kt);
            if node_depth < self.max_depth {
                let edge = ((path_bits >> (node_depth - depth)) & 1) as usize;
                self.arena.set_child(node, edge, built);
            }
            self.arena.recompute_node_weight(node);
            built = ChildRef::from_node(node);
        }
        built
    }

    #[inline(always)]
    fn child_to_existing_source(child: ChildRef) -> Option<ExistingSource> {
        if let Some(node) = child.as_node() {
            Some(ExistingSource::Node(node))
        } else {
            child
                .as_segment()
                .map(|segment| ExistingSource::Segment(segment, 0))
        }
    }

    #[inline(always)]
    fn update_source_state<L: CtLogAccess>(
        &mut self,
        logs: L,
        source: ExistingSource,
        sym_idx: usize,
    ) {
        match source {
            ExistingSource::Node(node_idx) => {
                let slot = node_idx.get();
                let mut counts = self.arena.nodes[slot].symbol_count;
                let mut log_prob_kt = self.arena.nodes[slot].log_prob_kt;
                apply_update_to_state_raw(logs, &mut counts, &mut log_prob_kt, sym_idx);
                self.arena.nodes[slot].symbol_count = counts;
                self.arena.nodes[slot].log_prob_kt = log_prob_kt;
            }
            ExistingSource::Segment(segment_idx, _) => {
                let slot = segment_idx.get();
                let mut counts = self.arena.segments[slot].symbol_count;
                let mut log_prob_kt = self.arena.segments[slot].log_prob_kt;
                apply_update_to_state_raw(logs, &mut counts, &mut log_prob_kt, sym_idx);
                self.arena.segments[slot].symbol_count = counts;
                self.arena.segments[slot].log_prob_kt = log_prob_kt;
            }
            ExistingSource::None => unreachable!("prepared update should never visit None"),
        }
    }

    #[inline(always)]
    fn recompute_source_weight(&mut self, source: ExistingSource) {
        match source {
            ExistingSource::Node(node_idx) => self.arena.recompute_node_weight(node_idx),
            ExistingSource::Segment(segment_idx, _) => self.recompute_segment_head(segment_idx),
            ExistingSource::None => unreachable!("prepared update should never visit None"),
        }
    }

    #[inline(always)]
    fn recompute_segment_head(&mut self, segment_idx: SegmentIndex) {
        let segment = self.arena.segments[segment_idx.get()];
        let head = if segment.tail.is_some() {
            let (alpha, log_alpha, log_one_minus_alpha) = self.segment_constants(segment.len());
            unary_chain_log_weight_precomputed(
                segment.log_prob_kt,
                self.arena.child_ref_weighted(segment.tail),
                alpha,
                log_alpha,
                log_one_minus_alpha,
            )
        } else {
            segment.log_prob_kt
        };
        self.arena.segments[segment_idx.get()].head_log_prob_weighted = head;
    }

    fn attach_missing_after_prepared_path(
        &mut self,
        history: &(impl HistoryAccess + ?Sized),
        sym_idx: usize,
        singleton_log_prob_kt: f64,
    ) {
        let Some(last_step) = self.prepared_steps.last().copied() else {
            return;
        };
        let depth = self.prepared_levels;
        match last_step.source {
            ExistingSource::Node(node_idx) => {
                debug_assert!(depth < self.max_depth);
                let path_edge = history_symbol(history, depth) as usize;
                debug_assert!(self.arena.child(node_idx, path_edge).is_none());
                let new_child =
                    self.build_missing_path(depth + 1, history, sym_idx, singleton_log_prob_kt);
                self.arena.set_child(node_idx, path_edge, new_child);
            }
            ExistingSource::Segment(segment_idx, offset) => {
                debug_assert!(depth < self.max_depth);
                debug_assert_eq!(offset + 1, self.arena.segment_len(segment_idx));
                debug_assert!(self.arena.segments[segment_idx.get()].tail.is_none());
                let new_tail =
                    self.build_missing_path(depth + 1, history, sym_idx, singleton_log_prob_kt);
                self.arena.set_segment_tail(segment_idx, new_tail);
            }
            ExistingSource::None => unreachable!("prepared path should never end in None source"),
        }
    }

    fn replace_prepared_child(
        &mut self,
        history: &(impl HistoryAccess + ?Sized),
        step_index: usize,
        current_start_depth: usize,
        new_child: ChildRef,
    ) {
        if step_index == 0 {
            let root_edge = history_symbol(history, 0) as usize;
            self.arena.set_child(self.root, root_edge, new_child);
            return;
        }

        match self.prepared_steps[step_index - 1].source {
            ExistingSource::Node(node_idx) => {
                let edge = history_symbol(history, current_start_depth - 1) as usize;
                self.arena.set_child(node_idx, edge, new_child);
            }
            ExistingSource::Segment(segment_idx, offset) => {
                debug_assert_eq!(offset + 1, self.arena.segment_len(segment_idx));
                self.arena.set_segment_tail(segment_idx, new_child);
            }
            ExistingSource::None => unreachable!("prepared path should never parent from None"),
        }
    }

    fn update_prepared_mismatch<L: CtLogAccess>(
        &mut self,
        logs: L,
        history: &(impl HistoryAccess + ?Sized),
        sym_idx: usize,
        singleton_log_prob_kt: f64,
    ) -> ChildRef {
        let last_index = self.prepared_steps.len() - 1;
        for idx in 0..last_index {
            self.update_source_state(logs, self.prepared_steps[idx].source, sym_idx);
        }

        let last_step = self.prepared_steps[last_index];
        let ExistingSource::Segment(segment_idx, offset_u32) = last_step.source else {
            unreachable!("prepared segment mismatch must end at a segment");
        };

        let original = self.arena.segments[segment_idx.get()];
        let offset = offset_u32 as usize;
        let seg_len = original.len() as usize;
        let history_len = history.len();
        let current_start_depth = self.prepared_levels - last_step.span as usize + 1;
        let node_depth = current_start_depth + offset;
        let path_edge = path_edge_at_depth(history, history_len, node_depth);
        let existing_edge = segment_edge_from_parts(original, offset, history, history_len);
        debug_assert_ne!(path_edge, existing_edge);

        let old_continuation = if offset + 1 < seg_len {
            if offset == 0 {
                let segment = &mut self.arena.segments[segment_idx.get()];
                segment.payload = original.payload.suffix_after(1);
                segment.tail = original.tail;
                segment.symbol_count = original.symbol_count;
                segment.log_prob_kt = original.log_prob_kt;
                self.recompute_segment_head(segment_idx);
                ChildRef::from_segment(segment_idx)
            } else {
                ChildRef::from_segment(self.arena.alloc_segment_with_parts(
                    original.symbol_count,
                    original.log_prob_kt,
                    original.tail,
                    original.payload.suffix_after(offset as u32 + 1),
                ))
            }
        } else {
            original.tail
        };

        let new_tail =
            self.build_missing_path(node_depth + 1, history, sym_idx, singleton_log_prob_kt);
        let mut updated_counts = original.symbol_count;
        let mut updated_log_prob_kt = original.log_prob_kt;
        apply_update_to_state_raw(logs, &mut updated_counts, &mut updated_log_prob_kt, sym_idx);

        let branch = self
            .arena
            .alloc_node_with_state(updated_counts, updated_log_prob_kt);
        self.arena
            .set_child(branch, existing_edge as usize, old_continuation);
        self.arena.set_child(branch, path_edge as usize, new_tail);
        self.arena.recompute_node_weight(branch);

        if offset == 0 {
            if seg_len == 1 {
                self.arena.free_segment(segment_idx);
            }
            self.replace_prepared_child(
                history,
                last_index,
                current_start_depth,
                ChildRef::from_node(branch),
            );
        } else {
            let segment = &mut self.arena.segments[segment_idx.get()];
            segment.payload = original.payload.prefix(offset as u32);
            segment.tail = ChildRef::from_node(branch);
            segment.symbol_count = updated_counts;
            segment.log_prob_kt = updated_log_prob_kt;
            self.recompute_segment_head(segment_idx);
        }

        for idx in (0..last_index).rev() {
            self.recompute_source_weight(self.prepared_steps[idx].source);
        }

        let root_edge = history_symbol(history, 0) as usize;
        self.arena.child(self.root, root_edge)
    }

    fn update_prepared_cached_path<L: CtLogAccess>(
        &mut self,
        logs: L,
        history: &(impl HistoryAccess + ?Sized),
        sym_idx: usize,
        singleton_log_prob_kt: f64,
    ) {
        debug_assert!(!self.prepared_steps.is_empty());
        debug_assert!(matches!(
            self.prepared_end,
            PreparedEnd::MaxDepth | PreparedEnd::MissingAfterCurrent
        ));

        if self.prepared_end == PreparedEnd::MissingAfterCurrent {
            self.attach_missing_after_prepared_path(history, sym_idx, singleton_log_prob_kt);
        }

        let last_index = self.prepared_steps.len() - 1;
        let mut child_weight = if self.prepared_end == PreparedEnd::MissingAfterCurrent {
            let last_step = self.prepared_steps[last_index];
            match last_step.source {
                ExistingSource::Node(node_idx) => {
                    let depth = self.prepared_levels;
                    let edge = history_symbol(history, depth) as usize;
                    self.arena
                        .child_ref_weighted(self.arena.child(node_idx, edge))
                }
                ExistingSource::Segment(segment_idx, offset) => {
                    debug_assert_eq!(offset + 1, self.arena.segment_len(segment_idx));
                    self.arena
                        .child_ref_weighted(self.arena.segments[segment_idx.get()].tail)
                }
                ExistingSource::None => unreachable!("prepared path should never end in None"),
            }
        } else {
            0.0
        };

        for idx in (0..=last_index).rev() {
            let step = self.prepared_steps[idx];
            match step.source {
                ExistingSource::Node(node_idx) => {
                    let (mut counts, mut log_prob_kt) =
                        self.source_counts_and_kt_log_prob(step.source);
                    apply_update_to_state_raw(logs, &mut counts, &mut log_prob_kt, sym_idx);
                    let weighted =
                        if idx == last_index && self.prepared_end == PreparedEnd::MaxDepth {
                            debug_assert_eq!(step.has_sibling, 0);
                            clamp_log_prob(log_prob_kt)
                        } else {
                            update_weighted_log_prob(
                                log_prob_kt,
                                child_weight,
                                step.sibling_weight,
                                false,
                            )
                        };
                    let slot = node_idx.get();
                    self.arena.nodes[slot].symbol_count = counts;
                    self.arena.nodes[slot].log_prob_kt = log_prob_kt;
                    self.arena.nodes[slot].log_prob_weighted = weighted;
                    child_weight = weighted;
                }
                ExistingSource::Segment(segment_idx, offset) => {
                    let (mut counts, mut log_prob_kt) =
                        self.source_counts_and_kt_log_prob(step.source);
                    apply_update_to_state_raw(logs, &mut counts, &mut log_prob_kt, sym_idx);
                    let slot = segment_idx.get();
                    let weighted =
                        if idx == last_index && self.prepared_end == PreparedEnd::MaxDepth {
                            debug_assert_eq!(offset + 1, self.arena.segment_len(segment_idx));
                            debug_assert!(self.arena.segments[slot].tail.is_none());
                            clamp_log_prob(log_prob_kt)
                        } else {
                            let (alpha, log_alpha, log_one_minus_alpha) =
                                self.segment_constants(step.span);
                            unary_chain_log_weight_precomputed(
                                log_prob_kt,
                                child_weight,
                                alpha,
                                log_alpha,
                                log_one_minus_alpha,
                            )
                        };
                    self.arena.segments[slot].symbol_count = counts;
                    self.arena.segments[slot].log_prob_kt = log_prob_kt;
                    self.arena.segments[slot].head_log_prob_weighted = weighted;
                    child_weight = weighted;
                }
                ExistingSource::None => unreachable!("prepared update should never visit None"),
            }
        }
    }

    fn update_child_fast<L: CtLogAccess>(
        &mut self,
        logs: L,
        child: ChildRef,
        depth: usize,
        history: &(impl HistoryAccess + ?Sized),
        sym_idx: usize,
        singleton_log_prob_kt: f64,
    ) -> ChildRef {
        if depth > self.max_depth {
            return child;
        }
        if child.is_none() {
            return self.build_missing_path(depth, history, sym_idx, singleton_log_prob_kt);
        }

        if let Some(node_idx) = child.as_node() {
            if depth < self.max_depth {
                let path_edge = history_symbol(history, depth) as usize;
                let next = self.arena.child(node_idx, path_edge);
                let updated = self.update_child_fast(
                    logs,
                    next,
                    depth + 1,
                    history,
                    sym_idx,
                    singleton_log_prob_kt,
                );
                if updated != next {
                    self.arena.set_child(node_idx, path_edge, updated);
                }
            }
            let mut counts = self.arena.nodes[node_idx.get()].symbol_count;
            let mut log_prob_kt = self.arena.nodes[node_idx.get()].log_prob_kt;
            apply_update_to_state_raw(logs, &mut counts, &mut log_prob_kt, sym_idx);
            self.arena.nodes[node_idx.get()].symbol_count = counts;
            self.arena.nodes[node_idx.get()].log_prob_kt = log_prob_kt;
            self.arena.recompute_node_weight(node_idx);
            return ChildRef::from_node(node_idx);
        }

        let segment_idx = child.as_segment().unwrap();
        let original = self.arena.segments[segment_idx.get()];
        let seg_len = original.len() as usize;
        let mut updated_counts = original.symbol_count;
        let mut updated_log_prob_kt = original.log_prob_kt;
        apply_update_to_state_raw(logs, &mut updated_counts, &mut updated_log_prob_kt, sym_idx);

        let depth_budget = self.max_depth.saturating_sub(depth);
        let comparable_len = if original.tail.is_none() {
            seg_len.saturating_sub(1)
        } else {
            seg_len
        }
        .min(depth_budget);
        let mismatch = first_segment_mismatch(original, depth, history, comparable_len).map(
            |(offset, path_edge, existing_edge)| (offset, depth + offset, path_edge, existing_edge),
        );

        if let Some((offset, node_depth, path_edge, existing_edge)) = mismatch {
            let old_continuation = if offset + 1 < seg_len {
                if offset == 0 {
                    let segment = &mut self.arena.segments[segment_idx.get()];
                    segment.payload = original.payload.suffix_after(1);
                    segment.tail = original.tail;
                    segment.symbol_count = original.symbol_count;
                    segment.log_prob_kt = original.log_prob_kt;
                    self.recompute_segment_head(segment_idx);
                    ChildRef::from_segment(segment_idx)
                } else {
                    ChildRef::from_segment(self.arena.alloc_segment_with_parts(
                        original.symbol_count,
                        original.log_prob_kt,
                        original.tail,
                        original.payload.suffix_after(offset as u32 + 1),
                    ))
                }
            } else {
                original.tail
            };

            let new_tail =
                self.build_missing_path(node_depth + 1, history, sym_idx, singleton_log_prob_kt);
            let branch = self
                .arena
                .alloc_node_with_state(updated_counts, updated_log_prob_kt);
            self.arena
                .set_child(branch, existing_edge as usize, old_continuation);
            self.arena.set_child(branch, path_edge as usize, new_tail);
            self.arena.recompute_node_weight(branch);

            if offset == 0 {
                if offset + 1 >= seg_len {
                    self.arena.free_segment(segment_idx);
                }
                return ChildRef::from_node(branch);
            }

            let segment = &mut self.arena.segments[segment_idx.get()];
            segment.payload = original.payload.prefix(offset as u32);
            segment.tail = ChildRef::from_node(branch);
            segment.symbol_count = updated_counts;
            segment.log_prob_kt = updated_log_prob_kt;
            self.recompute_segment_head(segment_idx);
            return ChildRef::from_segment(segment_idx);
        }

        if depth_budget < seg_len {
            self.arena.segments[segment_idx.get()].symbol_count = updated_counts;
            self.arena.segments[segment_idx.get()].log_prob_kt = updated_log_prob_kt;
            self.recompute_segment_head(segment_idx);
            return ChildRef::from_segment(segment_idx);
        }

        if original.tail.is_none() {
            let new_tail =
                self.build_missing_path(depth + seg_len, history, sym_idx, singleton_log_prob_kt);
            self.arena.segments[segment_idx.get()].tail = new_tail;
            self.arena.segments[segment_idx.get()].symbol_count = updated_counts;
            self.arena.segments[segment_idx.get()].log_prob_kt = updated_log_prob_kt;
            self.recompute_segment_head(segment_idx);
            return ChildRef::from_segment(segment_idx);
        }

        let tail = original.tail;
        let updated_tail = self.update_child_fast(
            logs,
            tail,
            depth + seg_len,
            history,
            sym_idx,
            singleton_log_prob_kt,
        );
        if updated_tail != tail {
            self.arena.set_segment_tail(segment_idx, updated_tail);
        }
        self.arena.segments[segment_idx.get()].symbol_count = updated_counts;
        self.arena.segments[segment_idx.get()].log_prob_kt = updated_log_prob_kt;
        self.recompute_segment_head(segment_idx);
        ChildRef::from_segment(segment_idx)
    }

    // Exact-mode updates thread together the log table, path state, and KT
    // singleton value; a wrapper would only hide the data dependencies in this
    // inner CTW update kernel.
    #[allow(clippy::too_many_arguments)]
    fn update_child_fast_exact<L: CtLogAccess>(
        &mut self,
        logs: L,
        child: ChildRef,
        depth: usize,
        history: &(impl HistoryAccess + ?Sized),
        path_bits: u64,
        sym_idx: usize,
        singleton_log_prob_kt: f64,
    ) -> ChildRef {
        debug_assert!(self.max_depth <= SEG_EXACT_MAX_LEN as usize);
        if depth > self.max_depth {
            return child;
        }
        if child.is_none() {
            return self.build_missing_path_exact_bits(
                depth,
                path_bits,
                sym_idx,
                singleton_log_prob_kt,
            );
        }

        if let Some(node_idx) = child.as_node() {
            if depth < self.max_depth {
                let path_edge = (path_bits & 1) as usize;
                let next = self.arena.child(node_idx, path_edge);
                let updated = self.update_child_fast_exact(
                    logs,
                    next,
                    depth + 1,
                    history,
                    shift_path_bits(path_bits, 1),
                    sym_idx,
                    singleton_log_prob_kt,
                );
                if updated != next {
                    self.arena.set_child(node_idx, path_edge, updated);
                }
            }
            let slot = node_idx.get();
            let mut counts = self.arena.nodes[slot].symbol_count;
            let mut log_prob_kt = self.arena.nodes[slot].log_prob_kt;
            apply_update_to_state_raw(logs, &mut counts, &mut log_prob_kt, sym_idx);
            let [left, right] = self.arena.nodes[slot].children;
            let weighted = if left.is_none() && right.is_none() {
                clamp_log_prob(log_prob_kt)
            } else {
                // Safety: `left`/`right` come from this node's stored children, so any
                // non-none child index is arena-owned and in-bounds for this arena.
                let w0 = unsafe { self.arena.child_ref_weighted_unchecked(left) };
                let w1 = unsafe { self.arena.child_ref_weighted_unchecked(right) };
                update_weighted_log_prob_non_leaf(log_prob_kt, w0, w1)
            };
            self.arena.nodes[slot].symbol_count = counts;
            self.arena.nodes[slot].log_prob_kt = log_prob_kt;
            self.arena.nodes[slot].log_prob_weighted = weighted;
            return ChildRef::from_node(node_idx);
        }

        let segment_idx = child.as_segment().unwrap();
        let original = self.arena.segments[segment_idx.get()];
        if !original.payload.is_exact() {
            return self.update_child_fast(
                logs,
                child,
                depth,
                history,
                sym_idx,
                singleton_log_prob_kt,
            );
        }

        let seg_len = original.len() as usize;
        let mut updated_counts = original.symbol_count;
        let mut updated_log_prob_kt = original.log_prob_kt;
        apply_update_to_state_raw(logs, &mut updated_counts, &mut updated_log_prob_kt, sym_idx);

        let depth_budget = self.max_depth.saturating_sub(depth);
        let comparable_len = if original.tail.is_none() {
            seg_len.saturating_sub(1)
        } else {
            seg_len
        }
        .min(depth_budget);
        let mismatch =
            first_exact_segment_mismatch(original.payload.exact_bits(), path_bits, comparable_len)
                .map(|(offset, path_edge, existing_edge)| {
                    (offset, depth + offset, path_edge, existing_edge)
                });

        if let Some((offset, node_depth, path_edge, existing_edge)) = mismatch {
            let old_continuation = if offset + 1 < seg_len {
                if offset == 0 {
                    let segment = &mut self.arena.segments[segment_idx.get()];
                    segment.payload = original.payload.suffix_after(1);
                    segment.tail = original.tail;
                    segment.symbol_count = original.symbol_count;
                    segment.log_prob_kt = original.log_prob_kt;
                    self.recompute_segment_head(segment_idx);
                    ChildRef::from_segment(segment_idx)
                } else {
                    ChildRef::from_segment(self.arena.alloc_segment_with_parts(
                        original.symbol_count,
                        original.log_prob_kt,
                        original.tail,
                        original.payload.suffix_after(offset as u32 + 1),
                    ))
                }
            } else {
                original.tail
            };

            let new_tail = self.build_missing_path_exact_bits(
                node_depth + 1,
                shift_path_bits(path_bits, offset + 1),
                sym_idx,
                singleton_log_prob_kt,
            );
            let branch = self
                .arena
                .alloc_node_with_state(updated_counts, updated_log_prob_kt);
            self.arena
                .set_child(branch, existing_edge as usize, old_continuation);
            self.arena.set_child(branch, path_edge as usize, new_tail);
            self.arena.recompute_node_weight(branch);

            if offset == 0 {
                if offset + 1 >= seg_len {
                    self.arena.free_segment(segment_idx);
                }
                return ChildRef::from_node(branch);
            }

            let segment = &mut self.arena.segments[segment_idx.get()];
            segment.payload = original.payload.prefix(offset as u32);
            segment.tail = ChildRef::from_node(branch);
            segment.symbol_count = updated_counts;
            segment.log_prob_kt = updated_log_prob_kt;
            self.recompute_segment_head(segment_idx);
            return ChildRef::from_segment(segment_idx);
        }

        if depth_budget < seg_len {
            self.arena.segments[segment_idx.get()].symbol_count = updated_counts;
            self.arena.segments[segment_idx.get()].log_prob_kt = updated_log_prob_kt;
            self.recompute_segment_head(segment_idx);
            return ChildRef::from_segment(segment_idx);
        }

        if original.tail.is_none() {
            let new_tail = self.build_missing_path_exact_bits(
                depth + seg_len,
                shift_path_bits(path_bits, seg_len),
                sym_idx,
                singleton_log_prob_kt,
            );
            self.arena.segments[segment_idx.get()].tail = new_tail;
            self.arena.segments[segment_idx.get()].symbol_count = updated_counts;
            self.arena.segments[segment_idx.get()].log_prob_kt = updated_log_prob_kt;
            self.recompute_segment_head(segment_idx);
            return ChildRef::from_segment(segment_idx);
        }

        let tail = original.tail;
        let updated_tail = self.update_child_fast_exact(
            logs,
            tail,
            depth + seg_len,
            history,
            shift_path_bits(path_bits, seg_len),
            sym_idx,
            singleton_log_prob_kt,
        );
        if updated_tail != tail {
            self.arena.set_segment_tail(segment_idx, updated_tail);
        }
        self.arena.segments[segment_idx.get()].symbol_count = updated_counts;
        self.arena.segments[segment_idx.get()].log_prob_kt = updated_log_prob_kt;
        self.recompute_segment_head(segment_idx);
        ChildRef::from_segment(segment_idx)
    }

    #[inline(always)]
    fn update_root_child<L: CtLogAccess>(
        &mut self,
        logs: L,
        child: ChildRef,
        history: &(impl HistoryAccess + ?Sized),
        sym_idx: usize,
        singleton_log_prob_kt: f64,
    ) -> ChildRef {
        if self.max_depth <= SEG_EXACT_MAX_LEN as usize {
            let path_bits = path_bits_from_history(history, 1, self.max_depth);
            self.update_child_fast_exact(
                logs,
                child,
                1,
                history,
                path_bits,
                sym_idx,
                singleton_log_prob_kt,
            )
        } else {
            self.update_child_fast(logs, child, 1, history, sym_idx, singleton_log_prob_kt)
        }
    }

    fn collect_existing_levels(&mut self, history: &(impl HistoryAccess + ?Sized)) -> ChildRef {
        if self.max_depth == 0 {
            self.detaches.clear();
            return ChildRef::NONE;
        }

        self.detaches.clear();
        self.levels.fill(LevelState::default());

        let root_edge = history_symbol(history, 0) as usize;
        let old_child = self.arena.child(self.root, root_edge);
        let mut source = if let Some(node) = old_child.as_node() {
            ExistingSource::Node(node)
        } else if let Some(segment) = old_child.as_segment() {
            ExistingSource::Segment(segment, 0)
        } else {
            ExistingSource::None
        };

        for depth in 1..=self.max_depth {
            let slot = depth - 1;
            self.levels[slot] = LevelState::default();

            match source {
                ExistingSource::None => {}
                ExistingSource::Node(node_idx) => {
                    self.levels[slot].symbol_count = self.arena.counts(node_idx);
                    self.levels[slot].log_prob_kt = self.arena.log_prob_kt(node_idx);
                    if depth < self.max_depth {
                        let path_edge = history_symbol(history, depth) as usize;
                        let sibling_edge = path_edge ^ 1;
                        let sibling = self.arena.child(node_idx, sibling_edge);
                        self.levels[slot].sibling = sibling;
                        if sibling.is_some() {
                            self.detaches.push(Detach::NodeChild {
                                node: node_idx,
                                edge: sibling_edge,
                            });
                        }
                        let next = self.arena.child(node_idx, path_edge);
                        source = if let Some(next_node) = next.as_node() {
                            ExistingSource::Node(next_node)
                        } else if let Some(next_segment) = next.as_segment() {
                            ExistingSource::Segment(next_segment, 0)
                        } else {
                            ExistingSource::None
                        };
                    }
                }
                ExistingSource::Segment(segment_idx, offset) => {
                    self.levels[slot].symbol_count = self.arena.segment_symbol_count(segment_idx);
                    self.levels[slot].log_prob_kt = self.arena.segment_log_prob_kt(segment_idx);
                    if depth < self.max_depth {
                        let path_edge = history_symbol(history, depth) as usize;
                        if self.arena.segment_has_child(segment_idx, offset) {
                            let existing_edge =
                                self.arena.segment_edge(segment_idx, offset, history);
                            if path_edge == existing_edge {
                                let seg_len = self.arena.segment_len(segment_idx);
                                if offset + 1 < seg_len {
                                    source = ExistingSource::Segment(segment_idx, offset + 1);
                                } else {
                                    let tail = self.arena.segments[segment_idx.get()].tail;
                                    source = if let Some(next_node) = tail.as_node() {
                                        ExistingSource::Node(next_node)
                                    } else if let Some(next_segment) = tail.as_segment() {
                                        ExistingSource::Segment(next_segment, 0)
                                    } else {
                                        ExistingSource::None
                                    };
                                }
                            } else {
                                let continuation = self.arena.detach_segment_continuation(
                                    segment_idx,
                                    offset,
                                    &mut self.detaches,
                                );
                                self.levels[slot].sibling = continuation;
                                source = ExistingSource::None;
                            }
                        } else {
                            source = ExistingSource::None;
                        }
                    }
                }
            }
        }

        old_child
    }

    fn rebuild_path_subtree(&mut self, history: &(impl HistoryAccess + ?Sized)) -> ChildRef {
        let mut built = ChildRef::NONE;

        for depth in (1..=self.max_depth).rev() {
            let level = self.levels[depth - 1];
            let visits = level.symbol_count[0] + level.symbol_count[1];
            if visits == 0 {
                built = ChildRef::NONE;
                continue;
            }

            let path_edge = if depth < self.max_depth {
                history_symbol(history, depth) as usize
            } else {
                0
            };
            let has_path_child = built.is_some();
            let has_sibling = level.sibling.is_some();
            let force_node = depth <= self.hot_prefix_depth();

            if force_node || (has_path_child && has_sibling) {
                let node = self
                    .arena
                    .alloc_node_with_state(level.symbol_count, level.log_prob_kt);
                if has_path_child {
                    self.arena.set_child(node, path_edge, built);
                }
                if has_sibling {
                    self.arena.set_child(node, path_edge ^ 1, level.sibling);
                }
                self.arena.recompute_node_weight(node);
                built = ChildRef::from_node(node);
            } else {
                let (edge, child) = if has_path_child {
                    (path_edge, built)
                } else if has_sibling {
                    (path_edge ^ 1, level.sibling)
                } else {
                    (path_edge, ChildRef::NONE)
                };
                built = self.arena.prepend_or_alloc_segment(
                    history,
                    depth,
                    level.symbol_count,
                    level.log_prob_kt,
                    child,
                    edge,
                    false,
                );
            }
        }

        built
    }

    fn apply_detaches(&mut self) {
        for detach in self.detaches.drain(..) {
            match detach {
                Detach::NodeChild { node, edge } => {
                    self.arena.set_child(node, edge, ChildRef::NONE);
                }
                Detach::SegmentNext { segment, new_len } => {
                    self.arena.segments[segment.get()].set_len(new_len);
                    self.arena.set_segment_tail(segment, ChildRef::NONE);
                }
            }
        }
    }

    fn update_with_logs<L: CtLogAccess>(
        &mut self,
        logs: L,
        sym: Symbol,
        history: &(impl HistoryAccess + ?Sized),
    ) {
        let sym_idx = sym as usize;
        let singleton_log_prob_kt = logs.log_half(0) - logs.log_int(1);
        {
            let slot = self.root.get();
            let mut counts = self.arena.nodes[slot].symbol_count;
            let mut log_prob_kt = self.arena.nodes[slot].log_prob_kt;
            apply_update_to_state_raw(logs, &mut counts, &mut log_prob_kt, sym_idx);
            self.arena.nodes[slot].symbol_count = counts;
            self.arena.nodes[slot].log_prob_kt = log_prob_kt;
        }

        if self.max_depth > 0 {
            let root_edge = history_symbol(history, 0) as usize;
            let old_child = self.arena.child(self.root, root_edge);
            let new_child =
                self.update_root_child(logs, old_child, history, sym_idx, singleton_log_prob_kt);
            self.arena.set_child(self.root, root_edge, new_child);
        }

        self.arena.recompute_node_weight(self.root);
    }

    fn update(&mut self, sym: Symbol, history: &(impl HistoryAccess + ?Sized)) {
        let upto = self.root_visits() + 1;
        if upto <= ctw_log_cache_limit() {
            self.with_cached_logs(upto, |this, logs| {
                this.update_with_logs(logs, sym, history);
            });
        } else {
            self.with_bounded_logs(upto, |this, logs| {
                this.update_with_logs(logs, sym, history);
            });
        }
    }

    fn update_prepared_with_logs<L: CtLogAccess>(
        &mut self,
        logs: L,
        history: &(impl HistoryAccess + ?Sized),
        sym_idx: usize,
        use_prepared: bool,
    ) {
        let singleton_log_prob_kt = logs.log_half(0) - logs.log_int(1);
        {
            let slot = self.root.get();
            let mut counts = self.arena.nodes[slot].symbol_count;
            let mut log_prob_kt = self.arena.nodes[slot].log_prob_kt;
            apply_update_to_state_raw(logs, &mut counts, &mut log_prob_kt, sym_idx);
            self.arena.nodes[slot].symbol_count = counts;
            self.arena.nodes[slot].log_prob_kt = log_prob_kt;
        }

        if self.max_depth > 0 {
            let root_edge = history_symbol(history, 0) as usize;
            let old_child = self.arena.child(self.root, root_edge);
            let new_child = if use_prepared {
                match self.prepared_end {
                    PreparedEnd::MissingAtRoot => {
                        self.build_missing_path(1, history, sym_idx, singleton_log_prob_kt)
                    }
                    PreparedEnd::MaxDepth | PreparedEnd::MissingAfterCurrent => {
                        if !self.prepared_steps.is_empty() {
                            self.update_prepared_cached_path(
                                logs,
                                history,
                                sym_idx,
                                singleton_log_prob_kt,
                            );
                        }
                        old_child
                    }
                    PreparedEnd::MismatchAtCurrentSegment => {
                        self.update_prepared_mismatch(logs, history, sym_idx, singleton_log_prob_kt)
                    }
                }
            } else {
                self.update_root_child(logs, old_child, history, sym_idx, singleton_log_prob_kt)
            };
            self.arena.set_child(self.root, root_edge, new_child);
        }

        self.arena.recompute_node_weight(self.root);
    }

    fn update_prepared(
        &mut self,
        sym: Symbol,
        history: &(impl HistoryAccess + ?Sized),
        use_prepared: bool,
    ) {
        let upto = self.root_visits() + 1;
        let sym_idx = sym as usize;
        if upto <= ctw_log_cache_limit() {
            self.with_cached_logs(upto, |this, logs| {
                this.update_prepared_with_logs(logs, history, sym_idx, use_prepared);
            });
        } else {
            self.with_bounded_logs(upto, |this, logs| {
                this.update_prepared_with_logs(logs, history, sym_idx, use_prepared);
            });
        }
    }

    fn revert_with_logs<L: CtLogAccess>(
        &mut self,
        logs: L,
        history: &(impl HistoryAccess + ?Sized),
        sym_idx: usize,
    ) {
        let old_child = self.collect_existing_levels(history);

        {
            let slot = self.root.get();
            let mut counts = self.arena.nodes[slot].symbol_count;
            let mut log_prob_kt = self.arena.nodes[slot].log_prob_kt;
            apply_revert_to_state_raw(logs, &mut counts, &mut log_prob_kt, sym_idx);
            self.arena.nodes[slot].symbol_count = counts;
            self.arena.nodes[slot].log_prob_kt = log_prob_kt;
        }

        for level in &mut self.levels {
            let mut counts = level.symbol_count;
            let mut log_prob_kt = level.log_prob_kt;
            apply_revert_to_state_raw(logs, &mut counts, &mut log_prob_kt, sym_idx);
            level.symbol_count = counts;
            level.log_prob_kt = log_prob_kt;
        }

        if self.max_depth > 0 {
            let new_child = self.rebuild_path_subtree(history);
            let root_edge = history_symbol(history, 0) as usize;
            self.apply_detaches();
            self.arena.free_child_ref(old_child);
            self.arena.set_child(self.root, root_edge, new_child);
        }

        self.arena.recompute_node_weight(self.root);
    }

    fn revert(&mut self, sym: Symbol, history: &(impl HistoryAccess + ?Sized)) {
        let upto = self.root_visits();
        let sym_idx = sym as usize;
        if upto <= ctw_log_cache_limit() {
            self.with_cached_logs(upto, |this, logs| {
                this.revert_with_logs(logs, history, sym_idx);
            });
        } else {
            self.with_bounded_logs(upto, |this, logs| {
                this.revert_with_logs(logs, history, sym_idx);
            });
        }
    }

    fn predict(&mut self, sym: Symbol, history: &(impl HistoryAccess + ?Sized)) -> f64 {
        self.prepared_steps.clear();
        self.prepared_levels = 0;
        self.prepared_end = PreparedEnd::MaxDepth;

        let (root_sibling, root_has_sibling, mut source) = if self.max_depth > 0 {
            let root_edge = history_symbol(history, 0) as usize;
            let path_child = self.arena.child(self.root, root_edge);
            let sibling = self.arena.child(self.root, root_edge ^ 1);
            (
                self.arena.child_ref_weighted(sibling),
                sibling.is_some() as u8,
                Self::child_to_existing_source(path_child).unwrap_or(ExistingSource::None),
            )
        } else {
            (0.0, 0, ExistingSource::None)
        };
        if self.max_depth > 0 && matches!(source, ExistingSource::None) {
            self.prepared_end = PreparedEnd::MissingAtRoot;
        }

        let history_len = history.len();
        let mut depth = 1usize;
        'walk: while depth <= self.max_depth {
            match source {
                ExistingSource::None => break,
                ExistingSource::Node(node_idx) => {
                    if depth == self.max_depth {
                        self.prepared_steps.push(PreparedStep {
                            source: ExistingSource::Node(node_idx),
                            span: 1,
                            sibling_weight: 0.0,
                            has_sibling: 0,
                        });
                        self.prepared_levels += 1;
                        break;
                    }
                    let path_edge = history_symbol(history, depth) as usize;
                    let sibling = self.arena.child(node_idx, path_edge ^ 1);
                    self.prepared_steps.push(PreparedStep {
                        source: ExistingSource::Node(node_idx),
                        span: 1,
                        sibling_weight: self.arena.child_ref_weighted(sibling),
                        has_sibling: sibling.is_some() as u8,
                    });
                    self.prepared_levels += 1;
                    let next = self.arena.child(node_idx, path_edge);
                    source = Self::child_to_existing_source(next).unwrap_or(ExistingSource::None);
                    if matches!(source, ExistingSource::None) {
                        self.prepared_end = PreparedEnd::MissingAfterCurrent;
                        break;
                    }
                    depth += 1;
                }
                ExistingSource::Segment(segment_idx, _) => {
                    let segment = self.arena.segments[segment_idx.get()];
                    if segment.payload.is_exact() {
                        let path_bits = path_bits_from_history(
                            history,
                            depth,
                            self.max_depth.saturating_sub(depth),
                        );
                        if let Some((next_depth, next_source)) =
                            self.walk_prepared_exact_segment(segment_idx, segment, depth, path_bits)
                        {
                            source = next_source;
                            if matches!(source, ExistingSource::None) {
                                self.prepared_end = PreparedEnd::MissingAfterCurrent;
                                break 'walk;
                            }
                            depth = next_depth;
                            continue 'walk;
                        }
                        break 'walk;
                    }
                    let seg_len = segment.len() as usize;
                    for offset in 0..seg_len {
                        let node_depth = depth + offset;
                        let span = (offset + 1) as u32;

                        if node_depth == self.max_depth {
                            self.prepared_steps.push(PreparedStep {
                                source: ExistingSource::Segment(segment_idx, offset as u32),
                                span,
                                sibling_weight: 0.0,
                                has_sibling: 0,
                            });
                            self.prepared_levels += span as usize;
                            break 'walk;
                        }

                        if offset + 1 >= seg_len && segment.tail.is_none() {
                            self.prepared_steps.push(PreparedStep {
                                source: ExistingSource::Segment(segment_idx, offset as u32),
                                span,
                                sibling_weight: 0.0,
                                has_sibling: 0,
                            });
                            self.prepared_levels += span as usize;
                            self.prepared_end = PreparedEnd::MissingAfterCurrent;
                            break 'walk;
                        }

                        let path_edge = path_edge_at_depth(history, history_len, node_depth);
                        let existing_edge =
                            segment_edge_from_parts(segment, offset, history, history_len);
                        if path_edge != existing_edge {
                            self.prepared_steps.push(PreparedStep {
                                source: ExistingSource::Segment(segment_idx, offset as u32),
                                span,
                                sibling_weight: self
                                    .arena
                                    .segment_continuation_weight(segment_idx, offset as u32),
                                has_sibling: 1,
                            });
                            self.prepared_levels += span as usize;
                            self.prepared_end = PreparedEnd::MismatchAtCurrentSegment;
                            break 'walk;
                        }

                        if offset + 1 < seg_len {
                            continue;
                        }

                        self.prepared_steps.push(PreparedStep {
                            source: ExistingSource::Segment(segment_idx, offset as u32),
                            span,
                            sibling_weight: 0.0,
                            has_sibling: 0,
                        });
                        self.prepared_levels += span as usize;
                        let tail = segment.tail;
                        source =
                            Self::child_to_existing_source(tail).unwrap_or(ExistingSource::None);
                        if matches!(source, ExistingSource::None) {
                            self.prepared_end = PreparedEnd::MissingAfterCurrent;
                            break 'walk;
                        }
                        depth = node_depth + 1;
                        continue 'walk;
                    }
                }
            }
        }

        let sym_idx = sym as usize;
        if self.prepared_levels == 0 {
            let counts = self.arena.counts(self.root);
            let kt_log_prob = self.arena.log_prob_kt(self.root);
            return if self.prepared_end == PreparedEnd::MaxDepth || root_has_sibling == 0 {
                predict_ratio_kt(counts, sym_idx)
            } else {
                predict_ratio_internal(kt_log_prob, counts, 0.0, root_sibling, 0.5, sym_idx)
            };
        }

        let last_step = *self.prepared_steps.last().unwrap();
        let (last_counts, last_kt_log_prob) = self.source_counts_and_kt_log_prob(last_step.source);
        let (mut child_weight, mut ratio) = if (self.prepared_end == PreparedEnd::MaxDepth
            && self.prepared_levels == self.max_depth)
            || last_step.has_sibling == 0
        {
            (last_kt_log_prob, predict_ratio_kt(last_counts, sym_idx))
        } else {
            combined_weight_ratio_internal(
                last_kt_log_prob,
                last_counts,
                0.0,
                last_step.sibling_weight,
                0.5,
                sym_idx,
            )
        };

        if matches!(last_step.source, ExistingSource::Segment(_, _)) && last_step.span > 1 {
            let (alpha, log_alpha, log_one_minus_alpha) =
                self.segment_constants(last_step.span - 1);
            (child_weight, ratio) = unary_chain_ratio_transform_precomputed(
                last_kt_log_prob,
                last_counts,
                child_weight,
                ratio,
                alpha,
                log_alpha,
                log_one_minus_alpha,
                sym_idx,
            );
        }

        for idx in (0..self.prepared_steps.len() - 1).rev() {
            let step = self.prepared_steps[idx];
            let (step_counts, step_kt_log_prob) = self.source_counts_and_kt_log_prob(step.source);
            match step.source {
                ExistingSource::Node(_) => {
                    (child_weight, ratio) = combined_weight_ratio_internal(
                        step_kt_log_prob,
                        step_counts,
                        child_weight,
                        step.sibling_weight,
                        ratio,
                        sym_idx,
                    );
                }
                ExistingSource::Segment(_, _) => {
                    let (alpha, log_alpha, log_one_minus_alpha) = self.segment_constants(step.span);
                    (child_weight, ratio) = unary_chain_ratio_transform_precomputed(
                        step_kt_log_prob,
                        step_counts,
                        child_weight,
                        ratio,
                        alpha,
                        log_alpha,
                        log_one_minus_alpha,
                        sym_idx,
                    );
                }
                ExistingSource::None => unreachable!("prepared step should never store None"),
            }
        }

        let root_counts = self.arena.counts(self.root);
        let root_kt_log_prob = self.arena.log_prob_kt(self.root);
        predict_ratio_internal(
            root_kt_log_prob,
            root_counts,
            child_weight,
            root_sibling,
            ratio,
            sym_idx,
        )
    }

    fn predict_one(&mut self, history: &(impl HistoryAccess + ?Sized)) -> f64 {
        self.prepared_steps.clear();
        self.prepared_levels = 0;
        self.prepared_end = PreparedEnd::MaxDepth;

        let (root_sibling, root_has_sibling, mut source) = if self.max_depth > 0 {
            let root_edge = history_symbol(history, 0) as usize;
            let path_child = self.arena.child(self.root, root_edge);
            let sibling = self.arena.child(self.root, root_edge ^ 1);
            (
                self.arena.child_ref_weighted(sibling),
                sibling.is_some() as u8,
                Self::child_to_existing_source(path_child).unwrap_or(ExistingSource::None),
            )
        } else {
            (0.0, 0, ExistingSource::None)
        };
        if self.max_depth > 0 && matches!(source, ExistingSource::None) {
            self.prepared_end = PreparedEnd::MissingAtRoot;
        }

        let history_len = history.len();
        let mut depth = 1usize;
        'walk: while depth <= self.max_depth {
            match source {
                ExistingSource::None => break,
                ExistingSource::Node(node_idx) => {
                    if depth == self.max_depth {
                        self.prepared_steps.push(PreparedStep {
                            source: ExistingSource::Node(node_idx),
                            span: 1,
                            sibling_weight: 0.0,
                            has_sibling: 0,
                        });
                        self.prepared_levels += 1;
                        break;
                    }
                    let path_edge = history_symbol(history, depth) as usize;
                    let sibling = self.arena.child(node_idx, path_edge ^ 1);
                    self.prepared_steps.push(PreparedStep {
                        source: ExistingSource::Node(node_idx),
                        span: 1,
                        sibling_weight: self.arena.child_ref_weighted(sibling),
                        has_sibling: sibling.is_some() as u8,
                    });
                    self.prepared_levels += 1;
                    let next = self.arena.child(node_idx, path_edge);
                    source = Self::child_to_existing_source(next).unwrap_or(ExistingSource::None);
                    if matches!(source, ExistingSource::None) {
                        self.prepared_end = PreparedEnd::MissingAfterCurrent;
                        break;
                    }
                    depth += 1;
                }
                ExistingSource::Segment(segment_idx, _) => {
                    let segment = self.arena.segments[segment_idx.get()];
                    if segment.payload.is_exact() {
                        let path_bits = path_bits_from_history(
                            history,
                            depth,
                            self.max_depth.saturating_sub(depth),
                        );
                        if let Some((next_depth, next_source)) =
                            self.walk_prepared_exact_segment(segment_idx, segment, depth, path_bits)
                        {
                            source = next_source;
                            if matches!(source, ExistingSource::None) {
                                self.prepared_end = PreparedEnd::MissingAfterCurrent;
                                break 'walk;
                            }
                            depth = next_depth;
                            continue 'walk;
                        }
                        break 'walk;
                    }
                    let seg_len = segment.len() as usize;
                    for offset in 0..seg_len {
                        let node_depth = depth + offset;
                        let span = (offset + 1) as u32;

                        if node_depth == self.max_depth {
                            self.prepared_steps.push(PreparedStep {
                                source: ExistingSource::Segment(segment_idx, offset as u32),
                                span,
                                sibling_weight: 0.0,
                                has_sibling: 0,
                            });
                            self.prepared_levels += span as usize;
                            break 'walk;
                        }

                        if offset + 1 >= seg_len && segment.tail.is_none() {
                            self.prepared_steps.push(PreparedStep {
                                source: ExistingSource::Segment(segment_idx, offset as u32),
                                span,
                                sibling_weight: 0.0,
                                has_sibling: 0,
                            });
                            self.prepared_levels += span as usize;
                            self.prepared_end = PreparedEnd::MissingAfterCurrent;
                            break 'walk;
                        }

                        let path_edge = path_edge_at_depth(history, history_len, node_depth);
                        let existing_edge =
                            segment_edge_from_parts(segment, offset, history, history_len);
                        if path_edge != existing_edge {
                            self.prepared_steps.push(PreparedStep {
                                source: ExistingSource::Segment(segment_idx, offset as u32),
                                span,
                                sibling_weight: self
                                    .arena
                                    .segment_continuation_weight(segment_idx, offset as u32),
                                has_sibling: 1,
                            });
                            self.prepared_levels += span as usize;
                            self.prepared_end = PreparedEnd::MismatchAtCurrentSegment;
                            break 'walk;
                        }

                        if offset + 1 < seg_len {
                            continue;
                        }

                        self.prepared_steps.push(PreparedStep {
                            source: ExistingSource::Segment(segment_idx, offset as u32),
                            span,
                            sibling_weight: 0.0,
                            has_sibling: 0,
                        });
                        self.prepared_levels += span as usize;
                        let tail = segment.tail;
                        source =
                            Self::child_to_existing_source(tail).unwrap_or(ExistingSource::None);
                        if matches!(source, ExistingSource::None) {
                            self.prepared_end = PreparedEnd::MissingAfterCurrent;
                            break 'walk;
                        }
                        depth = node_depth + 1;
                        continue 'walk;
                    }
                }
            }
        }

        if self.prepared_levels == 0 {
            let counts = self.arena.counts(self.root);
            let kt_log_prob = self.arena.log_prob_kt(self.root);
            return if self.prepared_end == PreparedEnd::MaxDepth || root_has_sibling == 0 {
                predict_ratio_kt_one(counts)
            } else {
                predict_ratio_internal_one(kt_log_prob, counts, 0.0, root_sibling, 0.5)
            };
        }

        let last_step = *self.prepared_steps.last().unwrap();
        let (last_counts, last_kt_log_prob) = self.source_counts_and_kt_log_prob(last_step.source);
        let (mut child_weight, mut ratio) = if (self.prepared_end == PreparedEnd::MaxDepth
            && self.prepared_levels == self.max_depth)
            || last_step.has_sibling == 0
        {
            (last_kt_log_prob, predict_ratio_kt_one(last_counts))
        } else {
            combined_weight_ratio_internal_one(
                last_kt_log_prob,
                last_counts,
                0.0,
                last_step.sibling_weight,
                0.5,
            )
        };

        if let ExistingSource::Segment(_, _) = last_step.source
            && last_step.span > 1
        {
            let (alpha, log_alpha, log_one_minus_alpha) =
                self.segment_constants(last_step.span - 1);
            (child_weight, ratio) = unary_chain_ratio_transform_precomputed_one(
                last_kt_log_prob,
                last_counts,
                child_weight,
                ratio,
                alpha,
                log_alpha,
                log_one_minus_alpha,
            );
        }

        for idx in (0..self.prepared_steps.len() - 1).rev() {
            let step = self.prepared_steps[idx];
            let (step_counts, step_kt_log_prob) = self.source_counts_and_kt_log_prob(step.source);
            match step.source {
                ExistingSource::Node(_) => {
                    (child_weight, ratio) = combined_weight_ratio_internal_one(
                        step_kt_log_prob,
                        step_counts,
                        child_weight,
                        step.sibling_weight,
                        ratio,
                    );
                }
                ExistingSource::Segment(_, _) => {
                    let (alpha, log_alpha, log_one_minus_alpha) = self.segment_constants(step.span);
                    (child_weight, ratio) = unary_chain_ratio_transform_precomputed_one(
                        step_kt_log_prob,
                        step_counts,
                        child_weight,
                        ratio,
                        alpha,
                        log_alpha,
                        log_one_minus_alpha,
                    );
                }
                ExistingSource::None => unreachable!("prepared step should never store None"),
            }
        }

        let root_counts = self.arena.counts(self.root);
        let root_kt_log_prob = self.arena.log_prob_kt(self.root);
        predict_ratio_internal_one(
            root_kt_log_prob,
            root_counts,
            child_weight,
            root_sibling,
            ratio,
        )
    }

    fn memory_usage(&self) -> usize {
        self.arena.memory_usage()
            + self.segment_alpha.capacity() * size_of::<f64>()
            + self.segment_log_alpha.capacity() * size_of::<f64>()
            + self.segment_log_one_minus_alpha.capacity() * size_of::<f64>()
            + self.levels.capacity() * size_of::<LevelState>()
            + self.detaches.capacity() * size_of::<Detach>()
            + self.prepared_steps.capacity() * size_of::<PreparedStep>()
    }

    #[cfg(any(test, feature = "research-tooling"))]
    fn scratch_memory_usage(&self) -> usize {
        self.segment_alpha.capacity() * size_of::<f64>()
            + self.segment_log_alpha.capacity() * size_of::<f64>()
            + self.segment_log_one_minus_alpha.capacity() * size_of::<f64>()
            + self.levels.capacity() * size_of::<LevelState>()
            + self.detaches.capacity() * size_of::<Detach>()
            + self.prepared_steps.capacity() * size_of::<PreparedStep>()
    }

    #[cfg(any(test, feature = "research-tooling"))]
    fn telemetry(&self, bit_index: usize) -> FacContextTreeTreeTelemetry {
        let mut exact_segments: usize = 0;
        let mut history_segments: usize = 0;
        let mut history_invert_segments: usize = 0;
        let mut const_segments: usize = 0;
        let mut segment_bits: u64 = 0;
        let mut max_segment_len: u32 = 0;
        for segment in &self.arena.segments {
            let len = segment.len();
            segment_bits = segment_bits.saturating_add(len as u64);
            max_segment_len = max_segment_len.max(len);
            match segment.payload.mode() {
                SEG_MODE_EXACT => exact_segments = exact_segments.saturating_add(1),
                SEG_MODE_HISTORY => history_segments = history_segments.saturating_add(1),
                SEG_MODE_HISTORY_INVERT => {
                    history_invert_segments = history_invert_segments.saturating_add(1);
                }
                SEG_MODE_CONST => const_segments = const_segments.saturating_add(1),
                _ => unreachable!("invalid ctw segment payload mode"),
            }
        }

        let node_payload_bytes = self.arena.nodes.len() * size_of::<CtNode>();
        let node_bytes = self.arena.nodes.capacity() * size_of::<CtNode>();
        let segment_payload_bytes = self.arena.segments.len() * size_of::<CtSegment>();
        let segment_bytes = self.arena.segments.capacity() * size_of::<CtSegment>();
        let free_list_bytes = self.arena.free_nodes.capacity() * size_of::<NodeIndex>()
            + self.arena.free_segments.capacity() * size_of::<SegmentIndex>();
        let scratch_bytes = self.scratch_memory_usage();
        let arena_slack_bytes = node_bytes
            .saturating_sub(node_payload_bytes)
            .saturating_add(segment_bytes.saturating_sub(segment_payload_bytes));

        FacContextTreeTreeTelemetry {
            bit_index,
            max_depth: self.max_depth,
            root_visits: self.root_visits(),
            nodes_len: self.arena.nodes.len(),
            nodes_capacity: self.arena.nodes.capacity(),
            segments_len: self.arena.segments.len(),
            segments_capacity: self.arena.segments.capacity(),
            free_nodes_len: self.arena.free_nodes.len(),
            free_nodes_capacity: self.arena.free_nodes.capacity(),
            free_segments_len: self.arena.free_segments.len(),
            free_segments_capacity: self.arena.free_segments.capacity(),
            node_bytes,
            node_payload_bytes,
            segment_bytes,
            segment_payload_bytes,
            free_list_bytes,
            scratch_bytes,
            total_bytes: node_bytes
                .saturating_add(segment_bytes)
                .saturating_add(free_list_bytes)
                .saturating_add(scratch_bytes),
            arena_slack_bytes,
            exact_segments,
            history_segments,
            history_invert_segments,
            const_segments,
            segment_bits,
            max_segment_len,
        }
    }
}

/// A Context Tree for binary sequence prediction.
#[derive(Clone)]
pub struct ContextTree {
    engine: CtEngine,
    history: BitHistory,
    history_version: u64,
    prepared_valid: bool,
    prepared_history_len: usize,
    prepared_history_version: u64,
}

#[derive(Clone)]
pub(crate) struct ContextTreeLifecycleSnapshot {
    history: BitHistory,
    history_version: u64,
}

impl ContextTree {
    /// Construct a binary CTW predictor with maximum context depth `depth`.
    pub fn new(depth: usize) -> Self {
        Self {
            engine: CtEngine::new(depth),
            history: BitHistory::default(),
            history_version: 0,
            prepared_valid: false,
            prepared_history_len: 0,
            prepared_history_version: 0,
        }
    }

    #[inline]
    fn bump_history_version(&mut self) {
        self.history_version = self.history_version.wrapping_add(1);
    }

    #[inline]
    fn clear_prepared_prediction(&mut self) {
        self.prepared_valid = false;
    }

    #[inline]
    fn prepared_prediction_matches_history(&self) -> bool {
        self.prepared_valid
            && self.prepared_history_len == self.history.len()
            && self.prepared_history_version == self.history_version
    }

    /// Reset tree parameters and clear conditioning history.
    pub fn clear(&mut self) {
        self.history.clear();
        self.engine.clear();
        self.history_version = 0;
        self.clear_prepared_prediction();
        self.prepared_history_len = 0;
        self.prepared_history_version = 0;
    }

    #[inline]
    pub(crate) fn reserve_for_symbols(&mut self, total_symbols: usize) {
        if total_symbols == 0 {
            return;
        }
        self.engine.reserve_for_symbols(total_symbols);
        self.history.reserve_exact(total_symbols);
    }

    #[inline]
    /// Observe one binary symbol and update the model.
    pub fn update(&mut self, sym: Symbol) {
        let use_prepared = self.prepared_prediction_matches_history();
        self.clear_prepared_prediction();
        if use_prepared {
            self.engine.update_prepared(sym, &self.history, true);
        } else {
            self.engine.update(sym, &self.history);
        }
        self.history.push(sym);
        self.bump_history_version();
    }

    #[inline]
    /// Revert the last symbol update if history is non-empty.
    pub fn revert(&mut self) {
        let Some(last_sym) = self.history.pop() else {
            return;
        };
        self.clear_prepared_prediction();
        self.engine.revert(last_sym, &self.history);
        self.bump_history_version();
    }

    #[inline]
    /// Append external symbols to history without touching model state.
    pub fn update_history(&mut self, symbols: &[Symbol]) {
        if symbols.is_empty() {
            return;
        }
        self.clear_prepared_prediction();
        self.history.extend_from_slice(symbols);
        self.bump_history_version();
    }

    #[inline]
    /// Remove one history symbol without reverting model statistics.
    pub fn revert_history(&mut self) {
        if self.history.pop().is_some() {
            self.clear_prepared_prediction();
            self.bump_history_version();
        }
    }

    /// Truncate the stored history to `new_size` symbols.
    pub fn truncate_history(&mut self, new_size: usize) {
        if new_size < self.history.len() {
            self.clear_prepared_prediction();
            self.history.truncate(new_size);
            self.bump_history_version();
        }
    }

    /// Capture rollback state for stream lifecycle transactions.
    ///
    /// This clones the full conditioning history, so the allocation and copy are
    /// O(history length). It is intended for stream lifecycle boundaries rather
    /// than per-symbol speculative prediction.
    pub(crate) fn lifecycle_snapshot(&self) -> ContextTreeLifecycleSnapshot {
        ContextTreeLifecycleSnapshot {
            history: self.history.clone(),
            history_version: self.history_version,
        }
    }

    pub(crate) fn restore_lifecycle_snapshot(&mut self, snapshot: ContextTreeLifecycleSnapshot) {
        self.history = snapshot.history;
        self.history_version = snapshot.history_version;
        self.clear_prepared_prediction();
    }

    #[inline]
    /// Predict `P(sym | history)` under current weighted CTW model.
    pub fn predict(&mut self, sym: Symbol) -> f64 {
        let prob = self.engine.predict(sym, &self.history);
        self.prepared_valid = true;
        self.prepared_history_len = self.history.len();
        self.prepared_history_version = self.history_version;
        prob
    }

    #[inline]
    pub(crate) fn predict_one(&mut self) -> f64 {
        let prob = self.engine.predict_one(&self.history);
        self.prepared_valid = true;
        self.prepared_history_len = self.history.len();
        self.prepared_history_version = self.history_version;
        prob
    }

    #[inline]
    /// Predict probability of symbol `true`.
    pub fn predict_sym_prob(&mut self) -> f64 {
        self.predict_one()
    }

    #[inline]
    /// Return log block probability accumulated by the root model.
    pub fn get_log_block_probability(&self) -> f64 {
        self.engine.get_log_block_probability()
    }

    #[inline]
    /// Maximum context depth configured for this tree.
    pub fn depth(&self) -> usize {
        self.engine.max_depth
    }

    #[inline]
    /// Current number of stored history symbols.
    pub fn history_size(&self) -> usize {
        self.history.len()
    }
}

#[derive(Clone)]
struct ContextTreeCore {
    engine: CtEngine,
    prepared_valid: bool,
    prepared_history_len: usize,
    prepared_history_version: u64,
}

#[derive(Clone, Copy)]
struct ContextTreeCorePreparedSnapshot {
    prepared_valid: bool,
    prepared_history_len: usize,
    prepared_history_version: u64,
}

impl ContextTreeCore {
    fn new(depth: usize) -> Self {
        Self {
            engine: CtEngine::new(depth),
            prepared_valid: false,
            prepared_history_len: 0,
            prepared_history_version: 0,
        }
    }

    fn clear(&mut self) {
        self.engine.clear();
        self.prepared_valid = false;
        self.prepared_history_len = 0;
        self.prepared_history_version = 0;
    }

    #[inline]
    fn reserve_for_symbols(&mut self, total_symbols: usize) {
        self.engine.reserve_for_symbols(total_symbols);
    }

    #[inline]
    fn update_predicted(
        &mut self,
        sym: Symbol,
        shared_history: &(impl HistoryAccess + ?Sized),
        history_version: u64,
    ) {
        let use_prepared = self.prepared_valid
            && self.prepared_history_len == shared_history.len()
            && self.prepared_history_version == history_version;
        self.prepared_valid = false;
        self.engine
            .update_prepared(sym, shared_history, use_prepared);
    }

    #[inline]
    fn update_predicted_with_logs<L: CtLogAccess>(
        &mut self,
        logs: L,
        sym: Symbol,
        shared_history: &(impl HistoryAccess + ?Sized),
        history_version: u64,
    ) {
        let use_prepared = self.prepared_valid
            && self.prepared_history_len == shared_history.len()
            && self.prepared_history_version == history_version;
        self.prepared_valid = false;
        self.engine
            .update_prepared_with_logs(logs, shared_history, sym as usize, use_prepared);
    }

    #[inline]
    fn revert(&mut self, last_sym: Symbol, shared_history: &(impl HistoryAccess + ?Sized)) {
        self.prepared_valid = false;
        self.engine.revert(last_sym, shared_history);
    }

    #[inline]
    fn predict(
        &mut self,
        sym: Symbol,
        shared_history: &(impl HistoryAccess + ?Sized),
        history_version: u64,
    ) -> f64 {
        let prob = self.engine.predict(sym, shared_history);
        self.prepared_valid = true;
        self.prepared_history_len = shared_history.len();
        self.prepared_history_version = history_version;
        prob
    }

    #[inline]
    fn predict_one(
        &mut self,
        shared_history: &(impl HistoryAccess + ?Sized),
        history_version: u64,
    ) -> f64 {
        let prob = self.engine.predict_one(shared_history);
        self.prepared_valid = true;
        self.prepared_history_len = shared_history.len();
        self.prepared_history_version = history_version;
        prob
    }

    #[inline]
    fn get_log_block_probability(&self) -> f64 {
        self.engine.get_log_block_probability()
    }

    #[inline]
    fn prepared_snapshot(&self) -> ContextTreeCorePreparedSnapshot {
        ContextTreeCorePreparedSnapshot {
            prepared_valid: self.prepared_valid,
            prepared_history_len: self.prepared_history_len,
            prepared_history_version: self.prepared_history_version,
        }
    }

    #[inline]
    fn restore_prepared_snapshot(&mut self, snapshot: ContextTreeCorePreparedSnapshot) {
        self.prepared_valid = snapshot.prepared_valid;
        self.prepared_history_len = snapshot.prepared_history_len;
        self.prepared_history_version = snapshot.prepared_history_version;
    }
}

/// Factorized Action-Conditional Context Tree Weighting.
#[derive(Clone)]
pub struct FacContextTree {
    trees: Vec<ContextTreeCore>,
    shared_history: BitHistory,
    base_depth: usize,
    num_bits: usize,
    shared_history_version: u64,
}

#[derive(Clone)]
pub(crate) struct FacContextTreeLifecycleSnapshot {
    shared_history: BitHistory,
    shared_history_version: u64,
    prepared: Vec<ContextTreeCorePreparedSnapshot>,
}

/// Approximate heap-memory breakdown for a [`FacContextTree`].
#[cfg(any(test, feature = "research-tooling"))]
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct FacContextTreeMemoryUsage {
    /// Bytes owned by per-bit CTW tree engines, including arenas and scratch buffers.
    pub tree_bytes: usize,
    /// Bytes held by the thread-local shared CTW logarithm cache.
    pub shared_log_cache_bytes: usize,
    /// Bytes reserved for the factorized shared history buffer.
    pub shared_history_bytes: usize,
}

/// Per-tree CTW arena and scratch telemetry.
#[cfg(any(test, feature = "research-tooling"))]
#[derive(Clone, Debug, PartialEq, Eq)]
#[non_exhaustive]
pub struct FacContextTreeTreeTelemetry {
    /// Factorized bit position for this tree.
    pub bit_index: usize,
    /// Maximum context depth for this tree.
    pub max_depth: usize,
    /// Number of symbols observed by this tree root.
    pub root_visits: usize,
    /// Number of allocated explicit nodes.
    pub nodes_len: usize,
    /// Reserved explicit-node capacity.
    pub nodes_capacity: usize,
    /// Number of allocated unary path segments.
    pub segments_len: usize,
    /// Reserved unary-segment capacity.
    pub segments_capacity: usize,
    /// Number of node slots currently on the free list.
    pub free_nodes_len: usize,
    /// Reserved free-node list capacity.
    pub free_nodes_capacity: usize,
    /// Number of segment slots currently on the free list.
    pub free_segments_len: usize,
    /// Reserved free-segment list capacity.
    pub free_segments_capacity: usize,
    /// Reserved explicit-node bytes.
    pub node_bytes: usize,
    /// Explicit-node payload bytes at current length.
    pub node_payload_bytes: usize,
    /// Reserved segment bytes.
    pub segment_bytes: usize,
    /// Unary-segment payload bytes at current length.
    pub segment_payload_bytes: usize,
    /// Reserved free-list bytes.
    pub free_list_bytes: usize,
    /// Reserved engine scratch bytes.
    pub scratch_bytes: usize,
    /// Total reserved tree bytes for this tree.
    pub total_bytes: usize,
    /// Reserved tree bytes currently unused by arena capacity.
    pub arena_slack_bytes: usize,
    /// Number of exact-bit segment payloads.
    pub exact_segments: usize,
    /// Number of history-anchor segment payloads.
    pub history_segments: usize,
    /// Number of inverted history-anchor segment payloads.
    pub history_invert_segments: usize,
    /// Number of constant-bit segment payloads.
    pub const_segments: usize,
    /// Sum of segment lengths in bits.
    pub segment_bits: u64,
    /// Maximum segment length in bits.
    pub max_segment_len: u32,
}

/// Detailed FAC-CTW memory telemetry.
#[cfg(any(test, feature = "research-tooling"))]
#[derive(Clone, Debug, PartialEq, Eq)]
#[non_exhaustive]
pub struct FacContextTreeTelemetry {
    /// Base depth used to construct tree 0.
    pub base_depth: usize,
    /// Number of factorized bit trees.
    pub num_bits: usize,
    /// Shared history length in bits.
    pub shared_history_len_bits: usize,
    /// Shared history reserved capacity in bits.
    pub shared_history_capacity_bits: usize,
    /// Reserved bytes for shared history.
    pub shared_history_bytes: usize,
    /// Shared-history payload bytes at current length.
    pub shared_history_payload_bytes: usize,
    /// Shared-history reserved slack bytes.
    pub shared_history_slack_bytes: usize,
    /// Reserved bytes for shared log caches.
    pub shared_log_cache_bytes: usize,
    /// Sum of per-tree reserved bytes.
    pub tree_bytes: usize,
    /// Sum of live per-tree node and segment payload bytes.
    pub tree_payload_bytes: usize,
    /// Sum of per-tree arena slack bytes.
    pub tree_arena_slack_bytes: usize,
    /// Total reserved bytes reported by FAC memory accounting.
    pub total_bytes: usize,
    /// Total reserved slack bytes attributable to allocator headroom.
    pub total_slack_bytes: usize,
    /// Sum of allocated explicit nodes across trees.
    pub nodes_len: usize,
    /// Sum of explicit-node capacity across trees.
    pub nodes_capacity: usize,
    /// Sum of allocated unary path segments across trees.
    pub segments_len: usize,
    /// Sum of unary-segment capacity across trees.
    pub segments_capacity: usize,
    /// Sum of node slots currently on free lists.
    pub free_nodes_len: usize,
    /// Sum of segment slots currently on free lists.
    pub free_segments_len: usize,
    /// Sum of exact-bit segment payloads across trees.
    pub exact_segments: usize,
    /// Sum of history-anchor segment payloads across trees.
    pub history_segments: usize,
    /// Sum of inverted history-anchor segment payloads across trees.
    pub history_invert_segments: usize,
    /// Sum of constant-bit segment payloads across trees.
    pub const_segments: usize,
    /// Sum of segment lengths in bits across trees.
    pub segment_bits: u64,
    /// Maximum segment length observed across trees.
    pub max_segment_len: u32,
    /// Per-tree telemetry.
    pub trees: Vec<FacContextTreeTreeTelemetry>,
}

#[cfg(any(test, feature = "research-tooling"))]
impl FacContextTreeMemoryUsage {
    /// Total approximate heap memory in bytes.
    #[inline]
    pub fn total_bytes(self) -> usize {
        self.tree_bytes
            .saturating_add(self.shared_log_cache_bytes)
            .saturating_add(self.shared_history_bytes)
    }
}

impl FacContextTree {
    /// Create a factorized CTW stack over `num_percept_bits` bit positions.
    ///
    /// Tree `i` uses depth `base_depth + i`, matching FAC-CTW's increasing context.
    pub fn new(base_depth: usize, num_percept_bits: usize) -> Self {
        let trees = (0..num_percept_bits)
            .map(|i| ContextTreeCore::new(base_depth + i))
            .collect();
        Self {
            trees,
            shared_history: BitHistory::default(),
            base_depth,
            num_bits: num_percept_bits,
            shared_history_version: 0,
        }
    }

    #[inline(always)]
    fn bump_shared_history_version(&mut self) {
        self.shared_history_version = self.shared_history_version.wrapping_add(1);
    }

    #[inline]
    /// Reserve history/tree capacity for approximately `total_symbols` updates.
    pub fn reserve_for_symbols(&mut self, total_symbols: usize) {
        if total_symbols == 0 {
            return;
        }
        self.shared_history
            .reserve_exact(total_symbols.saturating_mul(self.num_bits));
        for tree in &mut self.trees {
            tree.reserve_for_symbols(total_symbols);
        }
    }

    #[inline]
    /// Number of bit positions currently modeled per symbol.
    pub fn num_bits(&self) -> usize {
        self.num_bits
    }

    #[inline]
    /// Base depth used to construct the first factorized tree.
    pub fn base_depth(&self) -> usize {
        self.base_depth
    }

    #[inline]
    /// Update one bit position with a binary symbol.
    pub fn update(&mut self, sym: Symbol, bit_index: usize) {
        debug_assert!(bit_index < self.num_bits);
        self.trees[bit_index].update_predicted(
            sym,
            &self.shared_history,
            self.shared_history_version,
        );
        self.shared_history.push(sym);
        self.bump_shared_history_version();
    }

    #[inline]
    /// Update all bit positions from one byte, most-significant bit first.
    pub fn update_byte_msb(&mut self, byte: u8) {
        if self.num_bits != 8 {
            for bit_idx in 0..self.num_bits {
                let bit = ((byte >> (7 - bit_idx)) & 1) == 1;
                self.update(bit, bit_idx);
            }
            return;
        }

        let upto = self.trees[0].engine.root_visits() + 1;
        debug_assert!(
            self.trees
                .iter()
                .all(|tree| tree.engine.root_visits() + 1 == upto)
        );
        if upto <= ctw_log_cache_limit() {
            with_shared_cached_logs(upto, |logs| {
                for bit_idx in 0..8usize {
                    let bit = ((byte >> (7 - bit_idx)) & 1) == 1;
                    let tree = &mut self.trees[bit_idx];
                    tree.prepared_valid = false;
                    tree.engine
                        .update_with_logs(logs, bit, &self.shared_history);
                    self.shared_history.push(bit);
                }
            });
        } else {
            with_shared_bounded_logs(upto, |logs| {
                for bit_idx in 0..8usize {
                    let bit = ((byte >> (7 - bit_idx)) & 1) == 1;
                    let tree = &mut self.trees[bit_idx];
                    tree.prepared_valid = false;
                    tree.engine
                        .update_with_logs(logs, bit, &self.shared_history);
                    self.shared_history.push(bit);
                }
            });
        }
        self.bump_shared_history_version();
    }

    #[inline]
    fn log_prob_update_byte_msb_with_logs<L: CtLogAccess>(&mut self, logs: L, byte: u8) -> f64 {
        let mut logp = 0.0;
        for bit_idx in 0..8usize {
            let bit = ((byte >> (7 - bit_idx)) & 1) == 1;
            let p =
                self.trees[bit_idx].predict(bit, &self.shared_history, self.shared_history_version);
            if p.is_finite() && p > 0.0 {
                logp += p.ln();
            } else {
                logp = f64::NEG_INFINITY;
            }
            self.trees[bit_idx].update_predicted_with_logs(
                logs,
                bit,
                &self.shared_history,
                self.shared_history_version,
            );
            self.shared_history.push(bit);
            self.bump_shared_history_version();
        }
        logp
    }

    #[inline]
    /// Return the MSB-first log probability of `byte` and then update the model.
    pub fn log_prob_update_byte_msb(&mut self, byte: u8) -> f64 {
        debug_assert_eq!(self.num_bits, 8);
        let upto = self.trees[0].engine.root_visits() + 1;
        debug_assert!(
            self.trees
                .iter()
                .all(|tree| tree.engine.root_visits() + 1 == upto)
        );
        if upto <= ctw_log_cache_limit() {
            with_shared_cached_logs(upto, |logs| {
                self.log_prob_update_byte_msb_with_logs(logs, byte)
            })
        } else {
            with_shared_bounded_logs(upto, |logs| {
                self.log_prob_update_byte_msb_with_logs(logs, byte)
            })
        }
    }

    #[inline]
    /// Update all active bit positions from one byte, least-significant bit first.
    pub fn update_byte_lsb(&mut self, byte: u8) {
        let bits = self.num_bits.clamp(1, 8);
        let upto = self.trees[0].engine.root_visits() + 1;
        debug_assert!(
            self.trees
                .iter()
                .take(bits)
                .all(|tree| tree.engine.root_visits() + 1 == upto)
        );
        if upto <= ctw_log_cache_limit() {
            with_shared_cached_logs(upto, |logs| {
                for bit_idx in 0..bits {
                    let bit = ((byte >> bit_idx) & 1) == 1;
                    let tree = &mut self.trees[bit_idx];
                    tree.prepared_valid = false;
                    tree.engine
                        .update_with_logs(logs, bit, &self.shared_history);
                    self.shared_history.push(bit);
                }
            });
        } else {
            with_shared_bounded_logs(upto, |logs| {
                for bit_idx in 0..bits {
                    let bit = ((byte >> bit_idx) & 1) == 1;
                    let tree = &mut self.trees[bit_idx];
                    tree.prepared_valid = false;
                    tree.engine
                        .update_with_logs(logs, bit, &self.shared_history);
                    self.shared_history.push(bit);
                }
            });
        }
        self.bump_shared_history_version();
    }

    #[inline]
    /// Commit an update using the fast-path prepared by a prior prediction call.
    pub fn update_predicted(&mut self, sym: Symbol, bit_index: usize) {
        debug_assert!(bit_index < self.num_bits);
        self.trees[bit_index].update_predicted(
            sym,
            &self.shared_history,
            self.shared_history_version,
        );
        self.shared_history.push(sym);
        self.bump_shared_history_version();
    }

    #[inline]
    /// Predict a single bit probability for `bit_index`.
    pub fn predict(&mut self, sym: Symbol, bit_index: usize) -> f64 {
        debug_assert!(bit_index < self.num_bits);
        self.trees[bit_index].predict(sym, &self.shared_history, self.shared_history_version)
    }

    #[inline]
    pub(crate) fn predict_one(&mut self, bit_index: usize) -> f64 {
        debug_assert!(bit_index < self.num_bits);
        self.trees[bit_index].predict_one(&self.shared_history, self.shared_history_version)
    }

    #[inline]
    /// Revert the most recent bit update for `bit_index`.
    pub fn revert(&mut self, bit_index: usize) {
        debug_assert!(bit_index < self.num_bits);
        let Some(last_sym) = self.shared_history.pop() else {
            return;
        };
        self.trees[bit_index].revert(last_sym, &self.shared_history);
        self.bump_shared_history_version();
    }

    #[inline]
    /// Append raw shared-history symbols without model updates.
    pub fn update_history(&mut self, symbols: &[Symbol]) {
        if symbols.is_empty() {
            return;
        }
        self.shared_history.extend_from_slice(symbols);
        self.bump_shared_history_version();
    }

    #[inline]
    /// Drop the last `count` symbols from shared history.
    pub fn revert_history(&mut self, count: usize) {
        let old_len = self.shared_history.len();
        let new_len = self.shared_history.len().saturating_sub(count);
        if new_len == old_len {
            return;
        }
        self.shared_history.truncate(new_len);
        self.bump_shared_history_version();
    }

    #[inline]
    /// Clear shared history while preserving learned tree parameters.
    pub fn reset_history_only(&mut self) {
        if self.shared_history.is_empty() {
            return;
        }
        self.shared_history.clear();
        self.bump_shared_history_version();
    }

    /// Capture rollback state for stream lifecycle transactions.
    ///
    /// This clones the shared conditioning history and per-tree prepared-prefix
    /// state, so the allocation and copy are O(shared history length + number of
    /// FAC component trees). It is intended for stream lifecycle boundaries
    /// rather than per-symbol speculative prediction.
    pub(crate) fn lifecycle_snapshot(&self) -> FacContextTreeLifecycleSnapshot {
        FacContextTreeLifecycleSnapshot {
            shared_history: self.shared_history.clone(),
            shared_history_version: self.shared_history_version,
            prepared: self
                .trees
                .iter()
                .map(ContextTreeCore::prepared_snapshot)
                .collect(),
        }
    }

    pub(crate) fn restore_lifecycle_snapshot(&mut self, snapshot: FacContextTreeLifecycleSnapshot) {
        debug_assert_eq!(self.trees.len(), snapshot.prepared.len());
        self.shared_history = snapshot.shared_history;
        self.shared_history_version = snapshot.shared_history_version;
        for (tree, prepared) in self.trees.iter_mut().zip(snapshot.prepared) {
            tree.restore_prepared_snapshot(prepared);
        }
    }

    #[inline]
    /// Sum of per-tree log block probabilities.
    pub fn get_log_block_probability(&self) -> f64 {
        self.trees
            .iter()
            .map(|t| t.get_log_block_probability())
            .sum()
    }

    /// Clear all trees and shared history.
    pub fn clear(&mut self) {
        for tree in &mut self.trees {
            tree.clear();
        }
        self.shared_history.clear();
        self.shared_history_version = 0;
    }

    /// Approximate heap-memory usage broken down by CTW component.
    #[cfg(any(test, feature = "research-tooling"))]
    pub fn memory_usage_breakdown(&self) -> FacContextTreeMemoryUsage {
        let tree_mem: usize = self.trees.iter().map(|t| t.engine.memory_usage()).sum();
        let log_cache_mem = self
            .trees
            .first()
            .map(|t| t.engine.log_cache_memory_usage())
            .unwrap_or(0);
        let history_mem = self.shared_history.memory_usage();
        FacContextTreeMemoryUsage {
            tree_bytes: tree_mem,
            shared_log_cache_bytes: log_cache_mem,
            shared_history_bytes: history_mem,
        }
    }

    /// Detailed CTW arena, segment, scratch, history, and log-cache telemetry.
    #[cfg(any(test, feature = "research-tooling"))]
    pub fn telemetry(&self) -> FacContextTreeTelemetry {
        let usage = self.memory_usage_breakdown();
        let trees: Vec<FacContextTreeTreeTelemetry> = self
            .trees
            .iter()
            .enumerate()
            .map(|(bit_index, tree)| tree.engine.telemetry(bit_index))
            .collect();
        let shared_history_payload_bytes =
            self.shared_history.len().div_ceil(HISTORY_WORD_BITS) * size_of::<u64>();
        let shared_history_slack_bytes = usage
            .shared_history_bytes
            .saturating_sub(shared_history_payload_bytes);
        let nodes_len = trees.iter().map(|tree| tree.nodes_len).sum();
        let nodes_capacity = trees.iter().map(|tree| tree.nodes_capacity).sum();
        let segments_len = trees.iter().map(|tree| tree.segments_len).sum();
        let segments_capacity = trees.iter().map(|tree| tree.segments_capacity).sum();
        let free_nodes_len = trees.iter().map(|tree| tree.free_nodes_len).sum();
        let free_segments_len = trees.iter().map(|tree| tree.free_segments_len).sum();
        let exact_segments = trees.iter().map(|tree| tree.exact_segments).sum();
        let history_segments = trees.iter().map(|tree| tree.history_segments).sum();
        let history_invert_segments = trees.iter().map(|tree| tree.history_invert_segments).sum();
        let const_segments = trees.iter().map(|tree| tree.const_segments).sum();
        let segment_bits = trees.iter().map(|tree| tree.segment_bits).sum();
        let max_segment_len = trees
            .iter()
            .map(|tree| tree.max_segment_len)
            .max()
            .unwrap_or(0);
        let tree_payload_bytes = trees
            .iter()
            .map(|tree| {
                tree.node_payload_bytes
                    .saturating_add(tree.segment_payload_bytes)
            })
            .sum();
        let tree_arena_slack_bytes: usize = trees.iter().map(|tree| tree.arena_slack_bytes).sum();
        let total_slack_bytes = tree_arena_slack_bytes.saturating_add(shared_history_slack_bytes);

        FacContextTreeTelemetry {
            base_depth: self.base_depth,
            num_bits: self.num_bits,
            shared_history_len_bits: self.shared_history.len(),
            shared_history_capacity_bits: self.shared_history.words.capacity() * HISTORY_WORD_BITS,
            shared_history_bytes: usage.shared_history_bytes,
            shared_history_payload_bytes,
            shared_history_slack_bytes,
            shared_log_cache_bytes: usage.shared_log_cache_bytes,
            tree_bytes: usage.tree_bytes,
            tree_payload_bytes,
            tree_arena_slack_bytes,
            total_bytes: usage.total_bytes(),
            total_slack_bytes,
            nodes_len,
            nodes_capacity,
            segments_len,
            segments_capacity,
            free_nodes_len,
            free_segments_len,
            exact_segments,
            history_segments,
            history_invert_segments,
            const_segments,
            segment_bits,
            max_segment_len,
            trees,
        }
    }

    /// Approximate heap memory usage in bytes.
    pub fn memory_usage(&self) -> usize {
        let tree_mem: usize = self.trees.iter().map(|t| t.engine.memory_usage()).sum();
        let log_cache_mem = self
            .trees
            .first()
            .map(|t| t.engine.log_cache_memory_usage())
            .unwrap_or(0);
        let history_mem = self.shared_history.memory_usage();
        tree_mem
            .saturating_add(log_cache_mem)
            .saturating_add(history_mem)
    }
}

#[inline]
fn compact_symbol_msb_shift(bits_per_symbol: usize, bit_idx: usize) -> usize {
    bits_per_symbol.saturating_sub(1).saturating_sub(bit_idx)
}

#[inline]
pub(crate) fn ctw_symbol_bit_msb(symbol: u8, bits_per_symbol: usize, bit_idx: usize) -> bool {
    let bits = bits_per_symbol.clamp(1, 8);
    let shift = compact_symbol_msb_shift(bits, bit_idx);
    ((symbol >> shift) & 1) == 1
}

#[inline]
pub(crate) fn ctw_log_prob_msb(
    tree: &mut ContextTree,
    symbol: u8,
    bits_per_symbol: usize,
    min_prob: f64,
) -> f64 {
    let bits = bits_per_symbol.clamp(1, 8);
    let mut logp = 0.0;
    for bit_idx in 0..bits {
        let bit = ctw_symbol_bit_msb(symbol, bits, bit_idx);
        let p = tree.predict(bit);
        if p.is_finite() && p > 0.0 {
            logp += p.ln();
        } else {
            logp = f64::NEG_INFINITY;
        }
        tree.update(bit);
    }
    for _ in 0..bits {
        tree.revert();
    }
    if logp.is_finite() {
        logp.max(min_prob.ln())
    } else {
        min_prob.ln()
    }
}

#[inline]
pub(crate) fn ctw_log_prob_update_msb(
    tree: &mut ContextTree,
    symbol: u8,
    bits_per_symbol: usize,
    min_prob: f64,
) -> f64 {
    let bits = bits_per_symbol.clamp(1, 8);
    let mut logp = 0.0;
    for bit_idx in 0..bits {
        let bit = ctw_symbol_bit_msb(symbol, bits, bit_idx);
        let p = tree.predict(bit);
        if p.is_finite() && p > 0.0 {
            logp += p.ln();
        } else {
            logp = f64::NEG_INFINITY;
        }
        tree.update(bit);
    }
    if logp.is_finite() {
        logp.max(min_prob.ln())
    } else {
        min_prob.ln()
    }
}

#[inline]
pub(crate) fn ctw_log_prob_update_lsb(
    tree: &mut FacContextTree,
    symbol: u8,
    bits_per_symbol: usize,
    min_prob: f64,
) -> f64 {
    let mut logp = 0.0;
    for bit_idx in 0..bits_per_symbol {
        let bit = ((symbol >> bit_idx) & 1) == 1;
        let p = tree.predict(bit, bit_idx);
        if p.is_finite() && p > 0.0 {
            logp += p.ln();
        } else {
            logp = f64::NEG_INFINITY;
        }
        tree.update_predicted(bit, bit_idx);
    }
    if logp.is_finite() {
        logp.max(min_prob.ln())
    } else {
        min_prob.ln()
    }
}

pub(crate) fn fill_ctw_tree_log_probs(
    tree: &mut ContextTree,
    bits_per_symbol: usize,
    min_logp: f64,
    out: &mut [f64; 256],
) {
    let bits = bits_per_symbol.clamp(1, 8);
    let patterns = 1usize << bits;
    let mut pattern_logps = [f64::NEG_INFINITY; 256];
    let log_before = tree.get_log_block_probability();

    fn rec(
        tree: &mut ContextTree,
        depth: usize,
        bits: usize,
        log_before: f64,
        min_logp: f64,
        symbol_acc: u8,
        pattern_logps: &mut [f64; 256],
    ) {
        if depth == bits {
            let pat = symbol_acc as usize;
            let logp = (tree.get_log_block_probability() - log_before).max(min_logp);
            pattern_logps[pat] = logp;
            return;
        }

        for bit in [false, true] {
            tree.update(bit);
            let shift = compact_symbol_msb_shift(bits, depth);
            let next_symbol = if bit {
                symbol_acc | (1u8 << shift)
            } else {
                symbol_acc
            };
            rec(
                tree,
                depth + 1,
                bits,
                log_before,
                min_logp,
                next_symbol,
                pattern_logps,
            );
            tree.revert();
        }
    }

    rec(tree, 0, bits, log_before, min_logp, 0, &mut pattern_logps);

    if bits == 8 {
        out.copy_from_slice(&pattern_logps);
    } else {
        let aliases = 1usize << (8 - bits);
        let alias_ln = (aliases as f64).ln();
        let mask = patterns - 1;
        for byte in 0..256usize {
            out[byte] = pattern_logps[byte & mask] - alias_ln;
        }
    }
}

pub(crate) fn fill_fac_tree_log_probs(
    tree: &mut FacContextTree,
    bits_per_symbol: usize,
    msb_first: bool,
    min_logp: f64,
    out: &mut [f64; 256],
) {
    struct RecParams {
        bits: usize,
        msb_first: bool,
        log_before: f64,
        min_logp: f64,
    }

    let bits = bits_per_symbol.clamp(1, 8);
    let patterns = 1usize << bits;
    let mut pattern_logps = [f64::NEG_INFINITY; 256];
    let params = RecParams {
        bits,
        msb_first,
        log_before: tree.get_log_block_probability(),
        min_logp,
    };

    fn rec(
        tree: &mut FacContextTree,
        depth: usize,
        params: &RecParams,
        symbol_acc: u8,
        pattern_logps: &mut [f64; 256],
    ) {
        if depth == params.bits {
            let pat = symbol_acc as usize;
            let logp = (tree.get_log_block_probability() - params.log_before).max(params.min_logp);
            pattern_logps[pat] = logp;
            return;
        }

        for bit in [false, true] {
            tree.update(bit, depth);
            let mut next_symbol = symbol_acc;
            if params.msb_first {
                // Keep sub-byte MSB symbols packed into the low pattern range so the
                // alias expansion below can index them via `byte & mask`.
                let shift = compact_symbol_msb_shift(params.bits, depth);
                if bit {
                    next_symbol |= 1u8 << shift;
                }
            } else if bit {
                next_symbol |= 1u8 << depth;
            }
            rec(tree, depth + 1, params, next_symbol, pattern_logps);
            tree.revert(depth);
        }
    }

    rec(tree, 0, &params, 0, &mut pattern_logps);

    if bits == 8 {
        out.copy_from_slice(&pattern_logps);
    } else {
        let aliases = 1usize << (8 - bits);
        let alias_ln = (aliases as f64).ln();
        let mask = patterns - 1;
        for byte in 0..256usize {
            out[byte] = pattern_logps[byte & mask] - alias_ln;
        }
    }
}

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

    struct ScopedLogCacheLimit {
        previous: Option<usize>,
    }

    struct ScopedLogOverflowCacheSlots {
        previous: Option<usize>,
    }

    impl ScopedLogCacheLimit {
        fn set(limit: usize) -> Self {
            CTW_TEST_LOG_CACHE_LIMIT.with(|cell| {
                let previous = cell.replace(Some(limit));
                Self { previous }
            })
        }
    }

    impl ScopedLogOverflowCacheSlots {
        fn set(slots: usize) -> Self {
            CTW_TEST_LOG_OVERFLOW_CACHE_SLOTS.with(|cell| {
                let previous = cell.replace(Some(slots));
                Self { previous }
            })
        }
    }

    impl Drop for ScopedLogCacheLimit {
        fn drop(&mut self) {
            CTW_TEST_LOG_CACHE_LIMIT.with(|cell| {
                cell.replace(self.previous);
            });
        }
    }

    impl Drop for ScopedLogOverflowCacheSlots {
        fn drop(&mut self) {
            CTW_TEST_LOG_OVERFLOW_CACHE_SLOTS.with(|cell| {
                cell.replace(self.previous);
            });
        }
    }

    #[derive(Clone)]
    struct RefNode {
        children: [Option<Box<RefNode>>; 2],
        log_prob_kt: f64,
        log_prob_weighted: f64,
        symbol_count: [u32; 2],
    }

    impl Default for RefNode {
        fn default() -> Self {
            Self {
                children: [None, None],
                log_prob_kt: 0.0,
                log_prob_weighted: 0.0,
                symbol_count: [0, 0],
            }
        }
    }

    #[derive(Clone)]
    struct RefContextTree {
        root: RefNode,
        history: Vec<Symbol>,
        max_depth: usize,
        log_int: Vec<f64>,
        log_half: Vec<f64>,
    }

    impl RefContextTree {
        fn new(depth: usize) -> Self {
            Self {
                root: RefNode::default(),
                history: Vec::new(),
                max_depth: depth,
                log_int: vec![f64::NEG_INFINITY],
                log_half: vec![(0.5f64).ln()],
            }
        }

        fn root_visits(&self) -> usize {
            (self.root.symbol_count[0] + self.root.symbol_count[1]) as usize
        }

        fn recompute(node: &mut RefNode) {
            let w0 = node.children[0]
                .as_ref()
                .map(|c| c.log_prob_weighted)
                .unwrap_or(0.0);
            let w1 = node.children[1]
                .as_ref()
                .map(|c| c.log_prob_weighted)
                .unwrap_or(0.0);
            let is_leaf = node.children[0].is_none() && node.children[1].is_none();
            node.log_prob_weighted = update_weighted_log_prob(node.log_prob_kt, w0, w1, is_leaf);
        }

        fn update(&mut self, sym: Symbol) {
            let upto = self.root_visits() + 1;
            ensure_log_caches(&mut self.log_int, &mut self.log_half, upto);
            let sym_idx = sym as usize;
            Self::update_node(
                &mut self.root,
                0,
                self.max_depth,
                &self.history,
                sym_idx,
                CachedLogs::new(&self.log_int, &self.log_half),
            );
            self.history.push(sym);
        }

        fn revert(&mut self) {
            let Some(last_sym) = self.history.pop() else {
                return;
            };
            let upto = self.root_visits();
            ensure_log_caches(&mut self.log_int, &mut self.log_half, upto);
            let sym_idx = last_sym as usize;
            let _ = Self::revert_node(
                &mut self.root,
                0,
                self.max_depth,
                &self.history,
                sym_idx,
                CachedLogs::new(&self.log_int, &self.log_half),
            );
        }

        fn predict(&mut self, sym: Symbol) -> f64 {
            let sym_idx = sym as usize;
            let mut entries = Vec::with_capacity(self.max_depth + 1);
            let reached_max_depth = Self::collect_predict_entries(
                &self.root,
                0,
                self.max_depth,
                &self.history,
                &mut entries,
            );

            let deepest = entries.len() - 1;
            let mut ratio = if reached_max_depth && deepest == self.max_depth {
                predict_ratio_kt(entries[deepest].symbol_count, sym_idx)
            } else {
                0.5
            };
            for idx in (0..=deepest).rev() {
                if reached_max_depth && idx == deepest {
                    continue;
                }
                let child_weight = if idx < deepest {
                    entries[idx + 1].log_prob_weighted
                } else {
                    0.0
                };
                ratio = predict_ratio_internal(
                    entries[idx].log_prob_kt,
                    entries[idx].symbol_count,
                    child_weight,
                    entries[idx].sibling_weight,
                    ratio,
                    sym_idx,
                );
            }
            ratio
        }

        fn get_log_block_probability(&self) -> f64 {
            self.root.log_prob_weighted
        }

        fn update_node<L: CtLogAccess>(
            node: &mut RefNode,
            depth: usize,
            max_depth: usize,
            history: &(impl HistoryAccess + ?Sized),
            sym_idx: usize,
            logs: L,
        ) {
            if depth < max_depth {
                let edge = history_symbol(history, depth) as usize;
                if node.children[edge].is_none() {
                    node.children[edge] = Some(Box::new(RefNode::default()));
                }
                Self::update_node(
                    node.children[edge].as_deref_mut().unwrap(),
                    depth + 1,
                    max_depth,
                    history,
                    sym_idx,
                    logs,
                );
            }
            apply_update_to_state_raw(logs, &mut node.symbol_count, &mut node.log_prob_kt, sym_idx);
            Self::recompute(node);
        }

        fn revert_node<L: CtLogAccess>(
            node: &mut RefNode,
            depth: usize,
            max_depth: usize,
            history: &(impl HistoryAccess + ?Sized),
            sym_idx: usize,
            logs: L,
        ) -> bool {
            if depth < max_depth {
                let edge = history_symbol(history, depth) as usize;
                let remove_child = if let Some(child) = node.children[edge].as_deref_mut() {
                    Self::revert_node(child, depth + 1, max_depth, history, sym_idx, logs)
                } else {
                    false
                };
                if remove_child {
                    node.children[edge] = None;
                }
            }
            apply_revert_to_state_raw(logs, &mut node.symbol_count, &mut node.log_prob_kt, sym_idx);
            Self::recompute(node);
            node.symbol_count[0] + node.symbol_count[1] == 0
        }

        fn collect_predict_entries(
            node: &RefNode,
            depth: usize,
            max_depth: usize,
            history: &(impl HistoryAccess + ?Sized),
            entries: &mut Vec<PredictEntry>,
        ) -> bool {
            let sibling_weight = if depth < max_depth {
                let path_edge = history_symbol(history, depth) as usize;
                node.children[path_edge ^ 1]
                    .as_ref()
                    .map(|c| c.log_prob_weighted)
                    .unwrap_or(0.0)
            } else {
                0.0
            };
            entries.push(PredictEntry {
                symbol_count: node.symbol_count,
                log_prob_kt: node.log_prob_kt,
                log_prob_weighted: node.log_prob_weighted,
                sibling_weight,
                has_sibling: depth < max_depth
                    && node.children[(history_symbol(history, depth) as usize) ^ 1].is_some(),
            });
            if depth == max_depth {
                return true;
            }
            let edge = history_symbol(history, depth) as usize;
            let Some(child) = node.children[edge].as_ref() else {
                return false;
            };
            Self::collect_predict_entries(child, depth + 1, max_depth, history, entries)
        }
    }

    #[derive(Clone)]
    struct RefFacContextTree {
        trees: Vec<RefContextTree>,
        history: Vec<Symbol>,
    }

    impl RefFacContextTree {
        fn new(base_depth: usize, num_bits: usize) -> Self {
            Self {
                trees: (0..num_bits)
                    .map(|i| RefContextTree::new(base_depth + i))
                    .collect(),
                history: Vec::new(),
            }
        }

        fn update(&mut self, sym: Symbol, bit_index: usize) {
            let tree = &mut self.trees[bit_index];
            tree.history = self.history.clone();
            tree.update(sym);
            self.history.push(sym);
        }

        fn predict(&mut self, sym: Symbol, bit_index: usize) -> f64 {
            let tree = &mut self.trees[bit_index];
            tree.history = self.history.clone();
            tree.predict(sym)
        }

        fn revert(&mut self, bit_index: usize) {
            let Some(last_sym) = self.history.pop() else {
                return;
            };
            let tree = &mut self.trees[bit_index];
            tree.history = self.history.clone();
            tree.history.push(last_sym);
            tree.revert();
        }

        fn get_log_block_probability(&self) -> f64 {
            self.trees
                .iter()
                .map(RefContextTree::get_log_block_probability)
                .sum()
        }
    }

    fn assert_close(a: f64, b: f64) {
        let diff = (a - b).abs();
        let scale = a.abs().max(b.abs()).max(1.0);
        assert!(diff <= 1e-12 * scale, "a={a} b={b} diff={diff}");
    }

    fn child_after_hot_prefix(
        tree: &ContextTree,
        history_before_update: &(impl HistoryAccess + ?Sized),
    ) -> ChildRef {
        let hot_prefix_depth = tree.engine.hot_prefix_depth();
        if hot_prefix_depth == 0 {
            return ChildRef::NONE;
        }

        let root_edge = history_symbol(history_before_update, 0) as usize;
        let mut current = tree
            .engine
            .arena
            .child(tree.engine.root, root_edge)
            .as_node()
            .expect("hot-prefix node");
        for node_depth in 1..hot_prefix_depth {
            let edge = history_symbol(history_before_update, node_depth) as usize;
            current = tree
                .engine
                .arena
                .child(current, edge)
                .as_node()
                .expect("next hot-prefix node");
        }
        let tail_edge = history_symbol(history_before_update, hot_prefix_depth) as usize;
        tree.engine.arena.child(current, tail_edge)
    }

    #[test]
    #[should_panic(expected = "ctw node index overflow")]
    fn node_index_from_usize_rejects_overflow() {
        let _ = NodeIndex::from_usize(INDEX_LIMIT);
    }

    #[test]
    #[should_panic(expected = "ctw node index overflow")]
    fn node_index_from_usize_rejects_large_values() {
        let _ = NodeIndex::from_usize(u32::MAX as usize);
    }

    #[test]
    fn ctw_count_lane_stays_packed() {
        assert_eq!(std::mem::size_of::<CtNode>(), 32);
    }

    #[test]
    fn ctw_segment_payload_stays_packed() {
        assert_eq!(std::mem::size_of::<CtSegment>(), 40);
    }

    #[test]
    fn log_lookup_matches_direct_log_formulas_past_cache_limit() {
        let log_int = vec![f64::NEG_INFINITY, 0.0, (2.0f64).ln()];
        let log_half = vec![(0.5f64).ln(), (1.5f64).ln(), (2.5f64).ln()];
        let lookup = BoundedLogs::new(&log_int, &log_half);

        for n in 0..16usize {
            let expected_int = if n == 0 {
                f64::NEG_INFINITY
            } else {
                (n as f64).ln()
            };
            assert_eq!(lookup.log_int(n).to_bits(), expected_int.to_bits());
            assert_eq!(
                lookup.log_half(n).to_bits(),
                (n as f64 + 0.5).ln().to_bits()
            );
        }
    }

    #[test]
    fn log_lookup_overflow_cache_reuses_hot_exact_values() {
        let slots = [LogCacheSlot::empty(), LogCacheSlot::empty()];
        let misses = Cell::new(0usize);
        let first = BoundedLogs::lookup_overflow(&slots, 17, || {
            misses.set(misses.get() + 1);
            (17.0f64).ln()
        });
        let second = BoundedLogs::lookup_overflow(&slots, 17, || {
            misses.set(misses.get() + 1);
            f64::NAN
        });

        assert_eq!(first.to_bits(), (17.0f64).ln().to_bits());
        assert_eq!(second.to_bits(), first.to_bits());
        assert_eq!(misses.get(), 1);
    }

    #[test]
    fn shared_log_cache_respects_test_limit() {
        reset_shared_log_cache_for_test();
        let _limit = ScopedLogCacheLimit::set(3);
        let _overflow_slots = ScopedLogOverflowCacheSlots::set(4);
        with_shared_bounded_logs(32, |lookup| {
            assert_eq!(lookup.log_int(32).to_bits(), (32.0f64).ln().to_bits());
            assert_eq!(lookup.log_half(32).to_bits(), (32.5f64).ln().to_bits());
        });

        let (log_int_len, log_half_len) = shared_log_cache_lens();
        let (overflow_log_int_len, overflow_log_half_len) = shared_log_overflow_cache_lens();
        assert!(log_int_len <= 4, "log_int_len={log_int_len}");
        assert!(log_half_len <= 4, "log_half_len={log_half_len}");
        assert_eq!(overflow_log_int_len, 4);
        assert_eq!(overflow_log_half_len, 4);
    }

    #[test]
    fn fac_ctw_memory_usage_breakdown_sums_to_existing_total() {
        let mut fac = FacContextTree::new(7, 8);
        let payload = b"ctw memory usage breakdown payload";
        for &byte in payload {
            fac.update_byte_msb(byte);
        }

        let usage = fac.memory_usage_breakdown();
        assert_eq!(fac.memory_usage(), usage.total_bytes());
        assert!(usage.tree_bytes > 0);
        assert!(usage.shared_log_cache_bytes > 0);
        assert!(usage.shared_history_bytes > 0);

        let telemetry = fac.telemetry();
        assert_eq!(telemetry.total_bytes, usage.total_bytes());
        assert_eq!(telemetry.tree_bytes, usage.tree_bytes);
        assert_eq!(
            telemetry.shared_log_cache_bytes,
            usage.shared_log_cache_bytes
        );
        assert_eq!(telemetry.shared_history_bytes, usage.shared_history_bytes);
        assert_eq!(telemetry.shared_history_len_bits, payload.len() * 8);
        assert_eq!(telemetry.trees.len(), 8);
        assert_eq!(
            telemetry.nodes_len,
            telemetry
                .trees
                .iter()
                .map(|tree| tree.nodes_len)
                .sum::<usize>()
        );
        assert_eq!(
            telemetry.segments_len,
            telemetry
                .trees
                .iter()
                .map(|tree| tree.segments_len)
                .sum::<usize>()
        );
        assert_eq!(
            telemetry.segments_len,
            telemetry.exact_segments
                + telemetry.history_segments
                + telemetry.history_invert_segments
                + telemetry.const_segments
        );
    }

    #[test]
    fn bit_history_preserves_logical_bits_and_packs_memory() {
        let mut history = BitHistory::default();
        let bits: Vec<Symbol> = (0..130usize).map(|idx| (idx * 17 + 5) % 7 < 3).collect();
        history.extend_from_slice(&bits);
        assert_eq!(history.len(), bits.len());
        assert_eq!(history.to_vec(), bits);
        assert_eq!(history.memory_usage(), 3 * size_of::<u64>());
        for depth in 0..160usize {
            assert_eq!(
                history_symbol(&history, depth),
                history_symbol(&bits, depth)
            );
            for len in [0usize, 1, 2, 7, 31, 32, 33, 63, 64] {
                assert_eq!(
                    path_bits_from_history(&history, depth, len),
                    path_bits_from_history(&bits, depth, len),
                    "depth={depth} len={len}",
                );
            }
        }

        let last = history.pop();
        assert_eq!(last, bits.last().copied());
        assert_eq!(history.len(), bits.len() - 1);
        let popped_bits = &bits[..bits.len() - 1];
        for depth in 0..160usize {
            assert_eq!(
                history_symbol(&history, depth),
                history_symbol(popped_bits, depth)
            );
            for len in [0usize, 1, 2, 7, 31, 32, 33, 63, 64] {
                assert_eq!(
                    path_bits_from_history(&history, depth, len),
                    path_bits_from_history(popped_bits, depth, len),
                    "after pop depth={depth} len={len}",
                );
            }
        }

        history.truncate(65);
        assert_eq!(history.to_vec(), bits[..65].to_vec());
        assert_eq!(history.memory_usage(), 3 * size_of::<u64>());

        history.clear();
        assert!(history.is_empty());
        assert_eq!(history.memory_usage(), 3 * size_of::<u64>());
    }

    #[test]
    fn fac_ctw_history_memory_is_bit_packed() {
        let mut fac = FacContextTree::new(4, 8);
        fac.reserve_for_symbols(1_000);
        let usage = fac.memory_usage_breakdown();
        assert_eq!(
            usage.shared_history_bytes,
            history_word_len(8_000) * size_of::<u64>()
        );
    }

    #[test]
    fn fac_ctw_memory_usage_breakdown_reports_bounded_log_cache_component() {
        reset_shared_log_cache_for_test();
        let _limit = ScopedLogCacheLimit::set(2);
        let _overflow_slots = ScopedLogOverflowCacheSlots::set(4);
        let mut fac = FacContextTree::new(7, 8);
        for &byte in b"bounded log cache memory component payload" {
            fac.update_byte_msb(byte);
        }

        let usage = fac.memory_usage_breakdown();
        let (log_int_len, log_half_len) = shared_log_cache_lens();
        let (overflow_log_int_len, overflow_log_half_len) = shared_log_overflow_cache_lens();
        assert!(log_int_len <= 3, "log_int_len={log_int_len}");
        assert!(log_half_len <= 3, "log_half_len={log_half_len}");
        assert_eq!(overflow_log_int_len, 4);
        assert_eq!(overflow_log_half_len, 4);
        assert!(
            usage.shared_log_cache_bytes <= 16 * size_of::<f64>() + 32 * size_of::<LogCacheSlot>()
        );
        assert_eq!(fac.memory_usage(), usage.total_bytes());
    }

    #[test]
    fn bounded_log_lookup_preserves_context_tree_updates_and_reverts() {
        let mut bounded = ContextTree::new(9);
        let mut unbounded = bounded.clone();
        let stream = b"bounded exact log lookup context-tree parity payload";

        {
            let _limit = ScopedLogCacheLimit::set(3);
            for &byte in stream {
                for bit_idx in 0..8usize {
                    let bit = ((byte >> (7 - bit_idx)) & 1) == 1;
                    bounded.update(bit);
                }
            }
        }

        for &byte in stream {
            for bit_idx in 0..8usize {
                let bit = ((byte >> (7 - bit_idx)) & 1) == 1;
                unbounded.update(bit);
            }
        }

        assert_eq!(
            bounded.get_log_block_probability().to_bits(),
            unbounded.get_log_block_probability().to_bits()
        );
        for &sym in &[false, true] {
            assert_eq!(
                bounded.predict(sym).to_bits(),
                unbounded.predict(sym).to_bits()
            );
        }

        {
            let _limit = ScopedLogCacheLimit::set(3);
            for _ in 0..16usize {
                bounded.revert();
            }
        }
        for _ in 0..16usize {
            unbounded.revert();
        }

        assert_eq!(
            bounded.get_log_block_probability().to_bits(),
            unbounded.get_log_block_probability().to_bits()
        );
        for &sym in &[false, true] {
            assert_eq!(
                bounded.predict(sym).to_bits(),
                unbounded.predict(sym).to_bits()
            );
        }
    }

    #[test]
    fn bounded_log_lookup_preserves_fac_byte_fast_paths() {
        let mut bounded_msb = FacContextTree::new(7, 8);
        let mut unbounded_msb = bounded_msb.clone();
        let mut bounded_lsb = FacContextTree::new(7, 5);
        let mut unbounded_lsb = bounded_lsb.clone();
        let stream = b"bounded exact log lookup fac byte parity payload";

        {
            let _limit = ScopedLogCacheLimit::set(2);
            for &byte in stream {
                bounded_msb.update_byte_msb(byte);
                bounded_lsb.update_byte_lsb(byte);
            }
        }
        for &byte in stream {
            unbounded_msb.update_byte_msb(byte);
            unbounded_lsb.update_byte_lsb(byte);
        }

        assert_eq!(
            bounded_msb.get_log_block_probability().to_bits(),
            unbounded_msb.get_log_block_probability().to_bits()
        );
        assert_eq!(
            bounded_lsb.get_log_block_probability().to_bits(),
            unbounded_lsb.get_log_block_probability().to_bits()
        );
        for bit_idx in 0..bounded_msb.num_bits() {
            assert_eq!(
                bounded_msb.predict(false, bit_idx).to_bits(),
                unbounded_msb.predict(false, bit_idx).to_bits()
            );
            assert_eq!(
                bounded_msb.predict_one(bit_idx).to_bits(),
                unbounded_msb.predict_one(bit_idx).to_bits()
            );
        }
        for bit_idx in 0..bounded_lsb.num_bits() {
            assert_eq!(
                bounded_lsb.predict(false, bit_idx).to_bits(),
                unbounded_lsb.predict(false, bit_idx).to_bits()
            );
            assert_eq!(
                bounded_lsb.predict_one(bit_idx).to_bits(),
                unbounded_lsb.predict_one(bit_idx).to_bits()
            );
        }
    }

    #[test]
    fn context_tree_singleton_paths_use_hot_prefix_nodes() {
        let mut tree = ContextTree::new(13);
        tree.update(false);

        let hot_prefix_depth = tree.engine.hot_prefix_depth();
        let child = tree.engine.arena.child(tree.engine.root, 0);
        let mut current = child.as_node().expect("hot-prefix node");
        let mut visited_hot_prefix_nodes = 1usize;
        for depth in 1..hot_prefix_depth {
            let next = tree.engine.arena.child(current, 0);
            current = next.as_node().expect("next hot-prefix node");
            visited_hot_prefix_nodes += 1;
            assert!(depth < hot_prefix_depth);
        }
        assert_eq!(visited_hot_prefix_nodes, hot_prefix_depth);
        let segment = tree
            .engine
            .arena
            .child(current, 0)
            .as_segment()
            .expect("segment tail");
        assert!(tree.engine.arena.child(current, 1).is_none());
        assert!(tree.engine.arena.segments[segment.get()].tail.is_none());
        assert_close(
            tree.engine.arena.segments[segment.get()].head_log_prob_weighted,
            -std::f64::consts::LN_2,
        );
        assert_close(tree.get_log_block_probability(), -std::f64::consts::LN_2);
    }

    #[test]
    fn context_tree_missing_path_tail_uses_exact_segment_payloads() {
        let mut tree = ContextTree::new(13);
        tree.update(true);
        let child = tree.engine.arena.child(tree.engine.root, 0);
        let mut current = child.as_node().expect("hot-prefix node");
        for _ in 1..tree.engine.hot_prefix_depth() {
            current = tree
                .engine
                .arena
                .child(current, 0)
                .as_node()
                .expect("next hot-prefix node");
        }
        let segment = tree
            .engine
            .arena
            .child(current, 0)
            .as_segment()
            .expect("segment tail");
        let payload = tree.engine.arena.segments[segment.get()].payload;
        assert!(payload.is_exact());
        assert_eq!(
            payload.len() as usize,
            tree.engine.max_depth - tree.engine.hot_prefix_depth()
        );
        assert_eq!(payload.exact_bits() & low_bits_mask_u64(payload.len()), 0);
    }

    #[test]
    fn context_tree_missing_path_tail_uses_const_payload_beyond_exact_limit() {
        let mut tree = ContextTree::new(80);
        let history_before = tree.history.clone();
        tree.update(false);

        let segment = child_after_hot_prefix(&tree, &history_before)
            .as_segment()
            .expect("segment tail");
        let segment = tree.engine.arena.segments[segment.get()];
        assert_eq!(segment.payload.mode(), SEG_MODE_CONST);
        assert_eq!(
            segment.payload.len() as usize,
            tree.engine.max_depth - tree.engine.hot_prefix_depth()
        );
        assert!(!segment.payload.const_bit());
        assert!(segment.tail.is_none());
    }

    #[test]
    fn context_tree_missing_path_tail_uses_history_and_const_payloads_beyond_exact_limit() {
        let mut tree = ContextTree::new(80);
        let seeded_history: Vec<Symbol> = (0..80).map(|i| (i & 1) == 1).collect();
        tree.update_history(&seeded_history);
        let history_before = tree.history.clone();
        tree.update(false);

        let first_segment = child_after_hot_prefix(&tree, &history_before)
            .as_segment()
            .expect("history-backed segment tail");
        let first_segment = tree.engine.arena.segments[first_segment.get()];
        assert_eq!(first_segment.payload.mode(), SEG_MODE_HISTORY);
        assert_eq!(
            first_segment.payload.len() as usize,
            tree.engine.max_depth - tree.engine.hot_prefix_depth() - 1
        );
        for offset in [0usize, 1, 7, 31, 66] {
            assert_eq!(
                segment_edge_from_parts(
                    first_segment,
                    offset,
                    &history_before,
                    history_before.len()
                ),
                history_symbol(&history_before, tree.engine.hot_prefix_depth() + 1 + offset)
            );
        }

        let tail_segment = first_segment
            .tail
            .as_segment()
            .expect("constant fallback tail");
        let tail_segment = tree.engine.arena.segments[tail_segment.get()];
        assert_eq!(tail_segment.payload.mode(), SEG_MODE_CONST);
        assert_eq!(tail_segment.payload.len(), 1);
        assert!(!tail_segment.payload.const_bit());
        assert!(tail_segment.tail.is_none());
    }

    #[test]
    fn context_tree_matches_reference_on_short_sequences() {
        for depth in 0..=6usize {
            for len in 0..=6usize {
                for mask in 0..(1usize << len) {
                    let mut prod = ContextTree::new(depth);
                    let mut reference = RefContextTree::new(depth);
                    for step in 0..len {
                        let p_prod_0 = prod.predict(false);
                        let p_ref_0 = reference.predict(false);
                        assert!(
                            (p_prod_0 - p_ref_0).abs()
                                <= 1e-12 * p_prod_0.abs().max(p_ref_0.abs()).max(1.0),
                            "predict0 mismatch depth={depth} len={len} mask={mask} step={step} prod={p_prod_0} ref={p_ref_0} history={:?}",
                            prod.history
                        );
                        let p_prod_1 = prod.predict(true);
                        let p_ref_1 = reference.predict(true);
                        assert!(
                            (p_prod_1 - p_ref_1).abs()
                                <= 1e-12 * p_prod_1.abs().max(p_ref_1.abs()).max(1.0),
                            "predict1 mismatch depth={depth} len={len} mask={mask} step={step} prod={p_prod_1} ref={p_ref_1} history={:?}",
                            prod.history
                        );
                        let log_prod = prod.get_log_block_probability();
                        let log_ref = reference.get_log_block_probability();
                        assert!(
                            (log_prod - log_ref).abs()
                                <= 1e-12 * log_prod.abs().max(log_ref.abs()).max(1.0),
                            "log mismatch before update depth={depth} len={len} mask={mask} step={step} prod={log_prod} ref={log_ref} history={:?}",
                            prod.history
                        );
                        let bit = ((mask >> step) & 1) == 1;
                        prod.update(bit);
                        reference.update(bit);
                        let log_prod = prod.get_log_block_probability();
                        let log_ref = reference.get_log_block_probability();
                        assert!(
                            (log_prod - log_ref).abs()
                                <= 1e-12 * log_prod.abs().max(log_ref.abs()).max(1.0),
                            "log mismatch after update depth={depth} len={len} mask={mask} step={step} bit={bit} prod={log_prod} ref={log_ref} history={:?}",
                            prod.history
                        );
                    }
                    while prod.history_size() > 0 {
                        let p_prod_0 = prod.predict(false);
                        let p_ref_0 = reference.predict(false);
                        assert!(
                            (p_prod_0 - p_ref_0).abs()
                                <= 1e-12 * p_prod_0.abs().max(p_ref_0.abs()).max(1.0),
                            "revert predict0 mismatch depth={depth} len={len} mask={mask} prod={p_prod_0} ref={p_ref_0} history={:?}",
                            prod.history
                        );
                        let p_prod_1 = prod.predict(true);
                        let p_ref_1 = reference.predict(true);
                        assert!(
                            (p_prod_1 - p_ref_1).abs()
                                <= 1e-12 * p_prod_1.abs().max(p_ref_1.abs()).max(1.0),
                            "revert predict1 mismatch depth={depth} len={len} mask={mask} prod={p_prod_1} ref={p_ref_1} history={:?}",
                            prod.history
                        );
                        prod.revert();
                        reference.revert();
                        let log_prod = prod.get_log_block_probability();
                        let log_ref = reference.get_log_block_probability();
                        assert!(
                            (log_prod - log_ref).abs()
                                <= 1e-12 * log_prod.abs().max(log_ref.abs()).max(1.0),
                            "revert log mismatch depth={depth} len={len} mask={mask} prod={log_prod} ref={log_ref} history={:?}",
                            prod.history
                        );
                    }
                }
            }
        }
    }

    #[test]
    fn context_tree_long_depth_matches_reference_on_short_sequences() {
        for &depth in &[65usize, 80usize] {
            for len in 0..=6usize {
                for mask in 0..(1usize << len) {
                    let mut prod = ContextTree::new(depth);
                    let mut reference = RefContextTree::new(depth);
                    for step in 0..len {
                        let p_prod_0 = prod.predict(false);
                        let p_ref_0 = reference.predict(false);
                        assert!(
                            (p_prod_0 - p_ref_0).abs()
                                <= 1e-12 * p_prod_0.abs().max(p_ref_0.abs()).max(1.0),
                            "long-depth predict0 mismatch depth={depth} len={len} mask={mask} step={step} prod={p_prod_0} ref={p_ref_0} history={:?}",
                            prod.history
                        );
                        let p_prod_1 = prod.predict(true);
                        let p_ref_1 = reference.predict(true);
                        assert!(
                            (p_prod_1 - p_ref_1).abs()
                                <= 1e-12 * p_prod_1.abs().max(p_ref_1.abs()).max(1.0),
                            "long-depth predict1 mismatch depth={depth} len={len} mask={mask} step={step} prod={p_prod_1} ref={p_ref_1} history={:?}",
                            prod.history
                        );
                        assert_close(
                            prod.get_log_block_probability(),
                            reference.get_log_block_probability(),
                        );
                        let bit = ((mask >> step) & 1) == 1;
                        prod.update(bit);
                        reference.update(bit);
                        assert_close(
                            prod.get_log_block_probability(),
                            reference.get_log_block_probability(),
                        );
                    }

                    while prod.history_size() > 0 {
                        assert_close(prod.predict(false), reference.predict(false));
                        assert_close(prod.predict(true), reference.predict(true));
                        prod.revert();
                        reference.revert();
                        assert_close(
                            prod.get_log_block_probability(),
                            reference.get_log_block_probability(),
                        );
                    }
                }
            }
        }
    }

    #[test]
    fn fac_ctw_matches_reference_on_short_sequences() {
        let mut fac = FacContextTree::new(4, 4);
        let mut reference = RefFacContextTree::new(4, 4);
        let stream = [
            (true, 0usize),
            (false, 1usize),
            (true, 2usize),
            (true, 3usize),
            (false, 0usize),
            (false, 1usize),
            (true, 2usize),
            (false, 3usize),
        ];

        for &(bit, idx) in &stream {
            assert_close(fac.predict(false, idx), reference.predict(false, idx));
            assert_close(fac.predict(true, idx), reference.predict(true, idx));
            fac.update(bit, idx);
            reference.update(bit, idx);
            assert_close(
                fac.get_log_block_probability(),
                reference.get_log_block_probability(),
            );
        }

        for &(_, idx) in stream.iter().rev() {
            fac.revert(idx);
            reference.revert(idx);
            assert_close(
                fac.get_log_block_probability(),
                reference.get_log_block_probability(),
            );
        }
    }

    #[test]
    fn fac_ctw_history_consistency() {
        let mut fac = FacContextTree::new(4, 4);

        fac.update_history(&[true, false, true]);
        assert_eq!(fac.shared_history.len(), 3);

        fac.update(true, 0);
        fac.update(false, 1);
        assert_eq!(fac.shared_history.len(), 5);

        fac.revert(1);
        assert_eq!(fac.shared_history.len(), 4);

        fac.revert(0);
        assert_eq!(fac.shared_history.len(), 3);
    }

    #[test]
    fn fac_ctw_predict_one_matches_predict_true() {
        let mut fac = FacContextTree::new(6, 8);
        for &byte in b"predict-one exactness regression payload" {
            for bit_idx in 0..8usize {
                let p_generic = fac.predict(true, bit_idx);
                let p_one = fac.predict_one(bit_idx);
                assert_close(p_generic, p_one);
                let bit = ((byte >> (7 - bit_idx)) & 1) == 1;
                fac.update_predicted(bit, bit_idx);
            }
        }
    }

    #[test]
    fn fac_ctw_long_depth_predict_one_matches_predict_true() {
        let mut fac = FacContextTree::new(78, 4);
        for step in 0..24usize {
            for bit_idx in 0..fac.num_bits() {
                let p_generic = fac.predict(true, bit_idx);
                let p_one = fac.predict_one(bit_idx);
                assert_close(p_generic, p_one);
                let bit = ((step * 5 + bit_idx * 3) & 1) == 1;
                fac.update_predicted(bit, bit_idx);
            }
        }
    }

    #[test]
    fn fac_ctw_update_byte_msb_matches_bit_updates() {
        let mut by_byte = FacContextTree::new(6, 8);
        let mut by_bits = FacContextTree::new(6, 8);
        for &byte in b"byte update msb regression payload" {
            by_byte.update_byte_msb(byte);
            for bit_idx in 0..8usize {
                let bit = ((byte >> (7 - bit_idx)) & 1) == 1;
                by_bits.update(bit, bit_idx);
            }
            assert_close(
                by_byte.get_log_block_probability(),
                by_bits.get_log_block_probability(),
            );
            assert_eq!(by_byte.shared_history, by_bits.shared_history);
        }
    }

    #[test]
    fn fac_ctw_log_prob_update_byte_msb_matches_manual_fast_path() {
        let mut batched = FacContextTree::new(6, 8);
        for &byte in b"log prob update byte msb regression payload" {
            let mut manual = batched.clone();
            let observed = batched.log_prob_update_byte_msb(byte);
            let mut expected = 0.0;
            for bit_idx in 0..8usize {
                let bit = ((byte >> (7 - bit_idx)) & 1) == 1;
                let p = manual.predict(bit, bit_idx);
                if p.is_finite() && p > 0.0 {
                    expected += p.ln();
                } else {
                    expected = f64::NEG_INFINITY;
                }
                manual.update_predicted(bit, bit_idx);
            }
            assert_eq!(observed.to_bits(), expected.to_bits());
            assert_eq!(batched.shared_history, manual.shared_history);
            assert_eq!(
                batched.get_log_block_probability().to_bits(),
                manual.get_log_block_probability().to_bits(),
            );
        }
    }

    #[test]
    fn fac_ctw_update_byte_lsb_matches_bit_updates() {
        let mut by_byte = FacContextTree::new(6, 5);
        let mut by_bits = FacContextTree::new(6, 5);
        for &byte in b"byte update lsb regression payload" {
            by_byte.update_byte_lsb(byte);
            for bit_idx in 0..5usize {
                let bit = ((byte >> bit_idx) & 1) == 1;
                by_bits.update(bit, bit_idx);
            }
            assert_close(
                by_byte.get_log_block_probability(),
                by_bits.get_log_block_probability(),
            );
            assert_eq!(by_byte.shared_history, by_bits.shared_history);
        }
    }

    #[test]
    fn fac_ctw_log_cache_tracks_tree_visits_not_shared_history() {
        let mut fac = FacContextTree::new(8, 8);
        let updates_per_tree = 512usize;
        let (log_int_before, log_half_before) = shared_log_cache_lens();

        for step in 0..updates_per_tree {
            let bit = (step & 1) == 1;
            for bit_idx in 0..8usize {
                fac.update(bit, bit_idx);
            }
        }

        assert_eq!(fac.shared_history.len(), updates_per_tree * 8);
        for tree in &fac.trees {
            let visits = tree.engine.arena.visits(tree.engine.root) as usize;
            assert_eq!(visits, updates_per_tree);
        }

        let (log_int_after, log_half_after) = shared_log_cache_lens();
        let expected_len = updates_per_tree + 1;
        assert!(
            log_int_after <= log_int_before.max(expected_len),
            "log_int grew to {log_int_after} (before={log_int_before}, expected_len={expected_len})"
        );
        assert!(
            log_half_after <= log_half_before.max(expected_len),
            "log_half grew to {log_half_after} (before={log_half_before}, expected_len={expected_len})"
        );
    }

    fn seed_fac_cache_regression_state(fac: &mut FacContextTree) {
        for step in 0..24usize {
            for bit_idx in 0..fac.num_bits() {
                let bit = ((step * 3 + bit_idx) & 1) == 1;
                fac.update(bit, bit_idx);
            }
        }
    }

    fn assert_update_predicted_matches_fresh_after_history_rewrite<F>(mut rewrite: F)
    where
        F: FnMut(&mut FacContextTree),
    {
        let mut predicted = FacContextTree::new(6, 4);
        seed_fac_cache_regression_state(&mut predicted);
        let mut fresh = predicted.clone();
        let original_history = predicted.shared_history.clone();
        let target_bit = 2usize;

        let _ = predicted.predict(true, target_bit);
        rewrite(&mut predicted);
        rewrite(&mut fresh);

        assert_eq!(predicted.shared_history.len(), original_history.len());
        assert_ne!(predicted.shared_history, original_history);
        assert_eq!(predicted.shared_history, fresh.shared_history);

        predicted.update_predicted(false, target_bit);
        fresh.update(false, target_bit);

        assert_eq!(predicted.shared_history, fresh.shared_history);
        assert_close(
            predicted.get_log_block_probability(),
            fresh.get_log_block_probability(),
        );
        for bit_idx in 0..predicted.num_bits() {
            assert_close(
                predicted.predict(false, bit_idx),
                fresh.predict(false, bit_idx),
            );
            assert_close(
                predicted.predict(true, bit_idx),
                fresh.predict(true, bit_idx),
            );
        }
    }

    #[test]
    fn fac_ctw_update_predicted_ignores_stale_cache_after_reset_and_rewrite() {
        assert_update_predicted_matches_fresh_after_history_rewrite(|fac| {
            let mut rewritten = fac.shared_history.to_vec();
            for bit in &mut rewritten {
                *bit = !*bit;
            }
            fac.reset_history_only();
            fac.update_history(&rewritten);
        });
    }

    #[test]
    fn fac_ctw_update_predicted_ignores_stale_cache_after_revert_and_rewrite() {
        assert_update_predicted_matches_fresh_after_history_rewrite(|fac| {
            let original = fac.shared_history.to_vec();
            let keep = original.len() / 3;
            let remove = original.len() - keep;
            let mut rewritten_suffix = original[keep..].to_vec();
            for bit in &mut rewritten_suffix {
                *bit = !*bit;
            }
            fac.revert_history(remove);
            fac.update_history(&rewritten_suffix);
        });
    }

    #[test]
    fn fac_ctw_shared_history_version_tracks_mutations() {
        let mut fac = FacContextTree::new(4, 2);
        let mut version = fac.shared_history_version;

        fac.update_history(&[]);
        assert_eq!(fac.shared_history_version, version);

        fac.update_history(&[true, false]);
        assert_ne!(fac.shared_history_version, version);
        version = fac.shared_history_version;

        fac.revert_history(0);
        assert_eq!(fac.shared_history_version, version);

        fac.revert_history(1);
        assert_ne!(fac.shared_history_version, version);
        version = fac.shared_history_version;

        let _ = fac.predict(true, 0);
        assert_eq!(fac.shared_history_version, version);

        fac.update_predicted(true, 0);
        assert_ne!(fac.shared_history_version, version);
        version = fac.shared_history_version;

        fac.reset_history_only();
        assert_ne!(fac.shared_history_version, version);
    }

    #[test]
    fn context_tree_predict_preserves_state() {
        let mut tree = ContextTree::new(6);
        for &bit in &[true, false, true, true, false, false, true, false] {
            tree.update(bit);
        }
        let p0_before = tree.predict(false);
        let p1_before = tree.predict(true);
        let log_before = tree.get_log_block_probability();
        let history_before = tree.history.clone();
        let _ = tree.predict(true);

        assert_eq!(tree.history, history_before);
        assert_close(tree.get_log_block_probability(), log_before);
        assert_close(tree.predict(false), p0_before);
        assert_close(tree.predict(true), p1_before);
    }

    #[test]
    fn context_tree_predict_matches_update_ratio() {
        let mut tree = ContextTree::new(7);
        for &bit in &[true, false, true, false, true, true, false, true, false] {
            tree.update(bit);
        }
        for &sym in &[false, true] {
            let predicted = tree.predict(sym);
            let mut reference = tree.clone();
            let before = reference.get_log_block_probability();
            reference.update(sym);
            let after = reference.get_log_block_probability();
            assert_close(predicted, (after - before).exp());
        }
    }

    #[test]
    fn fac_ctw_predict_preserves_state() {
        let mut fac = FacContextTree::new(5, 8);
        for &byte in b"fac ctw state preservation" {
            for bit_idx in 0..8usize {
                let bit = ((byte >> (7 - bit_idx)) & 1) == 1;
                fac.update(bit, bit_idx);
            }
        }
        let p0_before = fac.predict(false, 3);
        let p1_before = fac.predict(true, 3);
        let log_before = fac.get_log_block_probability();
        let history_before = fac.shared_history.clone();
        let _ = fac.predict(true, 3);

        assert_eq!(fac.shared_history, history_before);
        assert_close(fac.get_log_block_probability(), log_before);
        assert_close(fac.predict(false, 3), p0_before);
        assert_close(fac.predict(true, 3), p1_before);
    }

    #[test]
    fn ctw_inline_node_weight_recompute_matches_recompute_node_weight() {
        fn inline_weight(arena: &CtArena, idx: NodeIndex) -> f64 {
            let slot = idx.get();
            let node = arena.nodes[slot];
            let [left, right] = node.children;
            if left.is_none() && right.is_none() {
                clamp_log_prob(node.log_prob_kt)
            } else {
                let w0 = arena.child_ref_weighted(left);
                let w1 = arena.child_ref_weighted(right);
                update_weighted_log_prob_non_leaf(node.log_prob_kt, w0, w1)
            }
        }

        let mut arena = CtArena::new();
        let parent = arena.alloc_node_with_state([3, 5], -1.75);
        let left_node = arena.alloc_node_with_state([2, 1], -0.25);
        let right_node = arena.alloc_node_with_state([1, 2], -0.50);
        arena.nodes[left_node.get()].log_prob_weighted = -0.333_333_333_f64;
        arena.nodes[right_node.get()].log_prob_weighted = -0.777_777_777_f64;

        let seg = arena.alloc_segment();
        arena.segments[seg.get()].head_log_prob_weighted = -0.125_f64;

        let cases = [
            (ChildRef::NONE, ChildRef::NONE, -2.0_f64),
            (
                ChildRef::from_node(left_node),
                ChildRef::from_node(right_node),
                -1.0_f64,
            ),
            (
                ChildRef::from_node(left_node),
                ChildRef::from_segment(seg),
                -0.625_f64,
            ),
            (
                ChildRef::from_segment(seg),
                ChildRef::from_node(right_node),
                -0.3125_f64,
            ),
        ];

        for (left, right, kt) in cases {
            let slot = parent.get();
            arena.nodes[slot].children = [left, right];
            arena.nodes[slot].log_prob_kt = kt;
            arena.nodes[slot].log_prob_weighted = f64::NAN;

            let expected = inline_weight(&arena, parent);
            arena.recompute_node_weight(parent);
            let actual = arena.nodes[slot].log_prob_weighted;
            assert_eq!(
                actual.to_bits(),
                expected.to_bits(),
                "node-weight recompute mismatch for children=({left:?},{right:?}) kt={kt}"
            );
        }
    }

    #[test]
    fn fac_ctw_predict_matches_update_ratio() {
        let mut fac = FacContextTree::new(6, 8);
        for &byte in b"fac ctw exact predictive ratio" {
            for bit_idx in 0..8usize {
                let bit = ((byte >> (7 - bit_idx)) & 1) == 1;
                fac.update(bit, bit_idx);
            }
        }
        for &sym in &[false, true] {
            let predicted = fac.predict(sym, 4);
            let mut reference = fac.clone();
            let before = reference.get_log_block_probability();
            reference.update(sym, 4);
            let after = reference.get_log_block_probability();
            assert_close(predicted, (after - before).exp());
        }
    }

    #[test]
    fn fac_ctw_update_predicted_matches_fresh_update_on_byte_stream() {
        let mut predicted = FacContextTree::new(6, 8);
        let mut fresh = predicted.clone();
        let stream = b"fac-ctw prepared update exactness regression";

        for (byte_pos, &byte) in stream.iter().enumerate() {
            for bit_idx in 0..8usize {
                let bit = ((byte >> (7 - bit_idx)) & 1) == 1;
                let _ = predicted.predict(true, bit_idx);
                predicted.update_predicted(bit, bit_idx);
                fresh.update(bit, bit_idx);

                let predicted_log = predicted.get_log_block_probability();
                let fresh_log = fresh.get_log_block_probability();
                assert!(
                    (predicted_log - fresh_log).abs()
                        <= 1e-12 * predicted_log.abs().max(fresh_log.abs()).max(1.0),
                    "log mismatch byte_pos={byte_pos} bit_idx={bit_idx} bit={bit} predicted={predicted_log} fresh={fresh_log}\nshared_history={:?}\npredicted_arena={:#?}\nfresh_arena={:#?}\npredicted_steps={:?}\nfresh_steps={:?}",
                    predicted.shared_history,
                    predicted.trees[bit_idx].engine.arena,
                    fresh.trees[bit_idx].engine.arena,
                    predicted.trees[bit_idx].engine.prepared_steps,
                    fresh.trees[bit_idx].engine.prepared_steps,
                );
                for probe_idx in 0..8usize {
                    let p_pred_0 = predicted.predict(false, probe_idx);
                    let p_fresh_0 = fresh.predict(false, probe_idx);
                    assert!(
                        (p_pred_0 - p_fresh_0).abs()
                            <= 1e-12 * p_pred_0.abs().max(p_fresh_0.abs()).max(1.0),
                        "predict0 mismatch byte_pos={byte_pos} bit_idx={bit_idx} probe_idx={probe_idx} predicted={p_pred_0} fresh={p_fresh_0}",
                    );
                    let p_pred_1 = predicted.predict(true, probe_idx);
                    let p_fresh_1 = fresh.predict(true, probe_idx);
                    assert!(
                        (p_pred_1 - p_fresh_1).abs()
                            <= 1e-12 * p_pred_1.abs().max(p_fresh_1.abs()).max(1.0),
                        "predict1 mismatch byte_pos={byte_pos} bit_idx={bit_idx} probe_idx={probe_idx} predicted={p_pred_1} fresh={p_fresh_1}",
                    );
                }
            }
        }
    }

    #[test]
    fn fac_ctw_long_depth_update_predicted_matches_fresh_update_on_bit_stream() {
        let mut predicted = FacContextTree::new(78, 4);
        let mut fresh = predicted.clone();

        for step in 0..20usize {
            for bit_idx in 0..predicted.num_bits() {
                let bit = ((step * 7 + bit_idx * 11) & 1) == 1;
                let _ = predicted.predict(true, bit_idx);
                predicted.update_predicted(bit, bit_idx);
                fresh.update(bit, bit_idx);
                assert_eq!(predicted.shared_history, fresh.shared_history);
                assert_close(
                    predicted.get_log_block_probability(),
                    fresh.get_log_block_probability(),
                );
            }
        }

        for bit_idx in 0..predicted.num_bits() {
            assert_close(
                predicted.predict(false, bit_idx),
                fresh.predict(false, bit_idx),
            );
            assert_close(
                predicted.predict(true, bit_idx),
                fresh.predict(true, bit_idx),
            );
        }
    }

    fn scan_symbol_space(tree: &mut FacContextTree, bits: usize) {
        fn rec(tree: &mut FacContextTree, bits: usize, depth: usize) {
            if depth == bits {
                return;
            }
            for bit in [false, true] {
                let bit_idx = depth;
                tree.update(bit, bit_idx);
                rec(tree, bits, depth + 1);
                tree.revert(bit_idx);
            }
        }
        rec(tree, bits, 0);
    }

    fn fac_symbol_bit(symbol: u8, msb_first: bool, bits: usize, bit_idx: usize) -> bool {
        if msb_first {
            ctw_symbol_bit_msb(symbol, bits, bit_idx)
        } else {
            ((symbol >> bit_idx) & 1) == 1
        }
    }

    fn symbol_log_prob(tree: &mut FacContextTree, symbol: u8, msb_first: bool, bits: usize) -> f64 {
        let before = tree.get_log_block_probability();
        for bit_idx in 0..bits {
            let bit = fac_symbol_bit(symbol, msb_first, bits, bit_idx);
            tree.update(bit, bit_idx);
        }
        let after = tree.get_log_block_probability();
        for bit_idx in (0..bits).rev() {
            tree.revert(bit_idx);
        }
        after - before
    }

    fn assert_symbol_scan_then_update_matches_plain(msb_first: bool) {
        let bits = 8usize;
        let mut with_scan = FacContextTree::new(7, bits);
        let mut plain = with_scan.clone();
        for &byte in b"pdf then update parity payload" {
            for bit_idx in 0..bits {
                let bit = fac_symbol_bit(byte, msb_first, bits, bit_idx);
                with_scan.update(bit, bit_idx);
                plain.update(bit, bit_idx);
            }
        }

        scan_symbol_space(&mut with_scan, bits);

        let observed = b'n';
        for bit_idx in 0..bits {
            let bit = fac_symbol_bit(observed, msb_first, bits, bit_idx);
            with_scan.update(bit, bit_idx);
            plain.update(bit, bit_idx);
        }

        for sym in 0u8..=255u8 {
            let lp_scan = symbol_log_prob(&mut with_scan, sym, msb_first, bits);
            let lp_plain = symbol_log_prob(&mut plain, sym, msb_first, bits);
            let diff = (lp_scan - lp_plain).abs();
            assert!(
                diff < 1e-12,
                "symbol={sym} lp_scan={lp_scan} lp_plain={lp_plain} diff={diff}",
            );
        }
    }

    fn assert_fill_fac_tree_log_probs_matches_direct_symbol_probs(msb_first: bool, bits: usize) {
        let bits = bits.clamp(1, 8);
        let min_logp = 1e-12f64.ln();
        let patterns = 1usize << bits;
        let alias_ln = if bits == 8 {
            0.0
        } else {
            ((1usize << (8 - bits)) as f64).ln()
        };
        let training = [0x0u8, 0x3, 0x5, 0x6, 0x9, 0xA, 0xC, 0xF, 0x7, 0x1];
        let mut tree = FacContextTree::new(7, bits);
        for &symbol in &training {
            for bit_idx in 0..bits {
                tree.update(fac_symbol_bit(symbol, msb_first, bits, bit_idx), bit_idx);
            }
        }

        let log_before = tree.get_log_block_probability();
        let mut predict_zero_before = vec![0.0; bits];
        let mut predict_one_before = vec![0.0; bits];
        for bit_idx in 0..bits {
            predict_zero_before[bit_idx] = tree.predict(false, bit_idx);
            predict_one_before[bit_idx] = tree.predict(true, bit_idx);
        }

        let mut out = [0.0; 256];
        fill_fac_tree_log_probs(&mut tree, bits, msb_first, min_logp, &mut out);

        assert_close(tree.get_log_block_probability(), log_before);
        for bit_idx in 0..bits {
            assert_close(tree.predict(false, bit_idx), predict_zero_before[bit_idx]);
            assert_close(tree.predict(true, bit_idx), predict_one_before[bit_idx]);
        }

        let mask = patterns - 1;
        for (byte, &actual) in out.iter().enumerate() {
            let symbol = if bits == 8 {
                byte as u8
            } else {
                (byte & mask) as u8
            };
            let expected =
                symbol_log_prob(&mut tree, symbol, msb_first, bits).max(min_logp) - alias_ln;
            assert_close(actual, expected);
        }
    }

    #[test]
    fn fac_ctw_symbol_scan_then_update_matches_plain_msb() {
        assert_symbol_scan_then_update_matches_plain(true);
    }

    #[test]
    fn fac_ctw_symbol_scan_then_update_matches_plain_lsb() {
        assert_symbol_scan_then_update_matches_plain(false);
    }

    #[test]
    fn fill_fac_tree_log_probs_matches_direct_symbol_probs_for_subbyte_msb() {
        assert_fill_fac_tree_log_probs_matches_direct_symbol_probs(true, 4);
    }

    #[test]
    fn fill_fac_tree_log_probs_matches_direct_symbol_probs_for_subbyte_lsb() {
        assert_fill_fac_tree_log_probs_matches_direct_symbol_probs(false, 4);
    }
}