texform_transform/finalize_ast/
mod.rs1use 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}