Skip to main content

hiko_syntax/
pretty.rs

1use crate::ast::*;
2use crate::intern::StringInterner;
3use std::fmt::Write;
4
5pub fn pretty_program(prog: &Program) -> String {
6    let interner = &prog.interner;
7    let mut buf = String::new();
8    for (i, decl) in prog.decls.iter().enumerate() {
9        if i > 0 {
10            buf.push('\n');
11        }
12        pretty_decl(&mut buf, decl, 0, interner);
13    }
14    buf
15}
16
17// ── Declarations ─────────────────────────────────────────────────────
18
19fn pretty_decl(buf: &mut String, decl: &Decl, indent: usize, interner: &StringInterner) {
20    match &decl.kind {
21        DeclKind::Val(pat, expr) => {
22            write_indent(buf, indent);
23            buf.push_str("val ");
24            pretty_pat(buf, pat, interner);
25            buf.push_str(" = ");
26            pretty_expr(buf, expr, indent, interner);
27        }
28        DeclKind::ValRec(name, expr) => {
29            write_indent(buf, indent);
30            write!(buf, "val rec {} = ", interner.resolve(*name)).unwrap();
31            pretty_expr(buf, expr, indent, interner);
32        }
33        DeclKind::Fun(bindings) => {
34            for (i, binding) in bindings.iter().enumerate() {
35                write_indent(buf, indent);
36                if i == 0 {
37                    buf.push_str("fun ");
38                } else {
39                    buf.push_str("and ");
40                }
41                pretty_fun_binding(buf, binding, indent, interner);
42            }
43        }
44        DeclKind::Datatype(dt) => {
45            write_indent(buf, indent);
46            buf.push_str("datatype ");
47            pretty_tyvars(buf, &dt.tyvars, interner);
48            buf.push_str(interner.resolve(dt.name));
49            buf.push_str(" =");
50            for (i, con) in dt.constructors.iter().enumerate() {
51                if i > 0 {
52                    buf.push_str(" |");
53                }
54                write!(buf, " {}", interner.resolve(con.name)).unwrap();
55                if let Some(ref ty) = con.payload {
56                    buf.push_str(" of ");
57                    pretty_type(buf, ty, interner);
58                }
59            }
60        }
61        DeclKind::TypeAlias(ta) => {
62            write_indent(buf, indent);
63            buf.push_str("type ");
64            pretty_tyvars(buf, &ta.tyvars, interner);
65            write!(buf, "{} = ", interner.resolve(ta.name)).unwrap();
66            pretty_type(buf, &ta.ty, interner);
67        }
68        DeclKind::Local(locals, body) => {
69            write_indent(buf, indent);
70            buf.push_str("local\n");
71            for d in locals {
72                pretty_decl(buf, d, indent + 2, interner);
73                buf.push('\n');
74            }
75            write_indent(buf, indent);
76            buf.push_str("in\n");
77            for d in body {
78                pretty_decl(buf, d, indent + 2, interner);
79                buf.push('\n');
80            }
81            write_indent(buf, indent);
82            buf.push_str("end");
83        }
84        DeclKind::Use(path) => {
85            write_indent(buf, indent);
86            buf.push_str("use ");
87            write_escaped_string(buf, path);
88        }
89        DeclKind::Effect(name, payload) => {
90            write_indent(buf, indent);
91            write!(buf, "effect {}", interner.resolve(*name)).unwrap();
92            if let Some(ty) = payload {
93                buf.push_str(" of ");
94                pretty_type(buf, ty, interner);
95            }
96        }
97    }
98}
99
100fn pretty_fun_binding(
101    buf: &mut String,
102    binding: &FunBinding,
103    indent: usize,
104    interner: &StringInterner,
105) {
106    for (i, clause) in binding.clauses.iter().enumerate() {
107        if i > 0 {
108            buf.push('\n');
109            write_indent(buf, indent + 2);
110            buf.push_str("| ");
111        }
112        buf.push_str(interner.resolve(binding.name));
113        for pat in &clause.pats {
114            buf.push(' ');
115            pretty_atom_pat(buf, pat, interner);
116        }
117        buf.push_str(" = ");
118        pretty_expr(buf, &clause.body, indent + 2, interner);
119    }
120}
121
122fn pretty_tyvars(buf: &mut String, tyvars: &[crate::intern::Symbol], interner: &StringInterner) {
123    match tyvars.len() {
124        0 => {}
125        1 => write!(buf, "{} ", interner.resolve(tyvars[0])).unwrap(),
126        _ => {
127            buf.push('(');
128            for (i, tv) in tyvars.iter().enumerate() {
129                if i > 0 {
130                    buf.push_str(", ");
131                }
132                buf.push_str(interner.resolve(*tv));
133            }
134            buf.push_str(") ");
135        }
136    }
137}
138
139// ── Expressions ──────────────────────────────────────────────────────
140
141fn pretty_expr(buf: &mut String, expr: &Expr, indent: usize, interner: &StringInterner) {
142    match &expr.kind {
143        ExprKind::IntLit(n) => write!(buf, "{n}").unwrap(),
144        ExprKind::FloatLit(f) => pretty_float(buf, *f),
145        ExprKind::StringLit(s) => write_escaped_string(buf, s),
146        ExprKind::CharLit(c) => write_escaped_char(buf, *c),
147        ExprKind::BoolLit(true) => buf.push_str("true"),
148        ExprKind::BoolLit(false) => buf.push_str("false"),
149        ExprKind::Unit => buf.push_str("()"),
150        ExprKind::Var(sym) => buf.push_str(interner.resolve(*sym)),
151        ExprKind::Constructor(sym) => buf.push_str(interner.resolve(*sym)),
152
153        ExprKind::Tuple(elems) => {
154            buf.push('(');
155            for (i, e) in elems.iter().enumerate() {
156                if i > 0 {
157                    buf.push_str(", ");
158                }
159                pretty_expr(buf, e, indent, interner);
160            }
161            buf.push(')');
162        }
163        ExprKind::List(elems) => {
164            buf.push('[');
165            for (i, e) in elems.iter().enumerate() {
166                if i > 0 {
167                    buf.push_str(", ");
168                }
169                pretty_expr(buf, e, indent, interner);
170            }
171            buf.push(']');
172        }
173        ExprKind::Cons(hd, tl) => {
174            pretty_cons_operand(buf, hd, indent, interner);
175            buf.push_str(" :: ");
176            pretty_expr(buf, tl, indent, interner);
177        }
178        ExprKind::BinOp(op, lhs, rhs) => {
179            let needs_parens_lhs = binop_needs_parens_lhs(op, lhs);
180            let needs_parens_rhs = binop_needs_parens_rhs(op, rhs);
181            if needs_parens_lhs {
182                buf.push('(');
183            }
184            pretty_expr(buf, lhs, indent, interner);
185            if needs_parens_lhs {
186                buf.push(')');
187            }
188            write!(buf, " {} ", binop_str(op)).unwrap();
189            if needs_parens_rhs {
190                buf.push('(');
191            }
192            pretty_expr(buf, rhs, indent, interner);
193            if needs_parens_rhs {
194                buf.push(')');
195            }
196        }
197        ExprKind::UnaryNeg(e) => {
198            buf.push('~');
199            pretty_atom_expr(buf, e, indent, interner);
200        }
201        ExprKind::Not(e) => {
202            buf.push_str("not ");
203            pretty_atom_expr(buf, e, indent, interner);
204        }
205        ExprKind::App(func, arg) => {
206            pretty_app_func(buf, func, indent, interner);
207            buf.push(' ');
208            pretty_atom_expr(buf, arg, indent, interner);
209        }
210        ExprKind::Fn(pat, body) => {
211            buf.push_str("fn ");
212            pretty_pat(buf, pat, interner);
213            buf.push_str(" => ");
214            pretty_expr(buf, body, indent, interner);
215        }
216        ExprKind::If(cond, then_br, else_br) => {
217            buf.push_str("if ");
218            pretty_expr(buf, cond, indent, interner);
219            buf.push_str(" then ");
220            pretty_expr(buf, then_br, indent, interner);
221            buf.push_str(" else ");
222            pretty_expr(buf, else_br, indent, interner);
223        }
224        ExprKind::Let(decls, body) => {
225            buf.push_str("let\n");
226            for d in decls {
227                pretty_decl(buf, d, indent + 2, interner);
228                buf.push('\n');
229            }
230            write_indent(buf, indent);
231            buf.push_str("in\n");
232            write_indent(buf, indent + 2);
233            pretty_expr(buf, body, indent + 2, interner);
234            buf.push('\n');
235            write_indent(buf, indent);
236            buf.push_str("end");
237        }
238        ExprKind::Case(scrutinee, branches) => {
239            buf.push_str("case ");
240            pretty_expr(buf, scrutinee, indent, interner);
241            buf.push_str(" of\n");
242            for (i, (pat, body)) in branches.iter().enumerate() {
243                write_indent(buf, indent + 2);
244                if i > 0 {
245                    buf.push_str("| ");
246                } else {
247                    buf.push_str("  ");
248                }
249                pretty_pat(buf, pat, interner);
250                buf.push_str(" => ");
251                pretty_expr(buf, body, indent + 4, interner);
252                if i + 1 < branches.len() {
253                    buf.push('\n');
254                }
255            }
256        }
257        ExprKind::Ann(e, ty) => {
258            pretty_expr(buf, e, indent, interner);
259            buf.push_str(" : ");
260            pretty_type(buf, ty, interner);
261        }
262        ExprKind::Paren(e) => {
263            buf.push('(');
264            pretty_expr(buf, e, indent, interner);
265            buf.push(')');
266        }
267        ExprKind::Perform(sym, arg) => {
268            write!(buf, "perform {} ", interner.resolve(*sym)).unwrap();
269            pretty_atom_expr(buf, arg, indent, interner);
270        }
271        ExprKind::Handle {
272            body,
273            return_var,
274            return_body,
275            handlers,
276        } => {
277            buf.push_str("handle\n");
278            write_indent(buf, indent + 2);
279            pretty_expr(buf, body, indent + 2, interner);
280            buf.push('\n');
281            write_indent(buf, indent);
282            buf.push_str("with\n");
283            write_indent(buf, indent + 2);
284            write!(buf, "return {} => ", interner.resolve(*return_var)).unwrap();
285            pretty_expr(buf, return_body, indent + 2, interner);
286            for handler in handlers {
287                buf.push('\n');
288                write_indent(buf, indent);
289                write!(
290                    buf,
291                    "| {} {} {} => ",
292                    interner.resolve(handler.effect_name),
293                    interner.resolve(handler.payload_var),
294                    interner.resolve(handler.cont_var),
295                )
296                .unwrap();
297                pretty_expr(buf, &handler.body, indent + 2, interner);
298            }
299        }
300        ExprKind::Resume(cont, arg) => {
301            buf.push_str("resume ");
302            pretty_atom_expr(buf, cont, indent, interner);
303            buf.push(' ');
304            pretty_atom_expr(buf, arg, indent, interner);
305        }
306    }
307}
308
309fn pretty_atom_expr(buf: &mut String, expr: &Expr, indent: usize, interner: &StringInterner) {
310    if needs_parens_as_atom(expr) {
311        buf.push('(');
312        pretty_expr(buf, expr, indent, interner);
313        buf.push(')');
314    } else {
315        pretty_expr(buf, expr, indent, interner);
316    }
317}
318
319fn pretty_app_func(buf: &mut String, expr: &Expr, indent: usize, interner: &StringInterner) {
320    match &expr.kind {
321        ExprKind::App(_, _) | ExprKind::Var(_) | ExprKind::Constructor(_) | ExprKind::Paren(_) => {
322            pretty_expr(buf, expr, indent, interner)
323        }
324        _ => {
325            buf.push('(');
326            pretty_expr(buf, expr, indent, interner);
327            buf.push(')');
328        }
329    }
330}
331
332fn pretty_cons_operand(buf: &mut String, expr: &Expr, indent: usize, interner: &StringInterner) {
333    match &expr.kind {
334        ExprKind::BinOp(BinOp::Orelse | BinOp::Andalso, _, _) | ExprKind::Ann(_, _) => {
335            buf.push('(');
336            pretty_expr(buf, expr, indent, interner);
337            buf.push(')');
338        }
339        _ => pretty_expr(buf, expr, indent, interner),
340    }
341}
342
343fn needs_parens_as_atom(expr: &Expr) -> bool {
344    !matches!(
345        &expr.kind,
346        ExprKind::IntLit(_)
347            | ExprKind::FloatLit(_)
348            | ExprKind::StringLit(_)
349            | ExprKind::CharLit(_)
350            | ExprKind::BoolLit(_)
351            | ExprKind::Unit
352            | ExprKind::Var(_)
353            | ExprKind::Constructor(_)
354            | ExprKind::Tuple(_)
355            | ExprKind::List(_)
356            | ExprKind::Paren(_)
357    )
358}
359
360// ── Operator precedence for parenthesization ─────────────────────────
361
362fn binop_prec(op: &BinOp) -> u8 {
363    match op {
364        BinOp::Orelse => 0,
365        BinOp::Andalso => 1,
366        BinOp::Eq
367        | BinOp::Ne
368        | BinOp::LtInt
369        | BinOp::GtInt
370        | BinOp::LeInt
371        | BinOp::GeInt
372        | BinOp::LtFloat
373        | BinOp::GtFloat
374        | BinOp::LeFloat
375        | BinOp::GeFloat => 2,
376        BinOp::AddInt | BinOp::SubInt | BinOp::AddFloat | BinOp::SubFloat | BinOp::ConcatStr => 4,
377        BinOp::MulInt | BinOp::DivInt | BinOp::ModInt | BinOp::MulFloat | BinOp::DivFloat => 5,
378    }
379}
380
381fn is_right_assoc(op: &BinOp) -> bool {
382    matches!(op, BinOp::Orelse | BinOp::Andalso)
383}
384
385fn binop_needs_parens_lhs(op: &BinOp, lhs: &Expr) -> bool {
386    if let ExprKind::BinOp(lhs_op, _, _) = &lhs.kind {
387        let lp = binop_prec(lhs_op);
388        let rp = binop_prec(op);
389        if is_right_assoc(op) {
390            lp <= rp
391        } else {
392            lp < rp
393        }
394    } else {
395        false
396    }
397}
398
399fn binop_needs_parens_rhs(op: &BinOp, rhs: &Expr) -> bool {
400    if let ExprKind::BinOp(rhs_op, _, _) = &rhs.kind {
401        let lp = binop_prec(rhs_op);
402        let rp = binop_prec(op);
403        if is_right_assoc(op) {
404            lp < rp
405        } else {
406            lp <= rp
407        }
408    } else {
409        false
410    }
411}
412
413fn binop_str(op: &BinOp) -> &'static str {
414    match op {
415        BinOp::AddInt => "+",
416        BinOp::SubInt => "-",
417        BinOp::MulInt => "*",
418        BinOp::DivInt => "/",
419        BinOp::ModInt => "mod",
420        BinOp::AddFloat => "+.",
421        BinOp::SubFloat => "-.",
422        BinOp::MulFloat => "*.",
423        BinOp::DivFloat => "/.",
424        BinOp::ConcatStr => "^",
425        BinOp::LtInt => "<",
426        BinOp::GtInt => ">",
427        BinOp::LeInt => "<=",
428        BinOp::GeInt => ">=",
429        BinOp::LtFloat => "<.",
430        BinOp::GtFloat => ">.",
431        BinOp::LeFloat => "<=.",
432        BinOp::GeFloat => ">=.",
433        BinOp::Eq => "=",
434        BinOp::Ne => "<>",
435        BinOp::Andalso => "andalso",
436        BinOp::Orelse => "orelse",
437    }
438}
439
440// ── Patterns ─────────────────────────────────────────────────────────
441
442fn pretty_pat(buf: &mut String, pat: &Pat, interner: &StringInterner) {
443    match &pat.kind {
444        PatKind::Wildcard => buf.push('_'),
445        PatKind::Var(sym) => buf.push_str(interner.resolve(*sym)),
446        PatKind::IntLit(n) => {
447            if *n < 0 {
448                write!(buf, "~{}", -n).unwrap();
449            } else {
450                write!(buf, "{n}").unwrap();
451            }
452        }
453        PatKind::FloatLit(f) => {
454            if *f < 0.0 {
455                buf.push('~');
456                pretty_float(buf, -f);
457            } else {
458                pretty_float(buf, *f);
459            }
460        }
461        PatKind::StringLit(s) => write_escaped_string(buf, s),
462        PatKind::CharLit(c) => write_escaped_char(buf, *c),
463        PatKind::BoolLit(true) => buf.push_str("true"),
464        PatKind::BoolLit(false) => buf.push_str("false"),
465        PatKind::Unit => buf.push_str("()"),
466        PatKind::Tuple(elems) => {
467            buf.push('(');
468            for (i, p) in elems.iter().enumerate() {
469                if i > 0 {
470                    buf.push_str(", ");
471                }
472                pretty_pat(buf, p, interner);
473            }
474            buf.push(')');
475        }
476        PatKind::Constructor(sym, None) => buf.push_str(interner.resolve(*sym)),
477        PatKind::Constructor(sym, Some(payload)) => {
478            buf.push_str(interner.resolve(*sym));
479            buf.push(' ');
480            pretty_atom_pat(buf, payload, interner);
481        }
482        PatKind::Cons(hd, tl) => {
483            pretty_atom_pat(buf, hd, interner);
484            buf.push_str(" :: ");
485            pretty_pat(buf, tl, interner);
486        }
487        PatKind::List(elems) => {
488            buf.push('[');
489            for (i, p) in elems.iter().enumerate() {
490                if i > 0 {
491                    buf.push_str(", ");
492                }
493                pretty_pat(buf, p, interner);
494            }
495            buf.push(']');
496        }
497        PatKind::Ann(p, ty) => {
498            pretty_pat(buf, p, interner);
499            buf.push_str(" : ");
500            pretty_type(buf, ty, interner);
501        }
502        PatKind::As(sym, p) => {
503            buf.push_str(interner.resolve(*sym));
504            buf.push_str(" as ");
505            pretty_pat(buf, p, interner);
506        }
507        PatKind::Paren(p) => {
508            buf.push('(');
509            pretty_pat(buf, p, interner);
510            buf.push(')');
511        }
512    }
513}
514
515fn pretty_atom_pat(buf: &mut String, pat: &Pat, interner: &StringInterner) {
516    match &pat.kind {
517        PatKind::Constructor(_, Some(_))
518        | PatKind::Cons(_, _)
519        | PatKind::Ann(_, _)
520        | PatKind::As(_, _) => {
521            buf.push('(');
522            pretty_pat(buf, pat, interner);
523            buf.push(')');
524        }
525        _ => pretty_pat(buf, pat, interner),
526    }
527}
528
529// ── Type expressions ─────────────────────────────────────────────────
530
531fn pretty_type(buf: &mut String, ty: &TypeExpr, interner: &StringInterner) {
532    match &ty.kind {
533        TypeExprKind::Named(sym) => buf.push_str(interner.resolve(*sym)),
534        TypeExprKind::Var(sym) => buf.push_str(interner.resolve(*sym)),
535        TypeExprKind::App(sym, args) => {
536            let name = interner.resolve(*sym);
537            if args.len() == 1 {
538                pretty_atom_type(buf, &args[0], interner);
539                write!(buf, " {name}").unwrap();
540            } else {
541                buf.push('(');
542                for (i, arg) in args.iter().enumerate() {
543                    if i > 0 {
544                        buf.push_str(", ");
545                    }
546                    pretty_type(buf, arg, interner);
547                }
548                write!(buf, ") {name}").unwrap();
549            }
550        }
551        TypeExprKind::Arrow(lhs, rhs) => {
552            pretty_arrow_lhs(buf, lhs, interner);
553            buf.push_str(" -> ");
554            pretty_type(buf, rhs, interner);
555        }
556        TypeExprKind::Tuple(elems) => {
557            for (i, t) in elems.iter().enumerate() {
558                if i > 0 {
559                    buf.push_str(" * ");
560                }
561                pretty_atom_type(buf, t, interner);
562            }
563        }
564        TypeExprKind::Paren(t) => {
565            buf.push('(');
566            pretty_type(buf, t, interner);
567            buf.push(')');
568        }
569    }
570}
571
572fn pretty_atom_type(buf: &mut String, ty: &TypeExpr, interner: &StringInterner) {
573    match &ty.kind {
574        TypeExprKind::Arrow(_, _) | TypeExprKind::Tuple(_) => {
575            buf.push('(');
576            pretty_type(buf, ty, interner);
577            buf.push(')');
578        }
579        _ => pretty_type(buf, ty, interner),
580    }
581}
582
583fn pretty_arrow_lhs(buf: &mut String, ty: &TypeExpr, interner: &StringInterner) {
584    match &ty.kind {
585        TypeExprKind::Arrow(_, _) => {
586            buf.push('(');
587            pretty_type(buf, ty, interner);
588            buf.push(')');
589        }
590        _ => pretty_type(buf, ty, interner),
591    }
592}
593
594// ── Utilities ────────────────────────────────────────────────────────
595
596fn write_indent(buf: &mut String, n: usize) {
597    buf.extend(std::iter::repeat_n(' ', n));
598}
599
600fn write_escaped(buf: &mut String, c: char) {
601    match c {
602        '\n' => buf.push_str("\\n"),
603        '\t' => buf.push_str("\\t"),
604        '\\' => buf.push_str("\\\\"),
605        '"' => buf.push_str("\\\""),
606        c => buf.push(c),
607    }
608}
609
610fn write_escaped_string(buf: &mut String, s: &str) {
611    buf.push('"');
612    for c in s.chars() {
613        write_escaped(buf, c);
614    }
615    buf.push('"');
616}
617
618fn write_escaped_char(buf: &mut String, c: char) {
619    buf.push_str("#\"");
620    write_escaped(buf, c);
621    buf.push('"');
622}
623
624fn pretty_float(buf: &mut String, f: f64) {
625    if f.is_nan() {
626        buf.push_str("0.0"); // NaN has no literal form; use 0.0 as placeholder
627        return;
628    }
629    if f.is_infinite() {
630        // No literal form for infinity
631        if f.is_sign_negative() {
632            buf.push_str("~1.0e308");
633        } else {
634            buf.push_str("1.0e308");
635        }
636        return;
637    }
638    let start = buf.len();
639    write!(buf, "{f}").unwrap();
640    let written = &buf[start..];
641    if !written.contains('.') && !written.contains('e') && !written.contains('E') {
642        buf.push_str(".0");
643    }
644}