polydat_grammar/

pprint.rs

// Copyright 2024-2026 Jonathan Shook
// SPDX-License-Identifier: Apache-2.0

//! AST → `.polydat` source pretty-printer.
//!
//! The printer gives a compiled `for` body its source text for
//! diagnostics (`pp_file`), re-emits rewritten expressions in
//! module inlining and tile lowering (`pp_expr`), and prints
//! expressions in `polydat explain`. The runtime compiles from the
//! AST; the printer is a faithful AST → source round-trip beside it.
//!
//! ## Round-trip contract
//!
//! For every `Statement`/`Expr` produced by the parser,
//! `pp_statement` / `pp_expr` produces source text that re-parses
//! into a semantically equivalent AST. "Semantically equivalent"
//! means same node types and identical inner data (modulo
//! `Span`s, which capture parser position and are not preserved
//! across re-parse).
//!
//! ## Precedence and parens
//!
//! `BinOp` expressions are emitted with parens around the whole
//! expression. This is uniformly safe — re-parsing produces the
//! same tree structure — at the cost of extra parens. The output
//! is the canonical spelling; the parens are uniform for round-trip
//! safety.

use crate::ast::{
    Arg, BinOpKind, Binding, BindingModifier, CallExpr, CursorDecl, Expr, ExternPort, ForStmt,
    ModuleDef, PolydatFile, Statement, TileBodyKind, TileDef, TileOptions, WireModifier,
};

/// Pretty-print a full file: every statement, separated by
/// newlines.
pub fn pp_file(file: &PolydatFile) -> String {
    let mut out = String::new();
    for stmt in &file.statements {
        out.push_str(&pp_statement(stmt));
        out.push('\n');
    }
    out
}

/// Pretty-print a top-level statement.
pub fn pp_statement(stmt: &Statement) -> String {
    match stmt {
        Statement::InputDecl(d) => match &d.ty {
            Some(ty) => format!("input {}: {}", d.name, ty),
            None => format!("input {}", d.name),
        },
        Statement::Binding(b) => pp_binding(b),
        Statement::ModuleDef(m) => pp_module_def(m),
        Statement::ExternPort(p) => pp_extern_port(p),
        Statement::Cursor(c) => pp_cursor(c),
        Statement::Pragma { name, .. } => format!("pragma {name}"),
        Statement::For(f) => pp_for_stmt(f, 0),
        Statement::Tile(t) => pp_tile(t),
    }
}

fn pp_tile(t: &TileDef) -> String {
    let mut out = format!("tile {}", t.name);
    if let Some(enc) = &t.encoding {
        out.push_str(&format!(" : {enc}"));
    }
    let defaults = TileOptions::default();
    let mut opts = Vec::new();
    if t.options.open != defaults.open || t.options.close != defaults.close {
        opts.push(format!(
            "delims \"{}\" \"{}\"",
            escape_string(&t.options.open),
            escape_string(&t.options.close)
        ));
    }
    if t.options.sigil != defaults.sigil {
        opts.push(format!("sigil \"{}\"", escape_string(&t.options.sigil)));
    }
    if t.options.strict {
        opts.push("strict".to_string());
    }
    if t.options.in_string {
        opts.push("instring".to_string());
    }
    if !opts.is_empty() {
        out.push_str(&format!(" ({})", opts.join(", ")));
    }
    // A tile binds a wire: `:=` precedes every body form.
    match t.body_kind {
        TileBodyKind::Block => {
            out.push_str(" := ");
            out.push_str(&t.body_text());
        }
        TileBodyKind::Heredoc => {
            out.push_str(" := <<<\n");
            out.push_str(&t.body_text());
            out.push_str("\n>>>");
        }
        TileBodyKind::Literal => {
            out.push_str(&format!(" := \"{}\"", escape_string(&t.body_text())))
        }
    }
    out
}

fn pp_for_stmt(f: &ForStmt, indent: usize) -> String {
    let pad = "    ".repeat(indent + 1);
    let mut body = String::new();
    for s in &f.body {
        body.push_str(&pad);
        body.push_str(&match s {
            Statement::For(inner) => pp_for_stmt(inner, indent + 1),
            other => pp_statement(other),
        });
        body.push('\n');
    }
    format!(
        "for {} {{\n{}{}}}",
        f.source.to_text(),
        body,
        "    ".repeat(indent)
    )
}

