mathtex-editor-core 0.3.0

//! LaTeX export with byte spans for every node, slot, and caret gap.

use std::borrow::Cow;
use std::collections::HashMap;
use std::ops::Range;

use crate::model::{Deco, FracStyle, Kind, Mark, MatrixEnv, NodeId, SeqId, SeqRange, Tree, Variant};
use crate::path::{CaretPath, Step};

/// LaTeX for the typesetter plus the spans that map it back to the document.
#[derive(Debug, Clone)]
pub struct Source {
    /// The LaTeX, which assumes amsmath and a `\hostbox` macro.
    pub tex: String,
    /// Byte ranges of every node, slot, and caret gap in `tex`.
    pub spans: SpanMap,
    /// The editor revision this source was exported at.
    pub revision: u64,
}

impl Source {
    /// Byte offset into `tex` of a caret gap, `None` when the path is not in this source.
    pub fn caret_offset(&self, at: &CaretPath) -> Option<usize> {
        self.spans.gaps.get(&at.steps)?.get(at.index).copied()
    }
}

/// Opaque byte ranges from document elements into [`Source::tex`], in export order.
#[derive(Debug, Clone, Default)]
pub struct SpanMap {
    pub(crate) owner: u64,
    pub(crate) nodes: Vec<(NodeId, Range<usize>)>,
    pub(crate) seqs: Vec<(SeqId, Range<usize>)>,
    pub(crate) gaps: HashMap<Vec<Step>, Vec<usize>>,
}

impl SpanMap {
    pub(crate) fn push_node(&mut self, id: NodeId, range: Range<usize>) {
        self.nodes.push((id, range));
    }

    pub(crate) fn push_seq(&mut self, id: SeqId, range: Range<usize>) {
        self.seqs.push((id, range));
    }
}

/// Placeholder export for typesetting, empty slots become `\phantom{x}` boxes.
pub(crate) fn source(tree: &Tree, owner: u64, revision: u64, placeholders: bool) -> Source {
    let mut ex = Exporter::new(tree, placeholders, None);
    ex.spans = Some(SpanMap { owner, ..SpanMap::default() });
    ex.emit_seq(tree.root());
    Source { tex: ex.out, spans: ex.spans.unwrap_or_default(), revision }
}

/// Clean export of the whole tree without spans.
pub(crate) fn clean_tex<'a>(tree: &'a Tree, host_box: Option<&'a mut dyn FnMut(u32) -> Option<String>>) -> String {
    let mut ex = Exporter::new(tree, false, host_box);
    ex.emit_seq(tree.root());
    ex.out
}

/// Clean export of a selected run, wrapped in `\text{}` when it comes from a text slot.
pub(crate) fn range_tex(tree: &Tree, sel: SeqRange) -> String {
    let mut ex = Exporter::new(tree, false, None);
    ex.text = tree.is_text_slot(sel.seq);
    if ex.text {
        ex.push("\\text{");
    }
    let items = tree.items(sel.seq);
    let hi = sel.hi().min(items.len());
    for (i, &n) in items.iter().enumerate().take(hi).skip(sel.lo()) {
        ex.emit_node(n, i);
    }
    if ex.text {
        ex.push("}");
    }
    ex.out
}

struct Exporter<'a> {
    tree: &'a Tree,
    out: String,
    /// Whether empty slots render as `\phantom{x}` placeholders.
    placeholders: bool,
    /// Whether atoms are being written inside `\text{}`.
    text: bool,
    spans: Option<SpanMap>,
    path: Vec<Step>,
    /// Offsets of separator spaces, so spans can start after them.
    seps: Vec<usize>,
    host_box: Option<&'a mut dyn FnMut(u32) -> Option<String>>,
}

impl<'a> Exporter<'a> {
    fn new(tree: &'a Tree, placeholders: bool, host_box: Option<&'a mut dyn FnMut(u32) -> Option<String>>) -> Self {
        Self { tree, out: String::new(), placeholders, text: false, spans: None, path: Vec::new(), seps: Vec::new(), host_box }
    }

    /// The single separator rule: a control word followed by a letter gets one space between them.
    fn push(&mut self, s: &str) {
        if s.starts_with(char::is_alphabetic) && ends_in_control_word(&self.out) {
            self.seps.push(self.out.len());
            self.out.push(' ');
        }
        self.out.push_str(s);
    }

