docling 1.71.0

//! LaTeX backend (core) — a lightweight scanner for simple `.tex` documents.
//!
//! docling's backend drives the full `pylatexenc` parser plus macro/environment/
//! math handlers (and optionally a Tectonic engine); this core handles the
//! common structural subset: `\title`/`\author`/`\maketitle`, sectioning,
//! `itemize`/`enumerate`, `tabular`, display math, and paragraph text with the
//! usual font macros stripped. Multi-file projects (`\input`/`\include`),
//! custom macros and rich math/citations are out of scope.

use crate::backend::DeclarativeBackend;
use crate::error::ConversionError;
use crate::source::SourceDocument;
use docling_core::tree::{ItemTree, ListMeta, TreeKind};
use docling_core::{DoclingDocument, Node, Table, TableCell};

pub struct LatexBackend;

impl DeclarativeBackend for LatexBackend {
    fn convert(&self, source: &SourceDocument) -> Result<DoclingDocument, ConversionError> {
        let raw = source.text()?;
        let text = strip_comments(&raw);
        let title = braced_arg(&text, "\\title");
        let author = braced_arg(&text, "\\author");
        let body = between(&text, "\\begin{document}", "\\end{document}").unwrap_or(&text);

        let mut doc = DoclingDocument::new(&source.name);
        // docling's `LatexDocumentBackend` parents every item to the body
        // (headings are flat siblings of their text); the tree records
        // upstream's labels — `paragraph` vs `text`, `formula`, list items
        // with `enumerated=False` — that the flat stream cannot express.
        doc.tree = Some(ItemTree::default());
        let chars: Vec<char> = body.chars().collect();
        let mut p = Parser {
            chars: &chars,
            i: 0,
            title,
            author,
        };
        p.run(&mut doc);
        Ok(doc)
    }
}

struct Parser<'a> {
    chars: &'a [char],
    i: usize,
    title: Option<String>,
    author: Option<String>,
}

const HEADINGS: &[(&str, u8)] = &[
    ("\\subsubsection", 4),
    ("\\subsection", 3),
    ("\\section", 2),
    ("\\paragraph", 5),
    ("\\chapter", 1),
];

