subms 0.2.1

//! `subms` — a tiny, std-only perf-test harness for Rust programs.
//!
//! Records timed samples per stage, computes percentiles, and emits a
//! stable JSON shape suitable for upload to
//! [submillisecond.com](https://submillisecond.com) cookbook samples.
//!
//! Zero external dependencies.
//!
//! # Example
//!
//! ```
//! use subms::PerfHarness;
//!
//! let mut h = PerfHarness::new("lsm-tree", "rust");
//! h.input("entries", &50_000.to_string());
//! h.input("bloom_mode", "on");
//! h.meta("sstables", "46");
//!
//! let put = h.stage("put", 50_000);
//! for _ in 0..50_000 {
//!     put.time(|| { /* work under test */ });
//! }
//!
//! h.write_json(&mut std::io::stdout()).unwrap();
//! ```
//!
//! # JSON shape (stable; matches the Java harness in the cookbook)
//!
//! ```text
//! {
//!   "workload": "lsm-tree",
//!   "lang": "rust",
//!   "timestamp": "2026-05-13T20:24:38Z",
//!   "inputs":  { "<k>": "<v>", ... },
//!   "meta":    { "<k>": "<v>", ... },
//!   "stages": {
//!     "<name>": {
//!       "count": <int>,
//!       "p50_ns": <int>, "p99_ns": <int>, "p999_ns": <int>, "max_ns": <int>,
//!       "mean_ns": <int>,
//!       "samples_ns": [<int>, ...]
//!     }
//!   }
//! }
//! ```

// `recipe` ships the Recipe trait + BenchParams + `benchmark()` helper that
// drives a recipe through the harness; `util` carries small reproducible
// helpers (deterministic LCG, etc.). Both live in this crate so a recipe
// author can grab everything with a single `use subms::*;`.
pub mod recipe;
pub mod util;

pub use recipe::{benchmark, BenchParams, Recipe};
pub use util::Lcg;

use std::collections::BTreeMap;
use std::fmt::Write as _;
use std::io::{self, Write};
use std::time::{Instant, SystemTime, UNIX_EPOCH};

/// Per-stage sample buffer + recorder.
pub struct Stage {
    name: String,
    samples: Vec<u64>,
}

impl Stage {
    fn new(name: &str, capacity: usize) -> Self {
        Self {
            name: name.to_string(),
            samples: Vec::with_capacity(capacity),
        }
    }
    /// Record an explicit duration in nanoseconds.
    pub fn record(&mut self, ns: u64) {
        self.samples.push(ns);
    }
    /// Time a closure and record its duration.
    pub fn time<F: FnOnce() -> R, R>(&mut self, f: F) -> R {
        let t0 = Instant::now();
        let r = f();
        self.samples.push(t0.elapsed().as_nanos() as u64);
        r
    }
    pub fn name(&self) -> &str {
        &self.name
    }
    pub fn samples(&self) -> &[u64] {
        &self.samples
    }
}

/// A workload run: inputs, metadata, and one or more timed stages.
pub struct PerfHarness {
    workload: String,
    lang: String,
    inputs: BTreeMap<String, String>,
    meta: BTreeMap<String, String>,
    stages: Vec<Stage>,
}

impl PerfHarness {
    pub fn new(workload: &str, lang: &str) -> Self {
        Self {
            workload: workload.to_string(),
            lang: lang.to_string(),
            inputs: BTreeMap::new(),
            meta: BTreeMap::new(),
            stages: Vec::new(),
        }
    }

    /// Add an input parameter (rendered in the cookbook UI's "inputs" panel).
    pub fn input(&mut self, key: &str, value: &str) -> &mut Self {
        self.inputs.insert(key.to_string(), value.to_string());
        self
    }

    /// Add a meta value (sstable count, host info, etc).
    pub fn meta(&mut self, key: &str, value: &str) -> &mut Self {
        self.meta.insert(key.to_string(), value.to_string());
        self
    }

    /// Create a new stage with sample-buffer capacity. Returns a mutable
    /// reference to the stage; record samples via [`Stage::time`] or
    /// [`Stage::record`].
    pub fn stage(&mut self, name: &str, capacity: usize) -> &mut Stage {
        self.stages.push(Stage::new(name, capacity));
        self.stages.last_mut().unwrap()
    }

    /// Borrow a previously-created stage by name.
    pub fn stage_mut(&mut self, name: &str) -> Option<&mut Stage> {
        self.stages.iter_mut().find(|s| s.name == name)
    }

