Checkpoint

src/ir/ast.rs

@@ -1,6 +1,7 @@
 use crate::{
     eval::PrimitiveType,
     syntax::{self, ConstantType, Location},
+    util::pretty::{pretty_comma_separated, PrettySymbol},
 };
 use internment::ArcIntern;
 use pretty::{BoxAllocator, DocAllocator, Pretty};
@@ -87,20 +88,33 @@ where
 {
     fn pretty(self, allocator: &'a D) -> pretty::DocBuilder<'a, D, A> {
         match self {
-            TopLevel::Function(name, args, stmts, body) => allocator
-                .text("function")
-                .append(allocator.space())
-                .append(allocator.text(name.as_ref().to_string()))
-                .append(
-                    allocator
-                        .intersperse(
-                            args.iter().map(|x| allocator.text(x.as_ref().to_string())),
-                            ", ",
+            TopLevel::Function(name, args, stmts, expr) => {
+                let base = allocator
+                    .text("function")
+                    .append(allocator.space())
+                    .append(allocator.text(name.as_ref().to_string()))
+                    .append(allocator.space())
+                    .append(
+                        pretty_comma_separated(
+                            allocator,
+                            &args.iter().map(PrettySymbol::from).collect(),
                         )
                         .parens(),
-                )
-                .append(allocator.space())
-                .append(body.pretty(allocator)),
+                    )
+                    .append(allocator.space());
+
+                let mut body = allocator.nil();
+                for stmt in stmts {
+                    body = body
+                        .append(stmt.pretty(allocator))
+                        .append(allocator.text(";"))
+                        .append(allocator.hardline());
+                }
+                body = body.append(expr.pretty(allocator));
+                body = body.append(allocator.hardline());
+                body = body.braces();
+                base.append(body)
+            }
 
             TopLevel::Statement(stmt) => stmt.pretty(allocator),
         }
@@ -389,6 +403,14 @@ where
     fn pretty(self, allocator: &'a D) -> pretty::DocBuilder<'a, D, A> {
         match self {
             Type::Primitive(pt) => allocator.text(format!("{}", pt)),
+            Type::Function(args, rettype) => {
+                pretty_comma_separated(allocator, &args.iter().collect())
+                    .parens()
+                    .append(allocator.space())
+                    .append(allocator.text("->"))
+                    .append(allocator.space())
+                    .append(rettype.pretty(allocator))
+            }
         }
     }
 }
@@ -397,6 +419,18 @@ impl fmt::Display for Type {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         match self {
             Type::Primitive(pt) => pt.fmt(f),
+            Type::Function(args, ret) => {
+                write!(f, "(")?;
+                let mut argiter = args.iter().peekable();
+                while let Some(arg) = argiter.next() {
+                    arg.fmt(f)?;
+                    if argiter.peek().is_some() {
+                        write!(f, ",")?;
+                    }
+                }
+                write!(f, "->")?;
+                ret.fmt(f)
+            }
         }
     }
 }