AILang/crates/ailang-core/src/pretty.rs

//! Pretty-printer: AST → human-readable text form.
//!
//! The text form is intended as a diff and review tool. The
//! canonical source remains the JSON form. Every pretty output is
//! deterministic.

use crate::ast::*;
use std::fmt::Write;

pub fn module(m: &Module) -> String {
    let mut s = String::new();
    writeln!(s, "(module {}", m.name).unwrap();
    if !m.imports.is_empty() {
        for imp in &m.imports {
            match &imp.alias {
                Some(a) => writeln!(s, "  (import {} as {})", imp.module, a).unwrap(),
                None => writeln!(s, "  (import {})", imp.module).unwrap(),
            }
        }
    }
    for (i, def) in m.defs.iter().enumerate() {
        if i > 0 {
            s.push('\n');
        }
        let body = def_block(def, 2);
        s.push_str(&body);
        s.push('\n');
    }
    s.push(')');
    s.push('\n');
    s
}

pub fn manifest(m: &Module) -> String {
    let mut s = String::new();
    writeln!(s, "module {}", m.name).unwrap();
    let max_name = m.defs.iter().map(|d| d.name().len()).max().unwrap_or(0);
    for def in &m.defs {
        let h = crate::hash::def_hash(def);
        let (kw, ty) = match def {
            Def::Fn(f) => ("fn", type_to_string(&f.ty)),
            Def::Const(c) => ("const", type_to_string(&c.ty)),
            Def::Type(t) => {
                let ctors = t
                    .ctors
                    .iter()
                    .map(|c| {
                        if c.fields.is_empty() {
                            c.name.clone()
                        } else {
                            format!(
                                "{}({})",
                                c.name,
                                c.fields
                                    .iter()
                                    .map(type_to_string)
                                    .collect::<Vec<_>>()
                                    .join(", ")
                            )
                        }
                    })
                    .collect::<Vec<_>>()
                    .join(" | ");
                ("type", ctors)
            }
        };
        writeln!(
            s,
            "  {kw:5} {name:<width$} :: {ty}    [{h}]",
            kw = kw,
            name = def.name(),
            width = max_name,
            ty = ty,
            h = h,
        )
        .unwrap();
    }
    s
}

fn def_block(def: &Def, indent: usize) -> String {
    let pad = " ".repeat(indent);
    match def {
        Def::Fn(f) => {
            let params = if f.params.is_empty() {
                "[]".to_string()
            } else {
                format!("[{}]", f.params.join(" "))
            };
            let mut s = format!(
                "{pad}(fn {name} :: {ty} {params}\n",
                pad = pad,
                name = f.name,
                ty = type_to_string(&f.ty),
                params = params,
            );
            s.push_str(&term_block(&f.body, indent + 2));
            s.push(')');
            s
        }
        Def::Const(c) => {
            let mut s = format!(
                "{pad}(const {name} :: {ty}\n",
                pad = pad,
                name = c.name,
                ty = type_to_string(&c.ty),
            );
            s.push_str(&term_block(&c.value, indent + 2));
            s.push(')');
            s
        }
        Def::Type(t) => {
            let mut s = format!("{pad}(type {name}\n", pad = pad, name = t.name);
            let inner = " ".repeat(indent + 2);
            for ctor in &t.ctors {
                if ctor.fields.is_empty() {
                    s.push_str(&format!("{inner}(| {})\n", ctor.name));
                } else {
                    let fs = ctor
                        .fields
                        .iter()
                        .map(type_to_string)
                        .collect::<Vec<_>>()
                        .join(" ");
                    s.push_str(&format!("{inner}(| {name} {fs})\n", name = ctor.name));
                }
            }
            s.push_str(&pad);
            s.push(')');
            s
        }
    }
}

fn term_block(t: &Term, indent: usize) -> String {
    let pad = " ".repeat(indent);
    match t {
        Term::Lit { lit } => format!("{pad}{}", lit_to_string(lit)),
        Term::Var { name } => format!("{pad}{name}"),
        Term::App { callee, args } => {
            let mut s = format!("{pad}(");
            s.push_str(&term_inline(callee));
            for a in args {
                s.push(' ');
                s.push_str(&term_inline(a));
            }
            s.push(')');
            s
        }
        Term::Let { name, value, body } => {
            let mut s = format!("{pad}(let {name}\n");
            s.push_str(&term_block(value, indent + 2));
            s.push('\n');
            s.push_str(&term_block(body, indent + 2));
            s.push(')');
            s
        }
        Term::If { cond, then, else_ } => {
            let mut s = format!("{pad}(if\n");
            s.push_str(&term_block(cond, indent + 2));
            s.push('\n');
            s.push_str(&term_block(then, indent + 2));
            s.push('\n');
            s.push_str(&term_block(else_, indent + 2));
            s.push(')');
            s
        }
        Term::Do { op, args } => {
            let mut s = format!("{pad}(do {op}");
            for a in args {
                s.push(' ');
                s.push_str(&term_inline(a));
            }
            s.push(')');
            s
        }
        Term::Ctor { type_name, ctor, args } => {
            let mut s = format!("{pad}({type_name}/{ctor}");
            for a in args {
                s.push(' ');
                s.push_str(&term_inline(a));
            }
            s.push(')');
            s
        }
        Term::Match { scrutinee, arms } => {
            let mut s = format!("{pad}(match {}\n", term_inline(scrutinee));
            let inner = " ".repeat(indent + 2);
            for arm in arms {
                s.push_str(&format!(
                    "{inner}(case {} ->\n",
                    pattern_to_string(&arm.pat)
                ));
                s.push_str(&term_block(&arm.body, indent + 4));
                s.push_str(")\n");
            }
            s.push_str(&pad);
            s.push(')');
            s
        }
        Term::Lam { params, param_tys, ret_ty, effects, body } => {
            // (\ [params] :: typed-sig . body)
            let typed_params: Vec<String> = params
                .iter()
                .zip(param_tys.iter())
                .map(|(n, t)| format!("{n}: {}", type_to_string(t)))
                .collect();
            let eff = if effects.is_empty() {
                String::new()
            } else {
                format!(" !{}", effects.join(","))
            };
            let mut s = format!(
                "{pad}(\\ ({}) -> {}{}\n",
                typed_params.join(" "),
                type_to_string(ret_ty),
                eff,
            );
            s.push_str(&term_block(body, indent + 2));
            s.push(')');
            s
        }
    }
}

