retex-engine 1.3.0

//! The profiler behind `retex profile`, and the origins of macro
//! definitions behind `retex explain` (`docs/design/profiling.md`).
//!
//! **Sampling.** A timer thread adds the time elapsed to [`PENDING`] every
//! millisecond. The engine looks at it at each macro call and each step
//! of the main loop ([`Engine::profile_tick`]), and charges the time to
//! its current state: the files and macros on the input stack, and the
//! engine work under way ([`Work`], marked by [`Engine::work_begin`] and
//! [`Engine::work_end`]). Without a profiler the counter stays zero, and
//! the check is one atomic load.
//!
//! A macro leaves the input stack when its last token is read (TeX's
//! `back_list` tail calls), so the time including a macro's callees is a
//! lower bound.
//!
//! **Origins.** With [`Origins`] on, each macro definition records the
//! innermost file being read and its line ([`Macro::origin`]). A file
//! belongs to a package: on the input stack when the file is opened, the
//! first file above the document's own files that is not one of them.
//!
//! [`Macro::origin`]: crate::token::Macro::origin

use std::collections::HashMap;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::time::{Duration, Instant};

use crate::Engine;
use crate::hash::{CsId, NO_CS};
use crate::input::{Level, LineSource};

/// Microseconds elapsed and not yet charged to the engine's state.
pub static PENDING: AtomicU64 = AtomicU64::new(0);

/// The interval between two samples.
pub const INTERVAL: Duration = Duration::from_millis(1);

/// Engine work, charged on its own besides the macros that started it.
#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub enum Work {
    /// Lua code (`\directlua`, `\luafunction`, callbacks).
    Lua,
    /// Loading a font (`define_font`).
    Fonts,
    /// Breaking a paragraph into lines.
    LineBreak,
    /// Shipping a page out (the page's PDF).
    Shipout,
    /// Reading an image (`\saveimageresource`, `\pdfximage`).
    Images,
}

impl Work {
    /// A short description.
    pub fn label(self) -> &'static str {
        match self {
            Work::Lua => "Lua",
            Work::Fonts => "font loading",
            Work::LineBreak => "line breaking",
            Work::Shipout => "shipping pages (PDF)",
            Work::Images => "images",
        }
    }
}

/// A file read by the job.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct FileInfo {
    /// The name, as TeX prints it.
    pub name: String,
    /// The package it belongs to; `None` for the document's own files.
    pub package: Option<String>,
}

/// Where macros were defined.
#[derive(Default)]
pub struct Origins {
    /// The files, by id - 1 ([`crate::input::TextLevel::profile_file`]).
    pub files: Vec<FileInfo>,
    /// Origins (file id, line), by origin - 1.
    places: Vec<(u32, u32)>,
    /// Origin of each place.
    ids: HashMap<(u32, u32), u32>,
    /// For `retex explain`: the control sequences whose definitions are
    /// recorded, with the origins of their meanings, in order.
    pub history: HashMap<CsId, Vec<u32>>,
}

impl Origins {
    /// The file and line of `origin`, if it is one.
    pub fn place(&self, origin: u32) -> Option<(&FileInfo, u32)> {
        let &(file, line) = self.places.get(origin.checked_sub(1)? as usize)?;
        Some((self.files.get(file.checked_sub(1)? as usize)?, line))
    }

    /// The file id of `origin`.
    pub fn file_of(&self, origin: u32) -> Option<u32> {
        self.places
            .get(origin.checked_sub(1)? as usize)
            .map(|p| p.0)
    }
}

/// Whether a file is one of the document's own: TeX sources and the
/// auxiliary files LaTeX writes.
fn is_document_file(name: &str) -> bool {
    let ext = std::path::Path::new(name)
        .extension()
        .map(|e| e.to_string_lossy().to_ascii_lowercase());
    matches!(
        ext.as_deref(),
        None | Some(
            "tex"
                | "ltx"
                | "aux"
                | "toc"
                | "lof"
                | "lot"
                | "bbl"
                | "ind"
                | "gls"
                | "out"
                | "nav"
                | "snm"
                | "vrb"
                | "loa"
                | "thm"
        )
    )
}

