Skip to main content

trident/ast/
display.rs

1//! Pretty-printing utilities for AST nodes.
2//!
3//! This module is the single source of truth for converting AST types,
4//! function signatures, and constant values to display strings.
5
6use super::{Expr, File, FileKind, FnDef, Item, Literal, Type};
7use crate::format;
8use crate::span::Spanned;
9
10/// Pretty-print a single function definition by wrapping it in a
11/// minimal synthetic `File` and running the canonical formatter.
12pub fn format_function(func: &FnDef) -> String {
13    // Build a minimal File containing only this function.
14    let file = File {
15        kind: FileKind::Program,
16        name: Spanned::dummy("_view".to_string()),
17        uses: Vec::new(),
18        declarations: Vec::new(),
19        items: vec![Spanned::dummy(Item::Fn(func.clone()))],
20    };
21
22    let formatted = format::format_file(&file, &[]);
23
24    // The formatter emits "program _view\n\n<fn>\n".
25    // Strip the synthetic header to isolate the function text.
26    strip_synthetic_header(&formatted)
27}
28
29/// Strip the synthetic "program _view\n\n" header produced by the
30/// formatter when we wrap a single function in a dummy File.
31fn strip_synthetic_header(formatted: &str) -> String {
32    // The formatter produces: "program _view\n\n<items>\n"
33    // Find the first blank line and take everything after it.
34    if let Some(pos) = formatted.find("\n\n") {
35        let rest = &formatted[pos + 2..];
36        rest.to_string()
37    } else {
38        formatted.to_string()
39    }
40}
41
42// ─── Canonical formatting helpers ──────────────────────────────────
43
44/// Format an AST type for display (documentation, diagnostics, hover).
45pub fn format_ast_type(ty: &Type) -> String {
46    match ty {
47        Type::Field => "Field".to_string(),
48        Type::XField => "XField".to_string(),
49        Type::Bool => "Bool".to_string(),
50        Type::U32 => "U32".to_string(),
51        Type::Digest => "Digest".to_string(),
52        Type::Array(inner, size) => format!("[{}; {}]", format_ast_type(inner), size),
53        Type::Tuple(elems) => {
54            let parts: Vec<_> = elems.iter().map(format_ast_type).collect();
55            format!("({})", parts.join(", "))
56        }
57        Type::Named(path) => path.as_dotted(),
58    }
59}
60
61/// Format a function signature for display (documentation, diagnostics).
62///
63/// Includes type parameters, parameter names and types, and return type.
64pub fn format_fn_signature(func: &FnDef) -> String {
65    let mut sig = String::from("fn ");
66    sig.push_str(&func.name.node);
67
68    if !func.type_params.is_empty() {
69        let params: Vec<_> = func.type_params.iter().map(|p| p.node.clone()).collect();
70        sig.push_str(&format!("<{}>", params.join(", ")));
71    }
72
73    sig.push('(');
74    let params: Vec<String> = func
75        .params
76        .iter()
77        .map(|p| format!("{}: {}", p.name.node, format_ast_type(&p.ty.node)))
78        .collect();
79    sig.push_str(&params.join(", "));
80    sig.push(')');
81
82    if let Some(ref ret) = func.return_ty {
83        sig.push_str(&format!(" -> {}", format_ast_type(&ret.node)));
84    }
85
86    sig
87}
88
89/// Format a constant value expression for display (documentation).
90pub fn format_const_value(expr: &Expr) -> String {
91    match expr {
92        Expr::Literal(Literal::Integer(n)) => n.to_string(),
93        Expr::Literal(Literal::Bool(b)) => b.to_string(),
94        _ => "...".to_string(),
95    }
96}
97
98#[cfg(test)]
99mod tests {
100    use super::*;
101    use crate::ast::navigate::find_function;
102
103    fn parse_file(source: &str) -> File {
104        crate::parse_source_silent(source, "test.tri").unwrap()
105    }
106
107    #[test]
108    fn test_format_function_produces_valid_source() {
109        let source = "program test\n\nfn add(a: Field, b: Field) -> Field {\n    a + b\n}\n";
110        let file = parse_file(source);
111        let func = find_function(&file, "add").expect("add function should exist");
112        let formatted = format_function(func);
113
114        assert!(formatted.contains("fn add("));
115        assert!(formatted.contains("a: Field, b: Field"));
116        assert!(formatted.contains("-> Field"));
117        assert!(formatted.contains("a + b"));
118    }
119
120    #[test]
121    fn test_format_function_with_annotations() {
122        let source = "program test\n\n#[requires(a + b < 1000)]\n#[ensures(result == a + b)]\nfn add(a: Field, b: Field) -> Field {\n    a + b\n}\n";
123        let file = parse_file(source);
124        let func = find_function(&file, "add").expect("add function should exist");
125        let formatted = format_function(func);
126
127        assert!(formatted.contains("#[requires("));
128        assert!(formatted.contains("#[ensures("));
129        assert!(formatted.contains("fn add("));
130    }
131
132    #[test]
133    fn test_format_function_pub() {
134        let source = "module test\n\npub fn helper(x: Field) -> Field {\n    x + 1\n}\n";
135        let file = parse_file(source);
136        let func = find_function(&file, "helper").expect("helper function should exist");
137        let formatted = format_function(func);
138
139        assert!(formatted.contains("pub fn helper("));
140    }
141}