pub fn pattern_to_string(p: &Pattern) -> String {
    match p {
        Pattern::Wild => "_".into(),
        Pattern::Var { name } => name.clone(),
        Pattern::Lit { lit } => lit_to_string(lit),
        Pattern::Ctor { ctor, fields } => {
            if fields.is_empty() {
                ctor.clone()
            } else {
                let fs = fields
                    .iter()
                    .map(pattern_to_string)
                    .collect::<Vec<_>>()
                    .join(" ");
                format!("({ctor} {fs})")
            }
        }
    }
}

fn term_inline(t: &Term) -> String {
    match t {
        Term::Lit { lit } => lit_to_string(lit),
        Term::Var { name } => name.clone(),
        Term::App { callee, args } => {
            let mut s = String::from("(");
            s.push_str(&term_inline(callee));
            for a in args {
                s.push(' ');
                s.push_str(&term_inline(a));
            }
            s.push(')');
            s
        }
        Term::Do { op, args } => {
            let mut s = format!("(do {op}");
            for a in args {
                s.push(' ');
                s.push_str(&term_inline(a));
            }
            s.push(')');
            s
        }
        Term::Ctor { type_name, ctor, args } => {
            let mut s = format!("({type_name}/{ctor}");
            for a in args {
                s.push(' ');
                s.push_str(&term_inline(a));
            }
            s.push(')');
            s
        }
        Term::Match { scrutinee, .. } => {
            // Don't render match inline — use a marker.
            format!("(match {} ...)", term_inline(scrutinee))
        }
        // Structural terms are tricky to render inline recursively; fallback:
        Term::Let { name, value, body } => {
            format!(
                "(let {name} {} {})",
                term_inline(value),
                term_inline(body)
            )
        }
        Term::If { cond, then, else_ } => {
            format!(
                "(if {} {} {})",
                term_inline(cond),
                term_inline(then),
                term_inline(else_)
            )
        }
        Term::Lam { params, .. } => {
            format!("(\\ {} ...)", params.join(" "))
        }
    }
}

fn lit_to_string(l: &Literal) -> String {
    match l {
        Literal::Int { value } => value.to_string(),
        Literal::Bool { value } => value.to_string(),
        Literal::Str { value } => {
            // serde_json escapes for us; the result is a valid
            // JSON string literal, which is enough as our canonical form.
            serde_json::to_string(value).unwrap()
        }
        Literal::Unit => "()".to_string(),
    }
}

pub fn type_to_string(t: &Type) -> String {
    match t {
        Type::Con { name } => name.clone(),
        Type::Var { name } => name.clone(),
        Type::Fn { params, ret, effects } => {
            let p = params
                .iter()
                .map(type_to_string)
                .collect::<Vec<_>>()
                .join(", ");
            let eff = if effects.is_empty() {
                String::new()
            } else {
                format!(" !{}", effects.join(","))
            };
            format!("({p}) -> {ret}{eff}", ret = type_to_string(ret))
        }
        Type::Forall { vars, body } => {
            format!("forall {}. {}", vars.join(" "), type_to_string(body))
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn sample_module() -> Module {
        Module {
            schema: crate::SCHEMA.into(),
            name: "sample".into(),
            imports: vec![],
            defs: vec![
                Def::Fn(FnDef {
                    name: "add".into(),
                    ty: Type::Fn {
                        params: vec![Type::int(), Type::int()],
                        ret: Box::new(Type::int()),
                        effects: vec![],
                    },
                    params: vec!["a".into(), "b".into()],
                    body: Term::App {
                        callee: Box::new(Term::Var { name: "+".into() }),
                        args: vec![
                            Term::Var { name: "a".into() },
                            Term::Var { name: "b".into() },
                        ],
                    },
                    doc: None,
                }),
            ],
        }
    }

    #[test]
    fn pretty_print_does_not_panic() {
        let s = module(&sample_module());
        assert!(s.contains("(module sample"));
        assert!(s.contains("(fn add"));
        assert!(s.contains("(+ a b)"));
    }

    #[test]
    fn manifest_contains_type_and_hash() {
        let s = manifest(&sample_module());
        assert!(s.contains("add"));
        assert!(s.contains("(Int, Int) -> Int"));
    }
}