Skip to main content

texform_transform/finalize_ast/
mod.rs

1//! Final AST cleanup that does not depend on rewrite metadata.
2
3use crate::ast::{Ast, ContentMode, Node, NodeId};
4
5#[derive(Clone, Copy, Debug, PartialEq, Eq)]
6pub struct FinalizeAstConfig {
7    pub enabled: bool,
8}
9
10impl FinalizeAstConfig {
11    pub const ENABLED: Self = Self { enabled: true };
12    pub const DISABLED: Self = Self { enabled: false };
13    pub const DEFAULTS: Self = Self::ENABLED;
14}
15
16#[derive(Clone, Debug, Default, PartialEq, Eq)]
17pub struct FinalizeAstReport {
18    pub steps: FinalizeAstStepReports,
19}
20
21#[derive(Clone, Debug, Default, PartialEq, Eq)]
22pub struct FinalizeAstStepReports {
23    pub merge_adjacent_primes: FinalizeAstStepReport,
24}
25
26#[derive(Clone, Debug, Default, PartialEq, Eq)]
27pub struct FinalizeAstStepReport {
28    pub applied_count: usize,
29}
30
31pub fn run(ast: &mut Ast, config: &FinalizeAstConfig, report: &mut FinalizeAstReport) {
32    if !config.enabled {
33        return;
34    }
35
36    visit(ast, ast.root(), report);
37    ast.assert_invariants();
38}
39
40fn visit(ast: &mut Ast, node: NodeId, report: &mut FinalizeAstReport) {
41    if is_math_sequence_container(ast, node) {
42        merge_adjacent_primes(ast, node, report);
43    }
44
45    for (child, _) in ast.edges(node) {
46        if ast.contains(child) {
47            visit(ast, child, report);
48        }
49    }
50}
51
52fn is_math_sequence_container(ast: &Ast, node: NodeId) -> bool {
53    matches!(
54        ast.node(node),
55        Node::Root {
56            mode: ContentMode::Math,
57            ..
58        } | Node::Group {
59            mode: ContentMode::Math,
60            ..
61        }
62    )
63}
64
65fn merge_adjacent_primes(ast: &mut Ast, parent: NodeId, report: &mut FinalizeAstReport) {
66    let children = ast.children(parent).to_vec();
67    let mut next_children = Vec::with_capacity(children.len());
68    let mut index = 0;
69    let mut changed = false;
70
71    while index < children.len() {
72        let Some((count, next_index)) = collect_prime_run(ast, &children, index) else {
73            next_children.push(children[index]);
74            index += 1;
75            continue;
76        };
77
78        next_children.push(ast.new_node(Node::Prime { count }));
79        report.steps.merge_adjacent_primes.applied_count += 1;
80        changed = true;
81        index = next_index;
82    }
83
84    if !changed {
85        return;
86    }
87
88    for removed in ast.replace_children(parent, next_children) {
89        ast.remove_detached(removed);
90    }
91}
92
93fn collect_prime_run(ast: &Ast, children: &[NodeId], start: usize) -> Option<(usize, usize)> {
94    let mut index = start;
95    let mut count = 0;
96
97    while let Some(child) = children.get(index) {
98        match ast.node(*child) {
99            Node::Prime { count: prime_count } => {
100                count += prime_count;
101                index += 1;
102            }
103            _ => break,
104        }
105    }
106
107    (index > start + 1).then_some((count, index))
108}