smith

//! Shared fixtures for Smith benchmarks (`SMH-SPEC-SPEC0001`, Performance
//! and portability).
//!
//! Benches live in `benches/`: `instructions` counts callgrind instructions
//! against the pins in `budgets.toml` (`cargo x perf`), `wall_clock` runs
//! criterion for local A/B comparisons (`cargo x bench`). Builders here run
//! in bench setup, outside every measured region, and use only public API,
//! so a bench measures what a caller of the crate gets.

#![deny(unsafe_code)]

use futures::StreamExt;
use smith::error::Result;
use smith::id::{MessageId, ToolCallId};
use smith::message::{Message, Role};
use smith::stream::{StopReason, StreamEvent};
use smith::{StreamFn, ToolArguments, Usage};
use smith_core::agent::Agent;
use smith_core::frame::{DEFAULT_MAX_FRAME_BYTES, Frame, encode_entry};
use smith_core::session::{EntryContent, EntryFrame, Session};
use smith_core::tools::ToolSession;
use std::sync::Arc;

/// Plaintext the agent fixture registers as a secret.
pub const SECRET: &str = "sk-bench-0123456789abcdef";

/// A session of `len` chat messages, alternating user and assistant.
///
/// # Errors
///
/// Returns the session error of the first failed append.
pub fn linear_session(len: usize) -> Result<Session> {
    let mut session = Session::new();
    for turn in 0..len {
        let role = if turn % 2 == 0 {
            Role::User
        } else {
            Role::Assistant
        };
        let message = Message::with_text(role, format!("turn {turn}: a short chat message"));
        session.append(&EntryContent::Message(message))?;
    }
    Ok(session)
}

/// A session of `pairs` tool calls, each followed by its result.
///
/// # Errors
///
/// Returns the session error of the first failed append.
pub fn tool_call_session(pairs: usize) -> Result<Session> {
    let mut session = Session::new();
    for call in 0..pairs {
        let call_id = ToolCallId::new();
        session.append(&EntryContent::ToolCall {
            name: "read".to_string(),
            input: ToolArguments::from_serializable(
                &serde_json::json!({"path": format!("src/module_{call}.rs"), "limit": 200}),
            )
            .map_err(|error| smith::error::SmithError::Session {
                code: smith::error::SessionCode::EntryDecode,
                message: format!("cannot encode benchmark tool arguments: {error}"),
            })?,
            call_id,
        })?;
        session.append(&EntryContent::ToolResult {
            ok: true,
            output: format!("pub fn item_{call}() -> u32 {{ {call} }}\n"),
            call_id,
        })?;
    }
    Ok(session)
}

/// Session plus exact aggregate frame-output capacity, prepared outside a
/// measured encode invocation.
pub struct PreallocatedEncodeFixture {
    session: Session,
    output: Vec<u8>,
}

/// Encode `session` during setup to size aggregate frame output exactly.
#[must_use]
pub fn preallocated_encode_fixture(session: Session) -> PreallocatedEncodeFixture {
    let output = Vec::with_capacity(encode_frames(&session, Vec::new()).len());
    PreallocatedEncodeFixture { session, output }
}

/// Encode every active entry into aggregate output supplied by setup.
#[must_use]
pub fn encode_preallocated(mut fixture: PreallocatedEncodeFixture) -> (Session, Vec<u8>) {
    let initial_capacity = fixture.output.capacity();
    encode_frames_into(&fixture.session, &mut fixture.output);
    debug_assert_eq!(fixture.output.len(), initial_capacity);
    debug_assert_eq!(fixture.output.capacity(), initial_capacity);
    (fixture.session, fixture.output)
}

fn encode_frames(session: &Session, mut output: Vec<u8>) -> Vec<u8> {
    encode_frames_into(session, &mut output);
    output
}

#[expect(
    clippy::expect_used,
    reason = "bounded benchmark fixtures must encode; panic makes a fixture fault visible"
)]
fn encode_frames_into(session: &Session, output: &mut Vec<u8>) {
    for entry in session.active_branch().entries() {
        let frame = encode_entry(entry, session.bytes(entry), DEFAULT_MAX_FRAME_BYTES)
            .expect("entry encodes");
        output.extend_from_slice(&frame);
    }
}

/// An in-memory frame corpus for restoring `session`: its header, then one
/// entry frame per active-branch record.
#[must_use]
pub fn frames_of(session: &Session) -> Vec<Frame> {
    let entries = session.active_branch().entries();
    let mut frames = Vec::with_capacity(entries.len() + 1);
    frames.push(Frame::Header {
        header: session.header(),
    });
    frames.extend(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(),
        },
    }));
    frames
}

