Merge pull request #210 from aiken-lang/fix-if-expressions

crates/aiken-lang/src/format.rs

@@ -6,10 +6,10 @@ use vec1::Vec1;
 use crate::{
     ast::{
         Annotation, Arg, ArgName, AssignmentKind, BinOp, CallArg, ClauseGuard, Constant, DataType,
-        Definition, Function, ModuleConstant, Pattern, RecordConstructor, RecordConstructorArg,
+        Definition, Function, IfBranch, ModuleConstant, Pattern, RecordConstructor,
-        RecordUpdateSpread, Span, TypeAlias, TypedArg, TypedConstant, UnqualifiedImport,
+        RecordConstructorArg, RecordUpdateSpread, Span, TypeAlias, TypedArg, TypedConstant,
-        UntypedArg, UntypedClause, UntypedClauseGuard, UntypedDefinition, UntypedModule,
+        UnqualifiedImport, UntypedArg, UntypedClause, UntypedClauseGuard, UntypedDefinition,
-        UntypedPattern, UntypedRecordUpdateArg, Use, CAPTURE_VARIABLE,
+        UntypedModule, UntypedPattern, UntypedRecordUpdateArg, Use, CAPTURE_VARIABLE,
     },
     docvec,
     expr::UntypedExpr,
@@ -133,14 +133,14 @@ impl<'comments> Formatter<'comments> {
             let start = def.location().start;
 
             match def {
-                Definition::Use { .. } => {
+                Definition::Use(import) => {
                     has_imports = true;
 
                     let comments = self.pop_comments(start);
 
                     let def = self.definition(def);
 
-                    imports.push(commented(def, comments))
+                    imports.push((import, commented(def, comments)))
                 }
 
                 _other => {
@@ -155,7 +155,15 @@ impl<'comments> Formatter<'comments> {
             }
         }
 
-        let imports = join(imports.into_iter(), line());
+        let imports = join(
+            imports
+                .into_iter()
+                .sorted_by(|(import_a, _), (import_b, _)| {
+                    Ord::cmp(&import_a.module, &import_b.module)
+                })
+                .map(|(_, doc)| doc),
+            line(),
+        );
 
         let declarations = join(declarations.into_iter(), lines(2));
 
@@ -253,12 +261,35 @@ impl<'comments> Formatter<'comments> {
                 ..
             }) => self.data_type(*public, *opaque, name, parameters, constructors, location),
 
-            Definition::Use(Use {
+            Definition::Use(import) => self.import(import),
+
+            Definition::ModuleConstant(ModuleConstant {
+                public,
+                name,
+                annotation,
+                value,
+                ..
+            }) => {
+                let head = pub_(*public).append("const ").append(name.as_str());
+                let head = match annotation {
+                    None => head,
+                    Some(t) => head.append(": ").append(self.annotation(t)),
+                };
+                head.append(" = ").append(self.const_expr(value))
+            }
+        }
+    }
+
+    fn import<'a>(
+        &mut self,
+        Use {
             module,
             as_name,
             unqualified,
             ..
-            }) => "use "
+        }: &'a Use<()>,
+    ) -> Document<'a> {
+        "use "
             .to_doc()
             .append(Document::String(module.join("/")))
             .append(if unqualified.is_empty() {
@@ -282,23 +313,7 @@ impl<'comments> Formatter<'comments> {
                 docvec![" as ", name]
             } else {
                 nil()
-                }),
+            })
-
-            Definition::ModuleConstant(ModuleConstant {
-                public,
-                name,
-                annotation,
-                value,
-                ..
-            }) => {
-                let head = pub_(*public).append("const ").append(name.as_str());
-                let head = match annotation {
-                    None => head,
-                    Some(t) => head.append(": ").append(self.annotation(t)),
-                };
-                head.append(" = ").append(self.const_expr(value))
-            }
-        }
     }
 
     fn const_expr<'a, A, B>(&mut self, value: &'a Constant<A, B>) -> Document<'a> {
@@ -663,44 +678,7 @@ impl<'comments> Formatter<'comments> {
                 branches,
                 final_else,
                 ..
-            } => {
+            } => self.if_expr(branches, final_else),
-                let first = branches.first();
-
-                break_("if", "if ")
-                    .append(self.wrap_expr(&first.condition))
-                    .nest(INDENT)
-                    .append(break_("", " "))
-                    .append("{")
-                    .group()
-                    .append(line())
-                    .nest(INDENT)
-                    .append(self.expr(&first.body))
-                    .append(line())
-                    .append("} ")
-                    .append(join(
-                        branches[1..].iter().map(|branch| {
-                            break_("else if", "else if ")
-                                .append(self.wrap_expr(&branch.condition))
-                                .nest(INDENT)
-                                .append(break_("", " "))
-                                .append("{")
-                                .group()
-                                .append(line())
-                                .nest(INDENT)
-                                .append(self.expr(&branch.body))
-                                .append(line())
-                                .append("} ")
-                        }),
-                        nil(),
-                    ))
-                    .append("else {")
-                    .group()
-                    .append(line().nest(INDENT))
-                    .append(self.expr(final_else))
-                    .append(line())
-                    .append("}")
-                    .force_break()
-            }
             UntypedExpr::Todo { label: None, .. } => "todo".to_doc(),
 
             UntypedExpr::Todo { label: Some(l), .. } => docvec!["todo(\"", l, "\")"],
@@ -917,6 +895,48 @@ impl<'comments> Formatter<'comments> {
         }
     }
 
