fix(format): Preserve comments between let bindings and their bodies

And type bindings
gluon-lang · Jul 3, 2017 · f9836f9 · f9836f9
1 parent 7b53849
commit f9836f9
Show file tree

Hide file tree

Showing 2 changed files with 39 additions and 19 deletions.
diff --git a/base/src/pretty_print.rs b/base/src/pretty_print.rs
@@ -208,12 +208,33 @@ impl<'a> ExprPrinter<'a> {
         arena.concat(repeat(arena.newline()).take(empty_lines))
     }
 
+    fn comments(&'a self, span: Span<BytePos>) -> DocBuilder<'a, Arena<'a>> {
+        let arena = &self.arena;
+        arena
+            .concat(
+                self.source
+                    .comments_between(span)
+                    .map(|comment| {
+                        chain![arena;
+                            if comment.is_empty() {
+                                arena.nil()
+                            } else {
+                                arena.text("// ").append(comment)
+                            },
+                            arena.newline()
+                        ]
+                    }),
+            )
+    }
+
     pub fn pretty_expr<Id>(&'a self, expr: &'a SpannedExpr<Id>) -> DocBuilder<'a, Arena<'a>>
     where
         Id: AsRef<str>,
     {
         self.pretty_expr_(BytePos::from(0), expr)
+            .append(self.comments(Span::new(expr.span.end, BytePos::from(self.source.src().len()))))
     }
+
     pub fn pretty_expr_<Id>(
         &'a self,
         previous_end: BytePos,
@@ -324,8 +345,7 @@ impl<'a> ExprPrinter<'a> {
                     arena.concat(prefixes.zip(binds).map(|(prefix, bind)| {
                         binding(prefix, bind)
                     }).interleave(newlines_iter!(self, binds.iter().map(|bind| bind.span())))),
-                    self.newlines(binds.last().unwrap().expr.span.end, body.span.start),
-                    pretty(body).group()
+                    self.pretty_expr_(binds.last().unwrap().span().end, body).group()
                 ]
             }
             Expr::Literal(_) => {
@@ -413,27 +433,12 @@ impl<'a> ExprPrinter<'a> {
                                 .group()
                         ].group()
                     }).interleave(newlines_iter!(self, binds.iter().map(|bind| bind.span())))),
-                    self.newlines(binds.last().unwrap().alias.span.end, body.span.start),
-                    pretty(body)
+                    self.pretty_expr_(binds.last().unwrap().alias.span.end, body)
                 ].group()
             }
             Expr::Error => arena.text("<error>"),
         };
-        arena
-            .concat(
-                self.source
-                    .comments_between(Span::new(previous_end, expr.span.start))
-                    .map(|comment| {
-                        chain![arena;
-                            if comment.is_empty() {
-                                arena.nil()
-                            } else {
-                                arena.text("// ").append(comment)
-                            },
-                            arena.newline()
-                        ]
-                    }),
-            )
+        self.comments(Span::new(previous_end, expr.span.start))
             .append(doc)
     }
 

diff --git a/parser/tests/pretty_print.rs b/parser/tests/pretty_print.rs
@@ -106,3 +106,18 @@ fn dont_lose_information_in_literals() {
 "#;
     assert_eq!(&format_expr(expr).unwrap(), expr);
 }
+
+
+#[test]
+fn preserve_comment_between_let_in() {
+    let expr = r#"
+// test
+let x = 1
+// test
+type Test = Int
+// test
+1
+// test2
+"#;
+    assert_eq!(&format_expr(expr).unwrap(), expr);
+}