From 5f8b25278399a125902c6ec08e3e96d65fcb9e88 Mon Sep 17 00:00:00 2001
From: rvcas <x@rvcas.dev>
Date: Mon, 30 Dec 2024 22:33:13 -0500
Subject: [PATCH] fix: adjust comment formatting in pipelines closes #985

---
 CHANGELOG.md                                  |  3 +-
 crates/aiken-lang/src/format.rs               | 31 +++++++++----------
 crates/aiken-lang/src/tests/format.rs         | 16 ++++++++++
 .../tests/snapshots/comment_in_pipeline.snap  | 13 ++++++++
 .../tests/snapshots/format_preserve_pipe.snap |  7 +++--
 5 files changed, 51 insertions(+), 19 deletions(-)
 create mode 100644 crates/aiken-lang/src/tests/snapshots/comment_in_pipeline.snap

diff --git a/CHANGELOG.md b/CHANGELOG.md
index e3098a5fa..7931d7fad 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -13,7 +13,8 @@
 
 ### Fixed
 
-- **aiken**: Fixed the panic error when using `aiken uplc decode` on cbor encoded flat bytes. @rvcas
+- **aiken**: panic error when using `aiken uplc decode` on cbor encoded flat bytes. @rvcas
+- **aiken-lang**: comment formatting in pipelines leading to confusion. @rvcas
 
 ## v1.1.9 - 2024-12-13
 
diff --git a/crates/aiken-lang/src/format.rs b/crates/aiken-lang/src/format.rs
index aa0e35c5b..6adfe3826 100644
--- a/crates/aiken-lang/src/format.rs
+++ b/crates/aiken-lang/src/format.rs
@@ -14,9 +14,7 @@ use crate::{
         extra::{Comment, ModuleExtra},
         token::Base,
     },
-    pretty::{
-        break_, concat, flex_break, join, line, lines, nil, prebreak, Document, Documentable,
-    },
+    pretty::{break_, concat, flex_break, join, line, lines, nil, Document, Documentable},
     tipo::{self, Type},
 };
 use itertools::Itertools;
@@ -1472,13 +1470,18 @@ impl<'comments> Formatter<'comments> {
         one_liner: bool,
     ) -> Document<'a> {
         let mut docs = Vec::with_capacity(expressions.len() * 3);
+
         let first = expressions.first();
+
         let first_precedence = first.binop_precedence();
+
         let first = self.wrap_expr(first);
+
         docs.push(self.operator_side(first, 5, first_precedence));
 
         for expr in expressions.iter().skip(1) {
             let comments = self.pop_comments(expr.location().start);
+
             let doc = match expr {
                 UntypedExpr::Fn {
                     fn_style: FnStyle::Capture,
@@ -1489,23 +1492,19 @@ impl<'comments> Formatter<'comments> {
                 _ => self.wrap_expr(expr),
             };
 
+            let space = if one_liner { break_("", " ") } else { line() };
+
+            let pipe = space
+                .append(commented("|> ".to_doc(), comments))
+                .nest(INDENT);
+
+            docs.push(pipe);
+
             let expr = self
                 .operator_side(doc, 4, expr.binop_precedence())
                 .nest(2 * INDENT);
 
-            match printed_comments(comments, true) {
-                None => {
-                    let pipe = prebreak("|> ", " |> ").nest(INDENT);
-                    docs.push(pipe.append(expr));
-                }
-                Some(comments) => {
-                    let pipe = prebreak("|> ", "|> ");
-                    docs.push(
-                        " ".to_doc()
-                            .append(comments.nest(INDENT).append(pipe.append(expr).group())),
-                    );
-                }
-            }
+            docs.push(expr);
         }
 
         if one_liner {
diff --git a/crates/aiken-lang/src/tests/format.rs b/crates/aiken-lang/src/tests/format.rs
index 64f104492..7d115bb62 100644
--- a/crates/aiken-lang/src/tests/format.rs
+++ b/crates/aiken-lang/src/tests/format.rs
@@ -1420,3 +1420,19 @@ fn multiline_if_is_2() {
         "#
     );
 }
+
+#[test]
+fn comment_in_pipeline() {
+    assert_format!(
+        r#"
+        fn foo() {
+            a
+            // stuff
+            // warning: wow
+            |> b
+            // Comment
+            |> c
+        }
+        "#
+    );
+}
diff --git a/crates/aiken-lang/src/tests/snapshots/comment_in_pipeline.snap b/crates/aiken-lang/src/tests/snapshots/comment_in_pipeline.snap
new file mode 100644
index 000000000..1be862b12
--- /dev/null
+++ b/crates/aiken-lang/src/tests/snapshots/comment_in_pipeline.snap
@@ -0,0 +1,13 @@
+---
+source: crates/aiken-lang/src/tests/format.rs
+description: "Code:\n\nfn foo() {\n    a\n    // stuff\n    // warning: wow\n    |> b\n    // Comment\n    |> c\n}\n"
+snapshot_kind: text
+---
+fn foo() {
+  a
+    // stuff
+    // warning: wow
+    |> b
+    // Comment
+    |> c
+}
diff --git a/crates/aiken-lang/src/tests/snapshots/format_preserve_pipe.snap b/crates/aiken-lang/src/tests/snapshots/format_preserve_pipe.snap
index a2ef28f9d..6b2a4813b 100644
--- a/crates/aiken-lang/src/tests/snapshots/format_preserve_pipe.snap
+++ b/crates/aiken-lang/src/tests/snapshots/format_preserve_pipe.snap
@@ -1,14 +1,17 @@
 ---
 source: crates/aiken-lang/src/tests/format.rs
 description: "Code:\n\nfn foo() {\n  a |> b |> c |> d\n}\n\nfn foo() {\n  a\n    // Foo\n    |> b// Some comments\n    |> c\n    |> d\n}\n\nfn baz() {\n  // Commented\n  however |> it_automatically_breaks |> into_multiple_lines |> anytime_when |> it_is_too_long // What?\n}\n"
+snapshot_kind: text
 ---
 fn foo() {
   a |> b |> c |> d
 }
 
 fn foo() {
-  a // Foo
-    |> b // Some comments
+  a
+    // Foo
+    |> b
+    // Some comments
     |> c
     |> d
 }