mathtex-editor-core 0.3.0

//! Maps model spans onto mathtex IR boxes to place the caret, selection, and placeholders.

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

use mathtex_ir::{ByteSpan, Fragment, LayoutNodeKind, Length, Placed};

use crate::export::SpanMap;
use crate::geometry::{HostObject, Metrics, Point, Rect, RenderOutput};
use crate::model::{Cursor, Kind, NodeId, SeqId, SeqRange, Tree};

/// Geometry for each model element resolved from the IR.
#[derive(Debug, Default)]
pub(crate) struct BoxMap {
    pub(crate) node: HashMap<NodeId, Rect>,
    pub(crate) seq: HashMap<SeqId, Rect>,
}

/// Scaled points to points, the SVG `viewBox` unit.
fn pt(len: Length) -> f64 {
    len.to_pt()
}

/// Scaled points held as `i64` to points.
fn sp(v: i64) -> f64 {
    v as f64 / f64::from(Length::SP_PER_PT)
}

/// A span leaf or container with its rectangle.
#[derive(Clone, Copy)]
struct Piece {
    node: usize,
    span: (usize, usize),
    rect: Rect,
}

/// The IR indexed once: parents and the pieces that carry geometry.
struct Layout {
    parent: Vec<Option<usize>>,
    /// Ink carrying glyphs and rules sorted by span.
    leaves: Vec<Piece>,
    /// Containers with no leaf inside their span and no such container below them, sorted by span.
    leafless: Vec<Piece>,
}

impl Layout {
    fn new(fragment: &Fragment) -> Self {
        let len = fragment.nodes.len();
        let mut parent = vec![None; len];
        // Absolute reference points in scaled points, the root sits on the surface baseline as in `flatten`.
        let mut abs: Vec<Option<(i64, i64)>> = vec![None; len];
        let mut stack = vec![(fragment.root, (0i64, i64::from(fragment.surface.baseline.0)), None)];
        while let Some((id, (px, py), up)) = stack.pop() {
            let Some(node) = fragment.node(id) else { continue };
            let i = id.index();
            if abs[i].is_some() {
                continue;
            }
            let here = (px + i64::from(node.origin.x.0), py + i64::from(node.origin.y.0));
            abs[i] = Some(here);
            parent[i] = up;
            stack.extend(node.children().iter().map(|&c| (c, here, Some(i))));
        }
        let mut layout = Layout { parent, leaves: Vec::new(), leafless: Vec::new() };
        layout.collect_leaves(fragment, &abs);
        layout.collect_containers(fragment, &abs);
        layout
    }

    /// One piece per glyph, so a run typeset from several atoms splits along its clusters, and one per rule.
    fn collect_leaves(&mut self, fragment: &Fragment, abs: &[Option<(i64, i64)>]) {
        let placed = fragment.flatten();
        for (k, item) in placed.iter().enumerate() {
            match *item {
                Placed::Glyph { node, x, source: Some(source), .. } => {
                    let (Some(n), Some((ax, ay))) = (fragment.node(node), abs[node.index()]) else { continue };
                    // A glyph reaches the next glyph of its run, the last one reaches the end of the run.
                    let next = match placed.get(k + 1) {
                        Some(&Placed::Glyph { node: other, x: nx, .. }) if other == node && nx > x => i64::from(nx.0),
                        _ => ax + i64::from(n.width.0),
                    };
                    let x = i64::from(x.0);
                    let rect = Rect {
                        x: sp(x),
                        y: sp(ay) - pt(n.height),
                        width: sp((next - x).max(0)),
                        height: pt(n.height + n.depth),
                    };
                    self.push_leaf(node.index(), source.span, rect);
                }
                Placed::Rule { node, x, y, width, height, source: Some(source) } => {
                    let rect = Rect { x: pt(x), y: pt(y), width: pt(width), height: pt(height) };
                    self.push_leaf(node.index(), source.span, rect);
                }
                _ => {}
            }
        }
        self.leaves.sort_by(|a, b| (a.span, a.node).cmp(&(b.span, b.node)).then(a.rect.x.total_cmp(&b.rect.x)));
    }

    fn push_leaf(&mut self, node: usize, span: ByteSpan, rect: Rect) {
        let span = (span.start as usize, span.end as usize);
        if span.0 < span.1 {
            self.leaves.push(Piece { node, span, rect });
        }
    }

