retex-engine 1.1.0

//! A macro profiler for performance work: with `RETEX_MACRO_PROFILE=<file>`
//! in the environment, every macro call is counted per control sequence,
//! and every 64th call samples the macros being expanded on the input
//! stack (inclusive counts). The report, written at the end of the run,
//! groups both by package prefix (`\pgfmath…`, `\__fp_…`, `\@…`) to show
//! which TeX-coded libraries the time goes to; `<file>.tsv` lists every
//! macro (calls, inclusive samples, name).
//!
//! Inclusive counts are a lower bound: a macro whose last token calls
//! another macro leaves the input stack before the callee runs (TeX's
//! `back_list` tail-call behaviour), so callers through tail calls are not
//! seen.

use std::collections::HashMap;
use std::fmt::Write as _;

use crate::Engine;
use crate::hash::CsId;
use crate::input::Level;

/// Counts of the profiler.
#[derive(Default)]
pub struct MacroProfile {
    /// Where to write the report.
    path: String,
    /// Calls per macro.
    calls: HashMap<CsId, u64>,
    /// Samples in which the macro was on the input stack.
    inclusive: HashMap<CsId, u64>,
    /// All calls.
    total: u64,
    /// Samples taken.
    samples: u64,
}

/// Calls between two input-stack samples.
const SAMPLE_EVERY: u64 = 64;

impl MacroProfile {
    /// The profiler, if `RETEX_MACRO_PROFILE` is set.
    pub fn from_env() -> Option<Box<MacroProfile>> {
        let path = std::env::var("RETEX_MACRO_PROFILE").ok()?;
        Some(Box::new(MacroProfile {
            path,
            ..MacroProfile::default()
        }))
    }
}

/// The package a control sequence name belongs to, from its prefix.
fn package_of(name: &str) -> String {
    let n = name.trim_start_matches('_');
    // expl3: module before the first `_` (`fp_add:Nn`, `__fp_parse…`).
    if name.contains(':') || name.starts_with("__") {
        let module = n.split(['_', ':']).next().unwrap_or(n);
        return format!("expl3 {module}");
    }
    for p in [
        "pgfmath",
        "pgfplots",
        "pgfkeys",
        "pgfsys",
        "pgfutil",
        "pgfpoint",
        "pgfpath",
        "pgf@",
        "pgf",
        "tikz@",
        "tikz",
        "tcb@",
        "tcb",
        "sisetup",
        "si@",
        "SI",
        "num",
        "bbl@",
        "beamer@",
        "beamer",
        "Hy@",
        "HyPsd@",
        "hyper",
        "Gin@",
        "Gread@",
        "@float",
        "@makecaption",
        "@ifnextchar",
        "@ifstar",
        "@xifnextchar",
        "@tempa",
        "@for",
        "@tfor",
        "@tempb",
        "@protected@",
        "@expl@",
        "fontspec",
        "f@",
    ] {
        if n.starts_with(p) {
            return p.trim_end_matches('@').to_owned();
        }
    }
    if n.contains('@') {
        return "LaTeX internals (@)".to_owned();
    }
    "(other)".to_owned()
}

impl Engine {
    /// Records a call of `cs` (called from `macro_call`).
    #[cold]
    pub(crate) fn profile_macro_call(&mut self, cs: CsId) {
        let Some(p) = self.macro_profile.as_mut() else {
            return;
        };
        *p.calls.entry(cs).or_default() += 1;
        p.total += 1;
        if p.total % SAMPLE_EVERY != 0 {
            return;
        }
        p.samples += 1;
        let mut seen: Vec<CsId> = vec![cs];
        for level in &self.input.stack {
            if let Level::Tokens(t) = level
                && t.name != crate::hash::NO_CS
                && !seen.contains(&t.name)
            {
                seen.push(t.name);
            }
        }
        for c in seen {
            *p.inclusive.entry(c).or_default() += 1;
        }
    }

    /// Writes the profile report, if profiling.
    pub fn write_macro_profile(&mut self) {
        let Some(p) = self.macro_profile.take() else {
            return;
        };
        let name = |e: &Engine, c: CsId| String::from_utf8_lossy(e.hash.name(c)).into_owned();
        let mut out = String::new();
        let _ = writeln!(
            out,
            "macro calls: {}, input-stack samples: {} (every {SAMPLE_EVERY} calls)\n",
            p.total, p.samples
        );
        let mut by_pkg: HashMap<String, (u64, u64)> = HashMap::new();
        for (&c, &n) in &p.calls {
            by_pkg.entry(package_of(&name(self, c))).or_default().0 += n;
        }
        // Inclusive per package: a sample counts once per package.
        // Approximated from per-macro inclusive counts, capped by samples.
        for (&c, &n) in &p.inclusive {
            let e = by_pkg.entry(package_of(&name(self, c))).or_default();
            e.1 = e.1.max(n);
        }
        let mut pk: Vec<_> = by_pkg.into_iter().collect();
        pk.sort_by(|a, b| b.1.0.cmp(&a.1.0));
        let _ = writeln!(
            out,
            "{:>7} {:>7}  package (calls %, max inclusive %)",
            "calls", "incl"
        );
        for (k, (c, i)) in pk.iter().take(40) {
            let _ = writeln!(
                out,
                "{:>6.2}% {:>6.2}%  {k}",
                *c as f64 * 100.0 / p.total.max(1) as f64,
                *i as f64 * 100.0 / p.samples.max(1) as f64
            );
        }
        let mut calls: Vec<_> = p.calls.iter().map(|(&c, &n)| (n, c)).collect();
        calls.sort_by(|a, b| b.0.cmp(&a.0));
        let _ = writeln!(out, "\ntop macros by calls:");
        for (n, c) in calls.iter().take(60) {
            let _ = writeln!(
                out,
                "{:>6.2}% {:>10}  \\{}",
                *n as f64 * 100.0 / p.total.max(1) as f64,
                n,
                name(self, *c)
            );
        }
        let mut incl: Vec<_> = p.inclusive.iter().map(|(&c, &n)| (n, c)).collect();
        incl.sort_by(|a, b| b.0.cmp(&a.0));
        let _ = writeln!(out, "\ntop macros by inclusive samples:");
        for (n, c) in incl.iter().take(60) {
            let _ = writeln!(
                out,
                "{:>6.2}%  \\{}",
                *n as f64 * 100.0 / p.samples.max(1) as f64,
                name(self, *c)
            );
        }
        let _ = std::fs::write(&p.path, out);
        // Every macro, for further analysis: calls, inclusive samples, name.
        let mut tsv = String::from("calls\tinclusive\tname\n");
        for (n, c) in &calls {
            let incl = p.inclusive.get(c).copied().unwrap_or(0);
            let _ = writeln!(tsv, "{n}\t{incl}\t\\{}", name(self, *c));
        }
        let _ = std::fs::write(format!("{}.tsv", p.path), tsv);
    }
}