impl Parser<'_> {
    /// The body walk follows docling's `_process_nodes` — a text buffer
    /// flushed at headings, environments, display math and paragraph breaks.
    /// pylatexenc hands docling the plain text between two macros as one
    /// *chars node*; `_process_chars_node` (as fixed by docling#4340) splits a
    /// node holding a paragraph break, glues its first part to the buffer
    /// unstripped, flushes, adds every *middle* part as a `paragraph` item of
    /// its own, and keeps the *last* part in the buffer so the inline macros,
    /// math and citations that continue the sentence land in the same item.
    /// Inline math is appended to the buffer too (`_process_math_node`). So a
    /// stretch of text is a `paragraph` only when it began right after a
    /// blank line, ends at one, and nothing but plain characters contributed
    /// to it; everything else is flushed by the node that follows it as a
    /// `text` item.
    fn run(&mut self, doc: &mut DoclingDocument) {
        let mut para = String::new();
        // Whether `para` started right after a paragraph break.
        let mut after_break = false;
        // Whether a non-chars node (inline math) contributed to `para`; a
        // macro shows up as a backslash in the buffer itself.
        let mut has_math = false;
        while self.i < self.chars.len() {
            let rest: String = self.chars[self.i..].iter().collect();
            if rest.starts_with("\\maketitle") {
                self.i += "\\maketitle".len();
                if let Some(t) = self.title.take() {
                    add_text(doc, "title", &t, None, None);
                    doc.push(Node::Heading { level: 1, text: t });
                }
                if let Some(a) = self.author.take() {
                    add_text(doc, "text", &a, None, None);
                    doc.push(Node::Paragraph { text: a });
                }
                after_break = false;
            } else if let Some((cmd, level)) = HEADINGS.iter().find(|(c, _)| {
                rest.starts_with(*c) && !rest[c.len()..].starts_with(|ch: char| ch.is_alphabetic())
            }) {
                flush(&mut para, doc, false);
                after_break = false;
                has_math = false;
                self.i += cmd.len();
                self.skip_star();
                let text = clean_inline(&self.read_group());
                // docling's `_get_heading_level`: part/chapter/section → 1,
                // subsection → 2, …; the flat level is one deeper (1 = title).
                add_text(
                    doc,
                    "section_header",
                    &text,
                    Some(level.saturating_sub(1).max(1)),
                    None,
                );
                doc.push(Node::Heading {
                    level: *level,
                    text,
                });
            } else if rest.starts_with("\\begin{") {
                flush(&mut para, doc, false);
                after_break = false;
                has_math = false;
                self.read_environment(doc);
            } else if rest.starts_with("\\[") || rest.starts_with("$$") {
                flush(&mut para, doc, false);
                after_break = false;
                has_math = false;
                let close = if rest.starts_with("\\[") { "\\]" } else { "$$" };
                self.i += 2;
                let math = self.read_until(close);
                emit_formula(doc, &math.split_whitespace().collect::<Vec<_>>().join(" "));
            } else if self.chars[self.i] == '$' {
                // Inline math joins the text buffer, `$…$` included (upstream's
                // `_process_math_node` appends `latex_verbatim()` for non-display
                // math), so `Inline math: $E = mc^2$` stays one `text` item.
                self.i += 1;
                let math = self.read_until("$");
                para.push_str(&format!(
                    "${}$",
                    math.split_whitespace().collect::<Vec<_>>().join(" ")
                ));
                has_math = true;
            } else if self.chars[self.i] == '\n' && self.peek_blank_line() {
                let as_paragraph = after_break && !has_math && !para.contains('\\');
                flush(&mut para, doc, as_paragraph);
                after_break = true;
                has_math = false;
                self.consume_blank_line();
            } else {
                para.push(self.chars[self.i]);
                self.i += 1;
            }
        }
        flush(&mut para, doc, false);
    }

    /// `\begin{env} … \end{env}` — handle the structural environments, ignore others.
    fn read_environment(&mut self, doc: &mut DoclingDocument) {
        self.i += "\\begin".len();
        let env = self.read_group();
        self.skip_optional();
        let inner = self.read_until(&format!("\\end{{{env}}}"));
        match env.as_str() {
            "itemize" | "enumerate" => emit_list(&inner, doc),
            "tabular" => emit_table(&inner, doc),
            "equation" | "displaymath" | "align" | "equation*" | "align*" | "gather"
            | "gather*" => {
                emit_formula(
                    doc,
                    &clean_inline(&inner)
                        .split_whitespace()
                        .collect::<Vec<_>>()
                        .join(" "),
                );
            }
            // Containers: render inner content (nested tabular, caption, …) transparently.
            "table" | "figure" | "document" | "abstract" | "center" => {
                let chars: Vec<char> = inner.chars().collect();
                let mut sub = Parser {
                    chars: &chars,
                    i: 0,
                    title: None,
                    author: None,
                };
                sub.run(doc);
            }
            _ => {}
        }
    }

    /// Skip an optional `[…]` argument (e.g. `\begin{table}[h]`).
    fn skip_optional(&mut self) {
        let mut j = self.i;
        while j < self.chars.len() && self.chars[j].is_whitespace() {
            j += 1;
        }
        if j < self.chars.len() && self.chars[j] == '[' {
            self.i = j + 1;
            while self.i < self.chars.len() && self.chars[self.i] != ']' {
                self.i += 1;
            }
            if self.i < self.chars.len() {
                self.i += 1;
            }
        }
    }

    fn skip_star(&mut self) {
        if self.i < self.chars.len() && self.chars[self.i] == '*' {
            self.i += 1;
        }
    }

    /// Read a `{…}` group (skipping leading whitespace), returning its raw content.
    fn read_group(&mut self) -> String {
        while self.i < self.chars.len() && self.chars[self.i].is_whitespace() {
            self.i += 1;
        }
        if self.i >= self.chars.len() || self.chars[self.i] != '{' {
            return String::new();
        }
        self.i += 1;
        let mut depth = 1;
        let mut out = String::new();
        while self.i < self.chars.len() && depth > 0 {
            match self.chars[self.i] {
                '{' => {
                    depth += 1;
                    out.push('{');
                }
                '}' => {
                    depth -= 1;
                    if depth > 0 {
                        out.push('}');
                    }
                }
                c => out.push(c),
            }
            self.i += 1;
        }
        out
    }

    fn read_until(&mut self, marker: &str) -> String {
        let mut out = String::new();
        let mlen = marker.chars().count();
        while self.i < self.chars.len() {
            if self.chars[self.i..].len() >= mlen
                && self.chars[self.i..self.i + mlen].iter().collect::<String>() == marker
            {
                self.i += mlen;
                return out;
            }
            out.push(self.chars[self.i]);
            self.i += 1;
        }
        out
    }

    fn peek_blank_line(&self) -> bool {
        let mut j = self.i + 1;
        while j < self.chars.len() && (self.chars[j] == ' ' || self.chars[j] == '\t') {
            j += 1;
        }
        j < self.chars.len() && self.chars[j] == '\n'
    }

    fn consume_blank_line(&mut self) {
        while self.i < self.chars.len() && self.chars[self.i] != '\n' {
            self.i += 1;
        }
        while self.i < self.chars.len()
            && (self.chars[self.i] == '\n'
                || self.chars[self.i] == ' '
                || self.chars[self.i] == '\t')
        {
            self.i += 1;
        }
    }
}