+    pub fn if_expr<'a>(
+        &mut self,
+        branches: &'a Vec1<IfBranch<UntypedExpr>>,
+        final_else: &'a UntypedExpr,
+    ) -> Document<'a> {
+        let if_branches = self
+            .if_branch(break_("if", "if "), branches.first())
+            .append(join(
+                branches[1..]
+                    .iter()
+                    .map(|branch| self.if_branch("else if".to_doc(), branch)),
+                nil(),
+            ));
+
+        let else_begin = line().append("} else {");
+
+        let else_body = line().append(self.expr(final_else)).nest(INDENT);
+
+        let else_end = line().append("}");
+
+        if_branches
+            .append(else_begin)
+            .append(else_body)
+            .append(else_end)
+            .force_break()
+    }
+
+    pub fn if_branch<'a>(
+        &mut self,
+        if_keyword: Document<'a>,
+        branch: &'a IfBranch<UntypedExpr>,
+    ) -> Document<'a> {
+        let if_begin = if_keyword
+            .append(self.wrap_expr(&branch.condition))
+            .append(break_("{", " {"))
+            .group();
+
+        let if_body = line().append(self.expr(&branch.body)).nest(INDENT);
+
+        if_begin.append(if_body)
+    }
+
     pub fn when<'a>(
         &mut self,
         subjects: &'a [UntypedExpr],

crates/aiken-lang/src/tests/format.rs

@@ -0,0 +1,199 @@
+use crate::{ast::ModuleKind, format, parser};
+use indoc::indoc;
+use pretty_assertions::assert_eq;
+
+fn assert_fmt(src: &str, expected: &str) {
+    let (module, extra) = parser::module(src, ModuleKind::Lib).unwrap();
+    let mut out = String::new();
+    format::pretty(&mut out, module, extra, src);
+
+    // Output is what we expect
+    assert_eq!(out, expected);
+
+    // Formatting is idempotent
+    let (module2, extra2) = parser::module(&out, ModuleKind::Lib).unwrap();
+    let mut out2 = String::new();
+    format::pretty(&mut out2, module2, extra2, &out);
+    assert_eq!(out, out2);
+}
+
+#[test]
+fn test_format_if() {
+    let src = indoc! {r#"
+        pub fn foo(a) {
+          if a { 14 } else { 42 }
+          }
+
+        pub fn bar(xs) {
+            list.map(xs, fn (x) { if x > 0 { "foo" } else { "bar" } })
+        }
+    "#};
+
+    let expected = indoc! {r#"
+        pub fn foo(a) {
+          if a {
+            14
+          } else {
+            42
+          }
+        }
+
+        pub fn bar(xs) {
+          list.map(
+            xs,
+            fn(x) {
+              if x > 0 {
+                "foo"
+              } else {
+                "bar"
+              }
+            },
+          )
+        }
+    "#};
+
+    assert_fmt(src, expected)
+}
+
+#[test]
+fn test_format_when() {
+    let src = indoc! {r#"
+        pub fn foo( a) {
+          when  a   is{
+            True  -> 14
+            False ->
+             42}
+         }
+    "#};
+
+    let expected = indoc! {r#"
+        pub fn foo(a) {
+          when a is {
+            True -> 14
+            False -> 42
+          }
+        }
+    "#};
+
+    assert_fmt(src, expected)
+}
+
+#[test]
+fn test_format_nested_if() {
+    let src = indoc! {r#"
+        pub fn foo(n) {
+          if n > 0 {
+            if n > 1 {
+              if n > 2 {
+                "foo"
+              } else {
+                "foo"
+              }
+            } else {
+              "bar"
+            }
+          } else {
+            if n < -1 {
+              "baz"
+            } else {
+              "biz"
+            }
+          }
+        }
+    "#};
+
+    assert_fmt(src, src)
+}
+
+#[test]
+fn test_format_nested_when_if() {
+    let src = indoc! {r#"
+        pub fn drop(xs: List<a>, n: Int) -> List<a> {
+          if n <= 0 {
+            xs
+          } else {
+            when xs is {
+            [] -> []
+            [_x, ..rest] -> drop(rest, n - 1)
+          }
+          }
+        }
+    "#};
+
+    let expected = indoc! {r#"
+        pub fn drop(xs: List<a>, n: Int) -> List<a> {
+          if n <= 0 {
+            xs
+          } else {
+            when xs is {
+              [] -> []
+              [_x, ..rest] -> drop(rest, n - 1)
+            }
+          }
+        }
+    "#};
+
+    assert_fmt(src, expected)
+}
+
+#[test]
+fn test_format_nested_when() {
+    let src = indoc! {r#"
+        fn foo() {
+          when a is {
+            None -> "foo"
+            Some(b) -> when b is {
+              None -> "foo"
+              Some(c) -> when c is {
+                None -> "foo"
+                Some(_) -> "foo"
+              }
+            }
+          }
+        }
+    "#};
+
+    let expected = indoc! {r#"
+        fn foo() {
+          when a is {
+            None -> "foo"
+            Some(b) ->
+              when b is {
+                None -> "foo"
+                Some(c) ->
+                  when c is {
+                    None -> "foo"
+                    Some(_) -> "foo"
+                  }
+              }
+          }
+        }
+    "#};
+
+    assert_fmt(src, expected)
+}
+
+#[test]
+fn test_format_imports() {
+    let src = indoc! {r#"
+        use aiken/list
+        // foo
+        use aiken/bytearray
+        use aiken/transaction/certificate
+        // bar
+        use aiken/transaction
+        use aiken/transaction/value
+    "#};
+
+    let expected = indoc! {r#"
+        // foo
+        use aiken/bytearray
+        use aiken/list
+        // bar
+        use aiken/transaction
+        use aiken/transaction/certificate
+        use aiken/transaction/value
+    "#};
+
+    assert_fmt(src, expected)
+}

crates/aiken-lang/src/tests/mod.rs

@@ -1,2 +1,3 @@
+mod format;
 mod lexer;
 mod parser;