From de51967e4fc53d78966daaef3c99aabbfcb723e4 Mon Sep 17 00:00:00 2001
From: Cassie Jones <code@witchoflight.com>
Date: Thu, 27 Feb 2020 19:26:23 +0100
Subject: [PATCH] Add a pretty-printer for the AST

This uses the pretty-crate for wadler-style pretty-printers. I'm
printing them with a pretty syntax and not with s-expressions, because
it's easier to print pretty than it is to parse pretty.
---
 Cargo.lock  | 23 ++++++++++++++++++
 Cargo.toml  |  1 +
 src/ast.rs  | 67 ++++++++++++++++++++++++++++++++++++++++++++++++++++-
 src/main.rs |  6 +++--
 4 files changed, 94 insertions(+), 3 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index fb1b2ed..1037821 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1,5 +1,11 @@
 # This file is automatically @generated by Cargo.
 # It is not intended for manual editing.
+[[package]]
+name = "arrayvec"
+version = "0.5.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "cff77d8686867eceff3105329d4698d96c2391c176d5d03adc90c7389162b5b8"
+
 [[package]]
 name = "ess"
 version = "0.4.2"
@@ -11,4 +17,21 @@ name = "ivy"
 version = "0.1.0"
 dependencies = [
  "ess",
+ "pretty",
 ]
+
+[[package]]
+name = "pretty"
+version = "0.9.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1891757f8427ce41706957fde1fec1b86aee3e335bd50320257705061507a24c"
+dependencies = [
+ "arrayvec",
+ "typed-arena",
+]
+
+[[package]]
+name = "typed-arena"
+version = "2.0.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0685c84d5d54d1c26f7d3eb96cd41550adb97baed141a761cf335d3d33bcd0ae"
diff --git a/Cargo.toml b/Cargo.toml
index e3a65aa..d96e26a 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -8,3 +8,4 @@ edition = "2018"
 
 [dependencies]
 ess = "0.4.2"
+pretty = "0.9"
diff --git a/src/ast.rs b/src/ast.rs
index 4a0fe8f..29c0562 100644
--- a/src/ast.rs
+++ b/src/ast.rs
@@ -1,4 +1,5 @@
 use ess::{self, Sexp};
+use pretty::RcDoc;
 use std::{
     cell::RefCell,
     collections::HashMap,
@@ -20,6 +21,15 @@ impl Name {
     }
 }
 
+impl Name {
+    fn to_doc(&self) -> RcDoc {
+        match self {
+            Name::Local(n) => RcDoc::text(format!("x{}", n)),
+            Name::Global(n) => RcDoc::text(format!("g{}", n)),
+        }
+    }
+}
+
 #[derive(PartialEq, Clone, Debug)]
 pub enum Ast {
     Let(Vec<(Name, Ast)>, Box<Ast>),
@@ -33,6 +43,58 @@ pub enum Ast {
     Num(i64),
 }
 
+impl Ast {
+    pub fn to_doc(&self) -> RcDoc {
+        match self {
+            Ast::Let(binds, body) => RcDoc::text("let")
+                .append(
+                    RcDoc::line()
+                        .append(RcDoc::intersperse(
+                            binds.iter().map(|(n, v)| {
+                                n.to_doc()
+                                    .append(RcDoc::space())
+                                    .append(RcDoc::text("="))
+                                    .append(RcDoc::line().append(v.to_doc()).nest(2).group())
+                            }),
+                            RcDoc::text(";").append(RcDoc::line()),
+                        ))
+                        .nest(2),
+                )
+                .append(RcDoc::line())
+                .append(RcDoc::text("in"))
+                .append(RcDoc::line())
+                .append(body.to_doc()),
+            Ast::Lam { binds, body, .. } => RcDoc::text("lam").append(
+                RcDoc::line()
+                    .append(RcDoc::text("("))
+                    .append(
+                        RcDoc::intersperse(binds.iter().map(|x| x.to_doc()), RcDoc::line())
+                            .nest(2)
+                            .group(),
+                    )
+                    .append(RcDoc::text(") =>"))
+                    .append(RcDoc::line())
+                    .append(body.to_doc())
+                    .nest(2)
+                    .group(),
+            ),
+            Ast::App(terms) => RcDoc::intersperse(
+                terms.iter().map(|x| match x {
+                    Ast::App(..) | Ast::Lam { .. } => {
+                        RcDoc::text("(").append(x.to_doc()).append(RcDoc::text(")"))
+                    }
+                    _ => x.to_doc(),
+                }),
+                RcDoc::line(),
+            )
+            .nest(2)
+            .group(),
+            Ast::Var(name) => name.to_doc(),
+            Ast::Num(n) => RcDoc::as_string(n),
+        }
+    }
+}
+
 struct Env<'e> {
     fresh: AtomicU32,
     vars: HashMap<String, Name>,
@@ -109,7 +171,10 @@ impl Drop for Env<'_> {
 }
 
 impl Ast {
-    pub fn parse(builtins: &[impl AsRef<str>], sexp: &Sexp) -> Result<(Ast, HashMap<String, Name>), &'static str> {
+    pub fn parse(
+        builtins: &[impl AsRef<str>],
+        sexp: &Sexp,
+    ) -> Result<(Ast, HashMap<String, Name>), &'static str> {
         let mut env = Env::new();
         let mut builtin_map = HashMap::new();
         for builtin in builtins {
diff --git a/src/main.rs b/src/main.rs
index 6099fb6..6ab7bb1 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -12,7 +12,9 @@ fn parse(input: &str) -> Result<Sexp, ParseError> {
     Ok(sexp)
 }
 
-static BUILTINS: &[&str] = &["true", "false", "<", ">", "<=", ">=", "==", "!=", "+", "-", "*", "/", "%"];
+static BUILTINS: &[&str] = &[
+    "true", "false", "<", ">", "<=", ">=", "==", "!=", "+", "-", "*", "/", "%",
+];
 
 fn main() -> io::Result<()> {
     let stdin = std::io::stdin();
@@ -33,6 +35,6 @@ fn main() -> io::Result<()> {
         }
     };
     println!("builtins: {:?}", builtins);
-    println!("{:#?}", ast);
+    println!("{}", ast.to_doc().pretty(80));
     Ok(())
 }
-- 
2.43.2