Fixed 'Incorrect comment indent inside if/else' issue. (#4459)

* Added test cases

* Fixed if condition comment issue

* Fixed extern C issue

* Removed previous test case

* Removed tmp file

* honor the authors intent

* Changed the file name to its original name

* Removed extra whitespace

src/items.rs

@@ -161,6 +161,14 @@ enum BodyElement<'a> {
     ForeignItem(&'a ast::ForeignItem),
 }
 
+impl BodyElement<'_> {
+    pub(crate) fn span(&self) -> Span {
+        match self {
+            BodyElement::ForeignItem(fi) => fi.span(),
+        }
+    }
+}
+
 /// Represents a fn's signature.
 pub(crate) struct FnSig<'a> {
     decl: &'a ast::FnDecl,
@@ -268,19 +276,19 @@ impl<'a> FmtVisitor<'a> {
             self.last_pos = item.span.lo() + BytePos(brace_pos as u32 + 1);
             self.block_indent = self.block_indent.block_indent(self.config);
 
-            if item.body.is_empty() {
-                self.format_missing_no_indent(item.span.hi() - BytePos(1));
-                self.block_indent = self.block_indent.block_unindent(self.config);
-                let indent_str = self.block_indent.to_string(self.config);
-                self.push_str(&indent_str);
-            } else {
+            if !item.body.is_empty() {
+                // Advance to first item (statement or inner attribute)
+                // within the block.
+                self.last_pos = item.body[0].span().lo();
                 for item in &item.body {
                     self.format_body_element(item);
                 }
-
-                self.block_indent = self.block_indent.block_unindent(self.config);
-                self.format_missing_with_indent(item.span.hi() - BytePos(1));
             }
+
+            self.format_missing_no_indent(item.span.hi() - BytePos(1));
+            self.block_indent = self.block_indent.block_unindent(self.config);
+            let indent_str = self.block_indent.to_string(self.config);
+            self.push_str(&indent_str);
         }
 
         self.push_str("}");

src/visitor.rs

@@ -5,7 +5,7 @@ use rustc_ast::{ast, attr::HasAttrs, token::DelimToken, visit};
 use rustc_span::{symbol, BytePos, Pos, Span, DUMMY_SP};
 
 use crate::attr::*;
-use crate::comment::{rewrite_comment, CodeCharKind, CommentCodeSlices};
+use crate::comment::{contains_comment, rewrite_comment, CodeCharKind, CommentCodeSlices};
 use crate::config::Version;
 use crate::config::{BraceStyle, Config};
 use crate::coverage::transform_missing_snippet;
@@ -261,14 +261,23 @@ impl<'b, 'a: 'b> FmtVisitor<'a> {
             trimmed.is_empty() || trimmed.chars().all(|c| c == ';')
         };
 
-        for (kind, offset, sub_slice) in CommentCodeSlices::new(self.snippet(span)) {
+        let comment_snippet = self.snippet(span);
+
+        let align_to_right = if unindent_comment && contains_comment(&comment_snippet) {
+            let first_lines = comment_snippet.splitn(2, '/').next().unwrap_or("");
+            last_line_width(first_lines) > last_line_width(&comment_snippet)
+        } else {
+            false
+        };
+
+        for (kind, offset, sub_slice) in CommentCodeSlices::new(comment_snippet) {
             let sub_slice = transform_missing_snippet(config, sub_slice);
 
             debug!("close_block: {:?} {:?} {:?}", kind, offset, sub_slice);
 
             match kind {
                 CodeCharKind::Comment => {
-                    if !unindented && unindent_comment {
+                    if !unindented && unindent_comment && !align_to_right {
                         unindented = true;
                         self.block_indent = self.block_indent.block_unindent(config);
                     }

tests/source/issue-4120.rs

@@ -0,0 +1,85 @@
+fn main() {
+    let x = if true {
+        1
+            // In if
+    } else {
+        0
+            // In else
+    };
+
+    let x = if true {
+        1
+             /* In if */
+    } else {
+        0
+             /* In else */
+    };
+
+    let z = if true {
+        if true {
+            1
+
+                 // In if level 2
+        } else {
+            2
+        }
+    } else {
+        3
+    };
+
+    let a = if true {
+        1
+  // In if
+    } else {
+        0
+  // In else
+    };
+
+    let a = if true {
+        1
+
+    // In if
+    } else {
+        0
+    // In else
+    };
+
+    let b = if true {
+        1
+
+    // In if
+    } else {
+        0
+        // In else
+    };
+
+    let c = if true {
+        1
+
+        // In if
+    } else {
+        0
+        // In else
+    };
+    for i in 0..2 {
+        println!("Something");
+        // In for
+    }
+
+    for i in 0..2 {
+        println!("Something");
+        /* In for */
+    }
+
+    extern "C" {
+        fn first();
+
+        // In foreign mod
+    }
+
+    extern "C" {
+        fn first();
+
+        /* In foreign mod */
+    }
+}

tests/target/issue-4120.rs

@@ -0,0 +1,85 @@
+fn main() {
+    let x = if true {
+        1
+        // In if
+    } else {
+        0
+        // In else
+    };
+
+    let x = if true {
+        1
+        /* In if */
+    } else {
+        0
+        /* In else */
+    };
+
+    let z = if true {
+        if true {
+            1
+
+            // In if level 2
+        } else {
+            2
+        }
+    } else {
+        3
+    };
+
+    let a = if true {
+        1
+    // In if
+    } else {
+        0
+        // In else
+    };
+
+    let a = if true {
+        1
+
+    // In if
+    } else {
+        0
+        // In else
+    };
+
+    let b = if true {
+        1
+
+    // In if
+    } else {
+        0
+        // In else
+    };
+
+    let c = if true {
+        1
+
+        // In if
+    } else {
+        0
+        // In else
+    };
+    for i in 0..2 {
+        println!("Something");
+        // In for
+    }
+
+    for i in 0..2 {
+        println!("Something");
+        /* In for */
+    }
+
+    extern "C" {
+        fn first();
+
+        // In foreign mod
+    }
+
+    extern "C" {
+        fn first();
+
+        /* In foreign mod */
+    }
+}