From 15bdb6972d7ed5964753e4d555201bcbe3133b12 Mon Sep 17 00:00:00 2001
From: rvcas <x@rvcas.dev>
Date: Thu, 16 Mar 2023 19:34:16 -0400
Subject: [PATCH] feat(fmt): better validator formatting with double supported

---
 crates/aiken-lang/src/format.rs       | 87 +++++++++++++++------------
 crates/aiken-lang/src/tests/format.rs | 26 ++++++++
 2 files changed, 73 insertions(+), 40 deletions(-)

diff --git a/crates/aiken-lang/src/format.rs b/crates/aiken-lang/src/format.rs
index 521263f2..cf20497f 100644
--- a/crates/aiken-lang/src/format.rs
+++ b/crates/aiken-lang/src/format.rs
@@ -240,10 +240,11 @@ impl<'comments> Formatter<'comments> {
 
             Definition::Validator(Validator {
                 end_position,
-                fun: function,
+                fun,
+                other_fun,
                 params,
                 ..
-            }) => self.definition_validator(params, function, *end_position),
+            }) => self.definition_validator(params, fun, other_fun, *end_position),
 
             Definition::Test(Function {
                 name,
@@ -507,53 +508,59 @@ impl<'comments> Formatter<'comments> {
         other_fun: &'a Option<UntypedFunction>,
         end_position: usize,
     ) -> Document<'a> {
-        // Fn and args
-        let head = "fn"
-            .to_doc()
-            .append(" ")
-            .append(fun.name.to_doc())
-            .append(wrap_args(
-                fun.arguments.iter().map(|e| (self.fn_arg(e), false)),
-            ));
-
-        // Add return annotation
-        let head = match &fun.return_annotation {
-            Some(anno) => head.append(" -> ").append(self.annotation(anno)),
-            None => head,
-        }
-        .group();
-
-        // Format body
-        let body = self.expr(&fun.body);
-
-        // Add any trailing comments
-        let body = match printed_comments(self.pop_comments(fun.end_position), false) {
-            Some(comments) => body.append(line()).append(comments),
-            None => body,
-        };
-
-        // validator name(params)
-        let v_head = "validator".to_doc().append(if !params.is_empty() {
-            wrap_args(params.iter().map(|e| (self.fn_arg(e), false)))
-        } else {
-            "".to_doc()
-        });
-
         // Stick it all together
-        let inner_fn = head
-            .append(" {")
-            .append(line().append(body).nest(INDENT).group())
-            .append(line())
-            .append("}");
+        let inner_fn = line()
+            .append(self.definition_fn(
+                &false,
+                "fn",
+                &fun.name,
+                &fun.arguments,
+                &fun.return_annotation,
+                &fun.body,
+                fun.end_position,
+            ))
+            .nest(INDENT)
+            .group()
+            .append(if other_fun.is_some() {
+                docvec![line()]
+            } else {
+                nil()
+            })
+            .append(
+                other_fun
+                    .as_ref()
+                    .map(|other| {
+                        line()
+                            .append(self.definition_fn(
+                                &false,
+                                "fn",
+                                &other.name,
+                                &other.arguments,
+                                &other.return_annotation,
+                                &other.body,
+                                other.end_position,
+                            ))
+                            .nest(INDENT)
+                            .group()
+                    })
+                    .unwrap_or_else(nil),
+            );
 
         let inner_fn = match printed_comments(self.pop_comments(end_position), false) {
             Some(comments) => inner_fn.append(line()).append(comments),
             None => inner_fn,
         };
 
+        // validator(params)
+        let v_head = "validator".to_doc().append(if !params.is_empty() {
+            wrap_args(params.iter().map(|e| (self.fn_arg(e), false)))
+        } else {
+            nil()
+        });
+
         v_head
             .append(" {")
-            .append(line().append(inner_fn).nest(INDENT).group())
+            .append(inner_fn)
             .append(line())
             .append("}")
     }
diff --git a/crates/aiken-lang/src/tests/format.rs b/crates/aiken-lang/src/tests/format.rs
index 7f28c8c4..c91dc156 100644
--- a/crates/aiken-lang/src/tests/format.rs
+++ b/crates/aiken-lang/src/tests/format.rs
@@ -76,6 +76,32 @@ fn test_format_validator() {
     assert_fmt(src, expected)
 }
 
+#[test]
+fn test_format_double_validator() {
+    let src = indoc! {r#"
+        validator ( param1 : ByteArray ) {
+        fn foo (d: Datum, r: Redeemer, ctx: ScriptContext) -> Bool {
+        True
+        }
+    fn bar(r: Redeemer, ctx    : ScriptContext  )   ->   Bool { True }
+        }
+    "#};
+
+    let expected = indoc! {r#"
+      validator(param1: ByteArray) {
+        fn foo(d: Datum, r: Redeemer, ctx: ScriptContext) -> Bool {
+          True
+        }
+
+        fn bar(r: Redeemer, ctx: ScriptContext) -> Bool {
+          True
+        }
+      }
+    "#};
+
+    assert_fmt(src, expected)
+}
+
 #[test]
 fn test_format_when() {
     let src = indoc! {r#"