Support flexible pipe operator formatting

Rules are now as follows:

  - If a pipeline contains a newline, then the entire pipeline is formatted over multiple lines.
  - If it doesn't, then it's formatted as a single-line UNLESS it cannot fit; in which case, we fallback to multiline again.

crates/aiken-lang/src/expr.rs

@@ -456,6 +456,7 @@ pub enum UntypedExpr {
 
     PipeLine {
         expressions: Vec1<Self>,
+        one_liner: bool,
     },
 
     Assignment {

crates/aiken-lang/src/format.rs

@@ -15,7 +15,9 @@ use crate::{
     docvec,
     expr::{UntypedExpr, DEFAULT_ERROR_STR, DEFAULT_TODO_STR},
     parser::extra::{Comment, ModuleExtra},
-    pretty::{break_, concat, flex_break, join, line, lines, nil, Document, Documentable},
+    pretty::{
+        break_, concat, flex_break, join, line, lines, nil, prebreak, Document, Documentable,
+    },
     tipo::{self, Type},
 };
 
@@ -683,7 +685,10 @@ impl<'comments> Formatter<'comments> {
                 ..
             } => self.if_expr(branches, final_else),
 
-            UntypedExpr::PipeLine { expressions, .. } => self.pipeline(expressions),
+            UntypedExpr::PipeLine {
+                expressions,
+                one_liner,
+            } => self.pipeline(expressions, *one_liner),
 
             UntypedExpr::Int { value, .. } => value.to_doc(),
 
@@ -1013,7 +1018,11 @@ impl<'comments> Formatter<'comments> {
         }
     }
 
-    fn pipeline<'a>(&mut self, expressions: &'a Vec1<UntypedExpr>) -> Document<'a> {
+    fn pipeline<'a>(
+        &mut self,
+        expressions: &'a Vec1<UntypedExpr>,
+        one_liner: bool,
+    ) -> Document<'a> {
         let mut docs = Vec::with_capacity(expressions.len() * 3);
         let first = expressions.first();
         let first_precedence = first.binop_precedence();
@@ -1031,13 +1040,30 @@ impl<'comments> Formatter<'comments> {
 
                 _ => self.wrap_expr(expr),
             };
-            docs.push(line());
+
-            docs.push(commented("|> ".to_doc(), comments));
+            let expr = self.operator_side(doc, 4, expr.binop_precedence());
-            docs.push(self.operator_side(doc, 4, expr.binop_precedence()));
+
+            match printed_comments(comments, true) {
+                None => {
+                    let pipe = prebreak("|> ", " |> ").nest(2);
+                    docs.push(pipe.append(expr));
+                }
+                Some(comments) => {
+                    let pipe = prebreak("|> ", "|> ");
+                    docs.push(
+                        " ".to_doc()
+                            .append(comments.nest(2).append(pipe.append(expr).group())),
+                    );
+                }
+            }
         }
 
+        if one_liner {
+            docs.to_doc().group()
+        } else {
             docs.to_doc().force_break()
         }