/// docling's `add_text` on the body: every LaTeX item is a body child.
fn add_text(
    doc: &mut DoclingDocument,
    label: &str,
    text: &str,
    level: Option<u8>,
    list: Option<ListMeta>,
) -> usize {
    tree(doc).add(
        None,
        None,
        TreeKind::Text {
            label: label.into(),
            text: text.into(),
            orig: None,
            formatting: None,
            hyperlink: None,
            level,
            list,
        },
    )
}

fn tree(doc: &mut DoclingDocument) -> &mut ItemTree {
    doc.tree.get_or_insert_with(ItemTree::default)
}

/// A display equation: a `formula` item holding the bare math; the flat
/// stream carries it `$$`-fenced for Markdown.
fn emit_formula(doc: &mut DoclingDocument, math: &str) {
    add_text(doc, "formula", math, None, None);
    doc.push(Node::Paragraph {
        text: format!("$${math}$$"),
    });
}

/// Flush the paragraph buffer. `as_paragraph` says the text is a middle part
/// of a chars node — it began right after a blank line, ends at one, and no
/// macro or math contributed to it — which upstream's `_process_chars_node`
/// adds as a `paragraph` item; everything that went through the text buffer
/// is a `text` item.
fn flush(para: &mut String, doc: &mut DoclingDocument, as_paragraph: bool) {
    let label = if as_paragraph { "paragraph" } else { "text" };
    let text = clean_inline(para);
    let text = text.split_whitespace().collect::<Vec<_>>().join(" ");
    if !text.is_empty() {
        add_text(doc, label, &text, None, None);
        doc.push(Node::Paragraph { text });
    }
    para.clear();
}

/// Emit `\item` entries of an itemize/enumerate body as (unordered) list items.
fn emit_list(inner: &str, doc: &mut DoclingDocument) {
    let mut first = true;
    // docling's `_process_environment` opens one `add_list_group` per
    // itemize/enumerate and adds every `\item` with `enumerated=False` and
    // an empty marker — an enumerate is not numbered in its JSON either.
    let mut group = None;
    for part in inner.split("\\item").skip(1) {
        let text = clean_inline(part);
        let text = text.split_whitespace().collect::<Vec<_>>().join(" ");
        if !text.is_empty() {
            let g = *group.get_or_insert_with(|| {
                tree(doc).add(
                    None,
                    None,
                    TreeKind::Group {
                        label: "list".into(),
                        name: "list".into(),
                    },
                )
            });
            tree(doc).add(
                Some(g),
                None,
                TreeKind::Text {
                    label: "list_item".into(),
                    text: text.clone(),
                    orig: None,
                    formatting: None,
                    hyperlink: None,
                    level: None,
                    list: Some(ListMeta::default()),
                },
            );
            doc.push(Node::ListItem {
                ordered: false,
                number: 0,
                first_in_list: first,
                text,
                level: 0,
                marker: None,
                location: None,
                dclx: None,
                href: None,
                layer: None,
            });
            first = false;
        }
    }
}

