fix: imports_granularity module with path containing self

rust-lang · Mar 11, 2022 · a82b008 · a82b008
1 parent 272fb42
commit a82b008
Show file tree

Hide file tree

Showing 3 changed files with 41 additions and 12 deletions.
diff --git a/src/imports.rs b/src/imports.rs
@@ -190,13 +190,17 @@ pub(crate) fn merge_use_trees(use_trees: Vec<UseTree>, merge_by: SharedPrefix) -
             continue;
         }
 
-        for flattened in use_tree.flatten() {
+        for mut flattened in use_tree.flatten() {
             if let Some(tree) = result
                 .iter_mut()
                 .find(|tree| tree.share_prefix(&flattened, merge_by))
             {
                 tree.merge(&flattened, merge_by);
             } else {
+                // If this is the first tree with this prefix, handle potential trailing ::self
+                if merge_by == SharedPrefix::Module {
+                    flattened = flattened.nest_trailing_self();
+                }
                 result.push(flattened);
             }
         }
@@ -208,17 +212,7 @@ pub(crate) fn flatten_use_trees(use_trees: Vec<UseTree>) -> Vec<UseTree> {
     use_trees
         .into_iter()
         .flat_map(UseTree::flatten)
-        .map(|mut tree| {
-            // If a path ends in `::self`, rewrite it to `::{self}`.
-            if let Some(UseSegment::Slf(..)) = tree.path.last() {
-                let self_segment = tree.path.pop().unwrap();
-                tree.path.push(UseSegment::List(vec![UseTree::from_path(
-                    vec![self_segment],
-                    DUMMY_SP,
-                )]));
-            }
-            tree
-        })
+        .map(UseTree::nest_trailing_self)
         .collect()
 }
 
@@ -635,6 +629,18 @@ impl UseTree {
             self.span = self.span.to(other.span);
         }
     }
+
+    /// If this tree ends in `::self`, rewrite it to `::{self}`.
+    fn nest_trailing_self(mut self) -> UseTree {
+        if let Some(UseSegment::Slf(..)) = self.path.last() {
+            let self_segment = self.path.pop().unwrap();
+            self.path.push(UseSegment::List(vec![UseTree::from_path(
+                vec![self_segment],
+                DUMMY_SP,
+            )]));
+        }
+        self
+    }
 }
 
 fn merge_rest(
@@ -1311,4 +1317,24 @@ mod test {
                 < parse_use_tree("std::cmp::{b, e, g, f}").normalize()
         );
     }
+
+    #[test]
+    fn test_use_tree_nest_trailing_self() {
+        assert_eq!(
+            parse_use_tree("a::b::self").nest_trailing_self(),
+            parse_use_tree("a::b::{self}")
+        );
+        assert_eq!(
+            parse_use_tree("a::b::c").nest_trailing_self(),
+            parse_use_tree("a::b::c")
+        );
+        assert_eq!(
+            parse_use_tree("a::b::{c, d}").nest_trailing_self(),
+            parse_use_tree("a::b::{c, d}")
+        );
+        assert_eq!(
+            parse_use_tree("a::b::{self, c}").nest_trailing_self(),
+            parse_use_tree("a::b::{self, c}")
+        );
+    }
 }
diff --git a/tests/source/imports_granularity_module.rs b/tests/source/imports_granularity_module.rs
@@ -4,6 +4,7 @@ use a::{b::c, d::e};
 use a::{f, g::{h, i}};
 use a::{j::{self, k::{self, l}, m}, n::{o::p, q}};
 pub use a::{r::s, t};
+use b::{c::d, self};
 
 #[cfg(test)]
 use foo::{a::b, c::d};

diff --git a/tests/target/imports_granularity_module.rs b/tests/target/imports_granularity_module.rs
@@ -10,6 +10,8 @@ use a::n::o::p;
 use a::n::q;
 pub use a::r::s;
 pub use a::t;
+use b::c::d;
+use b::{self};
 
 use foo::e;
 #[cfg(test)]