muon/

tree.rs

//! The Muon tree and its JSON serialization.
//!
//! `to_json` emits a compact, tagged JSON dump used by the golden tests so they
//! read close to `muon-spec.md`. It is deliberately not Muon: other-language
//! implementations can compare against it with their own JSON parser (see the
//! language-agnostic-suite item in `docs/TODO.md`).

use alloc::boxed::Box;
use alloc::format;
use alloc::string::String;
use alloc::vec::Vec;
use core::mem;

/// A single visually-atomic entity in a chain (`<item>` in the spec).
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Item<'a> {
    /// A comment, stored raw *including* its `//` or `/* */` delimiters.
    Comm(&'a str),
    /// A string literal; the inner slice, *between* the quotes, undecoded
    /// (escapes are validated but not expanded — see the crate `open` note).
    Str(&'a str),
    /// A symbol: a maximal run of symchars (identifier, number, operator, …).
    Sym(&'a str),
    /// Punctuation: a maximal run of punctchars standing between items (`:`).
    Punct(&'a str),
    /// A parenthesis-wrapped sequence.
    Parens(Group<'a>),
    /// A bracket-wrapped sequence.
    Brackets(Group<'a>),
    /// A brace-wrapped sequence.
    Braces(Group<'a>),
    /// An item marked with a leading sigil (`&`, `.`, or `#`).
    Prefixed { sigil: char, item: Box<Item<'a>> },
    /// A **run**: two or more atoms written with no whitespace between them, so
    /// they form one chain item (`Main.A.S`, `f(x)`, `(f x).1`, `3.14`). Members
    /// are in written order; none is itself a `Run`, a `Punct`, or a `Comm` (a
    /// punct or comment ends a run). A lone atom is never wrapped in a `Run`.
    Run(Vec<Item<'a>>),
}

/// The location of a node in the source: `length` bytes from byte `offset`.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub struct Span {
    pub offset: u32,
    pub length: u32,
}

impl Span {
    /// The byte offset just past the node.
    pub fn end(self) -> u32 {
        self.offset + self.length
    }
}

/// The sequence inside a bracket group, with the group's location. A group is
/// the one node that stores its span: a leaf borrows its text from the source,
/// so its span follows from where the slice starts (see [`Item::span`]), but an
/// empty `()` has no slice to measure, and whitespace may stand between a
/// bracket and the first item.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Group<'a> {
    pub seq: Seq<'a>,
    /// Covers the brackets, but not the sigils of a prefixed group.
    pub span: Span,
}

/// A chain: one or more juxtaposed items (`<chain>`). Never empty.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Chain<'a>(pub Vec<Item<'a>>);

/// A sequence: zero or more chains delimited by separators (`<seq>`).
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Seq<'a>(pub Vec<Chain<'a>>);

/// Tear the tree down iteratively, so dropping a deeply-nested tree (the parser
/// bounds nesting only by memory) can't overflow the stack the way the default
/// recursive `Drop` would. `Item` itself has no `Drop`, so leaves — the vast
/// majority of nodes — drop trivially; the cost falls only on the nesting nodes
/// walked here. The two recursion points are containers (a nested `Seq`) and
/// `Prefixed` (a boxed inner `Item`); both are drained onto a heap worklist
/// before the node holding them drops, so no `Drop` call nests.
impl Drop for Seq<'_> {
    fn drop(&mut self) {
        // Leaves — the vast majority of nodes — drop in place as `collect_nesting`
        // iterates; only the nesting nodes reach `work`, so a flat tree never
        // allocates one and a deep tree is torn down a level at a time.
        let mut work: Vec<Item> = Vec::new();
        collect_nesting(self, &mut work);
        while let Some(mut item) = work.pop() {
            match &mut item {
                Item::Parens(g) | Item::Brackets(g) | Item::Braces(g) => {
                    collect_nesting(&mut g.seq, &mut work);
                }
                Item::Prefixed { item, .. } => {
                    let inner = mem::replace(item.as_mut(), Item::Sym(""));
                    if inner.nests() {
                        work.push(inner);
                    }
                }
                Item::Run(members) => {
                    for m in mem::take(members) {
                        if m.nests() {
                            work.push(m);
                        }
                    }
                }
                _ => {}
            }
            // `item` drops here with its nesting already moved out: a container's
            // inner `Seq` is empty and a `Prefixed`'s box now holds a leaf, so
            // neither recurses.
        }
    }
}

/// Drop `seq`'s leaves in place and move its nesting items onto `work`, leaving
/// `seq` empty.
fn collect_nesting<'a>(seq: &mut Seq<'a>, work: &mut Vec<Item<'a>>) {
    for chain in mem::take(&mut seq.0) {
        for item in chain.0 {
            if item.nests() {
                work.push(item);
            }
            // else: a leaf, dropped here as the loop advances.
        }
    }
}

impl Seq<'_> {
    /// Serialize the tree to compact, tagged JSON.
    pub fn to_json(&self) -> String {
        let mut out = String::new();
        self.write_json(&mut out);
        out
    }

    fn write_json(&self, out: &mut String) {
        out.push_str("{\"seq\":[");
        for (i, chain) in self.0.iter().enumerate() {
            if i > 0 {
                out.push(',');
            }
            chain.write_json(out);
        }
        out.push_str("]}");
    }
}

impl Chain<'_> {
    fn write_json(&self, out: &mut String) {
        out.push_str("{\"chain\":[");
        for (i, item) in self.0.iter().enumerate() {
            if i > 0 {
                out.push(',');
            }
            item.write_json(out);
        }
        out.push_str("]}");
    }
}

impl Item<'_> {
    /// Where this item stands in `src`, the text it was parsed from. A string's
    /// span covers its quotes, a comment's its delimiters, a prefixed item's its
    /// sigils, and a group's its brackets.
    ///
    /// # Panics
    ///
    /// If the item was not parsed from `src`.
    pub fn span(&self, src: &str) -> Span {
        let leaf = |s: &str| {
            let range = src
                .substr_range(s)
                .expect("the item was not parsed from `src`");
            Span {
                offset: range.start as u32,
                length: s.len() as u32,
            }
        };
        match self {
            Item::Comm(s) | Item::Sym(s) | Item::Punct(s) => leaf(s),
            Item::Str(s) => {
                let inner = leaf(s);
                Span {
                    offset: inner.offset - 1,
                    length: inner.length + 2,
                }
            }
            Item::Parens(g) | Item::Brackets(g) | Item::Braces(g) => g.span,
            // A sigil is one ASCII byte, glued to its item.
            Item::Prefixed { item, .. } => {
                let inner = item.span(src);
                Span {
                    offset: inner.offset - 1,
                    length: inner.length + 1,
                }
            }
            Item::Run(members) => {
                let first = members[0].span(src);
                let last = members[members.len() - 1].span(src);
                Span {
                    offset: first.offset,
                    length: last.end() - first.offset,
                }
            }
        }
    }

    /// Whether this item owns nested items (a container's `Seq` or a `Prefixed`'s
    /// boxed item) — the nodes the iterative `Drop for Seq` must walk.
    fn nests(&self) -> bool {
        matches!(
            self,
            Item::Parens(_)
                | Item::Brackets(_)
                | Item::Braces(_)
                | Item::Prefixed { .. }
                | Item::Run(_)
        )
    }

    fn write_json(&self, out: &mut String) {
        match self {
            Item::Comm(s) => {
                out.push_str("{\"comm\":");
                json_str(s, out);
                out.push('}');
            }
            Item::Str(s) => {
                out.push_str("{\"str\":");
                json_str(s, out);
                out.push('}');
            }
            Item::Sym(s) => {
                out.push_str("{\"sym\":");
                json_str(s, out);
                out.push('}');
            }
            Item::Punct(s) => {
                out.push_str("{\"punct\":");
                json_str(s, out);
                out.push('}');
            }
            Item::Parens(g) => {
                out.push_str("{\"parens\":");
                g.seq.write_json(out);
                out.push('}');
            }
            Item::Brackets(g) => {
                out.push_str("{\"brackets\":");
                g.seq.write_json(out);
                out.push('}');
            }
            Item::Braces(g) => {
                out.push_str("{\"braces\":");
                g.seq.write_json(out);
                out.push('}');
            }
            Item::Prefixed { sigil, item } => {
                out.push_str("{\"pre\":");
                let mut buf = [0u8; 4];
                json_str(sigil.encode_utf8(&mut buf), out);
                out.push_str(",\"item\":");
                item.write_json(out);
                out.push('}');
            }
            Item::Run(members) => {
                out.push_str("{\"run\":[");
                for (i, m) in members.iter().enumerate() {
                    if i > 0 {
                        out.push(',');
                    }
                    m.write_json(out);
                }
                out.push_str("]}");
            }
        }
    }
}

/// Write `s` as a JSON string literal (quoted and escaped).
fn json_str(s: &str, out: &mut String) {
    out.push('"');
    for ch in s.chars() {
        match ch {
            '"' => out.push_str("\\\""),
            '\\' => out.push_str("\\\\"),
            '\n' => out.push_str("\\n"),
            '\r' => out.push_str("\\r"),
            '\t' => out.push_str("\\t"),
            c if (c as u32) < 0x20 => out.push_str(&format!("\\u{:04x}", c as u32)),
            c => out.push(c),
        }
    }
    out.push('"');
}