    /// Boxes whose span holds no leaf, such as phantoms in empty slots and host boxes, keeping the deepest.
    fn collect_containers(&mut self, fragment: &Fragment, abs: &[Option<(i64, i64)>]) {
        let mut leafless = Vec::new();
        for (i, n) in fragment.nodes.iter().enumerate() {
            let (LayoutNodeKind::Box(_), Some(s), Some((ax, ay))) = (&n.kind, n.primary_source, abs[i]) else {
                continue;
            };
            let span = (s.span.start as usize, s.span.end as usize);
            if span.0 >= span.1 || pieces_in(&self.leaves, span.0..span.1).next().is_some() {
                continue;
            }
            let rect = Rect { x: sp(ax), y: sp(ay) - pt(n.height), width: pt(n.width), height: pt(n.height + n.depth) };
            leafless.push(Piece { node: i, span, rect });
        }
        let mut has_leafless_below = vec![false; self.parent.len()];
        for c in &leafless {
            let mut cur = c.node;
            while let Some(p) = self.parent[cur] {
                if std::mem::replace(&mut has_leafless_below[p], true) {
                    break;
                }
                cur = p;
            }
        }
        self.leafless = leafless.into_iter().filter(|c| !has_leafless_below[c.node]).collect();
        self.leafless.sort_by_key(|p| (p.span, p.node));
    }

    fn union_in(&self, range: &Range<usize>) -> Option<Rect> {
        let leaves = pieces_in(&self.leaves, range.clone());
        let containers = pieces_in(&self.leafless, range.clone());
        union(leaves.chain(containers).map(|p| p.rect))
    }
}

fn pieces_in(pieces: &[Piece], range: Range<usize>) -> impl Iterator<Item = &Piece> + '_ {
    let start = pieces.partition_point(|p| p.span.0 < range.start);
    pieces[start..].iter().take_while(move |p| p.span.0 < range.end).filter(move |p| p.span.1 <= range.end)
}

/// Map each model element to the union of IR boxes whose source span lies in its export range.
pub(crate) fn match_boxes(spans: &SpanMap, fragment: &Fragment) -> BoxMap {
    let layout = Layout::new(fragment);
    let mut map = BoxMap::default();
    for (node, range) in &spans.nodes {
        if let Some(r) = layout.union_in(range) {
            map.node.insert(*node, r);
        }
    }
    for (seq, range) in &spans.seqs {
        if let Some(r) = layout.union_in(range) {
            map.seq.insert(*seq, r);
        }
    }
    map
}

/// Caret, selection, placeholder, menu, and host box geometry for the current state.
pub(crate) fn render(
    tree: &Tree,
    cursor: Cursor,
    sel: Option<SeqRange>,
    spans: &SpanMap,
    fragment: &Fragment,
    menu_anchor: Option<NodeId>,
) -> RenderOutput {
    let boxes = match_boxes(spans, fragment);
    let placeholders = spans
        .seqs
        .iter()
        .filter(|(s, _)| tree.is_empty(*s))
        .filter_map(|(s, _)| boxes.seq.get(s).copied())
        .collect();
    let host_objects = spans
        .nodes
        .iter()
        .filter_map(|(n, _)| match tree.kind(*n) {
            Some(Kind::HostBox { token }) => boxes.node.get(n).map(|&rect| HostObject { token: *token, rect }),
            _ => None,
        })
        .collect();
    RenderOutput {
        caret: caret_rect(&boxes, tree, cursor),
        selection: sel.map(|s| selection_rects(&boxes, tree, s)).unwrap_or_default(),
        placeholders,
        metrics: Metrics {
            width: pt(fragment.surface.width),
            height: pt(fragment.surface.height),
            baseline: pt(fragment.surface.baseline),
        },
        menu: menu_anchor.map(|n| boxes.node.get(&n).copied().unwrap_or(ZERO)),
        host_objects,
    }
}

/// Caret rect at the right edge of the left neighbor, else the left edge of the next item.
fn caret_rect(boxes: &BoxMap, tree: &Tree, cursor: Cursor) -> Rect {
    let items = tree.items(cursor.seq);
    let placement = if cursor.index > 0 {
        items.get(cursor.index - 1).map(|&n| (n, true))
    } else {
        items.get(cursor.index).map(|&n| (n, false))
    };
    if let Some(v) = placement.and_then(|(n, right)| boxes.node.get(&n).map(|v| (v, right))) {
        let (v, right) = v;
        return caret_at(if right { v.x + v.width } else { v.x }, v);
    }
    boxes.seq.get(&cursor.seq).map_or(ZERO, |v| caret_at(v.x, v))
}

fn caret_at(x: f64, v: &Rect) -> Rect {
    Rect { x, y: v.y, width: 0.0, height: v.height }
}