    /// A span start recorded before a push, moved past a separator that push inserted.
    fn start_after_sep(&self, start: usize) -> usize {
        if self.seps.binary_search(&start).is_ok() { start + 1 } else { start }
    }

    fn emit_seq(&mut self, seq: SeqId) {
        let start = self.out.len();
        let items = self.tree.items(seq);
        let mut gaps = Vec::with_capacity(items.len() + 1);
        if items.is_empty() && self.placeholders {
            self.push("\\phantom{x}");
        }
        for (i, &n) in items.iter().enumerate() {
            gaps.push(self.emit_node(n, i));
        }
        let end = self.out.len();
        let start = self.start_after_sep(start).min(end);
        gaps.push(if items.is_empty() { start } else { end });
        if let Some(spans) = &mut self.spans {
            spans.push_seq(seq, start..end);
            spans.gaps.insert(self.path.clone(), gaps);
        }
    }

    /// Emits the slot `seq` of the node at `index`, tracking the path for gap lookups.
    fn emit_slot(&mut self, node: NodeId, index: usize, seq: SeqId) {
        let slot = self.tree.slot_of(node, seq);
        if let Some(slot) = slot {
            self.path.push(Step { node: index, slot });
        }
        self.emit_seq(seq);
        if slot.is_some() {
            self.path.pop();
        }
    }

    fn emit_braced(&mut self, node: NodeId, index: usize, seq: SeqId) {
        self.push("{");
        self.emit_slot(node, index, seq);
        self.push("}");
    }

    /// Script and limit arguments drop their braces only around a single one character atom.
    fn emit_arg(&mut self, node: NodeId, index: usize, seq: SeqId) {
        let items = self.tree.items(seq);
        let bare = items.len() == 1
            && matches!(self.tree.kind(items[0]), Some(Kind::Atom(s)) if s.latex.chars().count() == 1);
        if bare { self.emit_slot(node, index, seq) } else { self.emit_braced(node, index, seq) }
    }

    /// Emits an optional attachment such as `^{..}`, clean export drops it when empty.
    fn emit_attachment(&mut self, node: NodeId, index: usize, marker: &str, seq: Option<SeqId>) {
        let Some(seq) = seq else { return };
        if self.placeholders || !self.tree.is_empty(seq) {
            self.push(marker);
            self.emit_arg(node, index, seq);
        }
    }

    /// A base can carry scripts bare only as one item that TeX reads as a single atom.
    fn script_safe_base(&self, seq: SeqId) -> bool {
        match self.tree.items(seq) {
            [only] => matches!(
                self.tree.kind(*only),
                Some(Kind::Atom(_) | Kind::Frac { .. } | Kind::Sqrt { .. } | Kind::Delim { .. })
                    | Some(Kind::Accent { .. } | Kind::Styled { .. })
            ),
            _ => false,
        }
    }

