Separate comments with different offset into different groups

src/comment.rs

@@ -1162,14 +1162,6 @@ impl FullCodeCharKind {
         self == FullCodeCharKind::InStringCommented
             || self == FullCodeCharKind::StartStringCommented
     }
-
-    fn to_codecharkind(self) -> CodeCharKind {
-        if self.is_comment() {
-            CodeCharKind::Comment
-        } else {
-            CodeCharKind::Normal
-        }
-    }
 }
 
 impl<T> CharClasses<T>
@@ -1484,16 +1476,24 @@ impl<'a> Iterator for UngroupedCommentCodeSlices<'a> {
 /// functional text. Line style comments contain their ending newlines.
 pub(crate) struct CommentCodeSlices<'a> {
     slice: &'a str,
-    last_slice_kind: CodeCharKind,
-    last_slice_end: usize,
+    ungrouped_code_slices: MultiPeek<UngroupedCommentCodeSlices<'a>>,
+    offset: Option<usize>,
 }
 
 impl<'a> CommentCodeSlices<'a> {
     pub(crate) fn new(slice: &'a str) -> CommentCodeSlices<'a> {
         CommentCodeSlices {
             slice,
-            last_slice_kind: CodeCharKind::Comment,
-            last_slice_end: 0,
+            ungrouped_code_slices: multipeek(UngroupedCommentCodeSlices::new(slice)),
+            offset: None,
+        }
+    }
+
+    pub(crate) fn with_offset(slice: &'a str, offset: usize) -> CommentCodeSlices<'a> {
+        CommentCodeSlices {
+            slice,
+            ungrouped_code_slices: multipeek(UngroupedCommentCodeSlices::new(slice)),
+            offset: Some(offset).filter(|o| *o != 0),
         }
     }
 }
@@ -1502,59 +1502,50 @@ impl<'a> Iterator for CommentCodeSlices<'a> {
     type Item = (CodeCharKind, usize, &'a str);
 
     fn next(&mut self) -> Option<Self::Item> {
-        if self.last_slice_end == self.slice.len() {
-            return None;
-        }
-
-        let mut sub_slice_end = self.last_slice_end;
-        let mut first_whitespace = None;
-        let subslice = &self.slice[self.last_slice_end..];
-        let mut iter = CharClasses::new(subslice.char_indices());
-
-        for (kind, (i, c)) in &mut iter {
-            let is_comment_connector = self.last_slice_kind == CodeCharKind::Normal
-                && &subslice[..2] == "//"
-                && [' ', '\t'].contains(&c);
-
-            if is_comment_connector && first_whitespace.is_none() {
-                first_whitespace = Some(i);
+        let first_chunk = self.ungrouped_code_slices.next()?;
+        if first_chunk.0 == CodeCharKind::Normal {
+            if !first_chunk.2.trim().is_empty() {
+                self.offset = Some(last_line_width(first_chunk.2)).filter(|o| *o != 0);
             }
+            return Some(first_chunk);
+        }
 
-            if kind.to_codecharkind() == self.last_slice_kind && !is_comment_connector {
-                let last_index = match first_whitespace {
-                    Some(j) => j,
-                    None => i,
-                };
-                sub_slice_end = self.last_slice_end + last_index;
-                break;
-            }
+        let mut comment_end_index = first_chunk.1 + first_chunk.2.len();
+        while let Some(&(k, i, s)) = self.ungrouped_code_slices.peek() {
+            match k {
+                CodeCharKind::Comment if self.offset.is_none() => {
+                    comment_end_index = i + s.len();
+                    self.ungrouped_code_slices.next()?;
+                }
+                CodeCharKind::Comment => break,
+                CodeCharKind::Normal if s.trim().is_empty() && count_newlines(s) == 0 => {
+                    let indent_width = s.len();
+                    if self.offset.map_or(false, |comment_offset| {
+                        !(indent_width < comment_offset + 2 && comment_offset < indent_width + 2)
+                    }) {
+                        break;
+                    }
 
-            if !is_comment_connector {
-                first_whitespace = None;
+                    match self.ungrouped_code_slices.peek() {
+                        Some((CodeCharKind::Comment, index, s)) => {
+                            comment_end_index = index + s.len();
+                            // Advance twice.
+                            self.ungrouped_code_slices.next()?;
+                            self.ungrouped_code_slices.next()?;
+                        }
+                        _ => break,
+                    }
+                }
+                CodeCharKind::Normal => break,
             }
         }
 
-        if let (None, true) = (iter.next(), sub_slice_end == self.last_slice_end) {
-            // This was the last subslice.
-            sub_slice_end = match first_whitespace {
-                Some(i) => self.last_slice_end + i,
-                None => self.slice.len(),
-            };
-        }
-
-        let kind = match self.last_slice_kind {
-            CodeCharKind::Comment => CodeCharKind::Normal,
-            CodeCharKind::Normal => CodeCharKind::Comment,
-        };
-        let res = (
-            kind,
-            self.last_slice_end,
-            &self.slice[self.last_slice_end..sub_slice_end],
-        );
-        self.last_slice_end = sub_slice_end;
-        self.last_slice_kind = kind;
-
-        Some(res)
+        let comment_start_index = first_chunk.1;
+        Some((
+            CodeCharKind::Comment,
+            comment_start_index,
+            &self.slice[comment_start_index..comment_end_index],
+        ))
     }
 }
 
@@ -1728,7 +1719,6 @@ mod test {
         let input = "// comment\n    test();";
         let mut iter = CommentCodeSlices::new(input);
 
-        assert_eq!((CodeCharKind::Normal, 0, ""), iter.next().unwrap());
         assert_eq!(
             (CodeCharKind::Comment, 0, "// comment\n"),
             iter.next().unwrap()
@@ -1742,18 +1732,60 @@ mod test {
 
     #[test]
     fn comment_code_slices_three() {
-        let input = "1 // comment\n    // comment2\n\n";
+        let input = "1 // comment\n  // comment2\n\n";
         let mut iter = CommentCodeSlices::new(input);
 
         assert_eq!((CodeCharKind::Normal, 0, "1 "), iter.next().unwrap());
         assert_eq!(
-            (CodeCharKind::Comment, 2, "// comment\n    // comment2\n"),
+            (CodeCharKind::Comment, 2, "// comment\n  // comment2\n"),
+            iter.next().unwrap()
+        );
+        assert_eq!((CodeCharKind::Normal, 27, "\n"), iter.next().unwrap());
+        assert_eq!(None, iter.next());
+    }
+
+    #[test]
+    fn comment_code_slices_four() {
+        let input = r#"
+if x == 3 {
+    x = 4;
+} // if x == 3
+// end of block
+"#;
+        let mut iter = CommentCodeSlices::new(input);
+
+        assert_eq!(
+            (CodeCharKind::Normal, 0, r#"
+if x == 3 {
+    x = 4;
+} "#),
+            iter.next().unwrap()
+        );
+        assert_eq!((CodeCharKind::Comment, 26, "// if x == 3\n",), iter.next().unwrap());
+        assert_eq!(
+            (CodeCharKind::Comment, 39, "// end of block\n"),
             iter.next().unwrap()
         );
-        assert_eq!((CodeCharKind::Normal, 29, "\n"), iter.next().unwrap());
         assert_eq!(None, iter.next());
     }
 
+    #[test]
+    fn comment_code_slices_five() {
+        let input = "1 // comment\r\n\r\n  // comment2\r\n";
+        let mut iter = CommentCodeSlices::new(input);
+
+        assert_eq!((CodeCharKind::Normal, 0, "1 "), iter.next().unwrap());
+        assert_eq!(
+            (CodeCharKind::Comment, 2, "// comment\r\n"),
+            iter.next().unwrap()
+        );
+        assert_eq!((CodeCharKind::Normal, 14, "\r\n  ",), iter.next().unwrap());
+        assert_eq!(
+            (CodeCharKind::Comment, 18, "// comment2\r\n"),
+            iter.next().unwrap()
+        );
+        assert_eq!(None, iter.next());
+    }
     #[test]
     #[rustfmt::skip]
     fn format_doc_comments() {

src/missed_spans.rs

@@ -6,7 +6,9 @@ use crate::config::FileName;
 use crate::coverage::transform_missing_snippet;
 use crate::shape::{Indent, Shape};
 use crate::source_map::LineRangeUtils;
-use crate::utils::{count_lf_crlf, count_newlines, last_line_width, mk_sp};
+use crate::utils::{
+    count_lf_crlf, count_newlines, last_line_contains_single_line_comment, last_line_width, mk_sp,
+};
 use crate::visitor::FmtVisitor;
 
 struct SnippetStatus {
@@ -183,7 +185,12 @@ impl<'a> FmtVisitor<'a> {
                 );
                 (lf_count, crlf_count, within_file_lines_range)
             };
-        for (kind, offset, subslice) in CommentCodeSlices::new(snippet) {
+        let last_line_offset = if last_line_contains_single_line_comment(&self.buffer) {
+            0
+        } else {
+            last_line_width(&self.buffer)
+        };
+        for (kind, offset, subslice) in CommentCodeSlices::with_offset(snippet, last_line_offset) {
             debug!("{:?}: {:?}", kind, subslice);
 
             let (lf_count, crlf_count, within_file_lines_range) =

src/visitor.rs

@@ -23,8 +23,9 @@ use crate::spanned::Spanned;
 use crate::stmt::Stmt;
 use crate::syntux::session::ParseSess;
 use crate::utils::{
-    self, contains_skip, count_newlines, depr_skip_annotation, inner_attributes, last_line_width,
-    mk_sp, ptr_vec_to_ref_vec, rewrite_ident, stmt_expr,
+    self, contains_skip, count_newlines, depr_skip_annotation, inner_attributes,
+    last_line_contains_single_line_comment, last_line_width, mk_sp, ptr_vec_to_ref_vec,
+    rewrite_ident, stmt_expr,
 };
 use crate::{ErrorKind, FormatReport, FormattingError};
 
@@ -249,7 +250,14 @@ 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 last_line_offset = if last_line_contains_single_line_comment(&self.buffer) {
+            0
+        } else {
+            last_line_width(&self.buffer)
+        };
+        for (kind, offset, sub_slice) in
+            CommentCodeSlices::with_offset(self.snippet(span), last_line_offset)
+        {
             let sub_slice = transform_missing_snippet(config, sub_slice);
             debug!("close_block: {:?} {:?} {:?}", kind, offset, sub_slice);
 
@@ -261,7 +269,8 @@ impl<'b, 'a: 'b> FmtVisitor<'a> {
                     }
                     let span_in_between = mk_sp(last_hi, span.lo() + BytePos::from_usize(offset));
                     let snippet_in_between = self.snippet(span_in_between);
-                    let mut comment_on_same_line = !snippet_in_between.contains("\n");
+                    let mut comment_on_same_line = !snippet_in_between.contains("\n")
+                        && !last_line_contains_single_line_comment(&self.buffer);
 
                     let mut comment_shape =
                         Shape::indented(self.block_indent, config).comment(config);