plotnik-compiler 0.3.2

//! AST/CST pretty-printer for debugging and test snapshots.

use std::collections::HashMap;
use std::fmt::Write;

use indexmap::IndexSet;
use rowan::NodeOrToken;

use crate::parser::{self as ast, SyntaxNode};

use super::Query;
use super::SourceKind;
use crate::analyze::type_check::Arity;

/// Returns indentation string for the given level.
fn indent(level: usize) -> String {
    "  ".repeat(level)
}

pub struct QueryPrinter<'q> {
    query: &'q Query,
    raw: bool,
    trivia: bool,
    arities: bool,
    spans: bool,
    symbols: bool,
}

impl<'q> QueryPrinter<'q> {
    pub fn new(query: &'q Query) -> Self {
        Self {
            query,
            raw: false,
            trivia: false,
            arities: false,
            spans: false,
            symbols: false,
        }
    }

    pub fn raw(mut self, value: bool) -> Self {
        self.raw = value;
        self
    }

    pub fn with_trivia(mut self, value: bool) -> Self {
        self.trivia = value;
        self
    }

    pub fn with_arities(mut self, value: bool) -> Self {
        self.arities = value;
        self
    }

    pub fn with_spans(mut self, value: bool) -> Self {
        self.spans = value;
        self
    }

    pub fn only_symbols(mut self, value: bool) -> Self {
        self.symbols = value;
        self
    }

    pub fn dump(&self) -> String {
        let mut out = String::new();
        self.format(&mut out).expect("String write never fails");
        out
    }

    pub fn format(&self, w: &mut impl Write) -> std::fmt::Result {
        if self.symbols {
            return self.format_symbols(w);
        }

        let source_map = self.query.source_map();
        let ast_map = self.query.asts();
        let show_headers = self.should_show_headers(source_map);
        let mut first = true;

        for source in source_map.iter() {
            let Some(root) = ast_map.get(&source.id) else {
                continue;
            };

            if show_headers {
                if !first {
                    writeln!(w)?;
                }
                writeln!(w, "# {}", source.kind.display_name())?;
            }

            if self.raw {
                self.format_cst(root.as_cst(), 0, w)?;
            } else {
                self.format_root(root, w)?;
            }

            first = false;
        }

        Ok(())
    }

    fn should_show_headers(&self, source_map: &super::source_map::SourceMap) -> bool {
        source_map.len() > 1
            || source_map
                .iter()
                .next()
                .is_some_and(|s| !matches!(s.kind, SourceKind::OneLiner))
    }

    fn format_symbols(&self, w: &mut impl Write) -> std::fmt::Result {
        let symbols = &self.query.symbol_table;
        if symbols.is_empty() {
            return Ok(());
        }

        let defined: IndexSet<&str> = symbols.keys().collect();

        // Collect body nodes from all files
        let mut body_nodes: HashMap<String, SyntaxNode> = HashMap::new();
        for root in self.query.asts().values() {
            for def in root.defs() {
                if let (Some(name_tok), Some(body)) = (def.name(), def.body()) {
                    body_nodes.insert(name_tok.text().to_string(), body.as_cst().clone());
                }
            }
        }

        for name in symbols.keys() {
            let mut visited = IndexSet::new();
            self.format_symbol_tree(name, 0, &defined, &body_nodes, &mut visited, w)?;
        }
        Ok(())
    }

    fn format_symbol_tree(
        &self,
        name: &str,
        depth: usize,
        defined: &indexmap::IndexSet<&str>,
        body_nodes: &std::collections::HashMap<String, SyntaxNode>,
        visited: &mut indexmap::IndexSet<String>,
        w: &mut impl Write,
    ) -> std::fmt::Result {
        let prefix = indent(depth);

        if visited.contains(name) {
            writeln!(w, "{}{} (cycle)", prefix, name)?;
            return Ok(());
        }

        let is_broken = !defined.contains(name);
        if is_broken {
            writeln!(w, "{}{}?", prefix, name)?;
            return Ok(());
        }

        let card = body_nodes
            .get(name)
            .map(|n| self.arity_mark(n))
            .unwrap_or("");
        writeln!(w, "{}{}{}", prefix, name, card)?;
        visited.insert(name.to_string());

        if let Some(body) = self.query.symbol_table.get(name) {
            let refs_set = crate::analyze::refs::ref_names(body);
            let mut refs: Vec<_> = refs_set.iter().map(|s| s.as_str()).collect();
            refs.sort();
            for r in refs {
                self.format_symbol_tree(r, depth + 1, defined, body_nodes, visited, w)?;
            }
        }

        visited.shift_remove(name);
        Ok(())
    }