    /// Emits a node and returns where its span starts.
    fn emit_node(&mut self, node: NodeId, index: usize) -> usize {
        let start = self.out.len();
        let Some(kind) = self.tree.kind(node).cloned() else {
            return start;
        };
        // Structures inside a text slot, only reachable from legacy documents, switch back to math.
        let wrap_math = self.text && !matches!(kind, Kind::Atom(_) | Kind::HostBox { .. });
        if wrap_math {
            self.push("\\ensuremath{");
            self.text = false;
        }
        match kind {
            Kind::Atom(s) => {
                let latex = if self.text { text_latex(&s.latex) } else { Cow::Borrowed(s.latex.as_str()) };
                self.push(&latex);
            }
            // The editor emits only the token macro, the object's content stays host side.
            Kind::HostBox { token } => {
                let content = self.host_box.as_mut().and_then(|f| f(token));
                match content {
                    Some(c) => self.push(&c),
                    None => self.push(&format!("\\hostbox{{{token}}}")),
                }
            }
            Kind::Frac { num, den, style } => {
                self.push(frac_cmd(style));
                self.emit_braced(node, index, num);
                self.emit_braced(node, index, den);
            }
            Kind::Script { base, sub, sup } => {
                if self.script_safe_base(base) {
                    self.emit_slot(node, index, base);
                } else {
                    self.emit_braced(node, index, base);
                }
                self.emit_attachment(node, index, "_", sub);
                self.emit_attachment(node, index, "^", sup);
            }
            Kind::BigOp { op, lower, upper } => {
                self.push(&op.latex);
                self.emit_attachment(node, index, "_", Some(lower));
                self.emit_attachment(node, index, "^", Some(upper));
            }
            Kind::Sqrt { index: degree, radicand } => {
                self.push("\\sqrt");
                if self.placeholders || !self.tree.is_empty(degree) {
                    self.push("[{");
                    self.emit_slot(node, index, degree);
                    self.push("}]");
                }
                self.emit_braced(node, index, radicand);
            }
            Kind::Delim { open, close, body } => {
                self.push("\\left");
                self.push(delim_tex(open));
                self.emit_slot(node, index, body);
                self.push("\\right");
                self.push(delim_tex(close));
            }
            Kind::Accent { mark, base } => {
                self.push(accent_cmd(mark));
                self.emit_braced(node, index, base);
            }
            Kind::UnderOver { base, over, under, over_deco, under_deco } => {
                self.emit_under_over(node, index, base, [(over, over_deco, true), (under, under_deco, false)]);
            }
            Kind::Styled { variant: Variant::Text, content } => {
                self.push("\\text{");
                self.text = true;
                self.emit_slot(node, index, content);
                self.text = false;
                self.push("}");
            }
            Kind::Styled { variant, content } => {
                self.push(variant_cmd(variant));
                self.emit_braced(node, index, content);
            }
            Kind::Matrix { env, rows } => self.emit_matrix(node, index, env, &rows),
        }
        if wrap_math {
            self.text = true;
            self.push("}");
        }
        let start = self.start_after_sep(start);
        if let Some(spans) = &mut self.spans {
            spans.push_node(node, start..self.out.len());
        }
        start
    }

    /// Wraps the base in the under decoration, then the over one, clean export drops empty labels.
    fn emit_under_over(&mut self, node: NodeId, index: usize, base: SeqId, labels: [(Option<SeqId>, Deco, bool); 2]) {
        let shown = |ex: &Self, l: Option<SeqId>| l.filter(|&s| ex.placeholders || !ex.tree.is_empty(s));
        let [over, under] = labels;
        // Outer wrapper first, so the over decoration encloses the under one.
        let mut closers: Vec<(Option<SeqId>, Deco, bool)> = Vec::new();
        for (label, deco, is_over) in [over, under] {
            if label.is_none() {
                continue;
            }
            let label = shown(self, label);
            match (deco, label) {
                (Deco::Brace, _) => self.push(if is_over { "\\overbrace{" } else { "\\underbrace{" }),
                (_, Some(l)) => {
                    self.push(if is_over { "\\overset" } else { "\\underset" });
                    self.emit_braced(node, index, l);
                    self.push("{");
                }
                (_, None) => {}
            }
            match deco {
                Deco::Arrow => self.push(if is_over { "\\overrightarrow{" } else { "\\underrightarrow{" }),
                Deco::Line => self.push(if is_over { "\\overline{" } else { "\\underline{" }),
                Deco::None | Deco::Brace => {}
            }
            closers.push((label, deco, is_over));
        }
        self.emit_slot(node, index, base);
        for (label, deco, is_over) in closers.into_iter().rev() {
            if matches!(deco, Deco::Arrow | Deco::Line) {
                self.push("}");
            }
            match (deco, label) {
                (Deco::Brace, Some(l)) => {
                    self.push(if is_over { "}^" } else { "}_" });
                    self.emit_braced(node, index, l);
                }
                (Deco::Brace, None) | (_, Some(_)) => self.push("}"),
                (_, None) => {}
            }
        }
    }

    fn emit_matrix(&mut self, node: NodeId, index: usize, env: MatrixEnv, rows: &[Vec<SeqId>]) {
        let name = matrix_env_name(env);
        self.push(&format!("\\begin{{{name}}}"));
        if env == MatrixEnv::Array {
            let cols = rows.first().map_or(0, Vec::len);
            self.push(&format!("{{{}}}", "c".repeat(cols)));
        }
        for (ri, row) in rows.iter().enumerate() {
            if ri > 0 {
                self.push(" \\\\ ");
            }
            for (ci, &cell) in row.iter().enumerate() {
                if ci > 0 {
                    self.push(" & ");
                }
                self.emit_slot(node, index, cell);
            }
        }
        self.push(&format!("\\end{{{name}}}"));
    }
}