/// Build a table from a `tabular` environment's rows (`\\` separated, `&` cells).
fn emit_table(inner: &str, doc: &mut DoclingDocument) {
    // Drop the column spec that follows \begin{tabular}.
    let body = match inner.find('}') {
        Some(p) if inner.trim_start().starts_with('{') => &inner[p + 1..],
        _ => inner,
    };
    let mut rows = Vec::new();
    for line in body.split("\\\\") {
        // docling's `_parse_table` finishes a row at `\\` and once more at the
        // end of the environment; rule macros (`\hline`, `\toprule`, …) add
        // no cell content, but the bare newlines around them are chars nodes,
        // so a segment that still holds *anything* after removing the rules —
        // typically the `\n\hline\n` after the last `\\` — is a row of one
        // empty cell, padded to the table's width (upstream's trailing blank
        // row). Only a segment with nothing at all yields no row.
        let line = line
            .replace("\\hline", "")
            .replace("\\toprule", "")
            .replace("\\midrule", "")
            .replace("\\bottomrule", "");
        if line.is_empty() {
            continue;
        }
        rows.push(
            line.split('&')
                .map(|c| {
                    let t = clean_inline(c);
                    t.split_whitespace().collect::<Vec<_>>().join(" ")
                })
                .collect::<Vec<_>>(),
        );
    }
    if rows.is_empty() {
        return;
    }
    // docling's `_parse_table` pads every row to the widest and writes each
    // slot as a plain `TableCell` — no header roles (a `tabular` has no
    // header markup), the trailing blank row included.
    let num_cols = rows.iter().map(Vec::len).max().unwrap_or(0);
    for r in &mut rows {
        r.resize(num_cols, String::new());
    }
    let cells = rows
        .iter()
        .enumerate()
        .flat_map(|(ri, row)| {
            row.iter().enumerate().map(move |(ci, text)| TableCell {
                text: text.clone(),
                bbox: None,
                start_row: ri,
                start_col: ci,
                row_span: 1,
                col_span: 1,
                column_header: false,
                row_header: false,
                row_section: false,
            })
        })
        .collect();
    let table = Table {
        rows,
        location: None,
        structure: None,
        cell_blocks: None,
        cells: Some(cells),
        caption: None,
        caption_parent: Default::default(),
    };
    tree(doc).add(
        None,
        None,
        TreeKind::Table {
            table: table.clone(),
            rich_cells: Vec::new(),
            captions: Vec::new(),
        },
    );
    doc.push(Node::Table(table));
}

/// Strip line comments (`%` to end of line, unless escaped `\%`).
fn strip_comments(text: &str) -> String {
    let mut out = String::new();
    for line in text.split_inclusive('\n') {
        let bytes: Vec<char> = line.chars().collect();
        let mut k = 0;
        while k < bytes.len() {
            if bytes[k] == '%' && (k == 0 || bytes[k - 1] != '\\') {
                out.push('\n');
                break;
            }
            out.push(bytes[k]);
            k += 1;
        }
    }
    out
}

/// The braced argument of `\cmd{…}` anywhere in the text.
fn braced_arg(text: &str, cmd: &str) -> Option<String> {
    let start = text.find(cmd)? + cmd.len();
    let rest: Vec<char> = text[start..].chars().collect();
    let mut k = 0;
    while k < rest.len() && rest[k].is_whitespace() {
        k += 1;
    }
    if k >= rest.len() || rest[k] != '{' {
        return None;
    }
    k += 1;
    let mut depth = 1;
    let mut out = String::new();
    while k < rest.len() && depth > 0 {
        match rest[k] {
            '{' => depth += 1,
            '}' => {
                depth -= 1;
                if depth == 0 {
                    break;
                }
            }
            _ => {}
        }
        out.push(rest[k]);
        k += 1;
    }
    Some(clean_inline(&out))
}