    fn format_cst(&self, node: &SyntaxNode, depth: usize, w: &mut impl Write) -> std::fmt::Result {
        let prefix = indent(depth);
        let card = self.arity_mark(node);
        let span = self.span_str(node.text_range());

        writeln!(w, "{}{:?}{}{}", prefix, node.kind(), card, span)?;

        for child in node.children_with_tokens() {
            match child {
                NodeOrToken::Node(n) => self.format_cst(&n, depth + 1, w)?,
                NodeOrToken::Token(t) => {
                    if !self.trivia && t.kind().is_trivia() {
                        continue;
                    }
                    let child_prefix = indent(depth + 1);
                    let child_span = self.span_str(t.text_range());
                    writeln!(
                        w,
                        "{}{:?}{} {:?}",
                        child_prefix,
                        t.kind(),
                        child_span,
                        t.text()
                    )?;
                }
            }
        }
        Ok(())
    }

    fn format_root(&self, root: &ast::Root, w: &mut impl Write) -> std::fmt::Result {
        let card = self.arity_mark(root.as_cst());
        let span = self.span_str(root.text_range());
        writeln!(w, "Root{}{}", card, span)?;

        for def in root.defs() {
            self.format_def(&def, 1, w)?;
        }
        // Parser wraps all top-level exprs in Def nodes, so this should be empty
        assert!(
            root.exprs().next().is_none(),
            "printer: unexpected bare Expr in Root (parser should wrap in Def)"
        );
        Ok(())
    }

    fn format_def(&self, def: &ast::Def, depth: usize, w: &mut impl Write) -> std::fmt::Result {
        let prefix = indent(depth);
        let card = self.arity_mark(def.as_cst());
        let span = self.span_str(def.text_range());
        let name = def.name().map(|t| t.text().to_string());

        match name {
            Some(n) => writeln!(w, "{}Def{}{} {}", prefix, card, span, n)?,
            None => writeln!(w, "{}Def{}{}", prefix, card, span)?,
        }

        let Some(body) = def.body() else {
            return Ok(());
        };
        self.format_expr(&body, depth + 1, w)
    }

    fn format_expr(&self, expr: &ast::Expr, depth: usize, w: &mut impl Write) -> std::fmt::Result {
        let prefix = indent(depth);
        let card = self.arity_mark(expr.as_cst());
        let span = self.span_str(expr.text_range());

        match expr {
            ast::Expr::NamedNode(n) => {
                if n.is_any() {
                    writeln!(w, "{}NamedNode{}{} (any)", prefix, card, span)?;
                } else {
                    let node_type = n.node_type().map(|tok| tok.text().to_string());
                    match node_type {
                        Some(ty) => writeln!(w, "{}NamedNode{}{} {}", prefix, card, span, ty)?,
                        None => writeln!(w, "{}NamedNode{}{}", prefix, card, span)?,
                    }
                }
                self.format_tree_children(n.as_cst(), depth + 1, w)?;
            }
            ast::Expr::Ref(r) => {
                let name = r.name().map(|t| t.text().to_string()).unwrap_or_default();
                writeln!(w, "{}Ref{}{} {}", prefix, card, span, name)?;
            }
            ast::Expr::AnonymousNode(a) => {
                if a.is_any() {
                    writeln!(w, "{}AnonymousNode{}{} (any)", prefix, card, span)?;
                } else {
                    let value = a.value().map(|t| t.text().to_string()).unwrap_or_default();
                    writeln!(w, "{}AnonymousNode{}{} \"{}\"", prefix, card, span, value)?;
                }
            }
            ast::Expr::AltExpr(a) => {
                writeln!(w, "{}Alt{}{}", prefix, card, span)?;
                for branch in a.branches() {
                    self.format_branch(&branch, depth + 1, w)?;
                }
                for expr in a.exprs() {
                    self.format_expr(&expr, depth + 1, w)?;
                }
            }
            ast::Expr::SeqExpr(s) => {
                writeln!(w, "{}Seq{}{}", prefix, card, span)?;
                self.format_tree_children(s.as_cst(), depth + 1, w)?;
            }
            ast::Expr::CapturedExpr(c) => {
                let name = c
                    .name()
                    .map(|t| t.text()[1..].to_string()) // Strip @ prefix
                    .unwrap_or_default();
                let type_ann = c
                    .type_annotation()
                    .and_then(|t| t.name())
                    .map(|t| t.text().to_string());
                match type_ann {
                    Some(ty) => writeln!(
                        w,
                        "{}CapturedExpr{}{} @{} :: {}",
                        prefix, card, span, name, ty
                    )?,
                    None => writeln!(w, "{}CapturedExpr{}{} @{}", prefix, card, span, name)?,
                }
                let Some(inner) = c.inner() else {
                    return Ok(());
                };
                self.format_expr(&inner, depth + 1, w)?;
            }
            ast::Expr::QuantifiedExpr(q) => {
                let op = q
                    .operator()
                    .map(|t| t.text().to_string())
                    .unwrap_or_default();
                writeln!(w, "{}QuantifiedExpr{}{} {}", prefix, card, span, op)?;
                let Some(inner) = q.inner() else {
                    return Ok(());
                };
                self.format_expr(&inner, depth + 1, w)?;
            }
            ast::Expr::FieldExpr(f) => {
                let name = f.name().map(|t| t.text().to_string()).unwrap_or_default();
                writeln!(w, "{}FieldExpr{}{} {}:", prefix, card, span, name)?;
                let Some(value) = f.value() else {
                    return Ok(());
                };
                self.format_expr(&value, depth + 1, w)?;
            }
        }
        Ok(())
    }

