From b3eaa409027187cb4f365a229c613ceefe0139bd Mon Sep 17 00:00:00 2001
From: Joscha <joscha@plugh.de>
Date: Mon, 21 Nov 2022 00:16:15 +0100
Subject: [PATCH] Pretty print function definitions

---
 src/pretty.rs          |  1 +
 src/pretty/expr.rs     |  2 +-
 src/pretty/func_def.rs | 98 ++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 100 insertions(+), 1 deletion(-)
 create mode 100644 src/pretty/func_def.rs

diff --git a/src/pretty.rs b/src/pretty.rs
index 149eb1c..e15aa53 100644
--- a/src/pretty.rs
+++ b/src/pretty.rs
@@ -2,6 +2,7 @@ mod basic;
 mod call;
 mod expr;
 mod field;
+mod func_def;
 mod lit;
 mod program;
 mod table_constr;
diff --git a/src/pretty/expr.rs b/src/pretty/expr.rs
index 454651e..7919963 100644
--- a/src/pretty/expr.rs
+++ b/src/pretty/expr.rs
@@ -31,7 +31,7 @@ impl<'a, D: DocAllocator<'a>> Pretty<'a, D> for Expr {
             Self::Var(var) => var.pretty(allocator),
             Self::TableConstr(constr) => constr.pretty(allocator),
             Self::TableDestr(destr) => destr.pretty(allocator),
-            Self::FuncDef(def) => allocator.text("<def>"),
+            Self::FuncDef(def) => def.pretty(allocator),
             Self::Paren {
                 s0,
                 inner,
diff --git a/src/pretty/func_def.rs b/src/pretty/func_def.rs
new file mode 100644
index 0000000..6fd8140
--- /dev/null
+++ b/src/pretty/func_def.rs
@@ -0,0 +1,98 @@
+use pretty::{DocAllocator, DocBuilder, Pretty};
+
+use crate::ast::FuncDef;
+
+impl<'a, D: DocAllocator<'a>> Pretty<'a, D> for FuncDef {
+    fn pretty(self, allocator: &'a D) -> DocBuilder<'a, D> {
+        match self {
+            Self::AnonNoArg {
+                s0,
+                s1,
+                s2,
+                body,
+                span: _,
+            } => allocator.text("function() ").append(body.pretty(allocator)),
+
+            Self::AnonArg {
+                s0,
+                s1,
+                arg,
+                s2,
+                s3,
+                body,
+                span: _,
+            } => allocator
+                .text("function")
+                .append(arg.pretty(allocator).parens())
+                .append(allocator.space())
+                .append(body.pretty(allocator)),
+
+            Self::AnonDestr {
+                s0,
+                pattern,
+                s1,
+                body,
+                span: _,
+            } => allocator
+                .text("function")
+                .append(pattern.pretty(allocator))
+                .append(allocator.space())
+                .append(body.pretty(allocator)),
+
+            Self::NamedNoArg {
+                local,
+                s0,
+                name,
+                s1,
+                s2,
+                s3,
+                body,
+                span: _,
+            } => local
+                .map(|s| allocator.text("local "))
+                .unwrap_or_else(|| allocator.nil())
+                .append(allocator.text("function "))
+                .append(name)
+                .append(allocator.text("() "))
+                .append(body.pretty(allocator)),
+
+            Self::NamedArg {
+                local,
+                s0,
+                name,
+                s1,
+                s2,
+                arg,
+                s3,
+                s4,
+                body,
+                span: _,
+            } => local
+                .map(|s| allocator.text("local "))
+                .unwrap_or_else(|| allocator.nil())
+                .append(allocator.text("function "))
+                .append(name)
+                .append(arg.pretty(allocator).parens())
+                .append(allocator.space())
+                .append(body.pretty(allocator)),
+
+            Self::NamedDestr {
+                local,
+                s0,
+                name,
+                s1,
+                pattern,
+                s2,
+                body,
+                span: _,
+            } => local
+                .map(|s| allocator.text("local "))
+                .unwrap_or_else(|| allocator.nil())
+                .append(allocator.text("function "))
+                .append(name)
+                .append(pattern.pretty(allocator))
+                .append(allocator.space())
+                .append(body.pretty(allocator)),
+        }
+    }
+}