Skip to main content

texform_core/
ast.rs

1//! Transform-ready mutable AST for tree editing.
2//!
3//! The parser still exposes [`syntax_node::SyntaxNode`] as the public parse
4//! result. This module defines a separate arena-backed IR that keeps the same
5//! semantics while supporting:
6//!
7//! - fast node kind checks
8//! - bidirectional navigation through parent links
9//! - safe structural edits without exposing raw mutable node access
10//!
11//! The AST owns all nodes in a [`slotmap::HopSlotMap`]. Tree edits must go
12//! through [`Ast`] methods so `nodes`, `parent`, and detached subtree tracking
13//! stay consistent.
14//!
15//! # Detached Roots
16//!
17//! A node with no parent is valid only when it is either:
18//!
19//! - the main [`Ast::root`], or
20//! - a tracked detached subtree root
21//!
22//! Detached roots let transforms stage or preserve subtrees without silently
23//! creating invisible orphans in the arena.
24//!
25//! # Structural Invariants Only
26//!
27//! [`Ast`] validates tree ownership, parent links, detached roots, and slot
28//! shape. It does not validate command signatures or argument semantics such as
29//! whether an [`ArgumentKind`] matches a particular [`ArgumentValue`].
30
31use std::collections::HashSet;
32
33use slotmap::{HopSlotMap, SecondaryMap, new_key_type};
34use texform_interface::syntax_node::{self, SyntaxNode};
35
36/// Re-exported content mode shared with parser syntax nodes.
37///
38/// Keeping a single definition avoids duplicate `Math` / `Text` enums drifting
39/// apart across parsing and transform stages.
40pub use texform_interface::syntax_node::ContentMode;
41
42new_key_type! {
43    /// Stable arena key for a node owned by [`Ast`].
44    pub struct NodeId;
45}
46
47/// Upward edge metadata for a node that is currently attached to a parent.
48#[derive(Clone, Copy, Debug, PartialEq, Eq)]
49pub struct ParentLink {
50    /// Parent node that owns the child.
51    pub parent: NodeId,
52    /// Concrete position occupied by the child within the parent.
53    pub slot: Slot,
54}
55
56/// The direct attachment site a child occupies in its parent.
57#[derive(Clone, Copy, Debug, PartialEq, Eq)]
58pub enum Slot {
59    /// Child at an index inside [`Node::Root::children`] / [`Node::Group::children`]
60    GroupChild(usize),
61    /// Content argument stored in an argument list slot
62    Argument(usize),
63    /// Base child of [`Node::Scripted`]
64    ScriptBase,
65    /// Subscript child of [`Node::Scripted`]
66    ScriptSub,
67    /// Superscript child of [`Node::Scripted`]
68    ScriptSup,
69    /// Left operand of [`Node::Infix`]
70    InfixLeft,
71    /// Right operand of [`Node::Infix`]
72    InfixRight,
73    /// Body child of [`Node::Environment`]
74    EnvBody,
75}
76
77/// Cheap node discriminant for queries that do not need full pattern matching.
78#[derive(Clone, Copy, Debug, PartialEq, Eq)]
79pub enum NodeKind {
80    Root,
81    Group,
82    Command,
83    Infix,
84    Declarative,
85    Environment,
86    Scripted,
87    Prime,
88    Text,
89    Char,
90    ActiveSpace,
91    /// Parser-produced error placeholder; see [`Node::Error`].
92    Error,
93}
94
95/// Delimiter value used by groups and argument kinds.
96#[derive(Clone, Debug, PartialEq, Eq)]
97pub enum Delimiter {
98    /// No delimiter, corresponding to `.` in LaTeX
99    None,
100    /// Single-character delimiter such as `(`, `)` or `|`
101    Char(char),
102    // The AST keeps this owned so future transforms are not restricted to
103    // interned or static control names.
104    /// Control-sequence delimiter such as `\langle` or `\rbrace`
105    ///
106    /// This is intentionally owned `String` data even though the parser stores
107    /// `&'static str`, because transforms may synthesize non-static names.
108    Control(String),
109}
110
111/// Concrete grouping form preserved from parsing.
112#[derive(Clone, Debug, PartialEq, Eq)]
113pub enum GroupKind {
114    /// Explicit brace group: `{...}`
115    Explicit,
116    /// Implicit synthetic group used by normalization or parser folding
117    Implicit,
118    /// Delimited group such as `\left( ... \right)`
119    Delimited { left: Delimiter, right: Delimiter },
120    /// Inline math segment inside text mode: `$...$`
121    InlineMath,
122}
123
124/// Argument slot kind preserved from xparse-aware parsing.
125#[derive(Clone, Debug, PartialEq, Eq)]
126pub enum ArgumentKind {
127    /// Standard mandatory argument
128    Mandatory,
129    /// Standard optional bracket argument
130    Optional,
131    /// Boolean star slot
132    Star,
133    /// Required braced-group form
134    Group,
135    /// Delimited argument with a matched open / close pair
136    Delimited { open: Delimiter, close: Delimiter },
137    /// Paired-candidate argument that records the matched delimiters
138    Paired { open: Delimiter, close: Delimiter },
139}
140
141/// Parsed argument payload.
142///
143/// Only content-carrying argument variants contribute a tree edge. All scalar variants
144/// are stored inline on the owning node and are skipped by tree traversal.
145#[derive(Clone, Debug, PartialEq, Eq)]
146pub enum ArgumentValue {
147    /// Child subtree used as math-mode argument content
148    MathContent(NodeId),
149    /// Child subtree used as text-mode argument content
150    TextContent(NodeId),
151    /// Parsed delimiter value
152    Delimiter(Delimiter),
153    /// Control-sequence name without leading backslash
154    CSName(String),
155    /// Raw dimension string
156    Dimension(String),
157    /// Raw integer string
158    Integer(String),
159    /// Raw key-value string
160    KeyVal(String),
161    /// Parsed column specification
162    Column(String),
163    /// Boolean value, primarily used for star slots
164    Boolean(bool),
165}
166
167/// Parsed command or environment argument.
168#[derive(Clone, Debug, PartialEq, Eq)]
169pub struct Argument {
170    /// Slot kind recorded by the parser
171    pub kind: ArgumentKind,
172    /// Parsed value stored in the slot
173    pub value: ArgumentValue,
174}
175
176/// Optional argument slot in a command or environment signature.
177///
178/// `None` means the slot exists in the spec but was not supplied in source.
179pub type ArgumentSlot = Option<Argument>;
180
181/// Mutable AST node stored inside [`Ast`].
182#[derive(Clone, Debug, PartialEq, Eq)]
183pub enum Node {
184    /// Main document root containing top-level children and parse mode.
185    Root {
186        /// Ordered direct children of the root
187        children: Vec<NodeId>,
188        /// Content mode used to parse the formula
189        mode: ContentMode,
190    },
191    /// Group node containing ordered children and mode metadata.
192    Group {
193        /// Ordered direct children of the group
194        children: Vec<NodeId>,
195        /// Source-level grouping form
196        kind: GroupKind,
197        /// Content mode used to parse the group
198        mode: ContentMode,
199    },
200    /// Prefix command with argument slots.
201    Command {
202        /// Command name without leading backslash
203        name: String,
204        /// Slots defined by the matched command spec
205        args: Vec<ArgumentSlot>,
206        /// Whether this name is present in the knowledge base.
207        known: bool,
208    },
209    /// Infix command with explicit left and right operands.
210    Infix {
211        /// Command name without leading backslash
212        name: String,
213        /// Additional argument slots owned by the infix node
214        args: Vec<ArgumentSlot>,
215        /// Left operand subtree
216        left: NodeId,
217        /// Right operand subtree
218        right: NodeId,
219    },
220    /// Declarative command with explicit argument slots.
221    Declarative {
222        /// Command name without leading backslash
223        name: String,
224        /// Additional argument slots owned by the declarative node
225        args: Vec<ArgumentSlot>,
226    },
227    /// Environment node whose body must always be a group.
228    Environment {
229        /// Environment name without `begin` / `end`
230        name: String,
231        /// Parsed argument slots attached to the environment
232        args: Vec<ArgumentSlot>,
233        /// Whether this name is present in the knowledge base.
234        known: bool,
235        /// Environment body subtree. Must be a [`Node::Group`]
236        body: NodeId,
237    },
238    /// Scripted expression such as `x_i^2`.
239    Scripted {
240        /// Base expression
241        base: NodeId,
242        /// Optional subscript subtree
243        subscript: Option<NodeId>,
244        /// Optional superscript subtree
245        superscript: Option<NodeId>,
246    },
247    /// Math prime shorthand represented by one or more consecutive prime marks.
248    Prime {
249        /// Number of consecutive prime marks. Must be greater than zero.
250        count: usize,
251    },
252    /// Text-mode text chunk
253    Text(String),
254    /// Single character node
255    Char(char),
256    /// Active `~` space node
257    ActiveSpace,
258    /// Parser-produced error placeholder, mirroring
259    /// [`SyntaxNode::Error`].
260    ///
261    /// `Error` is a first-class structural leaf: a tree containing `Error`
262    /// nodes is still structurally valid. Semantic completeness (the absence of
263    /// `Error` nodes) is a separate property, not a structural invariant.
264    /// `Error` carries the original recovery message and the source snippet so
265    /// serialization can round-trip it losslessly.
266    Error {
267        /// Human-readable recovery message captured by the parser.
268        message: String,
269        /// Verbatim source slice the parser failed to interpret.
270        snippet: String,
271    },
272}
273
274impl Node {
275    /// Return the lightweight discriminant for this node.
276    pub const fn kind(&self) -> NodeKind {
277        match self {
278            Node::Root { .. } => NodeKind::Root,
279            Node::Group { .. } => NodeKind::Group,
280            Node::Command { .. } => NodeKind::Command,
281            Node::Infix { .. } => NodeKind::Infix,
282            Node::Declarative { .. } => NodeKind::Declarative,
283            Node::Environment { .. } => NodeKind::Environment,
284            Node::Scripted { .. } => NodeKind::Scripted,
285            Node::Prime { .. } => NodeKind::Prime,
286            Node::Text(_) => NodeKind::Text,
287            Node::Char(_) => NodeKind::Char,
288            Node::ActiveSpace => NodeKind::ActiveSpace,
289            Node::Error { .. } => NodeKind::Error,
290        }
291    }
292}
293
294/// Arena-backed mutable AST used by transform-oriented code.
295///
296/// Public mutation APIs intentionally avoid exposing `&mut Node` so parent links
297/// and detached subtree tracking cannot be bypassed accidentally.
298#[derive(Debug, Clone)]
299pub struct Ast {
300    nodes: HopSlotMap<NodeId, Node>,
301    parent: SecondaryMap<NodeId, ParentLink>,
302    // Detached roots are valid subtrees that currently live in the arena but
303    // are not attached to the main root. This lets transforms stage nodes
304    // before insertion without turning them into invisible orphans.
305    detached_roots: HashSet<NodeId>,
306    root: NodeId,
307}
308
309impl Ast {
310    /// Create an empty AST with a math-mode root node.
311    pub fn new() -> Self {
312        Self::with_root_mode(ContentMode::Math)
313    }
314
315    /// Create an empty AST whose root uses the given content mode.
316    pub fn with_root_mode(mode: ContentMode) -> Self {
317        let mut nodes = HopSlotMap::with_key();
318        let root = nodes.insert(Node::Root {
319            children: Vec::new(),
320            mode,
321        });
322
323        Ast {
324            nodes,
325            parent: SecondaryMap::new(),
326            detached_roots: HashSet::new(),
327            root,
328        }
329    }
330
331    /// Convert a parsed [`SyntaxNode`] tree into a mutable [`Ast`].
332    ///
333    /// Conversion preserves shape:
334    ///
335    /// - content argument variants are converted directly to the referenced node
336    /// - single-node content is not wrapped in an implicit group
337    /// - `Delimiter::Control(&'static str)` becomes [`Delimiter::Control`] with
338    ///   owned `String` data
339    /// - [`SyntaxNode::Error`] is preserved losslessly as [`Node::Error`]; the
340    ///   resulting tree is structurally valid but semantically incomplete
341    ///
342    /// # Panics
343    ///
344    /// Panics if the converted tree violates AST invariants, such as an
345    /// environment body that is not a group.
346    pub fn from_syntax_root(node: &SyntaxNode) -> Self {
347        let SyntaxNode::Root { mode, children } = node else {
348            panic!("Ast::from_syntax_root expects SyntaxNode::Root");
349        };
350
351        let mut nodes = HopSlotMap::with_key();
352        let mut parent = SecondaryMap::new();
353        let converted_children: Vec<NodeId> = children
354            .iter()
355            .map(|child| Self::convert_syntax_node(child, &mut nodes, &mut parent))
356            .collect();
357        let root = nodes.insert(Node::Root {
358            children: converted_children,
359            mode: *mode,
360        });
361
362        for (child, slot) in Self::node_edges(nodes.get(root).expect("Converted node must exist")) {
363            parent.insert(child, ParentLink { parent: root, slot });
364        }
365
366        let ast = Ast {
367            nodes,
368            parent,
369            detached_roots: HashSet::new(),
370            root,
371        };
372        ast.assert_invariants();
373        ast
374    }
375
376    /// Convert this AST back into a [`SyntaxNode`] tree.
377    pub fn to_syntax_root(&self) -> SyntaxNode {
378        let Node::Root { children, mode } = self.node(self.root) else {
379            unreachable!("root must be a Root node");
380        };
381        SyntaxNode::Root {
382            mode: *mode,
383            children: children
384                .iter()
385                .map(|child| self.to_syntax_node(*child))
386                .collect(),
387        }
388    }
389
390    /// Return the main root node of the AST.
391    pub fn root(&self) -> NodeId {
392        self.root
393    }
394
395    /// Check whether `id` still exists in the arena.
396    pub fn contains(&self, id: NodeId) -> bool {
397        self.nodes.contains_key(id)
398    }
399
400    /// Return the lightweight node kind for `id`.
401    ///
402    /// # Panics
403    ///
404    /// Panics if `id` is invalid.
405    pub fn kind(&self, id: NodeId) -> NodeKind {
406        self.node(id).kind()
407    }
408
409    /// Borrow the full node data for `id`.
410    ///
411    /// # Panics
412    ///
413    /// Panics if `id` is invalid.
414    pub fn node(&self, id: NodeId) -> &Node {
415        self.nodes.get(id).expect("Invalid NodeId")
416    }
417
418    /// Non-panicking node borrow; returns `None` if `id` is invalid.
419    pub fn node_opt(&self, id: NodeId) -> Option<&Node> {
420        self.nodes.get(id)
421    }
422
423    /// Non-panicking mutable node borrow; returns `None` if `id` is invalid.
424    ///
425    /// Direct mutation through this helper must not change a node's edges.
426    /// Structural changes must keep using the edge-aware AST methods.
427    pub fn node_opt_mut(&mut self, id: NodeId) -> Option<&mut Node> {
428        self.nodes.get_mut(id)
429    }
430
431    /// Whether `id` is currently tracked as a detached root.
432    pub fn is_detached_root(&self, id: NodeId) -> bool {
433        self.detached_roots.contains(&id)
434    }
435
436    /// Whether `target` lies in the subtree rooted at `root`, inclusive.
437    pub fn subtree_contains_node(&self, root: NodeId, target: NodeId) -> bool {
438        self.subtree_contains(root, target)
439    }
440
441    /// Return the parent link for `id`, if the node is attached.
442    ///
443    /// Root, detached roots, and invalid or removed IDs return `None`. Callers
444    /// that need to distinguish a valid detached root from an invalid ID should
445    /// check [`Ast::contains`] first.
446    pub fn parent(&self, id: NodeId) -> Option<ParentLink> {
447        self.parent.get(id).copied()
448    }
449
450    /// Return the parent node ID for `id`, if attached.
451    ///
452    /// See [`Ast::parent`] for the `None` cases.
453    pub fn parent_id(&self, id: NodeId) -> Option<NodeId> {
454        self.parent(id).map(|link| link.parent)
455    }
456
457    /// Return the slot occupied by `id` in its parent, if attached.
458    ///
459    /// See [`Ast::parent`] for the `None` cases.
460    pub fn slot(&self, id: NodeId) -> Option<Slot> {
461        self.parent(id).map(|link| link.slot)
462    }
463
464    /// Return the direct children of a root/group node.
465    ///
466    /// Other node kinds return an empty slice.
467    ///
468    /// # Panics
469    ///
470    /// Panics if `id` is invalid.
471    pub fn children(&self, id: NodeId) -> &[NodeId] {
472        match self.node(id) {
473            Node::Root { children, .. } | Node::Group { children, .. } => children,
474            _ => &[],
475        }
476    }
477
478    /// Return the argument slots owned by a command-like node.
479    ///
480    /// Non-command-like nodes return an empty slice.
481    ///
482    /// # Panics
483    ///
484    /// Panics if `id` is invalid.
485    pub fn arg_slots(&self, id: NodeId) -> &[ArgumentSlot] {
486        match self.node(id) {
487            Node::Command { args, .. }
488            | Node::Infix { args, .. }
489            | Node::Declarative { args, .. }
490            | Node::Environment { args, .. } => args,
491            _ => &[],
492        }
493    }
494
495    /// Return every direct tree edge of `id` as `(child, slot)` pairs.
496    ///
497    /// Returned order matches the AST's direct traversal order. Only
498    /// Content argument entries are exposed as argument edges.
499    ///
500    /// # Panics
501    ///
502    /// Panics if `id` is invalid.
503    pub fn edges(&self, id: NodeId) -> Vec<(NodeId, Slot)> {
504        Self::node_edges(self.node(id))
505    }
506
507    /// Deep-copy the subtree rooted at `id` and return the copied detached root.
508    ///
509    /// The cloned tree preserves node shape and scalar argument values, but every
510    /// copied node receives a fresh [`NodeId`]. The returned root is detached, so
511    /// callers can attach it with [`Ast::new_node`], [`Ast::replace_node`], or
512    /// group insertion helpers.
513    ///
514    /// # Panics
515    ///
516    /// Panics if `id` is invalid or points at the main AST root.
517    pub fn clone_subtree(&mut self, id: NodeId) -> NodeId {
518        if id == self.root {
519            panic!("Cannot clone root node as a detached subtree");
520        }
521        if !self.contains(id) {
522            panic!("Invalid NodeId");
523        }
524
525        self.clone_subtree_impl(id)
526    }
527
528    /// Return the next sibling of `id` when it is attached as a group child.
529    ///
530    /// Nodes in non-`GroupChild` slots and invalid or removed IDs return
531    /// `None`.
532    pub fn next_sibling(&self, id: NodeId) -> Option<NodeId> {
533        let parent_link = self.parent(id)?;
534        let Slot::GroupChild(index) = parent_link.slot else {
535            return None;
536        };
537
538        self.children(parent_link.parent).get(index + 1).copied()
539    }
540
541    /// Return the previous sibling of `id` when it is attached as a group child.
542    ///
543    /// Nodes in non-`GroupChild` slots and invalid or removed IDs return
544    /// `None`.
545    pub fn prev_sibling(&self, id: NodeId) -> Option<NodeId> {
546        let parent_link = self.parent(id)?;
547        let Slot::GroupChild(index) = parent_link.slot else {
548            return None;
549        };
550
551        index
552            .checked_sub(1)
553            .and_then(|prev| self.children(parent_link.parent).get(prev).copied())
554    }
555
556    /// Depth-first search starting at `start`, returning the first matching node.
557    ///
558    /// # Panics
559    ///
560    /// Panics if `start` is invalid.
561    pub fn find<F>(&self, start: NodeId, predicate: F) -> Option<NodeId>
562    where
563        F: Fn(&Node) -> bool,
564    {
565        self.find_impl(start, &predicate)
566    }
567
568    /// Collect all matching nodes reachable from `start` in depth-first order.
569    ///
570    /// The returned vector is a snapshot of `NodeId`s at collection time.
571    /// Later mutations may delete or move those nodes, so callers should use
572    /// [`Ast::contains`] and, when necessary, re-check parent / slot state
573    /// before mutating.
574    ///
575    /// # Panics
576    ///
577    /// Panics if `start` is invalid.
578    pub fn find_all<F>(&self, start: NodeId, predicate: F) -> Vec<NodeId>
579    where
580        F: Fn(&Node) -> bool,
581    {
582        let mut result = Vec::new();
583        self.find_all_impl(start, &predicate, &mut result);
584        result
585    }
586
587    /// Check whether `id` is a character node matching `expected`.
588    ///
589    /// # Panics
590    ///
591    /// Panics if `id` is invalid.
592    pub fn is_char(&self, id: NodeId, expected: char) -> bool {
593        matches!(self.node(id), Node::Char(ch) if *ch == expected)
594    }
595
596    /// Check whether the subtree rooted at `start` contains a command named `name`.
597    ///
598    /// # Panics
599    ///
600    /// Panics if `start` is invalid.
601    pub fn subtree_contains_command(&self, start: NodeId, name: &str) -> bool {
602        self.find(
603            start,
604            |node| matches!(node, Node::Command { name: command_name, .. } if command_name == name),
605        )
606        .is_some()
607    }
608
609    /// Insert a new detached node into the arena and return its [`NodeId`].
610    ///
611    /// If `node` references child IDs, those children must already exist as
612    /// detached roots. They are adopted by the new node immediately.
613    ///
614    /// # Panics
615    ///
616    /// These panics indicate a caller-side structural invariant violation. The
617    /// AST is not guaranteed to be reusable after such a panic.
618    ///
619    /// Panics if:
620    ///
621    /// - the same child is referenced more than once
622    /// - a referenced child does not exist
623    /// - a referenced child is not a detached root
624    /// - adopting a child would introduce a cycle
625    /// - an environment body is not a group
626    pub fn new_node(&mut self, node: Node) -> NodeId {
627        if matches!(node, Node::Root { .. }) {
628            panic!("Cannot create detached root node");
629        }
630
631        Self::assert_unique_direct_children(&node);
632        let direct_children = Self::node_edges(&node);
633        let id = self.nodes.insert(node);
634        self.detached_roots.insert(id);
635
636        for &(child, slot) in &direct_children {
637            self.assert_child_is_detached_root(child);
638            self.assert_no_cycle(child, id);
639            self.assert_slot_shape(slot, child);
640            self.adopt_child(id, child, slot);
641        }
642
643        id
644    }
645
646    /// Append `child` to the end of a root/group child list.
647    ///
648    /// `child` must currently be a detached root.
649    ///
650    /// # Panics
651    ///
652    /// Panics if:
653    ///
654    /// - `parent` is not a root/group node
655    /// - `child` is invalid
656    /// - `child` is already attached
657    /// - `child` is not tracked as a detached root
658    /// - attaching `child` would introduce a cycle
659    pub fn append_child(&mut self, parent: NodeId, child: NodeId) {
660        self.assert_child_is_detached_root(child);
661        self.assert_no_cycle(child, parent);
662
663        let index = match self.node_mut(parent) {
664            Node::Root { children, .. } | Node::Group { children, .. } => {
665                let index = children.len();
666                children.push(child);
667                index
668            }
669            _ => panic!("Parent is not a root/group node"),
670        };
671
672        self.adopt_child(parent, child, Slot::GroupChild(index));
673    }
674
675    /// Insert `child` at `index` within a root/group child list.
676    ///
677    /// `child` must currently be a detached root.
678    ///
679    /// # Panics
680    ///
681    /// Panics if:
682    ///
683    /// - `parent` is not a root/group node
684    /// - `index` is out of bounds for `Vec::insert`
685    /// - `child` is invalid
686    /// - `child` is already attached
687    /// - `child` is not tracked as a detached root
688    /// - attaching `child` would introduce a cycle
689    pub fn insert_child(&mut self, parent: NodeId, index: usize, child: NodeId) {
690        self.assert_child_is_detached_root(child);
691        self.assert_no_cycle(child, parent);
692
693        match self.node_mut(parent) {
694            Node::Root { children, .. } | Node::Group { children, .. } => {
695                children.insert(index, child);
696            }
697            _ => panic!("Parent is not a root/group node"),
698        }
699
700        self.detached_roots.remove(&child);
701        self.reindex_group_children(parent, index);
702    }
703
704    /// Detach `id` from its parent and return the same node ID.
705    ///
706    /// This version only supports detaching nodes that occupy [`Slot::GroupChild`].
707    /// Detached nodes remain in the arena and become detached roots until they
708    /// are reattached or explicitly removed with [`Ast::remove_detached`].
709    ///
710    /// # Panics
711    ///
712    /// Panics if:
713    ///
714    /// - `id` is the main root
715    /// - `id` is already detached
716    /// - `id` is attached through a non-`GroupChild` slot
717    pub fn detach(&mut self, id: NodeId) -> NodeId {
718        if id == self.root {
719            panic!("Cannot detach root node");
720        }
721
722        let parent_link = self
723            .parent(id)
724            .unwrap_or_else(|| panic!("Cannot detach node without a parent"));
725        let Slot::GroupChild(index) = parent_link.slot else {
726            panic!("Can only detach GroupChild nodes");
727        };
728
729        match self.node_mut(parent_link.parent) {
730            Node::Root { children, .. } | Node::Group { children, .. } => {
731                assert_eq!(
732                    children.get(index).copied(),
733                    Some(id),
734                    "Group child index must match detached node"
735                );
736                children.remove(index);
737            }
738            _ => panic!("Parent is not a root/group node"),
739        }
740
741        self.release_child_as_detached_root(id);
742        self.reindex_group_children(parent_link.parent, index);
743        id
744    }
745
746    /// Replace all direct children of a root/group node.
747    ///
748    /// New children may be detached roots or existing direct children of the
749    /// same parent. Removed children are preserved as detached roots.
750    pub fn replace_children(&mut self, parent: NodeId, children: Vec<NodeId>) -> Vec<NodeId> {
751        let old_children = match self.node(parent) {
752            Node::Root { children, .. } | Node::Group { children, .. } => children.clone(),
753            _ => panic!("Parent is not a root/group node"),
754        };
755        let old_child_set: HashSet<NodeId> = old_children.iter().copied().collect();
756
757        let mut seen = HashSet::new();
758        for child in &children {
759            assert!(
760                seen.insert(*child),
761                "Node cannot reference the same child twice"
762            );
763            assert!(self.contains(*child), "Invalid child NodeId");
764            assert!(*child != self.root, "Cannot attach the AST root as a child");
765            self.assert_no_cycle(*child, parent);
766
767            if old_child_set.contains(child) {
768                assert_eq!(
769                    self.parent(*child),
770                    Some(ParentLink {
771                        parent,
772                        slot: Slot::GroupChild(
773                            old_children
774                                .iter()
775                                .position(|old_child| old_child == child)
776                                .expect("old child should have an index")
777                        )
778                    }),
779                    "Can only reuse direct children of the replaced parent"
780                );
781            } else {
782                self.assert_child_is_detached_root(*child);
783            }
784        }
785
786        let new_child_set: HashSet<NodeId> = children.iter().copied().collect();
787        let removed: Vec<NodeId> = old_children
788            .iter()
789            .copied()
790            .filter(|child| !new_child_set.contains(child))
791            .collect();
792
793        for child in &removed {
794            self.release_child_as_detached_root(*child);
795        }
796        for child in &children {
797            self.detached_roots.remove(child);
798        }
799
800        match self.node_mut(parent) {
801            Node::Root {
802                children: old_children,
803                ..
804            }
805            | Node::Group {
806                children: old_children,
807                ..
808            } => *old_children = children,
809            _ => panic!("Parent is not a root/group node"),
810        }
811        self.reindex_group_children(parent, 0);
812        removed
813    }
814
815    /// Detach a contiguous range of direct children from a root/group node.
816    ///
817    /// # Panics
818    ///
819    /// Panics if `parent` is not a root/group node or `range` is out of bounds
820    /// for `Vec::drain`.
821    pub fn detach_children_range(
822        &mut self,
823        parent: NodeId,
824        range: std::ops::Range<usize>,
825    ) -> Vec<NodeId> {
826        let removed = match self.node_mut(parent) {
827            Node::Root { children, .. } | Node::Group { children, .. } => {
828                children.drain(range.clone()).collect::<Vec<_>>()
829            }
830            _ => panic!("Parent is not a root/group node"),
831        };
832
833        for child in &removed {
834            self.release_child_as_detached_root(*child);
835        }
836        self.reindex_group_children(parent, range.start);
837        removed
838    }
839
840    /// Replace a content child or another single-child slot.
841    ///
842    /// The replacement must be a detached root. The old child becomes detached.
843    pub fn replace_content_child(&mut self, old: NodeId, replacement: NodeId) {
844        if old == self.root {
845            panic!("Cannot replace root node");
846        }
847        let parent_link = self
848            .parent(old)
849            .unwrap_or_else(|| panic!("Cannot replace node without a parent"));
850        self.assert_child_is_detached_root(replacement);
851        self.assert_no_cycle(replacement, parent_link.parent);
852        self.assert_slot_shape(parent_link.slot, replacement);
853
854        match self.node_mut(parent_link.parent) {
855            Node::Command { args, .. }
856            | Node::Infix { args, .. }
857            | Node::Declarative { args, .. }
858                if matches!(parent_link.slot, Slot::Argument(_)) =>
859            {
860                replace_argument_child(args, parent_link.slot, old, replacement);
861            }
862            Node::Infix { left, right, .. } => match parent_link.slot {
863                Slot::InfixLeft => {
864                    assert_eq!(*left, old, "Infix left operand must match old node");
865                    *left = replacement;
866                }
867                Slot::InfixRight => {
868                    assert_eq!(*right, old, "Infix right operand must match old node");
869                    *right = replacement;
870                }
871                _ => panic!("Expected infix operand slot"),
872            },
873            Node::Environment { args, body, .. } => match parent_link.slot {
874                Slot::Argument(_) => {
875                    replace_argument_child(args, parent_link.slot, old, replacement)
876                }
877                Slot::EnvBody => {
878                    assert_eq!(*body, old, "Environment body must match old node");
879                    *body = replacement;
880                }
881                _ => panic!("Expected environment body or argument slot"),
882            },
883            Node::Scripted {
884                base,
885                subscript,
886                superscript,
887            } => match parent_link.slot {
888                Slot::ScriptBase => {
889                    assert_eq!(*base, old, "Script base must match old node");
890                    *base = replacement;
891                }
892                Slot::ScriptSub => {
893                    assert_eq!(
894                        subscript,
895                        &Some(old),
896                        "Script subscript must match old node"
897                    );
898                    *subscript = Some(replacement);
899                }
900                Slot::ScriptSup => {
901                    assert_eq!(
902                        superscript,
903                        &Some(old),
904                        "Script superscript must match old node"
905                    );
906                    *superscript = Some(replacement);
907                }
908                _ => panic!("Expected script slot"),
909            },
910            _ => panic!("Parent does not have replaceable content children"),
911        }
912
913        self.release_child_as_detached_root(old);
914        self.adopt_child(parent_link.parent, replacement, parent_link.slot);
915    }
916
917    /// Remove an attached node and its entire subtree from the arena.
918    ///
919    /// This is implemented as [`Ast::detach`] followed by
920    /// [`Ast::remove_detached`].
921    ///
922    /// # Panics
923    ///
924    /// Panics under the same conditions as [`Ast::detach`] or
925    /// [`Ast::remove_detached`].
926    pub fn remove_node(&mut self, id: NodeId) {
927        let detached = self.detach(id);
928        self.remove_detached(detached);
929    }
930
931    /// Replace the node data stored at `id` while preserving the same [`NodeId`].
932    ///
933    /// The new node may:
934    ///
935    /// - reuse direct children already owned by `id`
936    /// - adopt detached roots
937    ///
938    /// Direct children removed by the replacement are not deleted. They become
939    /// detached roots so transforms can keep or inspect them explicitly.
940    ///
941    /// # Panics
942    ///
943    /// These panics indicate a caller-side structural invariant violation. The
944    /// AST is not guaranteed to be reusable after such a panic.
945    ///
946    /// Panics if:
947    ///
948    /// - `id` is the main root or invalid
949    /// - the new node references the same child more than once
950    /// - the new node is a root variant
951    /// - a reused child is not a direct child of `id`
952    /// - a newly introduced child is not a detached root
953    /// - adopting a child would introduce a cycle
954    /// - an environment body is not a group
955    pub fn replace_node(&mut self, id: NodeId, new_node: Node) -> Node {
956        if id == self.root {
957            panic!("Cannot replace root node");
958        }
959        if !self.contains(id) {
960            panic!("Invalid NodeId");
961        }
962        if matches!(new_node, Node::Root { .. }) {
963            panic!("Cannot replace node with root variant");
964        }
965
966        Self::assert_unique_direct_children(&new_node);
967
968        let old_edges = self.edges(id);
969        let old_children: HashSet<NodeId> = old_edges.iter().map(|(child, _)| *child).collect();
970        let new_edges = Self::node_edges(&new_node);
971        let was_detached = self.detached_roots.contains(&id);
972
973        for &(child, slot) in &new_edges {
974            self.assert_slot_shape(slot, child);
975            if old_children.contains(&child) {
976                let link = self
977                    .parent(child)
978                    .unwrap_or_else(|| panic!("Existing child is missing parent link"));
979                assert_eq!(
980                    link.parent, id,
981                    "Can only reuse direct children of the replaced node"
982                );
983            } else {
984                self.assert_child_is_detached_root(child);
985                self.assert_no_cycle(child, id);
986            }
987        }
988
989        // Children that disappear from the new node are preserved as detached
990        // roots instead of being silently orphaned or recursively deleted.
991        for &(child, _) in &old_edges {
992            if !new_edges.iter().any(|(new_child, _)| *new_child == child) {
993                self.release_child_as_detached_root(child);
994            }
995        }
996
997        for &(child, _) in &new_edges {
998            if !old_children.contains(&child) {
999                self.detached_roots.remove(&child);
1000            }
1001        }
1002
1003        let old_node = std::mem::replace(self.node_mut(id), new_node);
1004        for (child, slot) in new_edges {
1005            self.parent.insert(child, ParentLink { parent: id, slot });
1006        }
1007
1008        if was_detached {
1009            self.detached_roots.insert(id);
1010        } else {
1011            self.detached_roots.remove(&id);
1012        }
1013
1014        old_node
1015    }
1016
1017    /// Appends cloned math content into `out`, flattening implicit math groups.
1018    ///
1019    /// Parser-created content arguments often wrap multiple items in an
1020    /// implicit math group. Flattening that wrapper lets transforms compose
1021    /// output such as `\partial f` without introducing extra braces around `f`.
1022    pub fn append_cloned_math_content(&mut self, out: &mut Vec<NodeId>, source: NodeId) {
1023        match self.node(source) {
1024            Node::Group {
1025                children,
1026                kind: GroupKind::Implicit,
1027                mode: ContentMode::Math,
1028            } => {
1029                let children = children.clone();
1030                out.extend(children.into_iter().map(|child| self.clone_subtree(child)));
1031            }
1032            _ => out.push(self.clone_subtree(source)),
1033        }
1034    }
1035
1036    /// Creates an implicit math group containing `children`.
1037    pub fn implicit_math_group(&mut self, children: Vec<NodeId>) -> NodeId {
1038        self.new_node(Node::Group {
1039            children,
1040            kind: GroupKind::Implicit,
1041            mode: ContentMode::Math,
1042        })
1043    }
1044
1045    /// Creates a scripted node with only a superscript.
1046    pub fn superscript(&mut self, base: NodeId, superscript: NodeId) -> NodeId {
1047        self.new_node(Node::Scripted {
1048            base,
1049            subscript: None,
1050            superscript: Some(superscript),
1051        })
1052    }
1053
1054    /// Replaces `id` and removes any old child subtrees detached by the replacement.
1055    pub fn replace_node_drop_detached_children(&mut self, id: NodeId, replacement: Node) {
1056        let old_children: Vec<NodeId> =
1057            self.edges(id).into_iter().map(|(child, _)| child).collect();
1058        self.replace_node(id, replacement);
1059        for child in old_children {
1060            if self.parent(child).is_none() {
1061                self.remove_detached(child);
1062            }
1063        }
1064    }
1065
1066    /// Replaces a node with a math-mode sequence.
1067    ///
1068    /// Every node in `before`, `replacement`, and `after` must be a unique
1069    /// detached root. If `id` is a group child, `before` and `after` are
1070    /// inserted as real siblings around the replacement payload, and
1071    /// `replacement` is consumed into `id`. In single-child slots, the sequence
1072    /// is wrapped in an implicit math group because those slots cannot hold
1073    /// siblings.
1074    pub fn replace_with_math_sequence(
1075        &mut self,
1076        id: NodeId,
1077        before: Vec<NodeId>,
1078        replacement: NodeId,
1079        after: Vec<NodeId>,
1080    ) {
1081        self.assert_detached_root_sequence(
1082            before
1083                .iter()
1084                .copied()
1085                .chain(std::iter::once(replacement))
1086                .chain(after.iter().copied()),
1087        );
1088
1089        match self.parent(id).map(|link| link.slot) {
1090            Some(Slot::GroupChild(index)) => {
1091                let parent = self
1092                    .parent_id(id)
1093                    .expect("group child should have a parent");
1094                let before_len = before.len();
1095                let replacement_node = self.take_detached_root_node(replacement);
1096
1097                self.replace_node_drop_detached_children(id, replacement_node);
1098                for (offset, child) in before.into_iter().enumerate() {
1099                    self.insert_child(parent, index + offset, child);
1100                }
1101                for (offset, child) in after.into_iter().enumerate() {
1102                    self.insert_child(parent, index + before_len + 1 + offset, child);
1103                }
1104            }
1105            _ => {
1106                let mut children = before;
1107                children.push(replacement);
1108                children.extend(after);
1109                self.replace_node_drop_detached_children(
1110                    id,
1111                    Node::Group {
1112                        children,
1113                        kind: GroupKind::Implicit,
1114                        mode: ContentMode::Math,
1115                    },
1116                );
1117            }
1118        }
1119    }
1120
1121    /// Replaces a node with a math-mode sequence, moving scripts from a
1122    /// parent [`Node::Scripted`] onto the final emitted node when `id` is the
1123    /// scripted base.
1124    ///
1125    /// Every node in `before`, `first`, and `after` must be a unique detached
1126    /// root.
1127    pub fn replace_with_math_sequence_preserving_scripts(
1128        &mut self,
1129        id: NodeId,
1130        before: Vec<NodeId>,
1131        first: NodeId,
1132        mut after: Vec<NodeId>,
1133    ) {
1134        self.assert_detached_root_sequence(
1135            before
1136                .iter()
1137                .copied()
1138                .chain(std::iter::once(first))
1139                .chain(after.iter().copied()),
1140        );
1141
1142        if self.slot(id) != Some(Slot::ScriptBase) {
1143            self.replace_with_math_sequence(id, before, first, after);
1144            return;
1145        }
1146
1147        let Some(parent) = self.parent_id(id) else {
1148            self.replace_with_math_sequence(id, before, first, after);
1149            return;
1150        };
1151        let Node::Scripted {
1152            subscript,
1153            superscript,
1154            ..
1155        } = self.node(parent)
1156        else {
1157            self.replace_with_math_sequence(id, before, first, after);
1158            return;
1159        };
1160        let subscript = *subscript;
1161        let superscript = *superscript;
1162        let subscript = subscript.map(|node_id| self.clone_subtree(node_id));
1163        let superscript = superscript.map(|node_id| self.clone_subtree(node_id));
1164
1165        // Fixed fences carry scripts on the closing token, which is the final
1166        // node of these replacement sequences.
1167        let last = after.pop().unwrap_or(first);
1168        let scripted_last = self.new_node(Node::Scripted {
1169            base: last,
1170            subscript,
1171            superscript,
1172        });
1173        if after.is_empty() && last == first {
1174            self.replace_with_math_sequence(parent, before, scripted_last, after);
1175        } else {
1176            after.push(scripted_last);
1177            self.replace_with_math_sequence(parent, before, first, after);
1178        }
1179    }
1180
1181    /// Destroy a detached subtree and return the removed root node value.
1182    ///
1183    /// The subtree must already be detached from the main tree. All descendants
1184    /// are removed from the arena as well.
1185    ///
1186    /// # Panics
1187    ///
1188    /// Panics if:
1189    ///
1190    /// - `id` is the main root or invalid
1191    /// - `id` is still attached
1192    /// - `id` is not tracked as a detached root
1193    pub fn remove_detached(&mut self, id: NodeId) -> Node {
1194        if id == self.root {
1195            panic!("Cannot remove root node");
1196        }
1197        if !self.contains(id) {
1198            panic!("Invalid NodeId");
1199        }
1200        if self.parent(id).is_some() {
1201            panic!("Can only remove detached roots");
1202        }
1203        if !self.detached_roots.remove(&id) {
1204            panic!("Node is not a detached root");
1205        }
1206
1207        let mut postorder = Vec::new();
1208        self.collect_postorder(id, &mut postorder);
1209
1210        for node_id in postorder
1211            .iter()
1212            .copied()
1213            .take(postorder.len().saturating_sub(1))
1214        {
1215            self.parent.remove(node_id);
1216            self.detached_roots.remove(&node_id);
1217            self.nodes.remove(node_id);
1218        }
1219
1220        self.parent.remove(id);
1221        self.nodes.remove(id).expect("Detached root must exist")
1222    }
1223
1224    /// Assert all structural invariants of the AST.
1225    ///
1226    /// Checked conditions include:
1227    ///
1228    /// - the main root exists and has no parent
1229    /// - every parent link corresponds to a real direct edge
1230    /// - every node is reachable from either the root or a detached root
1231    /// - every parentless non-root node is tracked in detached roots
1232    /// - environment bodies are always groups
1233    ///
1234    /// This method is intended for tests and debug-time validation.
1235    ///
1236    /// # Panics
1237    ///
1238    /// Panics with a descriptive message when any invariant is violated.
1239    pub fn assert_invariants(&self) {
1240        assert!(self.contains(self.root), "Root node must exist");
1241        assert!(
1242            matches!(self.node(self.root), Node::Root { .. }),
1243            "ast.root() must be Node::Root"
1244        );
1245        assert!(
1246            self.parent(self.root).is_none(),
1247            "Root node must not have a parent"
1248        );
1249        assert!(
1250            !self.detached_roots.contains(&self.root),
1251            "Root node cannot be a detached root"
1252        );
1253
1254        for (id, node) in self.nodes.iter() {
1255            if matches!(node, Node::Root { .. }) {
1256                assert_eq!(id, self.root, "Only ast.root() may be Node::Root");
1257            }
1258        }
1259
1260        for (id, link) in self.parent.iter() {
1261            assert!(self.contains(id), "Parent map contains non-existent child");
1262            assert!(
1263                self.contains(link.parent),
1264                "Parent map points to non-existent parent"
1265            );
1266            let has_edge = self
1267                .edges(link.parent)
1268                .into_iter()
1269                .any(|(child, slot)| child == id && slot == link.slot);
1270            assert!(has_edge, "Parent link must match a direct edge");
1271        }
1272
1273        let mut visited = HashSet::new();
1274        self.assert_subtree(self.root, None, &mut visited);
1275
1276        for detached_root in &self.detached_roots {
1277            assert!(*detached_root != self.root, "Root cannot be detached");
1278            assert!(self.contains(*detached_root), "Detached root must exist");
1279            assert!(
1280                self.parent(*detached_root).is_none(),
1281                "Detached root must not have a parent"
1282            );
1283            self.assert_subtree(*detached_root, None, &mut visited);
1284        }
1285
1286        for (id, _) in self.nodes.iter() {
1287            assert!(visited.contains(&id), "Node is orphaned or unreachable");
1288            if id != self.root && self.parent(id).is_none() {
1289                assert!(
1290                    self.detached_roots.contains(&id),
1291                    "Rootless nodes must be tracked as detached roots"
1292                );
1293            }
1294        }
1295    }
1296
1297    fn node_mut(&mut self, id: NodeId) -> &mut Node {
1298        self.nodes.get_mut(id).expect("Invalid NodeId")
1299    }
1300
1301    fn assert_child_is_detached_root(&self, child: NodeId) {
1302        assert!(self.contains(child), "Invalid child NodeId");
1303        assert!(child != self.root, "Cannot attach the AST root as a child");
1304        assert!(
1305            self.parent(child).is_none(),
1306            "Cannot attach child that already has a parent"
1307        );
1308        assert!(
1309            self.detached_roots.contains(&child),
1310            "Can only attach detached roots"
1311        );
1312    }
1313
1314    fn assert_no_cycle(&self, child: NodeId, new_parent: NodeId) {
1315        assert!(
1316            !self.subtree_contains(child, new_parent),
1317            "Cannot create an ancestor cycle"
1318        );
1319    }
1320
1321    fn assert_slot_shape(&self, slot: Slot, child: NodeId) {
1322        if matches!(slot, Slot::EnvBody) {
1323            assert!(
1324                matches!(self.node(child), Node::Group { .. }),
1325                "Environment body must be a Group node"
1326            );
1327        }
1328    }
1329
1330    fn assert_detached_root_sequence(&self, nodes: impl IntoIterator<Item = NodeId>) {
1331        let mut seen = HashSet::new();
1332        for node in nodes {
1333            assert!(
1334                seen.insert(node),
1335                "Node cannot appear in a replacement sequence twice"
1336            );
1337            self.assert_child_is_detached_root(node);
1338        }
1339    }
1340
1341    fn adopt_child(&mut self, parent: NodeId, child: NodeId, slot: Slot) {
1342        self.detached_roots.remove(&child);
1343        self.parent.insert(child, ParentLink { parent, slot });
1344    }
1345
1346    fn release_child_as_detached_root(&mut self, child: NodeId) {
1347        self.parent.remove(child);
1348        self.detached_roots.insert(child);
1349    }
1350
1351    fn take_detached_root_node(&mut self, id: NodeId) -> Node {
1352        if id == self.root {
1353            panic!("Cannot consume root node as detached replacement");
1354        }
1355        if !self.contains(id) {
1356            panic!("Invalid NodeId");
1357        }
1358        if self.parent(id).is_some() {
1359            panic!("Can only consume detached roots");
1360        }
1361        if !self.detached_roots.remove(&id) {
1362            panic!("Node is not a detached root");
1363        }
1364
1365        let node = self.nodes.remove(id).expect("Detached root must exist");
1366        for (child, _) in Self::node_edges(&node) {
1367            let link = self
1368                .parent
1369                .remove(child)
1370                .expect("Detached root child should have a parent link");
1371            assert_eq!(
1372                link.parent, id,
1373                "Detached replacement child must point at the consumed root"
1374            );
1375            self.release_child_as_detached_root(child);
1376        }
1377        node
1378    }
1379
1380    fn reindex_group_children(&mut self, parent: NodeId, start: usize) {
1381        let children = match self.node(parent) {
1382            Node::Root { children, .. } | Node::Group { children, .. } => children.clone(),
1383            _ => panic!("Parent is not a root/group node"),
1384        };
1385
1386        // Group child slots store indices, so every insertion / removal requires
1387        // rewriting the suffix of the parent link table.
1388        for (index, child) in children.into_iter().enumerate().skip(start) {
1389            self.parent.insert(
1390                child,
1391                ParentLink {
1392                    parent,
1393                    slot: Slot::GroupChild(index),
1394                },
1395            );
1396        }
1397    }
1398
1399    fn collect_postorder(&self, id: NodeId, out: &mut Vec<NodeId>) {
1400        for (child, _) in self.edges(id) {
1401            self.collect_postorder(child, out);
1402        }
1403        out.push(id);
1404    }
1405
1406    fn assert_subtree(
1407        &self,
1408        id: NodeId,
1409        expected_parent: Option<ParentLink>,
1410        visited: &mut HashSet<NodeId>,
1411    ) {
1412        if let Node::Prime { count } = self.node(id) {
1413            assert!(*count > 0, "Prime count must be greater than zero");
1414        }
1415
1416        assert!(visited.insert(id), "Node is reachable from multiple roots");
1417        assert_eq!(
1418            self.parent(id),
1419            expected_parent,
1420            "Parent link must match traversal path"
1421        );
1422
1423        for (child, slot) in self.edges(id) {
1424            assert!(self.contains(child), "Direct edge points to invalid child");
1425            self.assert_slot_shape(slot, child);
1426            self.assert_subtree(child, Some(ParentLink { parent: id, slot }), visited);
1427        }
1428    }
1429
1430    fn subtree_contains(&self, root: NodeId, target: NodeId) -> bool {
1431        if root == target {
1432            return true;
1433        }
1434
1435        self.edges(root)
1436            .into_iter()
1437            .any(|(child, _)| self.subtree_contains(child, target))
1438    }
1439
1440    fn find_impl(&self, start: NodeId, predicate: &dyn Fn(&Node) -> bool) -> Option<NodeId> {
1441        if predicate(self.node(start)) {
1442            return Some(start);
1443        }
1444
1445        for (child, _) in self.edges(start) {
1446            if let Some(found) = self.find_impl(child, predicate) {
1447                return Some(found);
1448            }
1449        }
1450
1451        None
1452    }
1453
1454    fn find_all_impl(
1455        &self,
1456        start: NodeId,
1457        predicate: &dyn Fn(&Node) -> bool,
1458        out: &mut Vec<NodeId>,
1459    ) {
1460        if predicate(self.node(start)) {
1461            out.push(start);
1462        }
1463
1464        for (child, _) in self.edges(start) {
1465            self.find_all_impl(child, predicate, out);
1466        }
1467    }
1468
1469    fn to_syntax_node(&self, id: NodeId) -> SyntaxNode {
1470        match self.node(id) {
1471            Node::Root { .. } => unreachable!("nested Root is impossible"),
1472            Node::Group {
1473                children,
1474                kind,
1475                mode,
1476            } => SyntaxNode::Group {
1477                mode: *mode,
1478                kind: self.to_syntax_group_kind(kind),
1479                children: children
1480                    .iter()
1481                    .map(|child| self.to_syntax_node(*child))
1482                    .collect(),
1483            },
1484            Node::Command { name, args, known } => SyntaxNode::Command {
1485                name: name.clone(),
1486                args: args
1487                    .iter()
1488                    .map(|slot| self.to_syntax_arg_slot(slot))
1489                    .collect(),
1490                known: *known,
1491            },
1492            Node::Infix {
1493                name,
1494                args,
1495                left,
1496                right,
1497            } => SyntaxNode::Infix {
1498                name: name.clone(),
1499                args: args
1500                    .iter()
1501                    .map(|slot| self.to_syntax_arg_slot(slot))
1502                    .collect(),
1503                left: Box::new(self.to_syntax_node(*left)),
1504                right: Box::new(self.to_syntax_node(*right)),
1505            },
1506            Node::Declarative { name, args } => SyntaxNode::Declarative {
1507                name: name.clone(),
1508                args: args
1509                    .iter()
1510                    .map(|slot| self.to_syntax_arg_slot(slot))
1511                    .collect(),
1512            },
1513            Node::Environment {
1514                name,
1515                args,
1516                known,
1517                body,
1518            } => SyntaxNode::Environment {
1519                name: name.clone(),
1520                args: args
1521                    .iter()
1522                    .map(|slot| self.to_syntax_arg_slot(slot))
1523                    .collect(),
1524                known: *known,
1525                body: Box::new(self.to_syntax_node(*body)),
1526            },
1527            Node::Scripted {
1528                base,
1529                subscript,
1530                superscript,
1531            } => SyntaxNode::Scripted {
1532                base: Box::new(self.to_syntax_node(*base)),
1533                subscript: subscript.map(|node| Box::new(self.to_syntax_node(node))),
1534                superscript: superscript.map(|node| Box::new(self.to_syntax_node(node))),
1535            },
1536            Node::Prime { count } => SyntaxNode::Prime { count: *count },
1537            Node::Text(text) => SyntaxNode::Text(text.clone()),
1538            Node::Char(ch) => SyntaxNode::Char(*ch),
1539            Node::ActiveSpace => SyntaxNode::ActiveSpace,
1540            Node::Error { message, snippet } => SyntaxNode::Error {
1541                message: message.clone(),
1542                snippet: snippet.clone(),
1543            },
1544        }
1545    }
1546
1547    fn to_syntax_arg_slot(&self, slot: &ArgumentSlot) -> syntax_node::ArgumentSlot {
1548        slot.as_ref().map(|argument| syntax_node::Argument {
1549            kind: self.to_syntax_arg_kind(&argument.kind),
1550            value: self.to_syntax_arg_value(&argument.value),
1551        })
1552    }
1553
1554    fn to_syntax_arg_kind(&self, kind: &ArgumentKind) -> syntax_node::ArgumentKind {
1555        match kind {
1556            ArgumentKind::Mandatory => syntax_node::ArgumentKind::Mandatory,
1557            ArgumentKind::Optional => syntax_node::ArgumentKind::Optional,
1558            ArgumentKind::Star => syntax_node::ArgumentKind::Star,
1559            ArgumentKind::Group => syntax_node::ArgumentKind::Group,
1560            ArgumentKind::Delimited { open, close } => syntax_node::ArgumentKind::Delimited {
1561                open: self.to_syntax_delimiter(open),
1562                close: self.to_syntax_delimiter(close),
1563            },
1564            ArgumentKind::Paired { open, close } => syntax_node::ArgumentKind::Paired {
1565                open: self.to_syntax_delimiter(open),
1566                close: self.to_syntax_delimiter(close),
1567            },
1568        }
1569    }
1570
1571    fn to_syntax_arg_value(&self, value: &ArgumentValue) -> syntax_node::ArgumentValue {
1572        match value {
1573            ArgumentValue::MathContent(id) => {
1574                syntax_node::ArgumentValue::MathContent(self.to_syntax_node(*id))
1575            }
1576            ArgumentValue::TextContent(id) => {
1577                syntax_node::ArgumentValue::TextContent(self.to_syntax_node(*id))
1578            }
1579            ArgumentValue::Delimiter(delimiter) => {
1580                syntax_node::ArgumentValue::Delimiter(self.to_syntax_delimiter(delimiter))
1581            }
1582            ArgumentValue::CSName(value) => syntax_node::ArgumentValue::CSName(value.clone()),
1583            ArgumentValue::Dimension(value) => syntax_node::ArgumentValue::Dimension(value.clone()),
1584            ArgumentValue::Integer(value) => syntax_node::ArgumentValue::Integer(value.clone()),
1585            ArgumentValue::KeyVal(value) => syntax_node::ArgumentValue::KeyVal(value.clone()),
1586            ArgumentValue::Column(value) => syntax_node::ArgumentValue::Column(value.clone()),
1587            ArgumentValue::Boolean(value) => syntax_node::ArgumentValue::Boolean(*value),
1588        }
1589    }
1590
1591    fn to_syntax_group_kind(&self, kind: &GroupKind) -> syntax_node::GroupKind {
1592        match kind {
1593            GroupKind::Explicit => syntax_node::GroupKind::Explicit,
1594            GroupKind::Implicit => syntax_node::GroupKind::Implicit,
1595            GroupKind::Delimited { left, right } => syntax_node::GroupKind::Delimited {
1596                left: self.to_syntax_delimiter(left),
1597                right: self.to_syntax_delimiter(right),
1598            },
1599            GroupKind::InlineMath => syntax_node::GroupKind::InlineMath,
1600        }
1601    }
1602
1603    fn to_syntax_delimiter(&self, delimiter: &Delimiter) -> syntax_node::Delimiter {
1604        match delimiter {
1605            Delimiter::None => syntax_node::Delimiter::None,
1606            Delimiter::Char(ch) => syntax_node::Delimiter::Char(*ch),
1607            Delimiter::Control(name) => {
1608                syntax_node::Delimiter::Control(Box::leak(name.clone().into_boxed_str()))
1609            }
1610        }
1611    }
1612
1613    fn clone_subtree_impl(&mut self, id: NodeId) -> NodeId {
1614        let cloned = match self.node(id).clone() {
1615            Node::Root { .. } => panic!("Cannot clone root node as a detached subtree"),
1616            Node::Group {
1617                children,
1618                kind,
1619                mode,
1620            } => Node::Group {
1621                children: children
1622                    .into_iter()
1623                    .map(|child| self.clone_subtree_impl(child))
1624                    .collect(),
1625                kind,
1626                mode,
1627            },
1628            Node::Command { name, args, known } => Node::Command {
1629                name,
1630                args: self.clone_argument_slots(args),
1631                known,
1632            },
1633            Node::Infix {
1634                name,
1635                args,
1636                left,
1637                right,
1638            } => Node::Infix {
1639                name,
1640                args: self.clone_argument_slots(args),
1641                left: self.clone_subtree_impl(left),
1642                right: self.clone_subtree_impl(right),
1643            },
1644            Node::Declarative { name, args } => Node::Declarative {
1645                name,
1646                args: self.clone_argument_slots(args),
1647            },
1648            Node::Environment {
1649                name,
1650                args,
1651                known,
1652                body,
1653            } => Node::Environment {
1654                name,
1655                args: self.clone_argument_slots(args),
1656                known,
1657                body: self.clone_subtree_impl(body),
1658            },
1659            Node::Scripted {
1660                base,
1661                subscript,
1662                superscript,
1663            } => Node::Scripted {
1664                base: self.clone_subtree_impl(base),
1665                subscript: subscript.map(|child| self.clone_subtree_impl(child)),
1666                superscript: superscript.map(|child| self.clone_subtree_impl(child)),
1667            },
1668            Node::Prime { count } => Node::Prime { count },
1669            Node::Text(text) => Node::Text(text),
1670            Node::Char(ch) => Node::Char(ch),
1671            Node::ActiveSpace => Node::ActiveSpace,
1672            Node::Error { message, snippet } => Node::Error { message, snippet },
1673        };
1674
1675        self.new_node(cloned)
1676    }
1677
1678    fn clone_argument_slots(&mut self, args: Vec<ArgumentSlot>) -> Vec<ArgumentSlot> {
1679        args.into_iter()
1680            .map(|slot| {
1681                slot.map(|arg| Argument {
1682                    kind: arg.kind,
1683                    value: self.clone_argument_value(arg.value),
1684                })
1685            })
1686            .collect()
1687    }
1688
1689    fn clone_argument_value(&mut self, value: ArgumentValue) -> ArgumentValue {
1690        match value {
1691            ArgumentValue::MathContent(child) => {
1692                ArgumentValue::MathContent(self.clone_subtree_impl(child))
1693            }
1694            ArgumentValue::TextContent(child) => {
1695                ArgumentValue::TextContent(self.clone_subtree_impl(child))
1696            }
1697            ArgumentValue::Delimiter(delimiter) => ArgumentValue::Delimiter(delimiter),
1698            ArgumentValue::CSName(name) => ArgumentValue::CSName(name),
1699            ArgumentValue::Dimension(value) => ArgumentValue::Dimension(value),
1700            ArgumentValue::Integer(value) => ArgumentValue::Integer(value),
1701            ArgumentValue::KeyVal(value) => ArgumentValue::KeyVal(value),
1702            ArgumentValue::Column(value) => ArgumentValue::Column(value),
1703            ArgumentValue::Boolean(value) => ArgumentValue::Boolean(value),
1704        }
1705    }
1706
1707    fn assert_unique_direct_children(node: &Node) {
1708        let mut seen = HashSet::new();
1709        for (child, _) in Self::node_edges(node) {
1710            assert!(
1711                seen.insert(child),
1712                "Node cannot reference the same child twice"
1713            );
1714        }
1715    }
1716
1717    fn node_edges(node: &Node) -> Vec<(NodeId, Slot)> {
1718        let mut edges = Vec::new();
1719
1720        // This function is the single source of truth for direct traversal
1721        // order. Public `edges()` and internal mutation helpers both rely on it
1722        // so read and write paths stay aligned.
1723        match node {
1724            Node::Root { children, .. } | Node::Group { children, .. } => {
1725                for (index, child) in children.iter().copied().enumerate() {
1726                    edges.push((child, Slot::GroupChild(index)));
1727                }
1728            }
1729            Node::Command { args, .. } => {
1730                Self::push_argument_edges(args, &mut edges);
1731            }
1732            Node::Infix {
1733                args, left, right, ..
1734            } => {
1735                edges.push((*left, Slot::InfixLeft));
1736                Self::push_argument_edges(args, &mut edges);
1737                edges.push((*right, Slot::InfixRight));
1738            }
1739            Node::Declarative { args, .. } => {
1740                Self::push_argument_edges(args, &mut edges);
1741            }
1742            Node::Environment { args, body, .. } => {
1743                Self::push_argument_edges(args, &mut edges);
1744                edges.push((*body, Slot::EnvBody));
1745            }
1746            Node::Scripted {
1747                base,
1748                subscript,
1749                superscript,
1750            } => {
1751                edges.push((*base, Slot::ScriptBase));
1752                if let Some(subscript) = subscript {
1753                    edges.push((*subscript, Slot::ScriptSub));
1754                }
1755                if let Some(superscript) = superscript {
1756                    edges.push((*superscript, Slot::ScriptSup));
1757                }
1758            }
1759            Node::Prime { .. }
1760            | Node::Text(_)
1761            | Node::Char(_)
1762            | Node::ActiveSpace
1763            | Node::Error { .. } => {}
1764        }
1765
1766        edges
1767    }
1768
1769    fn push_argument_edges(args: &[ArgumentSlot], edges: &mut Vec<(NodeId, Slot)>) {
1770        for (index, slot) in args.iter().enumerate() {
1771            let Some(argument) = slot else {
1772                continue;
1773            };
1774            // Only content arguments participate in tree traversal. Scalar
1775            // values still live on the node, but they do not become AST edges.
1776            match &argument.value {
1777                ArgumentValue::MathContent(child) | ArgumentValue::TextContent(child) => {
1778                    edges.push((*child, Slot::Argument(index)));
1779                }
1780                _ => {}
1781            }
1782        }
1783    }
1784
1785    fn convert_syntax_node(
1786        node: &SyntaxNode,
1787        nodes: &mut HopSlotMap<NodeId, Node>,
1788        parent: &mut SecondaryMap<NodeId, ParentLink>,
1789    ) -> NodeId {
1790        // Conversion constructs child nodes first, then inserts the current
1791        // node, then wires direct parent links immediately. That keeps the
1792        // transformation local and avoids a second global rebuild pass over the
1793        // finished tree.
1794        let converted_node = match node {
1795            SyntaxNode::Root { .. } => {
1796                panic!("Ast::from_syntax_root does not accept nested SyntaxNode::Root")
1797            }
1798            SyntaxNode::Group {
1799                mode,
1800                kind,
1801                children,
1802            } => {
1803                let converted_children: Vec<NodeId> = children
1804                    .iter()
1805                    .map(|child| Self::convert_syntax_node(child, nodes, parent))
1806                    .collect();
1807                Node::Group {
1808                    children: converted_children,
1809                    kind: Self::convert_group_kind(kind),
1810                    mode: *mode,
1811                }
1812            }
1813            SyntaxNode::Command { name, args, known } => Node::Command {
1814                name: name.clone(),
1815                args: args
1816                    .iter()
1817                    .map(|slot| Self::convert_argument_slot(slot, nodes, parent))
1818                    .collect(),
1819                known: *known,
1820            },
1821            SyntaxNode::Infix {
1822                name,
1823                args,
1824                left,
1825                right,
1826            } => Node::Infix {
1827                name: name.clone(),
1828                args: args
1829                    .iter()
1830                    .map(|slot| Self::convert_argument_slot(slot, nodes, parent))
1831                    .collect(),
1832                left: Self::convert_syntax_node(left, nodes, parent),
1833                right: Self::convert_syntax_node(right, nodes, parent),
1834            },
1835            SyntaxNode::Declarative { name, args } => Node::Declarative {
1836                name: name.clone(),
1837                args: args
1838                    .iter()
1839                    .map(|slot| Self::convert_argument_slot(slot, nodes, parent))
1840                    .collect(),
1841            },
1842            SyntaxNode::Environment {
1843                name,
1844                args,
1845                known,
1846                body,
1847            } => Node::Environment {
1848                name: name.clone(),
1849                args: args
1850                    .iter()
1851                    .map(|slot| Self::convert_argument_slot(slot, nodes, parent))
1852                    .collect(),
1853                known: *known,
1854                body: Self::convert_syntax_node(body, nodes, parent),
1855            },
1856            SyntaxNode::Scripted {
1857                base,
1858                subscript,
1859                superscript,
1860            } => Node::Scripted {
1861                base: Self::convert_syntax_node(base, nodes, parent),
1862                subscript: subscript
1863                    .as_ref()
1864                    .map(|node| Self::convert_syntax_node(node, nodes, parent)),
1865                superscript: superscript
1866                    .as_ref()
1867                    .map(|node| Self::convert_syntax_node(node, nodes, parent)),
1868            },
1869            SyntaxNode::Error { message, snippet } => Node::Error {
1870                message: message.clone(),
1871                snippet: snippet.clone(),
1872            },
1873            SyntaxNode::Prime { count } => Node::Prime { count: *count },
1874            SyntaxNode::Text(text) => Node::Text(text.clone()),
1875            SyntaxNode::Char(ch) => Node::Char(*ch),
1876            SyntaxNode::ActiveSpace => Node::ActiveSpace,
1877        };
1878
1879        let id = nodes.insert(converted_node);
1880        let edges = Self::node_edges(nodes.get(id).expect("Converted node must exist"));
1881        for (child, slot) in edges {
1882            if matches!(slot, Slot::EnvBody) {
1883                assert!(
1884                    matches!(nodes.get(child), Some(Node::Group { .. })),
1885                    "Environment body must convert to a Group node"
1886                );
1887            }
1888            parent.insert(child, ParentLink { parent: id, slot });
1889        }
1890
1891        id
1892    }
1893
1894    fn convert_argument_slot(
1895        slot: &syntax_node::ArgumentSlot,
1896        nodes: &mut HopSlotMap<NodeId, Node>,
1897        parent: &mut SecondaryMap<NodeId, ParentLink>,
1898    ) -> ArgumentSlot {
1899        slot.as_ref().map(|arg| Argument {
1900            kind: Self::convert_argument_kind(arg.kind),
1901            value: Self::convert_argument_value(&arg.value, nodes, parent),
1902        })
1903    }
1904
1905    fn convert_argument_kind(kind: syntax_node::ArgumentKind) -> ArgumentKind {
1906        match kind {
1907            syntax_node::ArgumentKind::Mandatory => ArgumentKind::Mandatory,
1908            syntax_node::ArgumentKind::Optional => ArgumentKind::Optional,
1909            syntax_node::ArgumentKind::Star => ArgumentKind::Star,
1910            syntax_node::ArgumentKind::Group => ArgumentKind::Group,
1911            syntax_node::ArgumentKind::Delimited { open, close } => ArgumentKind::Delimited {
1912                open: Self::convert_delimiter(open),
1913                close: Self::convert_delimiter(close),
1914            },
1915            syntax_node::ArgumentKind::Paired { open, close } => ArgumentKind::Paired {
1916                open: Self::convert_delimiter(open),
1917                close: Self::convert_delimiter(close),
1918            },
1919        }
1920    }
1921
1922    fn convert_argument_value(
1923        value: &syntax_node::ArgumentValue,
1924        nodes: &mut HopSlotMap<NodeId, Node>,
1925        parent: &mut SecondaryMap<NodeId, ParentLink>,
1926    ) -> ArgumentValue {
1927        match value {
1928            // Conversion keeps the original shape instead of wrapping single
1929            // content nodes in an implicit group.
1930            syntax_node::ArgumentValue::MathContent(node) => {
1931                ArgumentValue::MathContent(Self::convert_syntax_node(node, nodes, parent))
1932            }
1933            syntax_node::ArgumentValue::TextContent(node) => {
1934                ArgumentValue::TextContent(Self::convert_syntax_node(node, nodes, parent))
1935            }
1936            syntax_node::ArgumentValue::Delimiter(delimiter) => {
1937                ArgumentValue::Delimiter(Self::convert_delimiter(*delimiter))
1938            }
1939            syntax_node::ArgumentValue::CSName(name) => ArgumentValue::CSName(name.clone()),
1940            syntax_node::ArgumentValue::Dimension(value) => ArgumentValue::Dimension(value.clone()),
1941            syntax_node::ArgumentValue::Integer(value) => ArgumentValue::Integer(value.clone()),
1942            syntax_node::ArgumentValue::KeyVal(value) => ArgumentValue::KeyVal(value.clone()),
1943            syntax_node::ArgumentValue::Column(value) => ArgumentValue::Column(value.clone()),
1944            syntax_node::ArgumentValue::Boolean(value) => ArgumentValue::Boolean(*value),
1945        }
1946    }
1947
1948    fn convert_group_kind(kind: &syntax_node::GroupKind) -> GroupKind {
1949        match kind {
1950            syntax_node::GroupKind::Explicit => GroupKind::Explicit,
1951            syntax_node::GroupKind::Implicit => GroupKind::Implicit,
1952            syntax_node::GroupKind::Delimited { left, right } => GroupKind::Delimited {
1953                left: Self::convert_delimiter(*left),
1954                right: Self::convert_delimiter(*right),
1955            },
1956            syntax_node::GroupKind::InlineMath => GroupKind::InlineMath,
1957        }
1958    }
1959
1960    fn convert_delimiter(delimiter: syntax_node::Delimiter) -> Delimiter {
1961        match delimiter {
1962            syntax_node::Delimiter::None => Delimiter::None,
1963            syntax_node::Delimiter::Char(ch) => Delimiter::Char(ch),
1964            syntax_node::Delimiter::Control(name) => Delimiter::Control(name.to_string()),
1965        }
1966    }
1967}
1968
1969fn replace_argument_child(args: &mut [ArgumentSlot], slot: Slot, old: NodeId, replacement: NodeId) {
1970    let Slot::Argument(index) = slot else {
1971        panic!("Expected argument slot");
1972    };
1973    let argument = args
1974        .get_mut(index)
1975        .and_then(Option::as_mut)
1976        .unwrap_or_else(|| panic!("Argument slot is missing"));
1977    match &mut argument.value {
1978        ArgumentValue::MathContent(child) | ArgumentValue::TextContent(child) => {
1979            assert_eq!(*child, old, "Argument child must match old node");
1980            *child = replacement;
1981        }
1982        _ => panic!("Argument slot is not content"),
1983    }
1984}
1985
1986impl Default for Ast {
1987    fn default() -> Self {
1988        Self::new()
1989    }
1990}
1991
1992#[cfg(test)]
1993mod tests {
1994    use super::{
1995        Argument, ArgumentKind, ArgumentValue, Ast, ContentMode, GroupKind, Node, NodeKind, Slot,
1996    };
1997
1998    fn panic_message(payload: Box<dyn std::any::Any + Send>) -> String {
1999        match payload.downcast::<String>() {
2000            Ok(message) => *message,
2001            Err(payload) => match payload.downcast::<&'static str>() {
2002                Ok(message) => message.to_string(),
2003                Err(_) => "non-string panic payload".to_string(),
2004            },
2005        }
2006    }
2007
2008    #[test]
2009    fn error_node_reports_error_kind() {
2010        let node = Node::Error {
2011            message: "unexpected token".to_string(),
2012            snippet: r"\frac".to_string(),
2013        };
2014
2015        assert_eq!(node.kind(), NodeKind::Error);
2016    }
2017
2018    #[test]
2019    fn error_node_is_a_leaf() {
2020        let mut ast = Ast::new();
2021        let error = ast.new_node(Node::Error {
2022            message: "bad".to_string(),
2023            snippet: "x".to_string(),
2024        });
2025        ast.append_child(ast.root(), error);
2026
2027        assert!(ast.edges(error).is_empty());
2028        ast.assert_invariants();
2029    }
2030
2031    #[test]
2032    fn clone_subtree_clones_error_node() {
2033        let mut ast = Ast::new();
2034        let error = ast.new_node(Node::Error {
2035            message: "bad".to_string(),
2036            snippet: "x".to_string(),
2037        });
2038
2039        let cloned = ast.clone_subtree(error);
2040
2041        assert_ne!(cloned, error);
2042        assert_eq!(
2043            ast.node(cloned),
2044            &Node::Error {
2045                message: "bad".to_string(),
2046                snippet: "x".to_string(),
2047            }
2048        );
2049        ast.assert_invariants();
2050    }
2051
2052    #[test]
2053    fn from_syntax_root_converts_error_nodes() {
2054        use texform_interface::syntax_node::{ContentMode as SynMode, SyntaxNode};
2055
2056        let syntax = SyntaxNode::Root {
2057            mode: SynMode::Math,
2058            children: vec![
2059                SyntaxNode::Char('a'),
2060                SyntaxNode::Error {
2061                    message: "unexpected".to_string(),
2062                    snippet: r"\bad".to_string(),
2063                },
2064            ],
2065        };
2066
2067        let ast = Ast::from_syntax_root(&syntax);
2068        let children = ast.children(ast.root());
2069
2070        assert_eq!(children.len(), 2);
2071        assert_eq!(ast.node(children[0]), &Node::Char('a'));
2072        assert_eq!(
2073            ast.node(children[1]),
2074            &Node::Error {
2075                message: "unexpected".to_string(),
2076                snippet: r"\bad".to_string(),
2077            }
2078        );
2079        ast.assert_invariants();
2080    }
2081
2082    #[test]
2083    #[should_panic(expected = "Only ast.root() may be Node::Root")]
2084    fn assert_invariants_rejects_additional_root_nodes() {
2085        let mut ast = Ast::new();
2086        let extra_root = ast.nodes.insert(Node::Root {
2087            children: Vec::new(),
2088            mode: ContentMode::Math,
2089        });
2090        ast.detached_roots.insert(extra_root);
2091
2092        ast.assert_invariants();
2093    }
2094
2095    #[test]
2096    fn clone_subtree_creates_detached_copy_with_rewired_children() {
2097        let mut ast = Ast::new();
2098        let numerator_child = ast.new_node(Node::Char('x'));
2099        let numerator = ast.new_node(Node::Group {
2100            children: vec![numerator_child],
2101            kind: GroupKind::Implicit,
2102            mode: ContentMode::Math,
2103        });
2104        let denominator = ast.new_node(Node::Char('t'));
2105        let frac = ast.new_node(Node::Command {
2106            name: "frac".to_string(),
2107            args: vec![
2108                Some(Argument {
2109                    kind: ArgumentKind::Mandatory,
2110                    value: ArgumentValue::MathContent(numerator),
2111                }),
2112                Some(Argument {
2113                    kind: ArgumentKind::Mandatory,
2114                    value: ArgumentValue::MathContent(denominator),
2115                }),
2116            ],
2117            known: true,
2118        });
2119
2120        let cloned = ast.clone_subtree(frac);
2121
2122        assert_ne!(cloned, frac);
2123        assert_eq!(ast.parent(cloned), None);
2124        assert!(ast.detached_roots.contains(&cloned));
2125
2126        let Node::Command { args, .. } = ast.node(cloned) else {
2127            panic!("cloned root should be a command");
2128        };
2129        let ArgumentValue::MathContent(cloned_numerator) =
2130            args[0].as_ref().expect("first argument should exist").value
2131        else {
2132            panic!("first argument should be math content");
2133        };
2134        let ArgumentValue::MathContent(cloned_denominator) = args[1]
2135            .as_ref()
2136            .expect("second argument should exist")
2137            .value
2138        else {
2139            panic!("second argument should be math content");
2140        };
2141
2142        assert_ne!(cloned_numerator, numerator);
2143        assert_ne!(cloned_denominator, denominator);
2144        assert_eq!(
2145            ast.parent(cloned_numerator),
2146            Some(super::ParentLink {
2147                parent: cloned,
2148                slot: Slot::Argument(0),
2149            })
2150        );
2151        assert_eq!(
2152            ast.parent(cloned_denominator),
2153            Some(super::ParentLink {
2154                parent: cloned,
2155                slot: Slot::Argument(1),
2156            })
2157        );
2158
2159        let cloned_numerator_children = ast.children(cloned_numerator);
2160        assert_eq!(cloned_numerator_children.len(), 1);
2161        assert_ne!(cloned_numerator_children[0], numerator_child);
2162        assert_eq!(
2163            ast.parent(cloned_numerator_children[0]),
2164            Some(super::ParentLink {
2165                parent: cloned_numerator,
2166                slot: Slot::GroupChild(0),
2167            })
2168        );
2169
2170        ast.assert_invariants();
2171    }
2172
2173    #[test]
2174    fn append_cloned_math_content_flattens_implicit_groups() {
2175        let mut ast = Ast::new();
2176        let x = ast.new_node(Node::Char('x'));
2177        let y = ast.new_node(Node::Char('y'));
2178        let source = ast.new_node(Node::Group {
2179            children: vec![x, y],
2180            kind: GroupKind::Implicit,
2181            mode: ContentMode::Math,
2182        });
2183        let mut out = Vec::new();
2184
2185        ast.append_cloned_math_content(&mut out, source);
2186
2187        assert_eq!(out.len(), 2);
2188        assert_ne!(out[0], x);
2189        assert_ne!(out[1], y);
2190        assert_eq!(ast.node(out[0]), &Node::Char('x'));
2191        assert_eq!(ast.node(out[1]), &Node::Char('y'));
2192        assert_eq!(ast.parent(out[0]), None);
2193        assert_eq!(ast.parent(out[1]), None);
2194        ast.assert_invariants();
2195    }
2196
2197    #[test]
2198    fn constructs_implicit_math_group() {
2199        let mut ast = Ast::new();
2200        let x = ast.new_node(Node::Char('x'));
2201        let y = ast.new_node(Node::Char('y'));
2202
2203        let group = ast.implicit_math_group(vec![x, y]);
2204
2205        assert_eq!(
2206            ast.node(group),
2207            &Node::Group {
2208                children: vec![x, y],
2209                kind: GroupKind::Implicit,
2210                mode: ContentMode::Math,
2211            }
2212        );
2213        assert_eq!(ast.parent_id(x), Some(group));
2214        assert_eq!(ast.parent_id(y), Some(group));
2215        ast.assert_invariants();
2216    }
2217
2218    #[test]
2219    fn constructs_superscript_node() {
2220        let mut ast = Ast::new();
2221        let base = ast.new_node(Node::Char('a'));
2222        let power = ast.new_node(Node::Char('2'));
2223
2224        let scripted = ast.superscript(base, power);
2225
2226        assert_eq!(
2227            ast.node(scripted),
2228            &Node::Scripted {
2229                base,
2230                subscript: None,
2231                superscript: Some(power),
2232            }
2233        );
2234        assert_eq!(ast.parent_id(base), Some(scripted));
2235        assert_eq!(ast.parent_id(power), Some(scripted));
2236        ast.assert_invariants();
2237    }
2238
2239    #[test]
2240    fn replace_children_detaches_removed_children_and_adopts_new_children() {
2241        let mut ast = Ast::new();
2242        let root = ast.root();
2243        let a = ast.new_node(Node::Char('a'));
2244        let b = ast.new_node(Node::Char('b'));
2245        let c = ast.new_node(Node::Char('c'));
2246        ast.append_child(root, a);
2247        ast.append_child(root, b);
2248
2249        let removed = ast.replace_children(root, vec![b, c]);
2250
2251        assert_eq!(removed, vec![a]);
2252        assert_eq!(ast.children(root), &[b, c]);
2253        assert_eq!(ast.parent(a), None);
2254        assert!(ast.detached_roots.contains(&a));
2255        assert_eq!(
2256            ast.parent(b),
2257            Some(super::ParentLink {
2258                parent: root,
2259                slot: Slot::GroupChild(0),
2260            })
2261        );
2262        assert_eq!(
2263            ast.parent(c),
2264            Some(super::ParentLink {
2265                parent: root,
2266                slot: Slot::GroupChild(1),
2267            })
2268        );
2269        ast.assert_invariants();
2270    }
2271
2272    #[test]
2273    fn detach_children_range_detaches_ordered_segment() {
2274        let mut ast = Ast::new();
2275        let root = ast.root();
2276        let a = ast.new_node(Node::Char('a'));
2277        let b = ast.new_node(Node::Char('b'));
2278        let c = ast.new_node(Node::Char('c'));
2279        let d = ast.new_node(Node::Char('d'));
2280        for child in [a, b, c, d] {
2281            ast.append_child(root, child);
2282        }
2283
2284        let removed = ast.detach_children_range(root, 1..3);
2285
2286        assert_eq!(removed, vec![b, c]);
2287        assert_eq!(ast.children(root), &[a, d]);
2288        assert_eq!(ast.parent(b), None);
2289        assert_eq!(ast.parent(c), None);
2290        assert!(ast.detached_roots.contains(&b));
2291        assert!(ast.detached_roots.contains(&c));
2292        assert_eq!(ast.slot(d), Some(Slot::GroupChild(1)));
2293        ast.assert_invariants();
2294    }
2295
2296    #[test]
2297    fn detach_panics_without_removing_wrong_child_when_parent_link_index_is_stale() {
2298        let mut ast = Ast::new();
2299        let root = ast.root();
2300        let a = ast.new_node(Node::Char('a'));
2301        let b = ast.new_node(Node::Char('b'));
2302        ast.append_child(root, a);
2303        ast.append_child(root, b);
2304        ast.parent.insert(
2305            b,
2306            super::ParentLink {
2307                parent: root,
2308                slot: Slot::GroupChild(0),
2309            },
2310        );
2311
2312        let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| ast.detach(b)));
2313
2314        let message = panic_message(result.expect_err("detach should reject stale child index"));
2315        assert!(
2316            message.contains("Group child index must match detached node"),
2317            "unexpected panic: {message}"
2318        );
2319        assert_eq!(ast.children(root), &[a, b]);
2320    }
2321
2322    #[test]
2323    fn replace_content_child_replaces_script_slot() {
2324        let mut ast = Ast::new();
2325        let base = ast.new_node(Node::Char('x'));
2326        let superscript = ast.new_node(Node::Char('y'));
2327        let scripted = ast.new_node(Node::Scripted {
2328            base,
2329            subscript: None,
2330            superscript: Some(superscript),
2331        });
2332        ast.append_child(ast.root(), scripted);
2333        let replacement = ast.new_node(Node::Group {
2334            children: Vec::new(),
2335            kind: GroupKind::Implicit,
2336            mode: ContentMode::Math,
2337        });
2338
2339        ast.replace_content_child(superscript, replacement);
2340
2341        assert_eq!(ast.parent(superscript), None);
2342        assert!(ast.detached_roots.contains(&superscript));
2343        assert_eq!(
2344            ast.parent(replacement),
2345            Some(super::ParentLink {
2346                parent: scripted,
2347                slot: Slot::ScriptSup,
2348            })
2349        );
2350        assert!(matches!(
2351            ast.node(scripted),
2352            Node::Scripted {
2353                superscript: Some(child),
2354                ..
2355            } if *child == replacement
2356        ));
2357        ast.assert_invariants();
2358    }
2359
2360    #[test]
2361    fn replace_content_child_replaces_infix_operand() {
2362        let mut ast = Ast::new();
2363        let left = ast.new_node(Node::Char('a'));
2364        let right = ast.new_node(Node::Char('b'));
2365        let infix = ast.new_node(Node::Infix {
2366            name: "over".to_string(),
2367            args: Vec::new(),
2368            left,
2369            right,
2370        });
2371        ast.append_child(ast.root(), infix);
2372        let replacement = ast.new_node(Node::Char('x'));
2373
2374        ast.replace_content_child(right, replacement);
2375
2376        assert_eq!(ast.parent(right), None);
2377        assert!(ast.detached_roots.contains(&right));
2378        assert_eq!(
2379            ast.parent(replacement),
2380            Some(super::ParentLink {
2381                parent: infix,
2382                slot: Slot::InfixRight,
2383            })
2384        );
2385        assert!(matches!(
2386            ast.node(infix),
2387            Node::Infix {
2388                right: child,
2389                ..
2390            } if *child == replacement
2391        ));
2392        ast.assert_invariants();
2393    }
2394
2395    #[test]
2396    fn replace_node_drop_detached_children_removes_old_subtree() {
2397        let mut ast = Ast::new();
2398        let old_child = ast.new_node(Node::Char('x'));
2399        let old_grandchild = ast.new_node(Node::Char('y'));
2400        let old_child = ast.new_node(Node::Group {
2401            children: vec![old_child, old_grandchild],
2402            kind: GroupKind::Implicit,
2403            mode: ContentMode::Math,
2404        });
2405        let target = ast.new_node(Node::Group {
2406            children: vec![old_child],
2407            kind: GroupKind::Implicit,
2408            mode: ContentMode::Math,
2409        });
2410        ast.append_child(ast.root(), target);
2411        let new_child = ast.new_node(Node::Char('z'));
2412
2413        ast.replace_node_drop_detached_children(
2414            target,
2415            Node::Group {
2416                children: vec![new_child],
2417                kind: GroupKind::Implicit,
2418                mode: ContentMode::Math,
2419            },
2420        );
2421
2422        assert!(!ast.contains(old_child));
2423        assert!(!ast.contains(old_grandchild));
2424        assert_eq!(ast.parent_id(new_child), Some(target));
2425        assert_eq!(ast.children(target), &[new_child]);
2426        ast.assert_invariants();
2427    }
2428
2429    #[test]
2430    fn replace_with_math_sequence_splices_group_children() {
2431        let mut ast = Ast::new();
2432        let target = ast.new_node(Node::Char('x'));
2433        let root = ast.root();
2434        ast.append_child(root, target);
2435        let before = ast.new_node(Node::Char('a'));
2436        let replacement = ast.new_node(Node::Char('b'));
2437        let after = ast.new_node(Node::Char('c'));
2438
2439        ast.replace_with_math_sequence(target, vec![before], replacement, vec![after]);
2440
2441        assert!(!ast.contains(replacement));
2442        assert_eq!(ast.children(root), &[before, target, after]);
2443        assert_eq!(ast.node(target), &Node::Char('b'));
2444        assert_eq!(ast.parent_id(before), Some(root));
2445        assert_eq!(ast.parent_id(after), Some(root));
2446        ast.assert_invariants();
2447    }
2448
2449    #[test]
2450    fn replace_with_math_sequence_wraps_single_child_slots() {
2451        let mut ast = Ast::new();
2452        let target = ast.new_node(Node::Char('x'));
2453        let command = ast.new_node(Node::Command {
2454            name: "sqrt".to_string(),
2455            args: vec![Some(Argument {
2456                kind: ArgumentKind::Mandatory,
2457                value: ArgumentValue::MathContent(target),
2458            })],
2459            known: true,
2460        });
2461        ast.append_child(ast.root(), command);
2462        let before = ast.new_node(Node::Char('a'));
2463        let replacement = ast.new_node(Node::Char('b'));
2464        let after = ast.new_node(Node::Char('c'));
2465
2466        ast.replace_with_math_sequence(target, vec![before], replacement, vec![after]);
2467
2468        assert_eq!(
2469            ast.node(target),
2470            &Node::Group {
2471                children: vec![before, replacement, after],
2472                kind: GroupKind::Implicit,
2473                mode: ContentMode::Math,
2474            }
2475        );
2476        assert_eq!(ast.parent_id(before), Some(target));
2477        assert_eq!(ast.parent_id(replacement), Some(target));
2478        assert_eq!(ast.parent_id(after), Some(target));
2479        ast.assert_invariants();
2480    }
2481
2482    #[test]
2483    fn replace_with_math_sequence_rejects_duplicate_sequence_node_before_replacement() {
2484        let mut ast = Ast::new();
2485        let root = ast.root();
2486        let target = ast.new_node(Node::Char('x'));
2487        ast.append_child(root, target);
2488        let replacement = ast.new_node(Node::Char('y'));
2489
2490        let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
2491            ast.replace_with_math_sequence(target, vec![replacement], replacement, Vec::new());
2492        }));
2493
2494        let message = panic_message(result.expect_err("duplicate sequence node should panic"));
2495        assert!(
2496            message.contains("Node cannot appear in a replacement sequence twice"),
2497            "unexpected panic: {message}"
2498        );
2499        assert_eq!(ast.node(target), &Node::Char('x'));
2500        assert_eq!(ast.children(root), &[target]);
2501        assert!(ast.contains(replacement));
2502        assert_eq!(ast.parent(replacement), None);
2503        assert!(ast.detached_roots.contains(&replacement));
2504    }
2505
2506    #[test]
2507    fn replace_with_math_sequence_rejects_attached_before_node_before_replacement() {
2508        let mut ast = Ast::new();
2509        let root = ast.root();
2510        let target = ast.new_node(Node::Char('x'));
2511        let attached = ast.new_node(Node::Char('a'));
2512        ast.append_child(root, attached);
2513        ast.append_child(root, target);
2514        let replacement = ast.new_node(Node::Char('y'));
2515
2516        let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
2517            ast.replace_with_math_sequence(target, vec![attached], replacement, Vec::new());
2518        }));
2519
2520        let message = panic_message(result.expect_err("attached sequence node should panic"));
2521        assert!(
2522            message.contains("Cannot attach child that already has a parent"),
2523            "unexpected panic: {message}"
2524        );
2525        assert_eq!(ast.node(target), &Node::Char('x'));
2526        assert_eq!(ast.children(root), &[attached, target]);
2527        assert!(ast.contains(replacement));
2528        assert_eq!(ast.parent(replacement), None);
2529        assert!(ast.detached_roots.contains(&replacement));
2530    }
2531
2532    #[test]
2533    fn replace_with_math_sequence_preserving_scripts_rejects_duplicate_before_staging() {
2534        let mut ast = Ast::new();
2535        let base = ast.new_node(Node::Char('x'));
2536        let superscript = ast.new_node(Node::Char('2'));
2537        let scripted = ast.new_node(Node::Scripted {
2538            base,
2539            subscript: None,
2540            superscript: Some(superscript),
2541        });
2542        ast.append_child(ast.root(), scripted);
2543        let first = ast.new_node(Node::Char('['));
2544
2545        let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
2546            ast.replace_with_math_sequence_preserving_scripts(base, vec![first], first, Vec::new());
2547        }));
2548
2549        let message = panic_message(result.expect_err("duplicate sequence node should panic"));
2550        assert!(
2551            message.contains("Node cannot appear in a replacement sequence twice"),
2552            "unexpected panic: {message}"
2553        );
2554        assert_eq!(
2555            ast.node(scripted),
2556            &Node::Scripted {
2557                base,
2558                subscript: None,
2559                superscript: Some(superscript),
2560            }
2561        );
2562        assert_eq!(ast.parent(first), None);
2563        assert!(ast.detached_roots.contains(&first));
2564    }
2565
2566    #[test]
2567    fn replace_with_math_sequence_preserving_scripts_moves_scripts_to_last_node() {
2568        let mut ast = Ast::new();
2569        let base = ast.new_node(Node::Char('x'));
2570        let subscript = ast.new_node(Node::Char('i'));
2571        let superscript = ast.new_node(Node::Char('2'));
2572        let scripted = ast.new_node(Node::Scripted {
2573            base,
2574            subscript: Some(subscript),
2575            superscript: Some(superscript),
2576        });
2577        ast.append_child(ast.root(), scripted);
2578        let open = ast.new_node(Node::Char('['));
2579        let body = ast.new_node(Node::Char('y'));
2580        let close = ast.new_node(Node::Char(']'));
2581
2582        ast.replace_with_math_sequence_preserving_scripts(
2583            base,
2584            Vec::new(),
2585            open,
2586            vec![body, close],
2587        );
2588
2589        let root_children = ast.children(ast.root()).to_vec();
2590        assert_eq!(root_children.len(), 3);
2591        assert_eq!(root_children[0], scripted);
2592        assert_eq!(ast.node(scripted), &Node::Char('['));
2593        assert_eq!(root_children[1], body);
2594        let Node::Scripted {
2595            base: scripted_close,
2596            subscript: moved_subscript,
2597            superscript: moved_superscript,
2598        } = ast.node(root_children[2])
2599        else {
2600            panic!("expected scripts to move to the close token");
2601        };
2602        assert_eq!(ast.node(*scripted_close), &Node::Char(']'));
2603        assert_eq!(
2604            moved_subscript.map(|id| ast.node(id)),
2605            Some(&Node::Char('i'))
2606        );
2607        assert_eq!(
2608            moved_superscript.map(|id| ast.node(id)),
2609            Some(&Node::Char('2'))
2610        );
2611        ast.assert_invariants();
2612    }
2613}