Fix formatting of if-expressions

Fixes #129.

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,
-        RecordUpdateSpread, Span, TypeAlias, TypedArg, TypedConstant, UnqualifiedImport,
-        UntypedArg, UntypedClause, UntypedClauseGuard, UntypedDefinition, UntypedModule,
-        UntypedPattern, UntypedRecordUpdateArg, Use, CAPTURE_VARIABLE,
+        Definition, Function, IfBranch, ModuleConstant, Pattern, RecordConstructor,
+        RecordConstructorArg, RecordUpdateSpread, Span, TypeAlias, TypedArg, TypedConstant,
+        UnqualifiedImport, UntypedArg, UntypedClause, UntypedClauseGuard, UntypedDefinition,
+        UntypedModule, UntypedPattern, UntypedRecordUpdateArg, Use, CAPTURE_VARIABLE,
     },
     docvec,
     expr::UntypedExpr,
@@ -663,44 +663,7 @@ impl<'comments> Formatter<'comments> {
                 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()
-            }
+            } => self.if_expr(branches, final_else),
             UntypedExpr::Todo { label: None, .. } => "todo".to_doc(),
 
             UntypedExpr::Todo { label: Some(l), .. } => docvec!["todo(\"", l, "\")"],
@@ -917,6 +880,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,173 @@
+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.clone(), extra.clone(), src);
+
+    // Output is what we expect
+    assert_eq!(out, expected);
+
+    // Formatting is idempotent
+    let mut out2 = String::new();
+    format::pretty(&mut out2, module, extra, &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)
+}

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

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