    fn format_tree_children(
        &self,
        node: &SyntaxNode,
        depth: usize,
        w: &mut impl Write,
    ) -> std::fmt::Result {
        use crate::parser::SyntaxKind;
        for child in node.children() {
            if child.kind() == SyntaxKind::Anchor {
                self.mark_anchor(depth, w)?;
            } else if child.kind() == SyntaxKind::NegatedField {
                self.format_negated_field(&ast::NegatedField::cast(child).unwrap(), depth, w)?;
            } else if let Some(expr) = ast::Expr::cast(child) {
                self.format_expr(&expr, depth, w)?;
            }
        }
        Ok(())
    }

    fn mark_anchor(&self, depth: usize, w: &mut impl Write) -> std::fmt::Result {
        writeln!(w, "{}.", indent(depth))
    }

    fn format_negated_field(
        &self,
        nf: &ast::NegatedField,
        depth: usize,
        w: &mut impl Write,
    ) -> std::fmt::Result {
        let prefix = indent(depth);
        let span = self.span_str(nf.text_range());
        let name = nf.name().map(|t| t.text().to_string()).unwrap_or_default();
        writeln!(w, "{}NegatedField{} -{}", prefix, span, name)
    }

    fn format_branch(
        &self,
        branch: &ast::Branch,
        depth: usize,
        w: &mut impl Write,
    ) -> std::fmt::Result {
        let prefix = indent(depth);
        let card = self.arity_mark(branch.as_cst());
        let span = self.span_str(branch.text_range());
        let label = branch.label().map(|t| t.text().to_string());

        match label {
            Some(l) => writeln!(w, "{}Branch{}{} {}:", prefix, card, span, l)?,
            None => writeln!(w, "{}Branch{}{}", prefix, card, span)?,
        }

        let Some(body) = branch.body() else {
            return Ok(());
        };
        self.format_expr(&body, depth + 1, w)
    }

    fn arity_mark(&self, node: &SyntaxNode) -> &'static str {
        if !self.arities {
            return "";
        }
        match self.query.get_arity(node) {
            Some(Arity::One) => "¹",
            Some(Arity::Many) => "⁺",
            None => "ˣ",
        }
    }

    fn span_str(&self, range: rowan::TextRange) -> String {
        if !self.spans {
            return String::new();
        }
        format!(
            " [{}..{}]",
            u32::from(range.start()),
            u32::from(range.end())
        )
    }
}

impl Query {
    pub fn printer(&self) -> QueryPrinter<'_> {
        QueryPrinter::new(self)
    }
}