/// Highlight for a single seq run, the union of the run's node boxes.
fn selection_rects(boxes: &BoxMap, tree: &Tree, sel: SeqRange) -> Vec<Rect> {
    let items = tree.items(sel.seq);
    let hi = sel.hi().min(items.len());
    union(items[sel.lo().min(hi)..hi].iter().filter_map(|n| boxes.node.get(n).copied())).into_iter().collect()
}

/// Reverse hit test from a point to a caret, `None` when the fragment has no matched geometry.
pub(crate) fn hit_test(tree: &Tree, spans: &SpanMap, fragment: &Fragment, point: Point) -> Option<Cursor> {
    if fragment.nodes.is_empty() {
        return None;
    }
    hit_test_boxes(&match_boxes(spans, fragment), spans, tree, point)
}

#[derive(Clone, Copy)]
enum Target {
    Node(NodeId),
    EmptySeq(SeqId),
}

/// Candidates in export order, so every tie resolves to the earliest element in the source.
fn candidates<'a>(boxes: &'a BoxMap, spans: &'a SpanMap, tree: &'a Tree) -> impl Iterator<Item = (Target, Rect, usize)> + 'a {
    let nodes = spans
        .nodes
        .iter()
        .filter_map(|(n, r)| boxes.node.get(n).map(|&rect| (Target::Node(*n), rect, r.len())));
    // Empty seqs have no node, so their placeholder box is the only direct target.
    let seqs = spans
        .seqs
        .iter()
        .filter(|(s, _)| tree.is_empty(*s))
        .filter_map(|(s, r)| boxes.seq.get(s).map(|&rect| (Target::EmptySeq(*s), rect, r.len())));
    nodes.chain(seqs)
}

fn hit_test_boxes(boxes: &BoxMap, spans: &SpanMap, tree: &Tree, point: Point) -> Option<Cursor> {
    // The smallest span containing the point wins, the first in export order on a tie.
    let mut best: Option<(Target, Rect, usize)> = None;
    for c in candidates(boxes, spans, tree).filter(|(_, r, _)| contains(r, point)) {
        if best.is_none_or(|b| c.2 < b.2) {
            best = Some(c);
        }
    }
    // Misses during a drag resolve to the nearest box instead of jumping to the document start.
    let chosen = best.map(|(t, r, _)| (t, r)).or_else(|| {
        let mut near: Option<(Target, Rect, f64)> = None;
        for (t, r, _) in candidates(boxes, spans, tree) {
            let d = rect_dist2(&r, point);
            if near.is_none_or(|n| d < n.2) {
                near = Some((t, r, d));
            }
        }
        near.map(|(t, r, _)| (t, r))
    })?;
    match chosen {
        (Target::Node(node), rect) => {
            let (seq, idx) = tree.index_in_parent(node)?;
            let index = if point.x > rect.x + rect.width / 2.0 { idx + 1 } else { idx };
            Some(Cursor { seq, index })
        }
        (Target::EmptySeq(seq), _) => Some(Cursor { seq, index: 0 }),
    }
}

/// Squared distance from `p` to the nearest point of `r`, zero inside.
fn rect_dist2(r: &Rect, p: Point) -> f64 {
    let dx = p.x - p.x.clamp(r.x, r.x + r.width);
    let dy = p.y - p.y.clamp(r.y, r.y + r.height);
    dx * dx + dy * dy
}

const ZERO: Rect = Rect { x: 0.0, y: 0.0, width: 0.0, height: 0.0 };

fn contains(r: &Rect, p: Point) -> bool {
    p.x >= r.x && p.x <= r.x + r.width && p.y >= r.y && p.y <= r.y + r.height
}

