From d8db4cfeb40ae1571cac7709d4ff3dbb212a7b45 Mon Sep 17 00:00:00 2001
From: Frank King <frankking1729@gmail.com>
Date: Mon, 4 Nov 2024 21:29:15 +0800
Subject: [PATCH] Fix rustfmt for attributes in `where` predicates

---
 src/tools/rustfmt/src/spanned.rs              |  6 +--
 src/tools/rustfmt/src/types.rs                | 39 +++++++++++++---
 .../tests/target/cfg_attribute_in_where.rs    | 45 ++++++++++++++-----
 3 files changed, 68 insertions(+), 22 deletions(-)

diff --git a/src/tools/rustfmt/src/spanned.rs b/src/tools/rustfmt/src/spanned.rs
index b10e2b3cd221a..db7c3486e7179 100644
--- a/src/tools/rustfmt/src/spanned.rs
+++ b/src/tools/rustfmt/src/spanned.rs
@@ -150,11 +150,7 @@ impl Spanned for ast::FieldDef {
 
 impl Spanned for ast::WherePredicate {
     fn span(&self) -> Span {
-        match self.kind {
-            ast::WherePredicateKind::BoundPredicate(ref p) => p.span,
-            ast::WherePredicateKind::RegionPredicate(ref p) => p.span,
-            ast::WherePredicateKind::EqPredicate(ref p) => p.span,
-        }
+        self.span
     }
 }
 
diff --git a/src/tools/rustfmt/src/types.rs b/src/tools/rustfmt/src/types.rs
index 44ca604f3d5c1..b79a62c24804f 100644
--- a/src/tools/rustfmt/src/types.rs
+++ b/src/tools/rustfmt/src/types.rs
@@ -462,12 +462,9 @@ impl Rewrite for ast::WherePredicate {
     }
 
     fn rewrite_result(&self, context: &RewriteContext<'_>, shape: Shape) -> RewriteResult {
-        let mut result = self.attrs.rewrite_result(context, shape)?;
-        if !self.attrs.is_empty() {
-            result.push(' ');
-        }
+        let attrs_str = self.attrs.rewrite_result(context, shape)?;
         // FIXME: dead spans?
-        result += &match self.kind {
+        let pred_str = &match self.kind {
             ast::WherePredicateKind::BoundPredicate(ast::WhereBoundPredicate {
                 ref bound_generic_params,
                 ref bounded_ty,
@@ -503,6 +500,38 @@ impl Rewrite for ast::WherePredicate {
             }
         };
 
+        let mut result = String::with_capacity(attrs_str.len() + pred_str.len() + 1);
+        result.push_str(&attrs_str);
+        let pred_start = self.kind.span().lo();
+        let line_len = last_line_width(&attrs_str) + 1 + first_line_width(&pred_str);
+        if let Some(last_attr) = self.attrs.last().filter(|last_attr| {
+            contains_comment(context.snippet(mk_sp(last_attr.span.hi(), pred_start)))
+        }) {
+            result = combine_strs_with_missing_comments(
+                context,
+                &result,
+                &pred_str,
+                mk_sp(last_attr.span.hi(), pred_start),
+                Shape {
+                    width: context.config.inline_attribute_width(),
+                    ..shape
+                },
+                !last_attr.is_doc_comment(),
+            )?;
+        } else {
+            if !self.attrs.is_empty() {
+                if context.config.inline_attribute_width() < line_len
+                    || self.attrs.len() > 1
+                    || self.attrs.last().is_some_and(|a| a.is_doc_comment())
+                {
+                    result.push_str(&shape.indent.to_string_with_newline(context.config));
+                } else {
+                    result.push(' ');
+                }
+            }
+            result.push_str(&pred_str);
+        }
+
         Ok(result)
     }
 }
diff --git a/src/tools/rustfmt/tests/target/cfg_attribute_in_where.rs b/src/tools/rustfmt/tests/target/cfg_attribute_in_where.rs
index 90c503f300e72..11f495b162995 100644
--- a/src/tools/rustfmt/tests/target/cfg_attribute_in_where.rs
+++ b/src/tools/rustfmt/tests/target/cfg_attribute_in_where.rs
@@ -1,3 +1,5 @@
+// rustfmt-inline_attribute_width: 40
+
 #![crate_type = "lib"]
 #![feature(cfg_attribute_in_where)]
 use std::marker::PhantomData;
@@ -10,35 +12,54 @@ trait TraitB {}
 
 trait A<T>
 where
-    #[cfg = a] T: TraitA,
-    #[cfg = b] T: TraitB,
+    #[cfg = a_very_long_attribute_name]
+    T: TraitA,
+    #[cfg = another_very_long_attribute_name]
+    T: TraitB,
 {
     type B<U>
     where
-        #[cfg = a] U: TraitA,
-        #[cfg = b] U: TraitB;
+        #[cfg = a]
+        // line comment after the attribute
+        U: TraitA,
+        #[cfg = b]
+        /* block comment after the attribute */
+        U: TraitB,
+        #[cfg = a] // short
+        U: TraitA,
+        #[cfg = b] /* short */ U: TraitB;
 
     fn foo<U>(&self)
     where
-        #[cfg = a] U: TraitA,
-        #[cfg = b] U: TraitB;
+        /// line doc comment before the attribute
+        U: TraitA,
+        /** line doc block comment before the attribute */
+        U: TraitB;
 }
 
 impl<T> A<T> for T
 where
-    #[cfg = a] T: TraitA,
-    #[cfg = b] T: TraitB,
+    #[doc = "line doc before the attribute"]
+    T: TraitA,
+    /** short doc */
+    T: TraitB,
 {
     type B<U>
         = ()
     where
-        #[cfg = a] U: TraitA,
-        #[cfg = b] U: TraitB;
+        #[doc = "short"] U: TraitA,
+        #[doc = "short"]
+        #[cfg = a]
+        U: TraitB;
 
     fn foo<U>(&self)
     where
-        #[cfg = a] U: TraitA,
-        #[cfg = b] U: TraitB,
+        #[cfg = a]
+        #[cfg = b]
+        U: TraitA,
+        /// line doc
+        #[cfg = c]
+        U: TraitB,
     {
     }
 }