From 731bf1cb98c50adae08c78ef8d90bd1f989acdaf Mon Sep 17 00:00:00 2001
From: Jeff <dev@jeffa.io>
Date: Sat, 6 Jan 2024 05:29:38 -0500
Subject: [PATCH] Write formatting tests; Improve formatting output

---
 src/abstract_tree/block.rs         |  8 +++---
 src/abstract_tree/function_node.rs | 31 +++++++++++++++++++++++
 src/abstract_tree/math.rs          |  2 +-
 src/value/function.rs              |  4 +--
 tests/format.rs                    | 40 ++++++++++++++++++++++++++++++
 5 files changed, 78 insertions(+), 7 deletions(-)
 create mode 100644 tests/format.rs

diff --git a/src/abstract_tree/block.rs b/src/abstract_tree/block.rs
index 44d9b01..3b080f0 100644
--- a/src/abstract_tree/block.rs
+++ b/src/abstract_tree/block.rs
@@ -134,15 +134,15 @@ impl AbstractTree for Block {
 impl Display for Block {
     fn fmt(&self, f: &mut Formatter) -> fmt::Result {
         if self.is_async {
-            write!(f, "async {{")?;
+            writeln!(f, "async {{")?;
         } else {
-            write!(f, "{{")?;
+            writeln!(f, "{{")?;
         }
 
         for statement in &self.statements {
-            write!(f, "    {statement}")?;
+            writeln!(f, "    {statement}")?;
         }
 
-        write!(f, "}}")
+        writeln!(f, "}}")
     }
 }
diff --git a/src/abstract_tree/function_node.rs b/src/abstract_tree/function_node.rs
index f2d8f1e..c61d531 100644
--- a/src/abstract_tree/function_node.rs
+++ b/src/abstract_tree/function_node.rs
@@ -1,3 +1,5 @@
+use std::fmt::{self, Display, Formatter};
+
 use serde::{Deserialize, Serialize};
 use tree_sitter::Node;
 
@@ -143,3 +145,32 @@ impl AbstractTree for FunctionNode {
         Ok(self.r#type().clone())
     }
 }
+
+impl Display for FunctionNode {
+    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
+        let FunctionNode {
+            parameters,
+            body,
+            r#type,
+            ..
+        } = self;
+
+        let (parameter_types, return_type) = if let Type::Function {
+            parameter_types,
+            return_type,
+        } = r#type
+        {
+            (parameter_types, return_type)
+        } else {
+            return Ok(());
+        };
+
+        write!(f, "(")?;
+
+        for (identifier, r#type) in parameters.iter().zip(parameter_types.iter()) {
+            write!(f, "{identifier} <{}>", r#type)?;
+        }
+
+        write!(f, ") <{return_type}> {body}")
+    }
+}
diff --git a/src/abstract_tree/math.rs b/src/abstract_tree/math.rs
index 308e86c..81ef5d5 100644
--- a/src/abstract_tree/math.rs
+++ b/src/abstract_tree/math.rs
@@ -96,7 +96,7 @@ impl Display for MathOperator {
             MathOperator::Add => write!(f, "+"),
             MathOperator::Subtract => write!(f, "-"),
             MathOperator::Multiply => write!(f, "*"),
-            MathOperator::Divide => write!(f, "\\"),
+            MathOperator::Divide => write!(f, "/"),
             MathOperator::Modulo => write!(f, "%"),
         }
     }
diff --git a/src/value/function.rs b/src/value/function.rs
index 038e136..d1154da 100644
--- a/src/value/function.rs
+++ b/src/value/function.rs
@@ -15,8 +15,8 @@ impl Display for Function {
     fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
         match self {
             Function::BuiltIn(built_in_function) => write!(f, "{}", built_in_function.r#type()),
-            Function::ContextDefined(context_defined_function) => {
-                write!(f, "{}", context_defined_function.r#type())
+            Function::ContextDefined(function_node) => {
+                write!(f, "{}", function_node)
             }
         }
     }
diff --git a/tests/format.rs b/tests/format.rs
new file mode 100644
index 0000000..f67f7f8
--- /dev/null
+++ b/tests/format.rs
@@ -0,0 +1,40 @@
+use dust_lang::*;
+
+#[test]
+fn format_simple_program() {
+    let mut interpreter = Interpreter::new(Map::new());
+
+    interpreter.run("x=1").unwrap();
+
+    assert_eq!(interpreter.format(), "x = 1");
+}
+
+const FORMATTED_BLOCK: &str = "{
+    1
+    2
+    3
+}
+";
+
+#[test]
+fn format_block() {
+    let mut interpreter = Interpreter::new(Map::new());
+
+    interpreter.run("{1 2 3}").unwrap();
+
+    assert_eq!(FORMATTED_BLOCK, interpreter.format());
+}
+
+const FORMATTED_FUNCTION: &str = "(x <int>) <num> {
+    x / 2
+}
+";
+
+#[test]
+fn format_function() {
+    let mut interpreter = Interpreter::new(Map::new());
+
+    interpreter.run("( x< int > )<num>{x/2}").unwrap();
+
+    assert_eq!(FORMATTED_FUNCTION, interpreter.format());
+}