/// The package name of file `name`: `tikz` for `tikz.sty`, `article
/// (class)` for `article.cls`, the file name otherwise.
fn package_name(name: &str) -> String {
    let path = std::path::Path::new(name);
    let file = path
        .file_name()
        .map_or_else(|| name.to_owned(), |f| f.to_string_lossy().into_owned());
    let stem = path
        .file_stem()
        .map_or_else(|| file.clone(), |f| f.to_string_lossy().into_owned());
    match path.extension().and_then(|e| e.to_str()) {
        Some("sty") => stem,
        Some("cls") => format!("{stem} (class)"),
        _ => file,
    }
}

/// A frame of a sample's stack.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub enum Frame {
    /// A file being read (file id).
    File(u32),
    /// A macro being expanded, with the origin of its meaning.
    Macro(CsId, u32),
    /// Engine work.
    Work(Work),
}

/// The sampling profiler.
pub struct Profiler {
    /// Stops the timer thread.
    stop: Arc<AtomicBool>,
    timer: Option<std::thread::JoinHandle<()>>,
    started: Instant,
    /// Time charged to each stack, in microseconds.
    stacks: HashMap<Box<[Frame]>, u64>,
    /// The work under way, with the input stack depth where it began.
    work: Vec<(Work, usize)>,
    /// Calls per control sequence, by id.
    calls: Vec<u64>,
    /// When the first file was read.
    first_file: Option<Duration>,
    /// When `\document` was first called.
    document: Option<Duration>,
    /// When each page was shipped out.
    pages: Vec<Duration>,
}

impl Profiler {
    /// Starts a profiler and its timer thread.
    fn start() -> Profiler {
        PENDING.store(0, Ordering::Relaxed);
        let stop = Arc::new(AtomicBool::new(false));
        let flag = Arc::clone(&stop);
        let timer = std::thread::Builder::new()
            .name("profiler".to_owned())
            .spawn(move || {
                let mut last = Instant::now();
                while !flag.load(Ordering::Relaxed) {
                    std::thread::sleep(INTERVAL);
                    let now = Instant::now();
                    let us = (now - last).as_micros() as u64;
                    last = now;
                    PENDING.fetch_add(us.max(1), Ordering::Relaxed);
                }
            })
            .ok();
        Profiler {
            stop,
            timer,
            started: Instant::now(),
            stacks: HashMap::new(),
            work: Vec::new(),
            calls: Vec::new(),
            first_file: None,
            document: None,
            pages: Vec::new(),
        }
    }

    fn stop_timer(&mut self) {
        self.stop.store(true, Ordering::Relaxed);
        if let Some(t) = self.timer.take() {
            let _ = t.join();
        }
    }
}

impl Drop for Profiler {
    fn drop(&mut self) {
        self.stop_timer();
        PENDING.store(0, Ordering::Relaxed);
    }
}

/// A frame of a [`Summary`], resolved.
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum SummaryFrame {
    /// A file.
    File {
        /// The file name, as TeX prints it.
        name: String,
        /// The package, `None` for the document's own files.
        package: Option<String>,
    },
    /// A macro.
    Macro {
        /// The name, without the backslash.
        name: String,
        /// The file and line of the definition; `None` if unknown (the
        /// format's macros).
        defined: Option<(String, u32)>,
        /// The package of that file; `None` for the document's own files
        /// and for unknown definitions.
        package: Option<String>,
    },
    /// Engine work.
    Work(Work),
}

/// What a profile found.
#[derive(Clone, Debug, Default)]
pub struct Summary {
    /// The time from the start of the profile to its end.
    pub total: Duration,
    /// When the first file was read.
    pub first_file: Option<Duration>,
    /// When the preamble ended (`\document`).
    pub document: Option<Duration>,
    /// When each page was shipped out.
    pub pages: Vec<Duration>,
    /// The frames of the stacks.
    pub frames: Vec<SummaryFrame>,
    /// The stacks, from the outermost frame (indices in `frames`), and
    /// the microseconds charged to each.
    pub stacks: Vec<(Vec<u32>, u64)>,
    /// Calls per macro name, for the macros called.
    pub calls: HashMap<String, u64>,
}

