avoid inserting newline between curly brace and comment

Author: RalfJung

src/parse/session.rs

@@ -1,3 +1,4 @@
+use std::ops::Range;
 use std::path::Path;
 use std::sync::atomic::{AtomicBool, Ordering};
 
@@ -237,6 +238,15 @@ impl ParseSess {
         }
     }
 
+    pub(crate) fn line_bounds(&self, pos: BytePos) -> Option<Range<BytePos>> {
+        let line = self.raw_psess.source_map().lookup_line(pos).ok();
+
+        match line {
+            Some(line_info) => Some(line_info.sf.line_bounds(line_info.line)),
+            None => None,
+        }
+    }
+
     pub(crate) fn line_of_byte_pos(&self, pos: BytePos) -> usize {
         self.raw_psess.source_map().lookup_char_pos(pos).line
     }

src/visitor.rs

@@ -223,6 +223,41 @@ impl<'b, 'a: 'b> FmtVisitor<'a> {
         self.push_str("{");
         self.trim_spaces_after_opening_brace(b, inner_attrs);
 
+        // Try to detect comments that refer to the block, not the first statement in the block.
+        if has_braces {
+            let block_line_range = self.psess.lookup_line_range(b.span);
+            if block_line_range.lo != block_line_range.hi {
+                // Skipping if a single line block
+                // Make sure there is no code on the first line we have to worry about.
+                // That would also be ambiguous: what should `if { /* comment */ statement;`
+                // even do -- should the comment be attached to the `if` or the `statement`?
+                // The current logic answers this with "statement".
+                let first_line_contains_stmt = if let Some(first_stmt) = b.stmts.first() {
+                    self.psess.lookup_line_range(first_stmt.span).lo == block_line_range.lo
+                } else {
+                    false
+                };
+                if !first_line_contains_stmt {
+                    let first_line_bounds = self.psess.line_bounds(self.last_pos).unwrap();
+                    let first_line_snip = self
+                        .snippet(mk_sp(self.last_pos, first_line_bounds.end))
+                        .trim();
+
+                    if contains_comment(first_line_snip) {
+                        if let Ok(comment) =
+                            rewrite_comment(first_line_snip, false, self.shape(), self.config)
+                        {
+                            self.push_str(" ");
+                            self.push_str(&comment);
+                            // This line has no statements, so we can jump to the end of it
+                            // (but *before* the newline).
+                            self.last_pos = first_line_bounds.end - BytePos::from_usize(1);
+                        }
+                    }
+                }
+            }
+        }
+
         // Format inner attributes if available.
         if let Some(attrs) = inner_attrs {
             self.visit_attrs(attrs, ast::AttrStyle::Inner);

tests/source/issue-3255.rs

@@ -0,0 +1,26 @@
+fn foo(){  
+    if true {      // Sample comment
+        1
+    }
+}
+
+
+fn foo(){  
+    if true {
+        // Sample comment
+        1
+    }
+}
+
+fn foo(){  
+    if true {    /* Sample comment */
+        1
+    }
+}
+
+fn foo(){  
+    if true {     /* Sample
+            comment */
+        1
+    }
+}

tests/target/async_fn.rs

@@ -13,8 +13,7 @@ async unsafe fn foo() {
 }
 
 async unsafe fn rust() {
-    async move {
-        // comment
+    async move { // comment
         Ok(())
     }
 }

tests/target/control-brace-style-always-next-line.rs

@@ -20,8 +20,7 @@ fn main() {
     }
 
     'while_label: while cond
-    {
-        // while comment
+    { // while comment
         ();
     }
 

tests/target/control-brace-style-always-same-line.rs

@@ -15,8 +15,7 @@ fn main() {
         ();
     }
 
-    'while_label: while cond {
-        // while comment
+    'while_label: while cond { // while comment
         ();
     }
 

tests/target/issue-3255.rs

@@ -0,0 +1,25 @@
+fn foo() {
+    if true { // Sample comment
+        1
+    }
+}
+
+fn foo() {
+    if true {
+        // Sample comment
+        1
+    }
+}
+
+fn foo() {
+    if true { /* Sample comment */
+        1
+    }
+}
+
+fn foo() {
+    if true { /* Sample
+            comment */
+        1
+    }
+}

tests/target/label_break.rs

@@ -19,8 +19,7 @@ fn main() {
             // comment
             break 'block 1;
         }
-        if bar() {
-            /* comment */
+        if bar() { /* comment */
             break 'block 2;
         }
         3

tests/target/match.rs

@@ -7,8 +7,7 @@ fn foo() {
         // Some comment.
         a => foo(),
         b if 0 < 42 => foo(),
-        c => {
-            // Another comment.
+        c => { // Another comment.
             // Comment.
             an_expression;
             foo()