From c7fc8584ffe0eb24b329573d2dd93d1cc4331a42 Mon Sep 17 00:00:00 2001
From: Joscha <joscha@plugh.de>
Date: Mon, 21 Nov 2022 00:02:10 +0100
Subject: [PATCH] Pretty print variable accesses and assignments

---
 src/pretty.rs      |  1 +
 src/pretty/expr.rs |  2 +-
 src/pretty/var.rs  | 45 +++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 47 insertions(+), 1 deletion(-)
 create mode 100644 src/pretty/var.rs

diff --git a/src/pretty.rs b/src/pretty.rs
index b291f7a..01a0dbb 100644
--- a/src/pretty.rs
+++ b/src/pretty.rs
@@ -5,5 +5,6 @@ mod field;
 mod lit;
 mod program;
 mod table_constr;
+mod var;
 
 const NEST_DEPTH: isize = 4;
diff --git a/src/pretty/expr.rs b/src/pretty/expr.rs
index 161a8c7..4869084 100644
--- a/src/pretty/expr.rs
+++ b/src/pretty/expr.rs
@@ -28,7 +28,7 @@ impl<'a, D: DocAllocator<'a>> Pretty<'a, D> for Expr {
             Self::Lit(lit) => lit.pretty(allocator),
             Self::Call(call) => call.pretty(allocator),
             Self::Field(field) => field.pretty(allocator),
-            Self::Var(var) => allocator.text("<var>"),
+            Self::Var(var) => var.pretty(allocator),
             Self::TableConstr(constr) => constr.pretty(allocator),
             Self::TableDestr(destr) => allocator.text("<destr>"),
             Self::FuncDef(def) => allocator.text("<def>"),
diff --git a/src/pretty/var.rs b/src/pretty/var.rs
new file mode 100644
index 0000000..a3b5418
--- /dev/null
+++ b/src/pretty/var.rs
@@ -0,0 +1,45 @@
+use pretty::{DocAllocator, DocBuilder, Pretty};
+
+use crate::ast::Var;
+
+impl<'a, D: DocAllocator<'a>> Pretty<'a, D> for Var {
+    fn pretty(self, allocator: &'a D) -> DocBuilder<'a, D> {
+        match self {
+            Self::Access {
+                s0,
+                index,
+                s1,
+                span: _,
+            } => index.pretty(allocator).brackets(),
+            Self::Assign {
+                local,
+                s0,
+                index,
+                s1,
+                s2,
+                s3,
+                value,
+                span: _,
+            } => local
+                .map(|s| allocator.text("local "))
+                .unwrap_or_else(|| allocator.nil())
+                .append(index.pretty(allocator).brackets())
+                .append(allocator.text(" = "))
+                .append(value.pretty(allocator)),
+            Self::AccessIdent(ident) => ident.pretty(allocator),
+            Self::AssignIdent {
+                local,
+                name,
+                s0,
+                s1,
+                value,
+                span: _,
+            } => local
+                .map(|s| allocator.text("local "))
+                .unwrap_or_else(|| allocator.nil())
+                .append(name.pretty(allocator))
+                .append(allocator.text(" = "))
+                .append(value.pretty(allocator)),
+        }
+    }
+}