Prevent duplicating comments and tokens when rewriting trait headers

src/header.rs

@@ -8,13 +8,14 @@
 use std::borrow::Cow;
 
 use rustc_ast as ast;
-use rustc_span::Span;
 use rustc_span::symbol::Ident;
+use rustc_span::{DUMMY_SP, Span};
 use tracing::debug;
 
 use crate::comment::{combine_strs_with_missing_comments, contains_comment};
 use crate::rewrite::RewriteContext;
 use crate::shape::Shape;
+use crate::source_map::SpanUtils;
 use crate::utils::rewrite_ident;
 
 pub(crate) fn format_header(
@@ -68,6 +69,23 @@ pub(crate) struct HeaderPart<'a> {
 }
 
 impl<'a> HeaderPart<'a> {
+    /// Create a new empty HeaderPart with a DUMMY_SP.
+    /// Because the snippet is empty these will be skipped when calling `format_header`.
+    fn empty() -> Self {
+        Self::new("", DUMMY_SP)
+    }
+
+    /// Find the `keyword` for the header within the given `span`
+    pub(crate) fn keyword(
+        context: &RewriteContext<'_>,
+        span: Span,
+        keyword: impl Into<Cow<'a, str>>,
+    ) -> Self {
+        let snippet = keyword.into();
+        let span = context.snippet_provider.subspan(span, &snippet);
+        Self::new(snippet, span)
+    }
+
     pub(crate) fn new(snippet: impl Into<Cow<'a, str>>, span: Span) -> Self {
         Self {
             snippet: snippet.into(),
@@ -103,4 +121,26 @@ impl<'a> HeaderPart<'a> {
 
         Self::new(snippet, vis.span)
     }
+
+    pub(crate) fn constness(constness: ast::Const) -> Self {
+        match constness {
+            ast::Const::No => Self::empty(),
+            ast::Const::Yes(span) => Self::new("const", span),
+        }
+    }
+
+    pub(crate) fn safety(safety: ast::Safety) -> Self {
+        match safety {
+            ast::Safety::Default => Self::empty(),
+            ast::Safety::Unsafe(span) => Self::new("unsafe", span),
+            ast::Safety::Safe(span) => Self::new("safe", span),
+        }
+    }
+
+    pub(crate) fn auto(context: &RewriteContext<'_>, span: Span, is_auto: ast::IsAuto) -> Self {
+        match is_auto {
+            ast::IsAuto::No => Self::empty(),
+            ast::IsAuto::Yes => Self::keyword(context, span, "auto"),
+        }
+    }
 }

src/items.rs

@@ -21,6 +21,7 @@ use crate::expr::{
     rewrite_assign_rhs_with, rewrite_assign_rhs_with_comments, rewrite_else_kw_with_comments,
     rewrite_let_else_block,
 };
+use crate::header::HeaderPart;
 use crate::lists::{ListFormatting, Separator, definitive_tactic, itemize_list, write_list};
 use crate::macros::{MacroPosition, rewrite_macro};
 use crate::overflow;
@@ -1165,20 +1166,39 @@ pub(crate) fn format_trait(
     } = *trait_;
 
     let mut result = String::with_capacity(128);
-    let header = format!(
-        "{}{}{}{}trait ",
-        format_visibility(context, &item.vis),
-        format_constness(constness),
-        format_safety(safety),
-        format_auto(is_auto),
-    );
-    result.push_str(&header);
+    let header = vec![
+        HeaderPart::visibility(context, &item.vis),
+        HeaderPart::constness(constness),
+        HeaderPart::safety(safety),
+        HeaderPart::auto(context, item.span, is_auto),
+        HeaderPart::keyword(context, item.span, "trait"),
+    ];
+
+    let shape = Shape::indented(offset, context.config);
+    let header_rewrite = crate::header::format_header(context, shape, header);
 
-    let body_lo = context.snippet_provider.span_after(item.span, "{");
+    result.push_str(&header_rewrite);
 
-    let shape = Shape::indented(offset, context.config).offset_left(result.len(), item.span)?;
+    // +1 to account for space between `trait` and the `ident`
+    let offset_left = last_line_width(&result) + 1;
+    let shape = Shape::indented(offset, context.config).offset_left(offset_left, item.span)?;
     let generics_str = rewrite_generics(context, rewrite_ident(context, ident), generics, shape)?;
-    result.push_str(&generics_str);
+
+    let comment_span_lo = context.snippet_provider.span_after(item.span, "trait");
+    let comment_span = mk_sp(comment_span_lo, ident.span.lo() - BytePos(1));
+
+    result = combine_strs_with_missing_comments(
+        context,
+        &result,
+        &generics_str,
+        comment_span,
+        // infinite_width so we don't put the `ident + generics_str` on the next line
+        shape.infinite_width(),
+        true,
+    )?;
+
+    // Keep track of the last position that we've written. This will help us recover comments
+    let mut current_rewite_lo = generics.span.hi();
 
     // FIXME(#2055): rustfmt fails to format when there are comments between trait bounds.
     if !bounds.is_empty() {
@@ -1199,18 +1219,15 @@ pub(crate) fn format_trait(
             &RhsAssignKind::Bounds,
             RhsTactics::ForceNextLineWithoutIndent,
         )?;
+
+        current_rewite_lo = bound_hi;
     }
 
     // Rewrite where-clause.
     if !generics.where_clause.predicates.is_empty() {
         let where_on_new_line = context.config.indent_style() != IndentStyle::Block;
 
         let where_budget = context.budget(last_line_width(&result));
-        let pos_before_where = if bounds.is_empty() {
-            generics.where_clause.span.lo()
-        } else {
-            bounds[bounds.len() - 1].span().hi()
-        };
         let option = WhereClauseOption::snuggled(&generics_str);
         let where_clause_str = rewrite_where_clause(
             context,
@@ -1220,7 +1237,7 @@ pub(crate) fn format_trait(
             where_on_new_line,
             "{",
             None,
-            pos_before_where,
+            current_rewite_lo,
             option,
         )?;
 
@@ -1236,27 +1253,23 @@ pub(crate) fn format_trait(
         }
         result.push_str(&where_clause_str);
     } else {
-        let item_snippet = context.snippet(item.span);
-        if let Some(lo) = item_snippet.find('/') {
-            // 1 = `{`
-            let comment_hi = if generics.params.len() > 0 {
-                generics.span.lo() - BytePos(1)
-            } else {
-                body_lo - BytePos(1)
-            };
-            let comment_lo = item.span.lo() + BytePos(lo as u32);
-            if comment_lo < comment_hi {
-                match recover_missing_comment_in_span(
-                    mk_sp(comment_lo, comment_hi),
-                    Shape::indented(offset, context.config),
-                    context,
-                    last_line_width(&result),
-                ) {
-                    Ok(ref missing_comment) if !missing_comment.is_empty() => {
-                        result.push_str(missing_comment);
-                    }
-                    _ => (),
+        let span_rest = mk_sp(current_rewite_lo, item.span.hi());
+        let body_lo = context.snippet_provider.span_before(span_rest, "{");
+        let comment_span = span_rest.with_hi(body_lo);
+        let snippet = context.snippet(comment_span);
+
+        if let Some(lo) = snippet.find('/') {
+            let comment_span = comment_span.with_lo(comment_span.lo() + BytePos(lo as u32));
+            match recover_missing_comment_in_span(
+                comment_span,
+                Shape::indented(offset, context.config),
+                context,
+                last_line_width(&result),
+            ) {
+                Ok(ref missing_comment) if !missing_comment.is_empty() => {
+                    result.push_str(missing_comment);
                 }
+                _ => (),
             }
         }
     }

src/macros.rs

@@ -435,11 +435,9 @@ pub(crate) fn rewrite_macro_def(
         let pos = context.snippet_provider.span_after(span, "macro_rules!");
         vec![HeaderPart::new("macro_rules!", span.with_hi(pos))]
     } else {
-        let macro_lo = context.snippet_provider.span_before(span, "macro");
-        let macro_hi = macro_lo + BytePos("macro".len() as u32);
         vec![
             HeaderPart::visibility(context, vis),
-            HeaderPart::new("macro", mk_sp(macro_lo, macro_hi)),
+            HeaderPart::keyword(context, span, "macro"),
         ]
     };
 

src/source_map.rs

@@ -5,9 +5,11 @@ use rustc_span::{BytePos, Span};
 
 use crate::comment::FindUncommented;
 use crate::config::file_lines::LineRange;
+use crate::utils::mk_sp;
 use crate::visitor::SnippetProvider;
 
 pub(crate) trait SpanUtils {
+    fn subspan(&self, original: Span, needle: &str) -> Span;
     fn span_after(&self, original: Span, needle: &str) -> BytePos;
     fn span_after_last(&self, original: Span, needle: &str) -> BytePos;
     fn span_before(&self, original: Span, needle: &str) -> BytePos;
@@ -26,6 +28,12 @@ pub(crate) trait LineRangeUtils {
 }
 
 impl SpanUtils for SnippetProvider {
+    fn subspan(&self, original: Span, needle: &str) -> Span {
+        let lo = self.span_before(original, needle);
+        let hi = lo + BytePos(needle.len() as u32);
+        mk_sp(lo, hi)
+    }
+
     fn span_after(&self, original: Span, needle: &str) -> BytePos {
         self.opt_span_after(original, needle).unwrap_or_else(|| {
             panic!(

src/utils.rs

@@ -117,14 +117,6 @@ pub(crate) fn format_safety(unsafety: ast::Safety) -> &'static str {
     }
 }
 
-#[inline]
-pub(crate) fn format_auto(is_auto: ast::IsAuto) -> &'static str {
-    match is_auto {
-        ast::IsAuto::Yes => "auto ",
-        ast::IsAuto::No => "",
-    }
-}
-
 #[inline]
 pub(crate) fn format_mutability(mutability: ast::Mutability) -> &'static str {
     match mutability {

tests/source/issue_5841.rs

@@ -0,0 +1,2 @@
+pub // should the following be private?
+trait Asso<'ciativity> {}

tests/source/trait.rs

@@ -102,9 +102,36 @@ trait FooBar = Foo
 
 // #2637
 auto trait Example {}
+
+auto /* post auto comment */ trait AutoExampleWithComments {}
+
 pub auto trait PubExample {}
+
+pub /* post pub comment */ auto /* post auto comment */ trait PubAutoExampleWithComments {}
+
 pub unsafe auto trait PubUnsafeExample {}
 
+pub // post pub comment
+unsafe
+/* post unsafe comment */
+auto /* post auto comment */ trait PubUnsafeAutoExampleWithComments {}
+
+pub // post pub comment
+const /* post const comment */
+unsafe
+/* post unsafe comment */
+auto /* post auto comment */ trait PubConstUnsafeAutoExampleWithComments {}
+
+pub /* first */ trait /* second */ MyTrait /* third */ {}
+
+pub
+   /* a */
+  trait
+    /* b */
+      Foo
+        /* c */
+{}
+
 // #3006
 trait Foo<'a> {
     type Bar<  'a  >;

tests/target/issue_5841.rs

@@ -0,0 +1,4 @@
+pub // should the following be private?
+trait Asso<'ciativity>
+{
+}

tests/target/issue_6160.rs

@@ -0,0 +1 @@
+pub/* (crate) */ trait MyTrait {}

tests/target/trait.rs

@@ -137,9 +137,40 @@ trait FooBar = Foo
 
 // #2637
 auto trait Example {}
+
+auto /* post auto comment */ trait AutoExampleWithComments {}
+
 pub auto trait PubExample {}
+
+pub /* post pub comment */ auto /* post auto comment */ trait PubAutoExampleWithComments {}
+
 pub unsafe auto trait PubUnsafeExample {}
 
+pub // post pub comment
+unsafe
+/* post unsafe comment */
+auto /* post auto comment */ trait PubUnsafeAutoExampleWithComments
+{
+}
+
+pub // post pub comment
+const /* post const comment */
+unsafe
+/* post unsafe comment */
+auto /* post auto comment */ trait PubConstUnsafeAutoExampleWithComments
+{
+}
+
+pub /* first */ trait /* second */ MyTrait /* third */ {}
+
+pub
+   /* a */
+  trait
+/* b */
+Foo /* c */
+{
+}
+
 // #3006
 trait Foo<'a> {
     type Bar<'a>;