keep comment appearing between parameter's name and its type (#3491)

5 files changed, 51 insertions, 18 deletions

diff --git a/src/expr.rs b/src/expr.rs
index f383eaca36f..1dbcf73618a 100644
--- a/src/expr.rs
+++ b/src/expr.rs
@@ -1691,7 +1691,7 @@ pub fn wrap_struct_field(
 }
 
 pub fn struct_lit_field_separator(config: &Config) -> &str {
-    colon_spaces(config.space_before_colon(), config.space_after_colon())
+    colon_spaces(config)
 }
 
 pub fn rewrite_field(
diff --git a/src/items.rs b/src/items.rs
index 1647ecfcc3b..c907d34b47c 100644
--- a/src/items.rs
+++ b/src/items.rs
@@ -36,7 +36,7 @@ const DEFAULT_VISIBILITY: ast::Visibility = source_map::Spanned {
 };
 
 fn type_annotation_separator(config: &Config) -> &str {
-    colon_spaces(config.space_before_colon(), config.space_after_colon())
+    colon_spaces(config)
 }
 
 // Statements of the form
@@ -1695,10 +1695,7 @@ fn rewrite_static(
     static_parts: &StaticParts<'_>,
     offset: Indent,
 ) -> Option<String> {
-    let colon = colon_spaces(
-        context.config.space_before_colon(),
-        context.config.space_after_colon(),
-    );
+    let colon = colon_spaces(context.config);
     let mut prefix = format!(
         "{}{}{} {}{}{}",
         format_visibility(context, static_parts.vis),
@@ -1828,6 +1825,42 @@ fn is_empty_infer(ty: &ast::Ty, pat_span: Span) -> bool {
     }
 }
 
+/// Recover any missing comments between the argument and the type.
+///
+/// # Returns
+///
+/// A 2-len tuple with the comment before the colon in first position, and the comment after the
+/// colon in second position.
+fn get_missing_arg_comments(
+    context: &RewriteContext<'_>,
+    pat_span: Span,
+    ty_span: Span,
+    shape: Shape,
+) -> (String, String) {
+    let missing_comment_span = mk_sp(pat_span.hi(), ty_span.lo());
+
+    let span_before_colon = {
+        let missing_comment_span_hi = context
+            .snippet_provider
+            .span_before(missing_comment_span, ":");
+        mk_sp(pat_span.hi(), missing_comment_span_hi)
+    };
+    let span_after_colon = {
+        let missing_comment_span_lo = context
+            .snippet_provider
+            .span_after(missing_comment_span, ":");
+        mk_sp(missing_comment_span_lo, ty_span.lo())
+    };
+
+    let comment_before_colon = rewrite_missing_comment(span_before_colon, shape, context)
+        .filter(|comment| !comment.is_empty())
+        .map_or(String::new(), |comment| format!(" {}", comment));
+    let comment_after_colon = rewrite_missing_comment(span_after_colon, shape, context)
+        .filter(|comment| !comment.is_empty())
+        .map_or(String::new(), |comment| format!("{} ", comment));
+    (comment_before_colon, comment_after_colon)
+}
+
 impl Rewrite for ast::Arg {
     fn rewrite(&self, context: &RewriteContext<'_>, shape: Shape) -> Option<String> {
         if let Some(ref explicit_self) = self.to_self() {
@@ -1838,13 +1871,11 @@ impl Rewrite for ast::Arg {
                 .rewrite(context, Shape::legacy(shape.width, shape.indent))?;
 
             if !is_empty_infer(&*self.ty, self.pat.span) {
-                if context.config.space_before_colon() {
-                    result.push_str(" ");
-                }
-                result.push_str(":");
-                if context.config.space_after_colon() {
-                    result.push_str(" ");
-                }
+                let (before_comment, after_comment) =
+                    get_missing_arg_comments(context, self.pat.span, self.ty.span, shape);
+                result.push_str(&before_comment);
+                result.push_str(colon_spaces(context.config));
+                result.push_str(&after_comment);
                 let overhead = last_line_width(&result);
                 let max_width = shape.width.checked_sub(overhead)?;
                 let ty_str = self
diff --git a/src/types.rs b/src/types.rs
index 81cfc20961d..8ed73b99913 100644
--- a/src/types.rs
+++ b/src/types.rs
@@ -378,10 +378,7 @@ where
 }
 
 fn type_bound_colon(context: &RewriteContext<'_>) -> &'static str {
-    colon_spaces(
-        context.config.space_before_colon(),
-        context.config.space_after_colon(),
-    )
+    colon_spaces(context.config)
 }
 
 impl Rewrite for ast::WherePredicate {
diff --git a/src/utils.rs b/src/utils.rs
index 743a8627612..a1ab5fd5aad 100644
--- a/src/utils.rs
+++ b/src/utils.rs
@@ -376,7 +376,9 @@ fn is_valid_str(snippet: &str, max_width: usize, shape: Shape) -> bool {
 }
 
 #[inline]
-pub fn colon_spaces(before: bool, after: bool) -> &'static str {
+pub fn colon_spaces(config: &Config) -> &'static str {
+    let before = config.space_before_colon();
+    let after = config.space_after_colon();
     match (before, after) {
         (true, true) => " : ",
         (true, false) => " :",
diff --git a/tests/target/issue-2976.rs b/tests/target/issue-2976.rs
new file mode 100644
index 00000000000..51c94a84ba2
--- /dev/null
+++ b/tests/target/issue-2976.rs
@@ -0,0 +1,3 @@
+fn a(_ /*comment*/: u8 /* toto */) {}
+fn b(/*comment*/ _: u8 /* tata */) {}
+fn c(_: /*comment*/ u8) {}