fn union(rects: impl IntoIterator<Item = Rect>) -> Option<Rect> {
    rects.into_iter().fold(None, |acc: Option<Rect>, r| {
        Some(match acc {
            None => r,
            Some(a) => {
                let (x0, y0) = (a.x.min(r.x), a.y.min(r.y));
                let (x1, y1) = ((a.x + a.width).max(r.x + r.width), (a.y + a.height).max(r.y + r.height));
                Rect { x: x0, y: y0, width: x1 - x0, height: y1 - y0 }
            }
        })
    })
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::model::{MathClass, MatrixEnv, Symbol};
    use mathtex_ir::{
        BoxKind, FontKey, FontRef, FragmentMetadata, GlyphId, GlyphRun, LayoutBox, LayoutNode, NodeId as IrId,
        Point as IrPoint, PositionedGlyph, SourceId, SourceMap, SourceRange,
    };

    fn atom(c: &str) -> Symbol {
        Symbol { latex: c.into(), class: MathClass::Ord }
    }

    fn rect(x: f64, y: f64, width: f64, height: f64) -> Rect {
        Rect { x, y, width, height }
    }

    /// `ab` with boxes for each atom and spans in export order.
    fn two_atoms() -> (Tree, SeqId, NodeId, NodeId) {
        let mut t = Tree::new();
        let root = t.root();
        t.insert_atom(Cursor { seq: root, index: 0 }, None, atom("a")).unwrap();
        t.insert_atom(Cursor { seq: root, index: 1 }, None, atom("b")).unwrap();
        let (a, b) = (t.items(root)[0], t.items(root)[1]);
        (t, root, a, b)
    }

    #[test]
    fn a_miss_falls_back_to_the_nearest_box() {
        let (t, root, a, b) = two_atoms();
        let mut boxes = BoxMap::default();
        boxes.node.insert(a, rect(0.0, 0.0, 1.0, 1.0));
        boxes.node.insert(b, rect(5.0, 0.0, 1.0, 1.0));
        let mut spans = SpanMap::default();
        spans.push_node(a, 0..1);
        spans.push_node(b, 1..2);
        let hit = |x, y| hit_test_boxes(&boxes, &spans, &t, Point { x, y });
        assert_eq!(hit(1.5, 0.5), Some(Cursor { seq: root, index: 1 }));
        assert_eq!(hit(4.0, 0.5), Some(Cursor { seq: root, index: 1 }));
        assert_eq!(hit(0.8, -5.0), Some(Cursor { seq: root, index: 1 }));
    }

    #[test]
    fn empty_geometry_is_no_hit() {
        let (t, _, _, _) = two_atoms();
        assert_eq!(hit_test_boxes(&BoxMap::default(), &SpanMap::default(), &t, Point { x: 3.0, y: 3.0 }), None);
        assert_eq!(hit_test(&t, &SpanMap::default(), &Fragment::default(), Point { x: 3.0, y: 3.0 }), None);
    }

    /// Equal boxes and spans resolve to the first element in export order, every time.
    #[test]
    fn ties_resolve_in_export_order() {
        let (t, root, a, b) = two_atoms();
        let mut boxes = BoxMap::default();
        boxes.node.insert(a, rect(0.0, 0.0, 2.0, 2.0));
        boxes.node.insert(b, rect(0.0, 0.0, 2.0, 2.0));
        let mut spans = SpanMap::default();
        spans.push_node(b, 1..2);
        spans.push_node(a, 0..1);
        for _ in 0..8 {
            assert_eq!(hit_test_boxes(&boxes, &spans, &t, Point { x: 0.5, y: 1.0 }), Some(Cursor { seq: root, index: 1 }));
        }
    }

    #[test]
    fn empty_matrix_cells_are_hit_through_their_placeholders() {
        let mut t = Tree::new();
        let root = t.root();
        t.insert_matrix(Cursor { seq: root, index: 0 }, None, MatrixEnv::Pmatrix, 2, 2).unwrap();
        let matrix = t.items(root)[0];
        let cells = t.child_seqs(matrix);
        let mut boxes = BoxMap::default();
        boxes.node.insert(matrix, rect(0.0, 0.0, 10.0, 10.0));
        let mut spans = SpanMap::default();
        spans.push_node(matrix, 0..40);
        for (i, &cell) in cells.iter().enumerate() {
            let (x, y) = (1.0 + 5.0 * (i % 2) as f64, 1.0 + 5.0 * (i / 2) as f64);
            boxes.seq.insert(cell, rect(x, y, 3.0, 3.0));
            spans.push_seq(cell, i * 11..i * 11 + 11);
        }
        let hit = |x, y| hit_test_boxes(&boxes, &spans, &t, Point { x, y });
        assert_eq!(hit(2.5, 2.5), Some(Cursor { seq: cells[0], index: 0 }));
        assert_eq!(hit(7.5, 2.5), Some(Cursor { seq: cells[1], index: 0 }));
        assert_eq!(hit(2.5, 7.5), Some(Cursor { seq: cells[2], index: 0 }));
        assert_eq!(hit(7.5, 7.5), Some(Cursor { seq: cells[3], index: 0 }));
    }

    const U: i32 = 65536;

    /// A node with extents `(width, height, depth)` in scaled points.
    fn node(id: u32, origin: (i32, i32), size: (i32, i32, i32), span: Option<(u32, u32)>, kind: LayoutNodeKind) -> LayoutNode {
        LayoutNode {
            id: IrId(id),
            origin: IrPoint::new(Length(origin.0), Length(origin.1)),
            width: Length(size.0),
            height: Length(size.1),
            depth: Length(size.2),
            primary_source: span.map(|(start, end)| SourceRange { source: SourceId(0), span: ByteSpan { start, end } }),
            kind,
        }
    }

    fn hbox(children: Vec<u32>) -> LayoutNodeKind {
        LayoutNodeKind::Box(LayoutBox { kind: BoxKind::Horizontal, children: children.into_iter().map(IrId).collect() })
    }

    /// A run whose glyphs sit at `(x, cluster)`.
    fn glyphs(at: &[(i32, (u32, u32))]) -> LayoutNodeKind {
        LayoutNodeKind::GlyphRun(GlyphRun {
            font: FontRef { key: Some(FontKey(1)), spec: String::new(), size: Length(10 * U) },
            glyphs: at
                .iter()
                .map(|&(x, (start, end))| PositionedGlyph {
                    glyph_id: GlyphId(1),
                    offset: IrPoint::new(Length(x), Length(0)),
                    cluster: Some(ByteSpan { start, end }),
                })
                .collect(),
        })
    }

    /// A validated fragment whose root is the last node.
    fn fragment(nodes: Vec<LayoutNode>) -> Fragment {
        let mut source_map = SourceMap::default();
        source_map.add_source("input");
        let root = IrId(nodes.len() as u32 - 1);
        Fragment::new(root, nodes, source_map, FragmentMetadata::default()).expect("a valid fragment")
    }

    /// Glyphs and boxes take their rects from the node extents around the baseline.
    #[test]
    fn rects_come_from_node_extents() {
        let fragment = fragment(vec![
            node(0, (0, 0), (5 * U, 4 * U, 2 * U), Some((0, 1)), glyphs(&[(0, (0, 1))])),
            node(1, (5 * U, 0), (3 * U, U, U), Some((1, 12)), hbox(vec![])),
            node(2, (0, 0), (8 * U, 4 * U, 2 * U), None, hbox(vec![0, 1])),
        ]);
        let layout = Layout::new(&fragment);
        assert_eq!(layout.union_in(&(0..1)), Some(rect(0.0, 0.0, 5.0, 6.0)));
        assert_eq!(layout.union_in(&(1..12)), Some(rect(5.0, 3.0, 3.0, 2.0)));
    }

    /// A run typeset from several atoms, as in a text slot, splits along its glyph clusters.
    #[test]
    fn a_run_splits_along_its_clusters() {
        let run = glyphs(&[(0, (0, 1)), (2 * U, (1, 2)), (3 * U, (2, 3))]);
        let fragment = fragment(vec![
            node(0, (U, 0), (6 * U, 2 * U, 0), Some((0, 3)), run),
            node(1, (0, 0), (7 * U, 2 * U, 0), None, hbox(vec![0])),
        ]);
        let layout = Layout::new(&fragment);
        assert_eq!(layout.union_in(&(0..1)), Some(rect(1.0, 0.0, 2.0, 2.0)));
        assert_eq!(layout.union_in(&(1..2)), Some(rect(3.0, 0.0, 1.0, 2.0)));
        assert_eq!(layout.union_in(&(2..3)), Some(rect(4.0, 0.0, 3.0, 2.0)));
    }

    /// Regression: matching looked up nodes linearly and rescanned every leaf per span.
    #[test]
    fn matching_scales_to_large_fragments() {
        let mut t = Tree::new();
        let root = t.root();
        for i in 0..5000 {
            t.insert_atom(Cursor { seq: root, index: i }, None, atom("x")).unwrap();
        }
        let src = crate::export::source(&t, 0, 0, true);
        let n = src.spans.nodes.len() as u32;
        let mut nodes = Vec::new();
        for (i, (_, r)) in src.spans.nodes.iter().enumerate() {
            let span = (r.start as u32, r.end as u32);
            nodes.push(node(i as u32, (i as i32 * U, 0), (U, U, 0), Some(span), glyphs(&[(0, span)])));
        }
        nodes.push(node(n, (0, 0), (n as i32 * U, U, 0), None, hbox((0..n).collect())));
        let fragment = fragment(nodes);
        let boxes = match_boxes(&src.spans, &fragment);
        assert_eq!(boxes.node.len(), 5000);
        assert_eq!(boxes.seq.get(&root).map(|r| r.width), Some(5000.0));
    }
}