/// Whether `s` ends in an unescaped control word such as `\alpha`, which a letter would extend.
fn ends_in_control_word(s: &str) -> bool {
    let letters = s.chars().rev().take_while(|c| c.is_alphabetic()).map(char::len_utf8).sum::<usize>();
    if letters == 0 {
        return false;
    }
    let slashes = s[..s.len() - letters].chars().rev().take_while(|&c| c == '\\').count();
    slashes % 2 == 1
}

/// The text mode spelling of an atom's math LaTeX, anything without one goes through `\ensuremath`.
fn text_latex(latex: &str) -> Cow<'_, str> {
    let mut chars = latex.chars();
    if let (Some(c), None) = (chars.next(), chars.next()) {
        return match c {
            '_' => Cow::Borrowed("\\_"),
            '^' => Cow::Borrowed("\\textasciicircum{}"),
            '~' => Cow::Borrowed("\\textasciitilde{}"),
            _ => Cow::Borrowed(latex),
        };
    }
    match latex {
        "\\%" | "\\#" | "\\&" | "\\$" | "\\_" | "\\{" | "\\}" | "\\ " => Cow::Borrowed(latex),
        "\\sim" => Cow::Borrowed("\\textasciitilde{}"),
        "\\backslash" => Cow::Borrowed("\\textbackslash{}"),
        "\\prime" => Cow::Borrowed("'"),
        _ => match latex.strip_prefix("\\text{").and_then(|l| l.strip_suffix('}')) {
            Some(inner) => Cow::Owned(inner.to_string()),
            None => Cow::Owned(format!("\\ensuremath{{{latex}}}")),
        },
    }
}

fn frac_cmd(style: FracStyle) -> &'static str {
    match style {
        FracStyle::Bar => "\\frac",
        FracStyle::Display => "\\dfrac",
        FracStyle::Text => "\\tfrac",
        FracStyle::Binom => "\\binom",
        // amsmath warns about `\atop`, the generalized fraction draws the same thing.
        FracStyle::Atop => "\\genfrac{}{}{0pt}{}",
    }
}

fn delim_tex(c: char) -> &'static str {
    match c {
        '(' => "(",
        ')' => ")",
        '[' => "[",
        ']' => "]",
        '{' => "\\{",
        '}' => "\\}",
        '|' => "|",
        '‖' => "\\|",
        '/' => "/",
        '⌈' => "\\lceil",
        '⌉' => "\\rceil",
        '⌊' => "\\lfloor",
        '⌋' => "\\rfloor",
        '⟨' => "\\langle",
        '⟩' => "\\rangle",
        _ => ".",
    }
}

fn accent_cmd(mark: Mark) -> &'static str {
    match mark {
        Mark::Hat => "\\hat",
        Mark::Vec => "\\vec",
        Mark::Bar => "\\bar",
        Mark::Tilde => "\\tilde",
        Mark::Dot => "\\dot",
        Mark::Ddot => "\\ddot",
        Mark::Widehat => "\\widehat",
        Mark::Widetilde => "\\widetilde",
        Mark::Overline => "\\overline",
        Mark::Underline => "\\underline",
        Mark::Check => "\\check",
        Mark::Breve => "\\breve",
    }
}

fn variant_cmd(v: Variant) -> &'static str {
    match v {
        Variant::Normal => "\\mathnormal",
        Variant::Bold => "\\mathbf",
        Variant::Blackboard => "\\mathbb",
        Variant::Calligraphic => "\\mathcal",
        Variant::Fraktur => "\\mathfrak",
        Variant::Roman => "\\mathrm",
        Variant::SansSerif => "\\mathsf",
        Variant::Typewriter => "\\mathtt",
        Variant::Text => "\\text",
        Variant::OperatorName => "\\operatorname",
    }
}

fn matrix_env_name(env: MatrixEnv) -> &'static str {
    match env {
        MatrixEnv::Matrix => "matrix",
        MatrixEnv::Pmatrix => "pmatrix",
        MatrixEnv::Bmatrix => "bmatrix",
        MatrixEnv::Vmatrix => "vmatrix",
        MatrixEnv::Cases => "cases",
        MatrixEnv::Aligned => "aligned",
        MatrixEnv::Array => "array",
    }
}