/// A JSON value `depth` levels deep, four keys per object, every object
/// in reverse key order so canonicalization has to reorder all of them.
#[must_use]
pub fn nested_json(depth: u32) -> serde_json::Value {
    let mut object = serde_json::Map::new();
    for key in (0..4u32).rev() {
        let value = match (depth, key) {
            (0, _) => serde_json::Value::from(format!("leaf {key}")),
            (_, 0) => serde_json::json!([key, "item", {"z": depth, "a": key}]),
            _ => nested_json(depth - 1),
        };
        object.insert(format!("key_{key}"), value);
    }
    serde_json::Value::Object(object)
}

/// An agent over a session of `history` messages with [`SECRET`]
/// registered, whose provider answers every request with one text delta
/// and a clean stop.
///
/// # Errors
///
/// Returns the session error of building the history.
pub fn masked_agent(history: usize) -> Result<Agent> {
    let mut tools = ToolSession::new(".");
    tools.session = linear_session(history)?;
    let mut agent = Agent::new(scripted_reply(), tools, "bench-model");
    agent.secrets_mut().register("api_key", SECRET);
    Ok(agent)
}

/// A provider that streams `ok` and stops, whatever the request.
fn scripted_reply() -> StreamFn {
    Arc::new(|_request, _cancel| {
        let reply = MessageId::new();
        futures::stream::iter([
            Ok(StreamEvent::text_delta(reply, "ok")),
            Ok(StreamEvent::stop(StopReason::EndTurn, Some(Usage::new()))),
        ])
        .boxed()
    })
}

/// SSE bytes: one `data:` event per payload.
fn sse(payloads: &[&str]) -> Vec<u8> {
    let mut out = String::new();
    for payload in payloads {
        out.push_str("data: ");
        out.push_str(payload);
        out.push_str("\n\n");
    }
    out.into_bytes()
}

/// An `OpenAI` chat-completions stream: text, fragmented tool arguments,
/// usage, done; the shapes of the adapter's conformance fixtures.
#[must_use]
pub fn openai_sse() -> Vec<u8> {
    sse(&[
        r#"{"choices":[{"delta":{"role":"assistant","content":""}}]}"#,
        r#"{"choices":[{"delta":{"content":"Hello"}}]}"#,
        r#"{"choices":[{"delta":{"content":" world"}}]}"#,
        r#"{"choices":[{"delta":{"tool_calls":[{"index":0,"id":"a","function":{"name":"write","arguments":""}}]}}]}"#,
        r#"{"choices":[{"delta":{"tool_calls":[{"index":0,"function":{"arguments":"{\"path\""}}]}}]}"#,
        r#"{"choices":[{"delta":{"tool_calls":[{"index":0,"function":{"arguments":":\"f.txt\",\"content\":\"hi\"}"}}]}}]}"#,
        r#"{"choices":[{"delta":{},"finish_reason":"tool_calls"}]}"#,
        r#"{"choices":[],"usage":{"prompt_tokens":9,"completion_tokens":2,"prompt_tokens_details":{"cached_tokens":4}}}"#,
        "[DONE]",
    ])
}

/// An Anthropic Messages stream: text, thinking, assembled tool input,
/// usage; the adapter's conformance fixture.
#[must_use]
pub fn anthropic_sse() -> Vec<u8> {
    sse(&[
        r#"{"type":"message_start","message":{"usage":{"input_tokens":12}}}"#,
        r#"{"type":"content_block_start","index":0,"content_block":{"type":"text"}}"#,
        r#"{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"Look"}}"#,
        r#"{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"ing"}}"#,
        r#"{"type":"content_block_stop","index":0}"#,
        r#"{"type":"content_block_start","index":1,"content_block":{"type":"thinking"}}"#,
        r#"{"type":"content_block_delta","index":1,"delta":{"type":"thinking_delta","thinking":"hmm"}}"#,
        r#"{"type":"content_block_stop","index":1}"#,
        r#"{"type":"content_block_start","index":2,"content_block":{"type":"tool_use","id":"t","name":"write"}}"#,
        r#"{"type":"content_block_delta","index":2,"delta":{"type":"input_json_delta","partial_json":"{\"path\""}}"#,
        r#"{"type":"content_block_delta","index":2,"delta":{"type":"input_json_delta","partial_json":":[1,2]}"}}"#,
        r#"{"type":"content_block_stop","index":2}"#,
        r#"{"type":"message_delta","delta":{"stop_reason":"tool_use"},"usage":{"output_tokens":7,"cache_read_input_tokens":3}}"#,
        r#"{"type":"message_stop"}"#,
    ])
}

