1use super::{Expr, File, FileKind, FnDef, Item, Literal, Type};
7use crate::format;
8use crate::span::Spanned;
9
10pub fn format_function(func: &FnDef) -> String {
13 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 strip_synthetic_header(&formatted)
27}
28
29fn strip_synthetic_header(formatted: &str) -> String {
32 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
42pub 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
61pub 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(¶ms.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
89pub 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}