smith

//! Allocation-count bounds on session hot paths (`SMH-SPEC-SPEC0001`,
//! memory accounting). Bounds only tighten.

use smith::message::{Message, Role};
use smith_alloc::{Counting, Scope, measure};
use smith_core::frame::Frame;
use smith_core::session::{EntryContent, EntryFrame, Session};
use std::sync::Mutex;

#[global_allocator]
static ALLOC: Counting = Counting;

/// Scope counters are process-wide; measurements must not overlap.
static MEASURING: Mutex<()> = Mutex::new(());

/// `fork_at` over 1000 entries: the branch list, the record `Vec`, and one
/// arena slice. Content bytes are copied, never decoded.
const FORK_AT_1000_ALLOCATIONS: u64 = 3;

/// Encoding a header frame: one buffer sized by the payload hint holds the
/// prefix and the whole payload.
const HEADER_ENCODE_ALLOCATIONS: u64 = 1;

/// `from_frames` over 1000 decoded frames: arena, stored records, the
/// parent-link table, the branch list, and one branch's records, each sized
/// up front. Independent of entry count.
const FROM_FRAMES_1000_ALLOCATIONS: u64 = 5;

fn linear_session(len: usize) -> Session {
    let mut session = Session::new();
    for turn in 0..len {
        let content = EntryContent::Message(Message::with_text(Role::User, format!("turn {turn}")));
        session.append(&content).unwrap();
    }
    session
}

/// Allocations and bytes charged to [`Scope::Session`] while `run` executes.
fn session_cost<T>(run: impl FnOnce() -> T) -> ((u64, u64), T) {
    let _serial = MEASURING
        .lock()
        .unwrap_or_else(std::sync::PoisonError::into_inner);
    let (out, cost) = measure(Scope::Session, run);
    ((cost.allocations, cost.bytes), out)
}

/// Cost of one `fork_at` of the leaf.
fn fork_cost(session: &Session) -> (u64, u64) {
    let leaf = session.selected_id().expect("non-empty session");
    let (cost, fork) = session_cost(|| session.fork_at(leaf).expect("leaf is known"));
    assert_eq!(
        fork.active_branch().entries().len(),
        session.active_branch().entries().len()
    );
    let last = fork.active_branch().entries()[fork.active_branch().entries().len() - 1];
    assert_eq!(fork.bytes(&last), session.bytes(&last));
    cost
}

/// Cost of restoring `session` from its decoded frames; frame decoding
/// itself happens before the measured scope.
fn from_frames_cost(session: &Session) -> (u64, u64) {
    let mut frames = vec![Frame::Header {
        header: session.header(),
    }];
    frames.extend(
        session
            .active_branch()
            .entries()
            .iter()
            .map(|entry| Frame::Known {
                entry: EntryFrame {
                    id: entry.id,
                    parent: entry.parent,
                    timestamp_ms: entry.timestamp_ms,
                    kind: entry.kind,
                    content: session.bytes(entry).to_vec(),
                },
            }),
    );
    let (cost, restored) = session_cost(|| Session::from_frames(&frames).expect("frames restore"));
    assert_eq!(&restored, session);
    cost
}

#[test]
fn fork_at_over_1000_entries_is_bounded_and_reproducible() {
    let session = linear_session(1000);
    let first = fork_cost(&session);
    let second = fork_cost(&session);
    assert_eq!(
        first.0, FORK_AT_1000_ALLOCATIONS,
        "fork_at allocated {} times, bound is {FORK_AT_1000_ALLOCATIONS}",
        first.0
    );
    assert_eq!(first, second);
}

#[test]
fn from_frames_allocates_independently_of_entry_count() {
    let small = from_frames_cost(&linear_session(100));
    let large = from_frames_cost(&linear_session(1000));
    assert_eq!(
        large.0, FROM_FRAMES_1000_ALLOCATIONS,
        "from_frames allocated {} times, bound is {FROM_FRAMES_1000_ALLOCATIONS}",
        large.0
    );
    assert_eq!(small.0, large.0);
    assert_eq!(large, from_frames_cost(&linear_session(1000)));
}

#[test]
fn header_frame_encodes_into_one_allocation() {
    let header = Frame::Header {
        header: Session::new().header(),
    };
    let (cost, encoded) =
        session_cost(|| header.encode(smith_core::frame::DEFAULT_MAX_FRAME_BYTES));
    assert!(encoded.is_ok());
    assert_eq!(
        cost.0, HEADER_ENCODE_ALLOCATIONS,
        "header encode allocated {} times, bound is {HEADER_ENCODE_ALLOCATIONS}",
        cost.0
    );
}