impl Engine {
    /// Turns on the recording of where macros are defined.
    pub fn record_origins(&mut self) {
        if self.origins.is_none() {
            self.origins = Some(Box::default());
        }
    }

    /// Starts profiling (and recording origins).
    pub fn start_profile(&mut self) {
        self.record_origins();
        self.profiler = Some(Box::new(Profiler::start()));
    }

    /// Charges the time elapsed since the last sample, if any, to the
    /// current state. Called often: cheap when there is nothing to do.
    #[inline]
    pub fn profile_tick(&mut self) {
        if PENDING.load(Ordering::Relaxed) != 0 {
            self.take_sample(None);
        }
    }

    /// Records a call of `cs` (from `macro_call`).
    pub(crate) fn profile_call(&mut self, cs: CsId) {
        // The caller's level may be gone already (its last token was
        // read), and the macro's is not there yet: the time goes to the
        // macro called.
        if PENDING.load(Ordering::Relaxed) != 0 {
            self.take_sample(Some(cs));
        }
        let Some(p) = self.profiler.as_mut() else {
            return;
        };
        let i = cs.0 as usize;
        if p.calls.len() <= i {
            p.calls.resize(i + 1, 0);
        }
        p.calls[i] += 1;
        // LaTeX's `\document` ends the preamble.
        if p.document.is_none() && self.hash.name(cs) == b"document" {
            p.document = Some(p.started.elapsed());
        }
    }

    /// Marks the beginning of engine work `w`.
    pub fn work_begin(&mut self, w: Work) {
        if self.profiler.is_none() {
            return;
        }
        self.profile_tick();
        let depth = self.input.stack.len();
        if let Some(p) = self.profiler.as_mut() {
            p.work.push((w, depth));
        }
    }

    /// Marks the end of the innermost engine work.
    pub fn work_end(&mut self) {
        if self.profiler.is_none() {
            return;
        }
        self.profile_tick();
        if let Some(p) = self.profiler.as_mut()
            && let Some((w, _)) = p.work.pop()
            && w == Work::Shipout
        {
            p.pages.push(p.started.elapsed());
        }
    }

    /// Charges the pending time to the current state, with the macro
    /// `called` innermost.
    #[cold]
    fn take_sample(&mut self, called: Option<CsId>) {
        let us = PENDING.swap(0, Ordering::Relaxed);
        if us == 0 || self.profiler.is_none() {
            return;
        }
        let mut frames = Vec::new();
        let work = self
            .profiler
            .as_ref()
            .map(|p| p.work.clone())
            .unwrap_or_default();
        let mut work = work.into_iter().peekable();
        let mut saw_file = false;
        for i in 0..self.input.stack.len() {
            while let Some(&(w, _)) = work.peek().filter(|(_, d)| *d <= i) {
                frames.push(Frame::Work(w));
                work.next();
            }
            match &self.input.stack[i] {
                Level::Text(t) if matches!(t.source, LineSource::File { .. }) => {
                    saw_file = true;
                    let id = self.file_id(i);
                    frames.push(Frame::File(id));
                }
                Level::Tokens(t) if t.name != NO_CS => {
                    let origin = self
                        .eqtb
                        .meaning(t.name)
                        .mac
                        .as_ref()
                        .map_or(0, |m| m.origin.get());
                    frames.push(Frame::Macro(t.name, origin));
                }
                _ => {}
            }
        }
        frames.extend(work.map(|(w, _)| Frame::Work(w)));
        if let Some(cs) = called {
            let origin = self
                .eqtb
                .meaning(cs)
                .mac
                .as_ref()
                .map_or(0, |m| m.origin.get());
            frames.push(Frame::Macro(cs, origin));
        }
        let Some(p) = self.profiler.as_mut() else {
            return;
        };
        if saw_file && p.first_file.is_none() {
            p.first_file = Some(p.started.elapsed());
        }
        *p.stacks.entry(frames.into_boxed_slice()).or_default() += us;
    }

