feat(fmt): better validator formatting with double supported

aiken-lang · Mar 17, 2023 · 15bdb69 · 15bdb69
1 parent ed92869
commit 15bdb69
Show file tree

Hide file tree

Showing 2 changed files with 72 additions and 39 deletions.
diff --git 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);
+        // Stick it all together
+        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),
+            );
 
-        // 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,
+        let inner_fn = match printed_comments(self.pop_comments(end_position), false) {
+            Some(comments) => inner_fn.append(line()).append(comments),
+            None => inner_fn,
         };
 
-        // validator name(params)
+        // 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 {
-            "".to_doc()
+            nil()
         });
 
-        // Stick it all together
-        let inner_fn = head
-            .append(" {")
-            .append(line().append(body).nest(INDENT).group())
-            .append(line())
-            .append("}");
-
-        let inner_fn = match printed_comments(self.pop_comments(end_position), false) {
-            Some(comments) => inner_fn.append(line()).append(comments),
-            None => inner_fn,
-        };
-
         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
@@ -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#"