fn between<'a>(text: &'a str, open: &str, close: &str) -> Option<&'a str> {
    let s = text.find(open)? + open.len();
    let e = text[s..].find(close)? + s;
    Some(&text[s..e])
}

/// Strip the common inline font/reference macros, keeping their content
/// (`\textbf{x}`→`x`, `\cite{k}`→`[k]`, `\ref{k}`→`[k]`), and unescape `\&`,`\%`,…
fn clean_inline(s: &str) -> String {
    let mut out = String::new();
    let chars: Vec<char> = s.chars().collect();
    let mut i = 0;
    while i < chars.len() {
        if chars[i] == '\\' && i + 1 < chars.len() {
            let rest: String = chars[i..].iter().collect();
            // \cite{k} / \ref{k} / \eqref{k} → [k]
            if let Some(cmd) = ["\\cite", "\\ref", "\\eqref", "\\citep", "\\citet"]
                .iter()
                .find(|c| rest.starts_with(**c))
            {
                i += cmd.len();
                let (arg, ni) = read_group_at(&chars, i);
                i = ni;
                out.push('[');
                out.push_str(&arg);
                out.push(']');
                continue;
            }
            // font/structure macros: keep the group content
            if let Some(cmd) = [
                "\\textbf",
                "\\textit",
                "\\emph",
                "\\texttt",
                "\\textrm",
                "\\textsc",
                "\\underline",
                "\\mbox",
                "\\caption",
                "\\text",
            ]
            .iter()
            .find(|c| rest.starts_with(**c))
            {
                i += cmd.len();
                let (arg, ni) = read_group_at(&chars, i);
                i = ni;
                out.push_str(&clean_inline(&arg));
                continue;
            }
            // escaped specials
            if matches!(chars[i + 1], '&' | '%' | '#' | '_' | '$' | '{' | '}') {
                out.push(chars[i + 1]);
                i += 2;
                continue;
            }
        }
        if chars[i] == '~' {
            out.push(' ');
            i += 1;
            continue;
        }
        out.push(chars[i]);
        i += 1;
    }
    out
}