/// A Gemini stream: text around a complete function call, usage, stop;
/// the adapter's conformance fixture.
#[must_use]
pub fn gemini_sse() -> Vec<u8> {
    sse(&[
        r#"{"candidates":[{"content":{"parts":[{"text":"Writing"}]}}]}"#,
        r#"{"candidates":[{"content":{"parts":[{"functionCall":{"name":"write","args":{"path":"f","content":"hi"}}}]}}],"usageMetadata":{"promptTokenCount":5,"candidatesTokenCount":9,"cachedContentTokenCount":2}}"#,
        r#"{"candidates":[{"content":{"parts":[{"text":" now"}]}}]}"#,
        r#"{"candidates":[{"content":{},"finishReason":"STOP"}]}"#,
    ])
}

/// Setup-owned resident bytes and equally sized output for byte-transfer models.
#[must_use]
pub fn byte_transfer_fixture(input: Vec<u8>) -> (Vec<u8>, Vec<u8>) {
    let output = vec![0; input.len()];
    (input, output)
}

/// Copy every resident input byte into setup-owned output.
#[must_use]
pub fn byte_transfer(mut fixture: (Vec<u8>, Vec<u8>)) -> (Vec<u8>, Vec<u8>) {
    fixture.1.copy_from_slice(std::hint::black_box(&fixture.0));
    fixture
}

/// Visit every JSON value and write its resident payload bytes without JSON
/// formatting, validation, sorting, or number conversion.
pub fn visit_value_payload(value: &serde_json::Value, output: &mut Vec<u8>) {
    match value {
        serde_json::Value::Null => output.push(0),
        serde_json::Value::Bool(value) => output.push(u8::from(*value)),
        serde_json::Value::Number(value) => {
            if let Some(value) = value.as_i64() {
                output.extend_from_slice(&value.to_ne_bytes());
            } else if let Some(value) = value.as_u64() {
                output.extend_from_slice(&value.to_ne_bytes());
            } else if let Some(value) = value.as_f64() {
                output.extend_from_slice(&value.to_bits().to_ne_bytes());
            }
        }
        serde_json::Value::String(value) => output.extend_from_slice(value.as_bytes()),
        serde_json::Value::Array(values) => {
            output.push(4);
            for value in values {
                visit_value_payload(value, output);
            }
        }
        serde_json::Value::Object(values) => {
            output.push(5);
            for (key, value) in values {
                output.extend_from_slice(key.as_bytes());
                visit_value_payload(value, output);
            }
        }
    }
}

/// Setup-owned canonical-JSON model input and sufficient output capacity.
#[must_use]
pub fn value_payload_fixture(input: serde_json::Value) -> (serde_json::Value, Vec<u8>) {
    let mut output = Vec::new();
    visit_value_payload(&input, &mut output);
    output.clear();
    (input, output)
}

/// Read all 128 input bits and write 47 bytes without alphabet mapping.
/// Repeating the input models output traffic, not a valid text encoding.
#[must_use]
pub fn proquint_encode_floor(mut fixture: (u128, Vec<u8>)) -> (u128, Vec<u8>) {
    let input = std::hint::black_box(fixture.0).to_le_bytes();
    assert_eq!(fixture.1.len(), 47);
    for output in fixture.1.chunks_mut(input.len()) {
        output.copy_from_slice(&input[..output.len()]);
    }
    fixture
}

/// Read all 47 input bytes and write 16 bytes without alphabet validation.
/// XOR reduction keeps every read observable; its work is model overhead.
#[must_use]
pub fn proquint_decode_floor(mut fixture: (String, Vec<u8>)) -> (String, Vec<u8>) {
    let input = std::hint::black_box(fixture.0.as_bytes());
    assert_eq!(input.len(), 47);
    assert_eq!(fixture.1.len(), std::mem::size_of::<u128>());
    let mut value = 0_u128;
    for chunk in input.chunks(std::mem::size_of::<u128>()) {
        let mut bytes = [0; std::mem::size_of::<u128>()];
        bytes[..chunk.len()].copy_from_slice(chunk);
        value ^= u128::from_le_bytes(bytes);
    }
    fixture.1.copy_from_slice(&value.to_le_bytes());
    fixture
}

#[cfg(test)]
#[expect(
    clippy::unwrap_used,
    reason = "tests may panic on invariant violations"
)]
mod tests {
    use super::*;
    use smith_ai::{AnthropicDecoder, GeminiDecoder, OpenAiDecoder, StreamDecoder};

    #[test]
    fn fixtures_have_the_sizes_benches_name() {
        let linear = linear_session(1000).unwrap();
        assert_eq!(linear.active_branch().entries().len(), 1000);
        assert_eq!(frames_of(&linear).len(), 1001);
        let calls = tool_call_session(500).unwrap();
        assert_eq!(calls.active_branch().entries().len(), 1000);
    }