    /// The file id of the text level at `i` of the input stack (a file),
    /// given on first use.
    fn file_id(&mut self, i: usize) -> u32 {
        let Some(Level::Text(t)) = self.input.stack.get(i) else {
            return 0;
        };
        if t.profile_file != 0 {
            return t.profile_file;
        }
        let Some(origins) = self.origins.as_mut() else {
            return 0;
        };
        // The files being read, from the outermost: the first is the
        // document; then the document's own files; the first other file
        // gives the package.
        let names: Vec<String> = self.input.stack[..=i]
            .iter()
            .filter_map(|l| match l {
                Level::Text(t) if matches!(t.source, LineSource::File { .. }) => {
                    Some(String::from_utf8_lossy(&t.name).into_owned())
                }
                _ => None,
            })
            .collect();
        let package = names
            .iter()
            .skip(1)
            .find(|n| !is_document_file(n))
            .map(|n| package_name(n));
        let name = names.last().cloned().unwrap_or_default();
        origins.files.push(FileInfo { name, package });
        let id = origins.files.len() as u32;
        if let Some(Level::Text(t)) = self.input.stack.get_mut(i) {
            t.profile_file = id;
        }
        id
    }

    /// The origin of a definition made now: the innermost file being read
    /// and its line (0 without [`Origins`] or outside files).
    pub fn origin_here(&mut self) -> u32 {
        if self.origins.is_none() {
            return 0;
        }
        let Some((i, line)) =
            self.input
                .stack
                .iter()
                .enumerate()
                .rev()
                .find_map(|(i, l)| match l {
                    Level::Text(t) if matches!(t.source, LineSource::File { .. }) => {
                        Some((i, t.line.max(0) as u32))
                    }
                    _ => None,
                })
        else {
            return 0;
        };
        let file = self.file_id(i);
        let Some(o) = self.origins.as_mut() else {
            return 0;
        };
        if let Some(&id) = o.ids.get(&(file, line)) {
            return id;
        }
        o.places.push((file, line));
        let id = o.places.len() as u32;
        o.ids.insert((file, line), id);
        id
    }

    /// Records a new meaning of `cs` for `retex explain`, if asked.
    pub(crate) fn record_definition(&mut self, cs: CsId) {
        let Some(o) = self.origins.as_ref() else {
            return;
        };
        if !o.history.contains_key(&cs) {
            return;
        }
        let origin = match &self.eqtb.meaning(cs).mac {
            Some(m) if m.origin.get() != 0 => m.origin.get(),
            _ => self.origin_here(),
        };
        if let Some(h) = self.origins.as_mut().and_then(|o| o.history.get_mut(&cs)) {
            h.push(origin);
        }
    }

    /// The `origins` section of a format, when origins are recorded: the
    /// files, the places, and the origin of each control sequence's
    /// macro.
    pub(crate) fn dump_origins(&self) -> Option<Vec<u8>> {
        let o = self.origins.as_ref()?;
        let mut w = crate::format::Writer::default();
        w.u(o.files.len() as u64);
        for f in &o.files {
            w.bytes(f.name.as_bytes());
        }
        w.u(o.places.len() as u64);
        for &(file, line) in &o.places {
            w.u(file.into());
            w.u(line.into());
        }
        let macros: Vec<(usize, u32)> = self
            .eqtb
            .cs
            .iter()
            .enumerate()
            .filter_map(|(i, s)| Some((i, s.val.mac.as_ref()?.origin.get())))
            .filter(|&(_, o)| o != 0)
            .collect();
        w.u(macros.len() as u64);
        for (cs, origin) in macros {
            w.u(cs as u64);
            w.u(origin.into());
        }
        Some(w.buf)
    }