/// Pretty-print an expression. Always emits parens around
/// `BinOp` for round-trip safety.
pub fn pp_expr(expr: &Expr) -> String {
    match expr {
        Expr::Ident(name, _) => name.clone(),
        Expr::IntLit(v, _) => v.to_string(),
        Expr::FloatLit(v, _) => format_float(*v),
        Expr::StringLit(s, _) => format!("\"{}\"", escape_string(s)),
        Expr::ArrayLit(elts, _) => {
            let parts: Vec<String> = elts.iter().map(pp_expr).collect();
            format!("[{}]", parts.join(", "))
        }
        Expr::Call(c) => pp_call(c),
        Expr::BinOp(lhs, op, rhs) => {
            format!("({} {} {})", pp_expr(lhs), pp_binop(*op), pp_expr(rhs))
        }
        Expr::UnaryNeg(e, _) => format!("(-{})", pp_expr(e)),
        Expr::UnaryBitNot(e, _) => format!("(!{})", pp_expr(e)),
        Expr::FieldAccess { source, field, .. } => format!("{source}.{field}"),
        Expr::Cast(e, ty, _) => format!("({} as {})", pp_expr(e), ty.to_keyword()),
        Expr::For(source) => format!("for {}", source.to_text()),
    }
}

fn pp_binding(b: &Binding) -> String {
    let mut target = if b.targets.len() == 1 {
        b.targets[0].clone()
    } else {
        format!("({})", b.targets.join(", "))
    };
    // `shared name: type := …` — cell type annotation
    // (scope_model.md §"Type stability").
    if let Some(ty) = &b.type_annotation {
        target = format!("{target}: {ty}");
    }
    let prefix = pp_modifier_prefix(b.modifier);
    if prefix.is_empty() {
        format!("{} := {}", target, pp_expr(&b.value))
    } else {
        format!("{} {} := {}", prefix, target, pp_expr(&b.value))
    }
}

fn pp_extern_port(p: &ExternPort) -> String {
    if let Some(default) = &p.default {
        format!("extern {}: {} = {}", p.name, p.typ, pp_expr(default))
    } else {
        format!("extern {}: {}", p.name, p.typ)
    }
}

fn pp_cursor(c: &CursorDecl) -> String {
    let mut out = format!("cursor {} = {}", c.name, pp_expr(&c.constructor));
    if let Some(over) = &c.over {
        out.push_str(" over ");
        out.push_str(&pp_expr(over));
    }
    out
}

fn pp_module_def(m: &ModuleDef) -> String {
    let params: Vec<String> = m
        .params
        .iter()
        .map(|p| format!("{}: {}", p.name, p.typ))
        .collect();
    let outputs: Vec<String> = m
        .outputs
        .iter()
        .map(|p| format!("{}: {}", p.name, p.typ))
        .collect();
    let mut body = String::new();
    for s in &m.body {
        body.push_str("    ");
        body.push_str(&pp_statement(s));
        body.push('\n');
    }
    format!(
        "{}({}) -> ({}) := {{\n{}}}",
        m.name,
        params.join(", "),
        outputs.join(", "),
        body
    )
}

fn pp_call(c: &CallExpr) -> String {
    let args: Vec<String> = c.args.iter().map(pp_arg).collect();
    format!("{}({})", c.func, args.join(", "))
}

fn pp_arg(arg: &Arg) -> String {
    match arg {
        Arg::Positional(e) => pp_expr(e),
        Arg::Named(name, e) => format!("{}: {}", name, pp_expr(e)),
    }
}

fn pp_modifier_prefix(m: BindingModifier) -> String {
    let mut parts: Vec<&str> = Vec::new();
    if m.has(WireModifier::Const) {
        parts.push("const");
    }
    if m.has(WireModifier::Shared) {
        parts.push("shared");
    }
    if m.has(WireModifier::Volatile) {
        parts.push("volatile");
    }
    parts.join(" ")
}