+    }
 
     fn pipe_capture_right_hand_side<'a>(&mut self, fun: &'a UntypedExpr) -> Document<'a> {
         let (fun, args) = match fun {

crates/aiken-lang/src/parser.rs

@@ -54,7 +54,6 @@ pub fn module(
                 extra.empty_lines.push(span.start);
                 None
             }
-            Token::NewLine => None,
             Token::LeftParen => {
                 if previous_is_newline {
                     Some((Token::NewLineLeftParen, *span))
@@ -62,6 +61,14 @@ pub fn module(
                     Some((Token::LeftParen, *span))
                 }
             }
+            Token::Pipe => {
+                if previous_is_newline {
+                    Some((Token::NewLinePipe, *span))
+                } else {
+                    Some((Token::Pipe, *span))
+                }
+            }
+            Token::NewLine => None,
             _ => Some((token, *span)),
         };
         previous_is_newline = current_is_newline;
@@ -1309,17 +1316,30 @@ pub fn expr_parser(
         // Pipeline
         disjunction
             .clone()
-            .then(just(Token::Pipe).ignore_then(disjunction).repeated())
+            .then(
-            .foldl(|l, r| {
+                choice((just(Token::Pipe), just(Token::NewLinePipe)))
-                let expressions = if let expr::UntypedExpr::PipeLine { mut expressions } = l {
+                    .then(disjunction)
+                    .repeated(),
+            )
+            .foldl(|l, (pipe, r)| {
+                if let expr::UntypedExpr::PipeLine {
+                    mut expressions,
+                    one_liner,
+                } = l
+                {
                     expressions.push(r);
-                    expressions
+                    return expr::UntypedExpr::PipeLine {
-                } else {
+                        expressions,
+                        one_liner,
+                    };
+                }
+
                 let mut expressions = Vec1::new(l);
                 expressions.push(r);
-                    expressions
+                expr::UntypedExpr::PipeLine {
-                };
+                    expressions,
-                expr::UntypedExpr::PipeLine { expressions }
+                    one_liner: pipe != Token::NewLinePipe,
+                }
             })
     })
 }

crates/aiken-lang/src/parser/token.rs

@@ -49,6 +49,7 @@ pub enum Token {
     Vbar,        // '|'
     VbarVbar,    // '||'
     AmperAmper,  // '&&'
+    NewLinePipe, // '↳|>'
     Pipe,        // '|>'
     Dot,         // '.'
     RArrow,      // '->'
@@ -134,6 +135,7 @@ impl fmt::Display for Token {
             Token::Vbar => "|",
             Token::VbarVbar => "||",
             Token::AmperAmper => "&&",
+            Token::NewLinePipe => "↳|>",
             Token::Pipe => "|>",
             Token::Dot => ".",
             Token::RArrow => "->",

crates/aiken-lang/src/pretty.rs

@@ -139,6 +139,7 @@ pub enum Document<'a> {
     Break {
         broken: &'a str,
         unbroken: &'a str,
+        break_first: bool,
         kind: BreakKind,
     },
 
@@ -207,6 +208,7 @@ fn fits(
             Document::Group(doc) => docs.push_front((indent, Mode::Unbroken, doc)),
 
             Document::Str(s) => limit -= s.len() as isize,
+
             Document::String(s) => limit -= s.len() as isize,
 
             Document::Break { unbroken, .. } => match mode {
@@ -255,31 +257,39 @@ fn format(
             Document::Break {
                 broken,
                 unbroken,
+                break_first,
                 kind: BreakKind::Flex,
             } => {
                 let unbroken_width = width + unbroken.len() as isize;
 
                 if fits(limit, unbroken_width, docs.clone()) {
                     writer.push_str(unbroken);
-
                     width = unbroken_width;
-                } else {
+                    continue;
-                    writer.push_str(broken);
+                }
 
+                if *break_first {
                     writer.push('\n');
-
                     for _ in 0..indent {
                         writer.push(' ');
                     }
+                    writer.push_str(broken);
+                } else {
+                    writer.push_str(broken);
+                    writer.push('\n');
+                    for _ in 0..indent {
+                        writer.push(' ');
+                    }
+                }
 
                 width = indent;
             }
-            }
 
             // Strict breaks are conditional to the mode
             Document::Break {
                 broken,
                 unbroken,
+                break_first,
                 kind: BreakKind::Strict,
             } => {
                 width = match mode {
@@ -289,6 +299,18 @@ fn format(
                         width + unbroken.len() as isize
                     }
 
+                    Mode::Broken | Mode::ForcedBroken if *break_first => {
+                        writer.push('\n');
+
+                        for _ in 0..indent {
+                            writer.push(' ');
+                        }
+
+                        writer.push_str(broken);
+
+                        indent
+                    }
+
                     Mode::Broken | Mode::ForcedBroken => {
                         writer.push_str(broken);
 
@@ -361,6 +383,16 @@ pub fn break_<'a>(broken: &'a str, unbroken: &'a str) -> Document<'a> {
         broken,
         unbroken,
         kind: BreakKind::Strict,
+        break_first: false,
+    }
+}
+
+pub fn prebreak<'a>(broken: &'a str, unbroken: &'a str) -> Document<'a> {
+    Document::Break {
+        broken,
+        unbroken,
+        kind: BreakKind::Strict,
+        break_first: true,
     }
 }
 
@@ -369,6 +401,16 @@ pub fn flex_break<'a>(broken: &'a str, unbroken: &'a str) -> Document<'a> {
         broken,
         unbroken,
         kind: BreakKind::Flex,
+        break_first: false,
+    }
+}
+
+pub fn flex_prebreak<'a>(broken: &'a str, unbroken: &'a str) -> Document<'a> {
+    Document::Break {
+        broken,
+        unbroken,
+        kind: BreakKind::Flex,
+        break_first: true,
     }
 }
 

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

@@ -560,3 +560,50 @@ fn test_unicode() {
 
     assert_fmt(src, src);
 }
+
+#[test]
+fn test_preserve_pipe() {
+    let src = indoc! { r#"
+        fn foo() {
+          a |> b |> c |> d
+        }
+
+        fn foo() {
+          a
+            // Foo
+            |> b// Some comments
+            |> c
+            |> d
+        }
+
+        fn baz() {
+          // Commented
+          however |> it_automatically_breaks |> into_multiple_lines |> anytime_when |> it_is_too_long // What?
+        }
+    "#};
+
+    let expected = indoc! { r#"
+        fn foo() {
+          a |> b |> c |> d
+        }
+
+        fn foo() {
+          a // Foo
+            |> b // Some comments
+            |> c
+            |> d
+        }
+
+        fn baz() {
+          // Commented
+          however
+            |> it_automatically_breaks
+            |> into_multiple_lines
+            |> anytime_when
+            |> it_is_too_long
+          // What?
+        }
+    "#};
+
+    assert_fmt(src, expected);
+}

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

@@ -545,6 +545,7 @@ fn pipeline() {
                 tipo: (),
             }],
             body: expr::UntypedExpr::PipeLine {
+                one_liner: false,
                 expressions: vec1::vec1![
                     expr::UntypedExpr::BinOp {
                         location: Span::new((), 31..36),
@@ -717,6 +718,7 @@ fn let_bindings() {
                     expr::UntypedExpr::Assignment {
                         location: Span::new((), 23..70),
                         value: Box::new(expr::UntypedExpr::PipeLine {
+                            one_liner: false,
                             expressions: vec1::vec1![
                                 expr::UntypedExpr::BinOp {
                                     location: Span::new((), 35..40),
@@ -1105,6 +1107,7 @@ fn anonymous_function() {
                         annotation: None,
                     },
                     expr::UntypedExpr::PipeLine {
+                        one_liner: true,
                         expressions: vec1::vec1![
                             expr::UntypedExpr::Int {
                                 location: Span::new((), 71..72),

crates/aiken-lang/src/tipo/expr.rs

@@ -274,7 +274,7 @@ impl<'a, 'b> ExprTyper<'a, 'b> {
                 location, value, ..
             } => Ok(self.infer_string(value, location)),
 
-            UntypedExpr::PipeLine { expressions } => self.infer_pipeline(expressions),
+            UntypedExpr::PipeLine { expressions, .. } => self.infer_pipeline(expressions),
 
             UntypedExpr::Fn {
                 location,