petr_ast/
pretty_print.rs

1//! Pretty-print the AST for tests and ease of development.
2
3use petr_utils::{PrettyPrint, SymbolInterner};
4
5use crate::*;
6
7impl PrettyPrint for Ast {
8    fn pretty_print(
9        &self,
10        interner: &SymbolInterner,
11        _indentation: usize,
12    ) -> String {
13        let mut buf = String::new();
14        for node in &self.modules {
15            buf.push_str(&node.pretty_print(interner, 0));
16        }
17        buf
18    }
19}
20
21impl PrettyPrint for Module {
22    fn pretty_print(
23        &self,
24        interner: &SymbolInterner,
25        indentation: usize,
26    ) -> String {
27        let mut buf = format!("{}module {} =\n", "  ".repeat(indentation), interner.get_path(&self.name).join("."));
28        for node in &self.nodes {
29            buf.push_str(&node.pretty_print(interner, 0));
30        }
31        buf
32    }
33}
34
35impl PrettyPrint for ImportStatement {
36    fn pretty_print(
37        &self,
38        interner: &SymbolInterner,
39        indentation: usize,
40    ) -> String {
41        let mut buf = format!(
42            "{}{} {}",
43            "  ".repeat(indentation),
44            if self.is_exported() { "export" } else { "import" },
45            self.path.iter().map(|id| interner.get(id.id)).collect::<Vec<_>>().join("."),
46        );
47        if let Some(alias) = self.alias {
48            buf.push_str(&format!(" as {}", interner.get(alias.id)));
49        }
50        buf.push('\n');
51        buf
52    }
53}
54
55impl PrettyPrint for AstNode {
56    fn pretty_print(
57        &self,
58        interner: &SymbolInterner,
59        indentation: usize,
60    ) -> String {
61        let mut string = match self {
62            AstNode::FunctionDeclaration(node) => node.pretty_print(interner, indentation),
63            AstNode::TypeDeclaration(ty) => ty.pretty_print(interner, indentation),
64            AstNode::ImportStatement(stmt) => stmt.pretty_print(interner, indentation),
65        };
66        let indentation_str = "  ".repeat(indentation);
67        string = format!("{indentation_str}{string}");
68        let indentation_str = format!("\n{indentation_str}");
69        string = string.replace('\n', &indentation_str);
70        string
71    }
72}
73
74impl PrettyPrint for TypeDeclaration {
75    fn pretty_print(
76        &self,
77        interner: &SymbolInterner,
78        indentation: usize,
79    ) -> String {
80        let TypeDeclaration { name, variants, visibility } = self;
81        format!(
82            "{}{}type {} =\n{}",
83            "  ".repeat(indentation),
84            if *visibility == Visibility::Exported { "exported " } else { "" },
85            name.pretty_print(interner, 0),
86            variants
87                .iter()
88                .map(|field| field.pretty_print(interner, indentation + 1))
89                .collect::<Vec<_>>()
90                .join(" |\n"),
91        )
92    }
93}
94
95impl PrettyPrint for TypeVariant {
96    fn pretty_print(
97        &self,
98        interner: &SymbolInterner,
99        indentation: usize,
100    ) -> String {
101        format!(
102            "{}{}({})",
103            "  ".repeat(indentation),
104            self.name.pretty_print(interner, 0),
105            self.fields
106                .iter()
107                .map(|field| {
108                    let name = field.item().name.pretty_print(interner, 0);
109                    let ty = field.item().ty.pretty_print(interner, 0);
110                    format!("{}: {}", name, ty)
111                })
112                .collect::<Vec<_>>()
113                .join(" ")
114        )
115    }
116}
117
118impl PrettyPrint for Ty {
119    fn pretty_print(
120        &self,
121        interner: &SymbolInterner,
122        _: usize,
123    ) -> String {
124        let name = match self {
125            Ty::Bool => "bool".to_string(),
126            Ty::Int => "int".to_string(),
127            Ty::String => "string".to_string(),
128            Ty::Unit => "unit".to_string(),
129            Ty::Named(name) => name.pretty_print(interner, 0),
130        };
131        format!("'{name}")
132    }
133}
134
135impl PrettyPrint for FunctionParameter {
136    fn pretty_print(
137        &self,
138        interner: &SymbolInterner,
139        indentation: usize,
140    ) -> String {
141        format!(
142            "{}{} ∈ {}",
143            "  ".repeat(indentation),
144            self.name.pretty_print(interner, 0),
145            self.ty.pretty_print(interner, 0)
146        )
147    }
148}
149
150impl PrettyPrint for Expression {
151    fn pretty_print(
152        &self,
153        interner: &SymbolInterner,
154        indentation: usize,
155    ) -> String {
156        match self {
157            Expression::Literal(Literal::Integer(i)) => i.to_string(),
158            Expression::Literal(Literal::Boolean(b)) => b.to_string(),
159            Expression::Literal(Literal::String(s)) => format!("\"{s}\""),
160            Expression::List(list) => list.pretty_print(interner, indentation),
161            Expression::Operator(op) => op.pretty_print(interner, indentation),
162            Expression::TypeConstructor(..) => "type constructor".to_string(),
163            Expression::FunctionCall(call) => call.pretty_print(interner, indentation),
164            Expression::Variable(v) => format!("var({})", interner.get(v.id)),
165            Expression::IntrinsicCall(call) => call.pretty_print(interner, indentation),
166            Expression::Binding(binding) => binding.pretty_print(interner, indentation + 1),
167        }
168    }
169}
170
171impl PrettyPrint for ExpressionWithBindings {
172    fn pretty_print(
173        &self,
174        interner: &SymbolInterner,
175        indentation: usize,
176    ) -> String {
177        let mut bindings = self.bindings.iter();
178        let Some(first_binding) = bindings.next() else {
179            return Default::default();
180        };
181        let mut buf = format!(
182            "\n{}let {} = {},",
183            "  ".repeat(indentation),
184            first_binding.name.pretty_print(interner, indentation + 1),
185            first_binding.val.pretty_print(interner, indentation + 1)
186        );
187        for (ix, binding) in bindings.enumerate() {
188            let is_last = ix == self.bindings.len() - 2;
189            buf.push_str(&format!(
190                "\n{}    {} = {}{}",
191                "  ".repeat(indentation),
192                binding.name.pretty_print(interner, indentation + 1),
193                binding.val.pretty_print(interner, indentation + 1),
194                if !is_last { "," } else { "" }
195            ));
196        }
197        buf.push_str(&format!("\n{}\n\n", self.expression.pretty_print(interner, indentation)));
198
199        buf
200    }
201}
202
203impl PrettyPrint for IntrinsicCall {
204    fn pretty_print(
205        &self,
206        interner: &SymbolInterner,
207        indentation: usize,
208    ) -> String {
209        format!(
210            "{}{}({})",
211            "  ".repeat(indentation),
212            self.intrinsic,
213            self.args
214                .iter()
215                .map(|arg| arg.pretty_print(interner, indentation))
216                .collect::<Vec<_>>()
217                .join(", ")
218        )
219    }
220}
221
222impl PrettyPrint for FunctionCall {
223    fn pretty_print(
224        &self,
225        interner: &SymbolInterner,
226        indentation: usize,
227    ) -> String {
228        format!(
229            "{}call {}({})",
230            "  ".repeat(indentation),
231            interner.get_path(&self.func_name).join("."),
232            self.args
233                .iter()
234                .map(|arg| arg.pretty_print(interner, indentation))
235                .collect::<Vec<_>>()
236                .join(", ")
237        )
238    }
239}
240
241impl PrettyPrint for List {
242    fn pretty_print(
243        &self,
244        interner: &SymbolInterner,
245        indentation: usize,
246    ) -> String {
247        format!(
248            "[{}]",
249            self.elements
250                .iter()
251                .map(|item| item.pretty_print(interner, indentation))
252                .collect::<Vec<_>>()
253                .join(", ")
254        )
255    }
256}
257
258impl PrettyPrint for OperatorExpression {
259    fn pretty_print(
260        &self,
261        interner: &SymbolInterner,
262        indentation: usize,
263    ) -> String {
264        let lhs = self.lhs.pretty_print(interner, indentation);
265        let rhs = self.rhs.pretty_print(interner, indentation);
266        let op = match self.op.item() {
267            Operator::Plus => "+",
268            Operator::Minus => "-",
269            Operator::Star => "*",
270            Operator::Slash => "/",
271        };
272        format!("{}({} {})", op, lhs, rhs)
273    }
274}
275
276impl PrettyPrint for FunctionDeclaration {
277    fn pretty_print(
278        &self,
279        interner: &SymbolInterner,
280        indentation: usize,
281    ) -> String {
282        let FunctionDeclaration {
283            name,
284            parameters,
285            return_type,
286            body,
287            visibility,
288        } = self;
289        format!(
290            "{}{}Func {}({}{}{}) -> {} {}\n",
291            "  ".repeat(indentation),
292            if *visibility == Visibility::Exported { "exported " } else { "" },
293            name.pretty_print(interner, 0),
294            if parameters.is_empty() { "" } else { "\n" },
295            parameters
296                .iter()
297                .map(|param| param.pretty_print(interner, indentation + 1))
298                .collect::<Vec<_>>()
299                .join(",\n"),
300            if parameters.is_empty() { "" } else { "\n" },
301            return_type.pretty_print(interner, indentation),
302            body.pretty_print(interner, indentation)
303        )
304    }
305}