Fix formatting of function expressions with traces

Fixes #471

crates/aiken-lang/src/format.rs

@@ -569,7 +569,7 @@ impl<'comments> Formatter<'comments> {
     ) -> Document<'a> {
         let args = wrap_args(args.iter().map(|e| (self.fn_arg(e), false))).group();
         let body = match body {
-            UntypedExpr::When { .. } => self.expr(body).force_break(),
+            UntypedExpr::Trace { .. } | UntypedExpr::When { .. } => self.expr(body).force_break(),
             _ => self.expr(body),
         };
 

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

@@ -698,3 +698,28 @@ fn weird_comments() {
 
     assert_fmt(src, expected);
 }
+
+#[test]
+fn format_trace_callback() {
+    let src = indoc! { r#"
+      fn foo() {
+        list.any([], fn (e) { trace @"foo"
+          e
+        })
+      }
+    "#};
+
+    let expected = indoc! { r#"
+      fn foo() {
+        list.any(
+          [],
+          fn(e) {
+            trace @"foo"
+            e
+          },
+        )
+      }
+    "#};
+
+    assert_fmt(src, expected);
+}