    #[test]
    fn preallocated_frame_output_matches_fixture_without_growth() {
        for session in [
            linear_session(1000).unwrap(),
            tool_call_session(500).unwrap(),
        ] {
            let expected = encode_frames(&session, Vec::new());
            let fixture = preallocated_encode_fixture(session);
            let output_capacity = fixture.output.capacity();
            assert_eq!(output_capacity, expected.len());

            let (session, output) = encode_preallocated(fixture);
            assert_eq!(output, expected);
            assert_eq!(output.len(), output_capacity);
            assert_eq!(output.capacity(), output_capacity);
            assert_eq!(session.active_branch().entries().len(), 1000);
        }
    }

    #[test]
    fn nested_json_is_not_canonical() {
        let value = nested_json(2);
        let keys: Vec<&String> = value.as_object().unwrap().keys().collect();
        assert_eq!(keys, ["key_3", "key_2", "key_1", "key_0"]);
    }

    #[test]
    fn every_sse_fixture_decodes_to_a_clean_stop() {
        let cases: [(Box<dyn StreamDecoder>, Vec<u8>); 3] = [
            (Box::new(OpenAiDecoder::new()), openai_sse()),
            (Box::new(AnthropicDecoder::new()), anthropic_sse()),
            (Box::new(GeminiDecoder::new()), gemini_sse()),
        ];
        for (mut decoder, bytes) in cases {
            let mut events = decoder.push(&bytes).unwrap();
            events.extend(decoder.finish().unwrap());
            assert!(
                matches!(events.last(), Some(StreamEvent::Stop { .. })),
                "{events:?}"
            );
        }
    }

    #[test]
    fn masked_agent_answers_a_turn() {
        let mut agent = masked_agent(10).unwrap();
        agent.run_turn(format!("use {SECRET}")).unwrap();
        let entries = agent.session().active_branch().entries();
        assert_eq!(entries.len(), 12);
        let input = agent.session().content(&entries[10]).unwrap();
        assert!(!format!("{input:?}").contains(SECRET), "{input:?}");
    }

    #[test]
    fn byte_transfer_model_reads_and_writes_every_resident_byte_without_growth() {
        let fixture = byte_transfer_fixture(vec![1, 2, 3, 4]);
        let capacity = fixture.1.capacity();
        let (input, output) = byte_transfer(fixture);
        assert_eq!(output, input);
        assert_eq!(output.capacity(), capacity);
    }

    #[test]
    fn value_payload_model_visits_number_and_container_payloads_without_growth() {
        let fixture = value_payload_fixture(serde_json::json!({"key": ["payload", 7, true, null]}));
        let capacity = fixture.1.capacity();
        let (value, mut output) = fixture;
        visit_value_payload(&value, &mut output);
        for payload in [b"key".as_slice(), b"payload"] {
            assert!(output.windows(payload.len()).any(|bytes| bytes == payload));
        }
        assert!(output.contains(&7));
        assert!(output.contains(&1));
        assert!(output.contains(&0));
        assert_eq!(output.capacity(), capacity);
    }

    #[test]
    fn proquint_models_match_codec_input_and_output_widths_without_growth() {
        let bits = 0x0123_4567_89ab_cdef_0123_4567_89ab_cdef_u128;
        let text = smith::id::proquint::encode(bits);
        let encode = (bits, vec![0; text.len()]);
        let encode_capacity = encode.1.capacity();
        let (_, encoded) = proquint_encode_floor(encode);
        assert_eq!(encoded.len(), text.len());
        assert_eq!(encoded.capacity(), encode_capacity);
        assert_eq!(&encoded[..16], &bits.to_le_bytes());

        let decode = (text, vec![0; std::mem::size_of::<u128>()]);
        let decode_capacity = decode.1.capacity();
        let (_, decoded) = proquint_decode_floor(decode);
        assert_eq!(decoded.len(), std::mem::size_of::<u128>());
        assert_eq!(decoded.capacity(), decode_capacity);
    }

    #[test]
    fn proquint_models_observe_every_input_bit_and_byte() {
        let (_, zero) = proquint_encode_floor((0, vec![0; 47]));
        for bit in 0..128 {
            let (_, output) = proquint_encode_floor((1_u128 << bit, vec![0; 47]));
            assert_ne!(output, zero, "input bit {bit} was not observed");
        }
        let (_, zero) = proquint_decode_floor(("a".repeat(47), vec![0; 16]));
        for position in 0..47 {
            let mut input = vec![b'a'; 47];
            input[position] = b'b';
            let (_, output) =
                proquint_decode_floor((String::from_utf8(input).unwrap(), vec![0; 16]));
            assert_ne!(output, zero, "input byte {position} was not observed");
        }
    }
}