From ac910f4b58ce12982abb876c65b77dc9304cffb7 Mon Sep 17 00:00:00 2001
From: Paul Gey <narpfel@gmx.de>
Date: Sun, 17 Oct 2021 16:21:53 +0200
Subject: [PATCH] Preserve attributes for `imports_granularity=Item`

Fixes #5030
---
 src/imports.rs             | 24 +++++++++++++++++-------
 tests/source/issue-5030.rs |  7 +++++++
 tests/target/issue-5030.rs |  6 ++++++
 3 files changed, 30 insertions(+), 7 deletions(-)
 create mode 100644 tests/source/issue-5030.rs
 create mode 100644 tests/target/issue-5030.rs

diff --git a/src/imports.rs b/src/imports.rs
index c60bec6d4a2..143f9e3daa9 100644
--- a/src/imports.rs
+++ b/src/imports.rs
@@ -108,6 +108,12 @@ pub(crate) struct UseTree {
     attrs: Option<Vec<ast::Attribute>>,
 }
 
+#[derive(Copy, Clone)]
+enum Attributes {
+    Preserve,
+    Discard,
+}
+
 impl PartialEq for UseTree {
     fn eq(&self, other: &UseTree) -> bool {
         self.path == other.path
@@ -190,7 +196,7 @@ pub(crate) fn merge_use_trees(use_trees: Vec<UseTree>, merge_by: SharedPrefix) -
             continue;
         }
 
-        for flattened in use_tree.flatten() {
+        for flattened in use_tree.flatten(Attributes::Discard) {
             if let Some(tree) = result
                 .iter_mut()
                 .find(|tree| tree.share_prefix(&flattened, merge_by))
@@ -207,7 +213,7 @@ pub(crate) fn merge_use_trees(use_trees: Vec<UseTree>, merge_by: SharedPrefix) -
 pub(crate) fn flatten_use_trees(use_trees: Vec<UseTree>) -> Vec<UseTree> {
     use_trees
         .into_iter()
-        .flat_map(UseTree::flatten)
+        .flat_map(|tree| tree.flatten(Attributes::Preserve))
         .map(|mut tree| {
             // If a path ends in `::self`, rewrite it to `::{self}`.
             if let Some(UseSegment::Slf(..)) = tree.path.last() {
@@ -587,7 +593,7 @@ impl UseTree {
         }
     }
 
-    fn flatten(self) -> Vec<UseTree> {
+    fn flatten(self, preserve_attrs: Attributes) -> Vec<UseTree> {
         if self.path.is_empty() {
             return vec![self];
         }
@@ -601,7 +607,7 @@ impl UseTree {
                 let prefix = &self.path[..self.path.len() - 1];
                 let mut result = vec![];
                 for nested_use_tree in list {
-                    for flattend in &mut nested_use_tree.clone().flatten() {
+                    for flattend in &mut nested_use_tree.clone().flatten(preserve_attrs) {
                         let mut new_path = prefix.to_vec();
                         new_path.append(&mut flattend.path);
                         result.push(UseTree {
@@ -609,7 +615,11 @@ impl UseTree {
                             span: self.span,
                             list_item: None,
                             visibility: self.visibility.clone(),
-                            attrs: None,
+                            // only retain attributes for `ImportGranularity::Item`
+                            attrs: match preserve_attrs {
+                                Attributes::Preserve => self.attrs.clone(),
+                                Attributes::Discard => None,
+                            },
                         });
                     }
                 }
@@ -1228,12 +1238,12 @@ mod test {
     #[test]
     fn test_use_tree_flatten() {
         assert_eq!(
-            parse_use_tree("a::b::{c, d, e, f}").flatten(),
+            parse_use_tree("a::b::{c, d, e, f}").flatten(Attributes::Discard),
             parse_use_trees!("a::b::c", "a::b::d", "a::b::e", "a::b::f",)
         );
 
         assert_eq!(
-            parse_use_tree("a::b::{c::{d, e, f}, g, h::{i, j, k}}").flatten(),
+            parse_use_tree("a::b::{c::{d, e, f}, g, h::{i, j, k}}").flatten(Attributes::Discard),
             parse_use_trees![
                 "a::b::c::d",
                 "a::b::c::e",
diff --git a/tests/source/issue-5030.rs b/tests/source/issue-5030.rs
new file mode 100644
index 00000000000..f367e79f01f
--- /dev/null
+++ b/tests/source/issue-5030.rs
@@ -0,0 +1,7 @@
+// rustfmt-imports_granularity: Item
+
+#[cfg(feature = "foo")]
+use std::collections::{
+    HashMap,
+    HashSet,
+};
diff --git a/tests/target/issue-5030.rs b/tests/target/issue-5030.rs
new file mode 100644
index 00000000000..b371331ed00
--- /dev/null
+++ b/tests/target/issue-5030.rs
@@ -0,0 +1,6 @@
+// rustfmt-imports_granularity: Item
+
+#[cfg(feature = "foo")]
+use std::collections::HashMap;
+#[cfg(feature = "foo")]
+use std::collections::HashSet;