From f9836f9c3e9c14b82f436d0dfe5034c3bc79c48d Mon Sep 17 00:00:00 2001
From: Markus Westerlind <marwes91@gmail.com>
Date: Mon, 3 Jul 2017 08:34:01 +0200
Subject: [PATCH] fix(format): Preserve comments between let bindings and their
 bodies

And type bindings
---
 base/src/pretty_print.rs     | 43 ++++++++++++++++++++----------------
 parser/tests/pretty_print.rs | 15 +++++++++++++
 2 files changed, 39 insertions(+), 19 deletions(-)

diff --git a/base/src/pretty_print.rs b/base/src/pretty_print.rs
index 9241df2f7e..8322f82c47 100644
--- 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
index 94f050462c..a28d4b29e8 100644
--- 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);
+}