From cc18e7cff28995e25e2f704ab80cbbfbe0428c96 Mon Sep 17 00:00:00 2001
From: KtorZ <matthias.benkort@gmail.com>
Date: Thu, 30 Mar 2023 09:21:46 +0200
Subject: [PATCH] Fix formatting of function expressions with traces

  Fixes #471
---
 crates/aiken-lang/src/format.rs       |  2 +-
 crates/aiken-lang/src/tests/format.rs | 25 +++++++++++++++++++++++++
 2 files changed, 26 insertions(+), 1 deletion(-)

diff --git a/crates/aiken-lang/src/format.rs b/crates/aiken-lang/src/format.rs
index 9881ff62..fa6d2d4d 100644
--- a/crates/aiken-lang/src/format.rs
+++ b/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),
         };
 
diff --git a/crates/aiken-lang/src/tests/format.rs b/crates/aiken-lang/src/tests/format.rs
index c91dc156..674331e1 100644
--- a/crates/aiken-lang/src/tests/format.rs
+++ b/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);
+}