Make code generic over allocator

This way, I can use the same code with any allocator I want. More
importantly, I can also use the nice DocBuilder-exclusive helper
functions!

src/main.rs

@@ -1,6 +1,7 @@
 use std::fs;
 use std::path::PathBuf;
 
+use ::pretty::{Pretty, RcAllocator};
 use chumsky::Parser as _;
 use clap::Parser;
 
@@ -48,7 +49,7 @@ fn main() -> anyhow::Result<()> {
             match parser::parser().parse(stream) {
                 Ok(program) => {
                     let mut out = vec![];
-                    program.to_doc().render(100, &mut out)?;
+                    program.pretty(&RcAllocator).render(100, &mut out)?;
                     println!("{}", String::from_utf8(out)?);
                 }
                 Err(errs) => {

src/pretty/basic.rs

@@ -1,37 +1,36 @@
-use pretty::RcDoc;
+use pretty::{DocAllocator, DocBuilder};
 
 use crate::ast::Separated;
 
 impl<E, S1, S2> Separated<E, S1, S2> {
-    pub fn to_doc<FE, FS1, FS2>(
-        &self,
+    pub fn pretty<'a, D, FE, FS1, FS2>(
+        self,
+        allocator: &'a D,
         elem_to_doc: FE,
         separator_to_doc: FS1,
         trailing_separator_to_doc: FS2,
-    ) -> RcDoc
+    ) -> DocBuilder<'a, D>
     where
-        FE: Fn(&E) -> RcDoc,
-        FS1: Fn(&S1) -> RcDoc,
-        FS2: Fn(&S2) -> RcDoc,
+        D: DocAllocator<'a>,
+        FE: Fn(&'a D, E) -> DocBuilder<'a, D>,
+        FS1: Fn(&'a D, S1) -> DocBuilder<'a, D>,
+        FS2: Fn(&'a D, S2) -> DocBuilder<'a, D>,
     {
         match self {
-            Separated::Empty(_) => RcDoc::nil(),
+            Separated::Empty(_) => allocator.nil(),
             Separated::NonEmpty {
                 first_elem,
                 last_elems,
                 trailing,
                 span: _span,
-            } => elem_to_doc(first_elem)
-                .append(RcDoc::concat(
-                    last_elems
-                        .iter()
-                        .map(|(s, e)| separator_to_doc(s).append(elem_to_doc(e))),
-                ))
+            } => elem_to_doc(allocator, first_elem)
+                .append(allocator.concat(last_elems.into_iter().map(|(s, e)| {
+                    separator_to_doc(allocator, s).append(elem_to_doc(allocator, e))
+                })))
                 .append(
                     trailing
-                        .as_ref()
-                        .map(trailing_separator_to_doc)
-                        .unwrap_or_else(RcDoc::nil),
+                        .map(|s| trailing_separator_to_doc(allocator, s))
+                        .unwrap_or_else(|| allocator.nil()),
                 ),
         }
     }

src/pretty/expr.rs

@@ -1,23 +1,23 @@
-use pretty::RcDoc;
+use pretty::{DocAllocator, DocBuilder, Pretty};
 
 use crate::ast::Expr;
 
-impl Expr {
-    pub fn to_doc(&self) -> RcDoc {
+impl<'a, D: DocAllocator<'a>> Pretty<'a, D> for Expr {
+    fn pretty(self, allocator: &'a D) -> DocBuilder<'a, D> {
         match self {
-            Expr::Lit(lit) => RcDoc::text("<lit>"),
-            Expr::Call(call) => RcDoc::text("<call>"),
-            Expr::Field(field) => RcDoc::text("<field>"),
-            Expr::Var(var) => RcDoc::text("<var>"),
-            Expr::TableConstr(constr) => RcDoc::text("<onstr>"),
-            Expr::TableDestr(destr) => RcDoc::text("<destr>"),
-            Expr::FuncDef(def) => RcDoc::text("<def>"),
+            Expr::Lit(lit) => allocator.text("<lit>"),
+            Expr::Call(call) => allocator.text("<call>"),
+            Expr::Field(field) => allocator.text("<field>"),
+            Expr::Var(var) => allocator.text("<var>"),
+            Expr::TableConstr(constr) => allocator.text("<onstr>"),
+            Expr::TableDestr(destr) => allocator.text("<destr>"),
+            Expr::FuncDef(def) => allocator.text("<def>"),
             Expr::Paren {
                 s0,
                 inner,
                 s1,
                 span: _,
-            } => RcDoc::text("(").append(inner.to_doc()).append(")"),
+            } => inner.pretty(allocator).parens(),
 
             // TODO Check whether parentheses are necessary
             Expr::Neg {
@@ -25,7 +25,7 @@ impl Expr {
                 s0,
                 expr,
                 span: _,
-            } => RcDoc::text("-").append(expr.to_doc()),
+            } => allocator.text("-").append(expr.pretty(allocator)),
 
             // TODO Check whether parentheses are necessary
             Expr::Not {
@@ -33,7 +33,7 @@ impl Expr {
                 s0,
                 expr,
                 span: _,
-            } => RcDoc::text("not ").append(expr.to_doc()),
+            } => allocator.text("not ").append(expr.pretty(allocator)),
 
             Expr::BinOp {
                 left,
@@ -42,7 +42,7 @@ impl Expr {
                 s1,
                 right,
                 span: _,
-            } => RcDoc::text("<binop>"),
+            } => allocator.text("<binop>"),
         }
     }
 }

src/pretty/program.rs

@@ -1,29 +1,31 @@
-use pretty::RcDoc;
+use pretty::{DocAllocator, DocBuilder, Pretty};
 
 use crate::ast::Program;
 
-impl Program {
-    pub fn to_doc(&self) -> RcDoc {
+impl<'a, D: DocAllocator<'a>> Pretty<'a, D> for Program {
+    fn pretty(self, allocator: &'a D) -> DocBuilder<'a, D> {
         match self {
             Program::Expr {
                 s0,
                 expr,
                 s1,
                 span: _,
-            } => expr.to_doc(),
+            } => expr.pretty(allocator),
             Program::Module {
                 s0,
                 s1,
                 elems,
                 s2,
                 span: _,
-            } => RcDoc::text("module")
-                .append(RcDoc::line())
-                .append(RcDoc::line())
-                .append(elems.to_doc(
-                    |e| RcDoc::text("<elem>"),
-                    |(s0, s1)| RcDoc::text(",").append(RcDoc::line()),
-                    |s| RcDoc::text(","),
+            } => allocator
+                .text("module")
+                .append(allocator.line())
+                .append(allocator.line())
+                .append(elems.pretty(
+                    allocator,
+                    |a, e| a.text("<elem>"),
+                    |a, (s0, s1)| a.text(",").append(a.line()),
+                    |a, s| a.text(","),
                 )),
         }
     }