    /// Serialise to the standard JSON shape.
    pub fn write_json<W: Write>(&self, out: &mut W) -> io::Result<()> {
        let mut s = String::with_capacity(64 * 1024);
        s.push('{');
        json_kv_str(&mut s, "workload", &self.workload);
        s.push(',');
        json_kv_str(&mut s, "lang", &self.lang);
        s.push(',');
        json_kv_str(&mut s, "timestamp", &iso8601_now());
        s.push(',');
        s.push_str("\"inputs\":");
        json_map(&mut s, &self.inputs);
        s.push(',');
        s.push_str("\"meta\":");
        json_map(&mut s, &self.meta);
        s.push(',');
        s.push_str("\"stages\":{");
        for (i, stage) in self.stages.iter().enumerate() {
            if i > 0 {
                s.push(',');
            }
            json_str(&mut s, &stage.name);
            s.push(':');
            json_stage(&mut s, &stage.samples);
        }
        s.push_str("}}");
        out.write_all(s.as_bytes())?;
        out.write_all(b"\n")?;
        Ok(())
    }

    /// Drop a stage if you never recorded into it.
    pub fn discard_stage(&mut self, name: &str) {
        self.stages.retain(|s| s.name != name);
    }
}

fn json_str(out: &mut String, s: &str) {
    out.push('"');
    for c in s.chars() {
        match c {
            '"' => out.push_str("\\\""),
            '\\' => out.push_str("\\\\"),
            '\n' => out.push_str("\\n"),
            '\r' => out.push_str("\\r"),
            '\t' => out.push_str("\\t"),
            c if (c as u32) < 0x20 => {
                let _ = write!(out, "\\u{:04x}", c as u32);
            }
            c => out.push(c),
        }
    }
    out.push('"');
}

fn json_kv_str(out: &mut String, k: &str, v: &str) {
    json_str(out, k);
    out.push(':');
    json_str(out, v);
}

fn json_map(out: &mut String, m: &BTreeMap<String, String>) {
    out.push('{');
    for (i, (k, v)) in m.iter().enumerate() {
        if i > 0 {
            out.push(',');
        }
        json_kv_str(out, k, v);
    }
    out.push('}');
}

fn json_stage(out: &mut String, samples: &[u64]) {
    let mut sorted = samples.to_vec();
    sorted.sort_unstable();
    let n = sorted.len();
    let p = |q: f64| -> u64 {
        if n == 0 {
            return 0;
        }
        let idx = ((q * n as f64) as usize).min(n - 1);
        sorted[idx]
    };
    let max = sorted.last().copied().unwrap_or(0);
    let mean = if n == 0 {
        0
    } else {
        sorted.iter().sum::<u64>() / n as u64
    };

    out.push('{');
    let _ = write!(out, "\"count\":{},", n);
    let _ = write!(out, "\"p50_ns\":{},", p(0.50));
    let _ = write!(out, "\"p99_ns\":{},", p(0.99));
    let _ = write!(out, "\"p999_ns\":{},", p(0.999));
    let _ = write!(out, "\"max_ns\":{},", max);
    let _ = write!(out, "\"mean_ns\":{},", mean);
    // Downsample samples_ns to at most 500 evenly-spaced points so the JSON
    // stays compact for the web layer. Order preserved (chronological).
    out.push_str("\"samples_ns\":[");
    let step = (n / 500).max(1);
    let mut first = true;
    for i in (0..n).step_by(step) {
        if !first {
            out.push(',');
        }
        first = false;
        let _ = write!(out, "{}", samples[i]);
    }
    out.push_str("]}");
}

fn iso8601_now() -> String {
    let d = SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap_or_default();
    let secs = d.as_secs() as i64;
    let mut year = 1970i64;
    let mut days = secs / 86_400;
    let rem = secs % 86_400;
    let hour = rem / 3600;
    let minute = (rem % 3600) / 60;
    let second = rem % 60;
    while days >= year_days(year) {
        days -= year_days(year);
        year += 1;
    }
    let mut month = 1u32;
    for m in 1..=12 {
        let dm = month_days(year, m);
        if days < dm as i64 {
            month = m;
            break;
        }
        days -= dm as i64;
    }
    let day = (days + 1) as u32;
    format!(
        "{:04}-{:02}-{:02}T{:02}:{:02}:{:02}Z",
        year, month, day, hour, minute, second
    )
}

fn year_days(y: i64) -> i64 {
    if (y % 4 == 0 && y % 100 != 0) || (y % 400 == 0) {
        366
    } else {
        365
    }
}
fn month_days(y: i64, m: u32) -> u32 {
    match m {
        1 | 3 | 5 | 7 | 8 | 10 | 12 => 31,
        4 | 6 | 9 | 11 => 30,
        2 => {
            if (y % 4 == 0 && y % 100 != 0) || (y % 400 == 0) {
                29
            } else {
                28
            }
        }
        _ => 0,
    }
}

/// Parse stdin `key=value` lines into a flat map. Skips blank lines and `#` comments.
pub fn read_stdin_kv() -> BTreeMap<String, String> {
    use std::io::BufRead;
    let mut m = BTreeMap::new();
    let stdin = io::stdin();
    for line in stdin.lock().lines().map_while(Result::ok) {
        let line = line.trim();
        if line.is_empty() || line.starts_with('#') {
            continue;
        }
        if let Some((k, v)) = line.split_once('=') {
            m.insert(k.trim().to_string(), v.trim().to_string());
        }
    }
    m
}