fn pp_binop(op: BinOpKind) -> &'static str {
    match op {
        BinOpKind::Add => "+",
        BinOpKind::Sub => "-",
        BinOpKind::Mul => "*",
        BinOpKind::Div => "/",
        BinOpKind::Mod => "%",
        BinOpKind::Pow => "**",
        BinOpKind::BitAnd => "&",
        BinOpKind::BitOr => "|",
        BinOpKind::BitXor => "^",
        BinOpKind::Shl => "<<",
        BinOpKind::Shr => ">>",
        BinOpKind::Eq => "==",
        BinOpKind::Ne => "!=",
        BinOpKind::Lt => "<",
        BinOpKind::Gt => ">",
        BinOpKind::Le => "<=",
        BinOpKind::Ge => ">=",
        BinOpKind::And => "&&",
        BinOpKind::Or => "||",
    }
}

fn escape_string(s: &str) -> String {
    let mut out = String::with_capacity(s.len());
    for c in s.chars() {
        match c {
            '\\' => out.push_str("\\\\"),
            '"' => out.push_str("\\\""),
            '\n' => out.push_str("\\n"),
            '\t' => out.push_str("\\t"),
            '\r' => out.push_str("\\r"),
            c => out.push(c),
        }
    }
    out
}

fn format_float(v: f64) -> String {
    if v.is_finite() && v == v.trunc() && v.abs() < 1e18 {
        format!("{v:.1}")
    } else {
        format!("{v}")
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{lexer, parser};

    fn parse(src: &str) -> PolydatFile {
        let tokens = lexer::lex(src).expect("lex");
        parser::parse(tokens).expect("parse")
    }

    fn round_trip(src: &str) {
        let ast1 = parse(src);
        let printed = pp_file(&ast1);
        let ast2 = parse(&printed);
        let printed2 = pp_file(&ast2);
        assert_eq!(
            printed, printed2,
            "second-pass print should be idempotent.\n\
             original source:\n{src}\n\n\
             first print:\n{printed}\n\n\
             second print:\n{printed2}"
        );
    }

    #[test]
    fn round_trip_simple_const() {
        round_trip("const x := 42\n");
    }

    #[test]
    fn round_trip_string_const() {
        round_trip("const dataset := \"sift1m\"\n");
    }

    #[test]
    fn round_trip_init_binding() {
        round_trip("const prebuffer := dataset_prebuffer(\"example\")\n");
    }

    #[test]
    fn round_trip_function_call() {
        round_trip("ratio := mod(cycle, 100)\n");
    }

    #[test]
    fn round_trip_named_args() {
        round_trip("v := dist_normal(mean: 72.0, stddev: 5.0)\n");
    }

    #[test]
    fn round_trip_binop() {
        round_trip("y := (x + 1)\n");
    }

    #[test]
    fn round_trip_inputs() {
        round_trip("input (cycle: u64, thread: u64)\n");
    }

    #[test]
    fn round_trip_extern() {
        round_trip("extern dataset: String\n");
    }

    #[test]
    fn round_trip_tuple_destructure() {
        round_trip("(a, b) := unpack(cycle)\n");
    }

    #[test]
    fn round_trip_workload_typical() {
        // Mirrors the shape of full_cql_vector workload bindings.
        let src = "\
const dataset := \"sift1m\"
const prefix := \"vec_default\"
profiles := matching_profiles(dataset, prefix)
table := first(profiles)
";
        round_trip(src);
    }

    #[test]
    fn round_trip_string_escapes() {
        round_trip("const s := \"hello \\\"world\\\"\"\n");
    }

    #[test]
    fn round_trip_array_literal() {
        round_trip("const weights := [60.0, 20.0, 15.0, 5.0]\n");
    }

    #[test]
    fn round_trip_cursor() {
        round_trip("cursor users = range(0, 1000000)\n");
    }

    #[test]
    fn round_trip_cursor_with_over() {
        // The `over <expr>` partition clause (cursor_partitions.md §7.2) must survive
        // projection — pp_cursor emits it so over-bearing cursors
        // round-trip faithfully.
        round_trip("cursor q = range(0, 100) over p\n");
    }

    #[test]
    fn pp_cursor_emits_over_clause() {
        let ast = parse("cursor q = range(0, 100) over p\n");
        let printed = pp_file(&ast);
        assert!(
            printed.contains(" over p"),
            "projected cursor must retain its `over` clause, got:\n{printed}"
        );
    }
}