    /// Loads the `origins` section of a format, when origins are
    /// recorded. The format's files belong to expl3 (`expl3-code.tex`,
    /// `l3….tex`) or to the LaTeX kernel (or the format).
    pub(crate) fn undump_origins(&mut self, data: &[u8]) -> Result<(), crate::format::FormatError> {
        if self.origins.is_none() {
            return Ok(());
        }
        let kernel = if self.latex_release().is_some() {
            "LaTeX kernel".to_owned()
        } else {
            format!("format {}", self.format_name)
        };
        let mut r = crate::format::Reader::new(data);
        let mut o = Origins::default();
        for _ in 0..r.u()? {
            let name = String::from_utf8_lossy(r.bytes()?).into_owned();
            let file = std::path::Path::new(&name)
                .file_name()
                .map(|f| f.to_string_lossy().into_owned())
                .unwrap_or_default();
            let package = if file.starts_with("expl3") || file.starts_with("l3") {
                "expl3".to_owned()
            } else {
                kernel.clone()
            };
            o.files.push(FileInfo {
                name,
                package: Some(package),
            });
        }
        for i in 0..r.u()? {
            let place = (r.u32()?, r.u32()?);
            o.places.push(place);
            o.ids.insert(place, i as u32 + 1);
        }
        for _ in 0..r.u()? {
            let (cs, origin) = (r.u()? as usize, r.u32()?);
            if let Some(m) = self.eqtb.cs.get(cs).and_then(|s| s.val.mac.as_ref()) {
                m.origin.set(origin);
            }
        }
        // What was recorded before (nothing is, before a format loads)
        // keeps its numbers after the format's.
        if let Some(old) = self.origins.take() {
            o.history = old.history;
        }
        self.origins = Some(Box::new(o));
        Ok(())
    }

    /// Stops the profiler and returns what it found.
    pub fn finish_profile(&mut self) -> Option<Summary> {
        self.profile_tick();
        let mut p = self.profiler.take()?;
        p.stop_timer();
        let total = p.started.elapsed();
        // The time since the last sample.
        let rest = PENDING.swap(0, Ordering::Relaxed);
        if rest > 0 {
            *p.stacks.entry(Box::default()).or_default() += rest;
        }
        let mut s = Summary {
            total,
            first_file: p.first_file,
            document: p.document,
            pages: std::mem::take(&mut p.pages),
            ..Summary::default()
        };
        let mut index: HashMap<Frame, u32> = HashMap::new();
        for (stack, us) in std::mem::take(&mut p.stacks) {
            let ids = stack
                .iter()
                .map(|f| match index.get(f) {
                    Some(&i) => i,
                    None => {
                        let i = s.frames.len() as u32;
                        s.frames.push(self.summary_frame(*f));
                        index.insert(*f, i);
                        i
                    }
                })
                .collect();
            s.stacks.push((ids, us));
        }
        s.stacks
            .sort_by(|a, b| b.1.cmp(&a.1).then_with(|| a.0.cmp(&b.0)));
        for (i, &n) in p.calls.iter().enumerate() {
            if n > 0 {
                let name = String::from_utf8_lossy(self.hash.name(CsId(i as u32))).into_owned();
                *s.calls.entry(name).or_default() += n;
            }
        }
        Some(s)
    }

    fn summary_frame(&self, f: Frame) -> SummaryFrame {
        let file = |id: u32| {
            self.origins
                .as_ref()
                .and_then(|o| o.files.get(id.checked_sub(1)? as usize))
        };
        match f {
            Frame::File(id) => match file(id) {
                Some(fi) => SummaryFrame::File {
                    name: fi.name.clone(),
                    package: fi.package.clone(),
                },
                None => SummaryFrame::File {
                    name: "?".to_owned(),
                    package: None,
                },
            },
            Frame::Macro(cs, origin) => {
                let name = String::from_utf8_lossy(self.hash.name(cs)).into_owned();
                let place = self.origins.as_ref().and_then(|o| o.place(origin));
                SummaryFrame::Macro {
                    name,
                    defined: place.map(|(fi, line)| (fi.name.clone(), line)),
                    package: place.and_then(|(fi, _)| fi.package.clone()),
                }
            }
            Frame::Work(w) => SummaryFrame::Work(w),
        }
    }
}

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

    #[test]
    fn packages_of_files() {
        assert!(is_document_file("chapter1.tex"));
        assert!(is_document_file("./doc.aux"));
        assert!(is_document_file("Makefile"));
        assert!(!is_document_file("tikz.sty"));
        assert_eq!(
            package_name("/usr/share/texmf/tex/latex/pgf/tikz.sty"),
            "tikz"
        );
        assert_eq!(package_name("article.cls"), "article (class)");
        assert_eq!(package_name("t1lmr.fd"), "t1lmr.fd");
    }
}