/// Read a `{…}` group starting at `i` (after the macro name); returns (content, next-index).
fn read_group_at(chars: &[char], mut i: usize) -> (String, usize) {
    while i < chars.len() && chars[i].is_whitespace() {
        i += 1;
    }
    if i >= chars.len() || chars[i] != '{' {
        return (String::new(), i);
    }
    i += 1;
    let mut depth = 1;
    let mut out = String::new();
    while i < chars.len() && depth > 0 {
        match chars[i] {
            '{' => {
                depth += 1;
                out.push('{');
            }
            '}' => {
                depth -= 1;
                if depth > 0 {
                    out.push('}');
                }
            }
            c => out.push(c),
        }
        i += 1;
    }
    (out, i)
}

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

    /// Deliberate deviation from docling (see [`Parser::run`]): the space
    /// after a paragraph's last formatting macro is kept (`bf more.`, not
    /// upstream's `bfmore.`) and a sentence continuing with a macro after a
    /// paragraph break stays one paragraph (`C emx and it.`, not `C` alone).
    #[test]
    fn paragraph_breaks_keep_the_space_after_the_last_macro() {
        let tex = "\\begin{document}\nA \\textit{it} text. B \\textbf{bf} more.\n\n\
            C \\emph{em}x and \\textit{it}.\n\\end{document}";
        let src = SourceDocument::from_bytes("d", InputFormat::Latex, tex.as_bytes().to_vec());
        let md = LatexBackend.convert(&src).unwrap().export_to_markdown();
        assert_eq!(md.trim(), "A it text. B bf more.\n\nC emx and it.");
    }

    /// docling's `_parse_table` finishes one more row after the last `\\`:
    /// the newline around a trailing `\hline` is a chars node, so every
    /// tabular ends with a blank row.
    #[test]
    fn tabular_keeps_the_trailing_blank_row() {
        let tex = "\\begin{document}\n\\begin{tabular}{|c|c|}\n\\hline\nH1 & H2 \\\\\n\\hline\n\
            a & b \\\\\n\\hline\n\\end{tabular}\n\\end{document}";
        let src = SourceDocument::from_bytes("d", InputFormat::Latex, tex.as_bytes().to_vec());
        let md = LatexBackend.convert(&src).unwrap().export_to_markdown();
        assert!(
            md.contains("| H1   | H2   |\n|------|------|\n| a    | b    |\n|      |      |"),
            "got:\n{md}"
        );
    }

    /// The item tree docling's backend builds: everything on the body, the
    /// `paragraph` label for a chars node's middle part vs `text` for
    /// buffered text (a sentence continued by inline math included,
    /// docling#4340), `formula` items, list items that are never
    /// `enumerated`, and `tabular` cells with no header roles down to the
    /// trailing blank row.
    #[test]
    fn tree_labels_and_list_meta_follow_docling() {
        let tex = "\\title{T}\\author{A}\n\\begin{document}\n\\maketitle\n\n            \\section{Math}\n\nInline math: $x$\n\nDisplay:\n$$y = 1$$\n\n            \\subsection{List}\nAfter a heading \\textbf{bold} text.\n\n            \\begin{enumerate}\n\\item one\n\\item two\n\\end{enumerate}\n\n            \\begin{tabular}{cc}\n\\hline\na & b \\\\\n\\hline\n\\end{tabular}\n\\end{document}";
        let src = SourceDocument::from_bytes("d", InputFormat::Latex, tex.as_bytes().to_vec());
        let json: serde_json::Value =
            serde_json::from_str(&LatexBackend.convert(&src).unwrap().export_to_json()).unwrap();
        let texts = json["texts"].as_array().unwrap();
        let labels: Vec<(&str, &str)> = texts
            .iter()
            .map(|t| (t["label"].as_str().unwrap(), t["text"].as_str().unwrap()))
            .collect();
        assert_eq!(
            labels,
            [
                ("title", "T"),
                ("text", "A"),
                ("section_header", "Math"),
                // Inline math stays in the sentence's buffer; the text before
                // display math is the buffer's last part — both `text`.
                ("text", "Inline math: $x$"),
                ("text", "Display:"),
                ("formula", "y = 1"),
                ("section_header", "List"),
                // Merged across the macro (docling#4339): the buffer's label.
                ("text", "After a heading bold text."),
                ("list_item", "one"),
                ("list_item", "two"),
            ]
        );
        assert_eq!(texts[2]["level"], 1);
        assert_eq!(texts[6]["level"], 2);
        assert_eq!(texts[9]["enumerated"], false);
        assert_eq!(texts[9]["marker"], "");
        assert_eq!(texts[9]["parent"]["$ref"], "#/groups/0");
        assert_eq!(json["groups"][0]["label"], "list");
        assert_eq!(json["groups"][0]["name"], "list");
        assert!(texts
            .iter()
            .all(|t| t["parent"]["$ref"] == "#/body" || t["parent"]["$ref"] == "#/groups/0"));
        let data = &json["tables"][0]["data"];
        assert_eq!(data["num_rows"], 2);
        assert_eq!(data["num_cols"], 2);
        let cells = data["table_cells"].as_array().unwrap();
        assert_eq!(cells.len(), 4);
        assert!(cells.iter().all(|c| c["column_header"] == false));
        assert_eq!(cells[3]["text"], "");
    }

    #[test]
    fn title_sections_lists_and_font_macros() {
        let tex = "\\title{T}\\author{A}\n\\begin{document}\n\\maketitle\n\
            \\section{Intro}\nHello \\textbf{world}.\n\
            \\begin{itemize}\n\\item one\n\\item two\n\\end{itemize}\n\\end{document}";
        let src = SourceDocument::from_bytes("d", InputFormat::Latex, tex.as_bytes().to_vec());
        let md = LatexBackend.convert(&src).unwrap().export_to_markdown();
        assert!(
            md.starts_with("# T\n\nA\n\n## Intro\n\nHello world.\n\n- one\n- two"),
            "got:\n{md}"
        );
    }
}