Skip to main content

lex_syntax/
printer.rs

1//! Pretty-printer for the syntax tree. Designed so
2//! `parse(text) → print → parse` round-trips on the canonical AST.
3
4use crate::syntax::*;
5use std::fmt::Write;
6
7pub fn print_program(program: &Program) -> String {
8    let mut p = Printer::new();
9    p.program(program);
10    p.out
11}
12
13struct Printer {
14    out: String,
15    indent: usize,
16}
17
18impl Printer {
19    fn new() -> Self { Self { out: String::new(), indent: 0 } }
20
21    fn write_indent(&mut self) {
22        for _ in 0..self.indent {
23            self.out.push_str("  ");
24        }
25    }
26
27    fn nl(&mut self) {
28        self.out.push('\n');
29    }
30
31    fn program(&mut self, p: &Program) {
32        for (i, item) in p.items.iter().enumerate() {
33            if i > 0 { self.nl(); }
34            self.item(item);
35        }
36        if !p.items.is_empty() {
37            self.nl();
38        }
39    }
40
41    fn item(&mut self, item: &Item) {
42        match item {
43            Item::Import(i) => {
44                writeln!(self.out, "import \"{}\" as {}", i.reference, i.alias).unwrap();
45            }
46            Item::TypeDecl(td) => self.type_decl(td),
47            Item::FnDecl(fd) => self.fn_decl(fd),
48        }
49    }
50
51    fn type_decl(&mut self, td: &TypeDecl) {
52        write!(self.out, "type {}", td.name).unwrap();
53        if !td.params.is_empty() {
54            write!(self.out, "[{}]", td.params.join(", ")).unwrap();
55        }
56        write!(self.out, " = ").unwrap();
57        self.type_expr(&td.definition);
58        self.nl();
59    }
60
61    fn fn_decl(&mut self, fd: &FnDecl) {
62        write!(self.out, "fn {}", fd.name).unwrap();
63        if !fd.type_params.is_empty() {
64            write!(self.out, "[{}]", fd.type_params.join(", ")).unwrap();
65        }
66        write!(self.out, "(").unwrap();
67        for (i, p) in fd.params.iter().enumerate() {
68            if i > 0 { write!(self.out, ", ").unwrap(); }
69            write!(self.out, "{} :: ", p.name).unwrap();
70            self.type_expr(&p.ty);
71        }
72        write!(self.out, ") -> ").unwrap();
73        self.effects(&fd.effects);
74        self.type_expr(&fd.return_type);
75        if fd.examples.is_empty() {
76            write!(self.out, " ").unwrap();
77        } else {
78            self.examples(&fd.name, &fd.examples);
79            self.nl();
80        }
81        self.block(&fd.body);
82        self.nl();
83    }
84
85    fn examples(&mut self, fn_name: &str, examples: &[Example]) {
86        if examples.is_empty() { return; }
87        self.nl();
88        write!(self.out, "  examples {{").unwrap();
89        self.indent += 2;
90        for (i, ex) in examples.iter().enumerate() {
91            self.nl();
92            self.write_indent();
93            write!(self.out, "{}(", fn_name).unwrap();
94            for (j, a) in ex.args.iter().enumerate() {
95                if j > 0 { write!(self.out, ", ").unwrap(); }
96                self.expr(a);
97            }
98            write!(self.out, ") => ").unwrap();
99            self.expr(&ex.expected);
100            if i + 1 < examples.len() {
101                write!(self.out, ",").unwrap();
102            }
103        }
104        self.indent -= 2;
105        self.nl();
106        write!(self.out, "  }}").unwrap();
107    }
108
109    fn effects(&mut self, effects: &[Effect]) {
110        if effects.is_empty() { return; }
111        write!(self.out, "[").unwrap();
112        for (i, e) in effects.iter().enumerate() {
113            if i > 0 { write!(self.out, ", ").unwrap(); }
114            write!(self.out, "{}", e.name).unwrap();
115            if let Some(arg) = &e.arg {
116                match arg {
117                    EffectArg::Str(s) => write!(self.out, "(\"{}\")", s).unwrap(),
118                    EffectArg::Int(n) => write!(self.out, "({})", n).unwrap(),
119                    EffectArg::Ident(s) => write!(self.out, "({})", s).unwrap(),
120                }
121            }
122        }
123        write!(self.out, "] ").unwrap();
124    }
125
126    fn type_expr(&mut self, t: &TypeExpr) {
127        match t {
128            TypeExpr::Named { name, args } => {
129                write!(self.out, "{}", name).unwrap();
130                if !args.is_empty() {
131                    write!(self.out, "[").unwrap();
132                    for (i, a) in args.iter().enumerate() {
133                        if i > 0 { write!(self.out, ", ").unwrap(); }
134                        self.type_expr(a);
135                    }
136                    write!(self.out, "]").unwrap();
137                }
138            }
139            TypeExpr::Record(fs) => {
140                write!(self.out, "{{ ").unwrap();
141                for (i, f) in fs.iter().enumerate() {
142                    if i > 0 { write!(self.out, ", ").unwrap(); }
143                    write!(self.out, "{} :: ", f.name).unwrap();
144                    self.type_expr(&f.ty);
145                }
146                write!(self.out, " }}").unwrap();
147            }
148            TypeExpr::RecordWithSpreads { spreads, fields } => {
149                write!(self.out, "{{ ").unwrap();
150                for (i, s) in spreads.iter().enumerate() {
151                    if i > 0 { write!(self.out, ", ").unwrap(); }
152                    write!(self.out, "...{}", s).unwrap();
153                }
154                for f in fields {
155                    write!(self.out, ", {} :: ", f.name).unwrap();
156                    self.type_expr(&f.ty);
157                }
158                write!(self.out, " }}").unwrap();
159            }
160            TypeExpr::Tuple(items) => {
161                write!(self.out, "(").unwrap();
162                for (i, it) in items.iter().enumerate() {
163                    if i > 0 { write!(self.out, ", ").unwrap(); }
164                    self.type_expr(it);
165                }
166                write!(self.out, ")").unwrap();
167            }
168            TypeExpr::Function { params, effects, ret } => {
169                write!(self.out, "(").unwrap();
170                for (i, p) in params.iter().enumerate() {
171                    if i > 0 { write!(self.out, ", ").unwrap(); }
172                    self.type_expr(p);
173                }
174                write!(self.out, ") -> ").unwrap();
175                self.effects(effects);
176                self.type_expr(ret);
177            }
178            TypeExpr::Union(variants) => {
179                for (i, v) in variants.iter().enumerate() {
180                    if i > 0 { write!(self.out, " | ").unwrap(); }
181                    write!(self.out, "{}", v.name).unwrap();
182                    if let Some(payload) = &v.payload {
183                        write!(self.out, "(").unwrap();
184                        self.type_expr(payload);
185                        write!(self.out, ")").unwrap();
186                    }
187                }
188            }
189            TypeExpr::Refined { base, binding, predicate } => {
190                self.type_expr(base);
191                write!(self.out, "{{{} | ", binding).unwrap();
192                self.expr(predicate);
193                write!(self.out, "}}").unwrap();
194            }
195        }
196    }
197
198    fn block(&mut self, b: &Block) {
199        write!(self.out, "{{").unwrap();
200        self.indent += 1;
201        for stmt in &b.statements {
202            self.nl();
203            self.write_indent();
204            self.statement(stmt);
205        }
206        self.nl();
207        self.write_indent();
208        self.expr(&b.result);
209        self.indent -= 1;
210        self.nl();
211        self.write_indent();
212        write!(self.out, "}}").unwrap();
213    }
214
215    fn statement(&mut self, s: &Statement) {
216        match s {
217            Statement::Let { name, ty, value } => {
218                write!(self.out, "let {}", name).unwrap();
219                if let Some(ty) = ty {
220                    write!(self.out, " :: ").unwrap();
221                    self.type_expr(ty);
222                }
223                write!(self.out, " := ").unwrap();
224                self.expr(value);
225            }
226            Statement::Expr(e) => self.expr(e),
227        }
228    }
229
230    fn expr(&mut self, e: &Expr) {
231        self.expr_prec(e, 0);
232    }
233
234    fn expr_prec(&mut self, e: &Expr, parent_prec: u8) {
235        match e {
236            Expr::Lit(l) => self.literal(l),
237            Expr::Var(n) => { write!(self.out, "{}", n).unwrap(); }
238            Expr::Block(b) => self.block(b),
239            Expr::Call { callee, args } => {
240                self.expr_prec(callee, 100);
241                write!(self.out, "(").unwrap();
242                for (i, a) in args.iter().enumerate() {
243                    if i > 0 { write!(self.out, ", ").unwrap(); }
244                    self.expr(a);
245                }
246                write!(self.out, ")").unwrap();
247            }
248            Expr::Pipe { left, right } => {
249                if parent_prec > 0 { write!(self.out, "(").unwrap(); }
250                self.expr_prec(left, 1);
251                write!(self.out, " |> ").unwrap();
252                self.expr_prec(right, 1);
253                if parent_prec > 0 { write!(self.out, ")").unwrap(); }
254            }
255            Expr::Try(inner) => {
256                self.expr_prec(inner, 100);
257                write!(self.out, "?").unwrap();
258            }
259            Expr::Field { value, field } => {
260                self.expr_prec(value, 100);
261                write!(self.out, ".{}", field).unwrap();
262            }
263            Expr::BinOp { op, lhs, rhs } => {
264                let prec = op.precedence() + 10;
265                if parent_prec > prec { write!(self.out, "(").unwrap(); }
266                self.expr_prec(lhs, prec);
267                write!(self.out, " {} ", op.as_str()).unwrap();
268                self.expr_prec(rhs, prec + 1);
269                if parent_prec > prec { write!(self.out, ")").unwrap(); }
270            }
271            Expr::UnaryOp { op, expr } => {
272                let s = match op { UnaryOp::Neg => "-", UnaryOp::Not => "not " };
273                write!(self.out, "{}", s).unwrap();
274                self.expr_prec(expr, 100);
275            }
276            Expr::If { cond, then_block, else_block } => {
277                write!(self.out, "if ").unwrap();
278                self.expr(cond);
279                write!(self.out, " ").unwrap();
280                self.block(then_block);
281                write!(self.out, " else ").unwrap();
282                self.block(else_block);
283            }
284            Expr::Match { scrutinee, arms } => {
285                write!(self.out, "match ").unwrap();
286                self.expr(scrutinee);
287                write!(self.out, " {{").unwrap();
288                self.indent += 1;
289                for arm in arms {
290                    self.nl();
291                    self.write_indent();
292                    self.pattern(&arm.pattern);
293                    write!(self.out, " => ").unwrap();
294                    self.expr(&arm.body);
295                    write!(self.out, ",").unwrap();
296                }
297                self.indent -= 1;
298                self.nl();
299                self.write_indent();
300                write!(self.out, "}}").unwrap();
301            }
302            Expr::RecordLit(fields) => {
303                write!(self.out, "{{ ").unwrap();
304                for (i, f) in fields.iter().enumerate() {
305                    if i > 0 { write!(self.out, ", ").unwrap(); }
306                    write!(self.out, "{}: ", f.name).unwrap();
307                    self.expr(&f.value);
308                }
309                write!(self.out, " }}").unwrap();
310            }
311            Expr::TupleLit(items) => {
312                write!(self.out, "(").unwrap();
313                for (i, it) in items.iter().enumerate() {
314                    if i > 0 { write!(self.out, ", ").unwrap(); }
315                    self.expr(it);
316                }
317                write!(self.out, ")").unwrap();
318            }
319            Expr::ListLit(items) => {
320                write!(self.out, "[").unwrap();
321                for (i, it) in items.iter().enumerate() {
322                    if i > 0 { write!(self.out, ", ").unwrap(); }
323                    self.expr(it);
324                }
325                write!(self.out, "]").unwrap();
326            }
327            Expr::Constructor { name, args } => {
328                write!(self.out, "{}", name).unwrap();
329                if !args.is_empty() {
330                    write!(self.out, "(").unwrap();
331                    for (i, a) in args.iter().enumerate() {
332                        if i > 0 { write!(self.out, ", ").unwrap(); }
333                        self.expr(a);
334                    }
335                    write!(self.out, ")").unwrap();
336                }
337            }
338            Expr::Ascription { value, ty } => {
339                write!(self.out, "(").unwrap();
340                self.expr(value);
341                write!(self.out, " :: ").unwrap();
342                self.type_expr(ty);
343                write!(self.out, ")").unwrap();
344            }
345            Expr::Lambda(l) => {
346                write!(self.out, "fn (").unwrap();
347                for (i, p) in l.params.iter().enumerate() {
348                    if i > 0 { write!(self.out, ", ").unwrap(); }
349                    write!(self.out, "{} :: ", p.name).unwrap();
350                    self.type_expr(&p.ty);
351                }
352                write!(self.out, ") -> ").unwrap();
353                self.effects(&l.effects);
354                self.type_expr(&l.return_type);
355                write!(self.out, " ").unwrap();
356                self.block(&l.body);
357            }
358        }
359    }
360
361    fn literal(&mut self, l: &Literal) {
362        match l {
363            Literal::Int(n) => write!(self.out, "{}", n).unwrap(),
364            Literal::Float(n) => write!(self.out, "{}", format_float(*n)).unwrap(),
365            Literal::Str(s) => write!(self.out, "\"{}\"", escape(s)).unwrap(),
366            Literal::Bytes(b) => {
367                write!(self.out, "b\"").unwrap();
368                for &c in b {
369                    if c.is_ascii() && (c as char).is_ascii_graphic() && c != b'"' && c != b'\\' {
370                        self.out.push(c as char);
371                    } else {
372                        write!(self.out, "\\x{:02x}", c).unwrap();
373                    }
374                }
375                write!(self.out, "\"").unwrap();
376            }
377            Literal::Bool(b) => write!(self.out, "{}", b).unwrap(),
378            Literal::Unit => write!(self.out, "()").unwrap(),
379        }
380    }
381
382    fn pattern(&mut self, p: &Pattern) {
383        match p {
384            Pattern::Lit(l) => self.literal(l),
385            Pattern::Var(n) => { write!(self.out, "{}", n).unwrap(); }
386            Pattern::Wild => { write!(self.out, "_").unwrap(); }
387            Pattern::Constructor { name, args } => {
388                write!(self.out, "{}", name).unwrap();
389                if !args.is_empty() {
390                    write!(self.out, "(").unwrap();
391                    for (i, a) in args.iter().enumerate() {
392                        if i > 0 { write!(self.out, ", ").unwrap(); }
393                        self.pattern(a);
394                    }
395                    write!(self.out, ")").unwrap();
396                }
397            }
398            Pattern::Record { fields, rest: _ } => {
399                write!(self.out, "{{ ").unwrap();
400                for (i, f) in fields.iter().enumerate() {
401                    if i > 0 { write!(self.out, ", ").unwrap(); }
402                    write!(self.out, "{}", f.name).unwrap();
403                    if let Some(p) = &f.pattern {
404                        write!(self.out, ": ").unwrap();
405                        self.pattern(p);
406                    }
407                }
408                write!(self.out, " }}").unwrap();
409            }
410            Pattern::Tuple(items) => {
411                write!(self.out, "(").unwrap();
412                for (i, it) in items.iter().enumerate() {
413                    if i > 0 { write!(self.out, ", ").unwrap(); }
414                    self.pattern(it);
415                }
416                write!(self.out, ")").unwrap();
417            }
418        }
419    }
420}
421
422fn escape(s: &str) -> String {
423    let mut out = String::with_capacity(s.len());
424    for c in s.chars() {
425        match c {
426            '\\' => out.push_str("\\\\"),
427            '"' => out.push_str("\\\""),
428            '\n' => out.push_str("\\n"),
429            '\t' => out.push_str("\\t"),
430            '\r' => out.push_str("\\r"),
431            c => out.push(c),
432        }
433    }
434    out
435}
436
437fn format_float(n: f64) -> String {
438    if n.is_finite() && n == n.trunc() {
439        format!("{:.1}", n)
440    } else {
441        format!("{}", n)
442    }
443}