Extract shared logic

2 files changed, 40 insertions, 29 deletions

diff --git a/crates/ide-assists/src/handlers/inline_call.rs b/crates/ide-assists/src/handlers/inline_call.rs
index 96890ad51a6..9f51cdaf8b1 100644
--- a/crates/ide-assists/src/handlers/inline_call.rs
+++ b/crates/ide-assists/src/handlers/inline_call.rs
@@ -7,6 +7,7 @@ use ide_db::{
     imports::insert_use::remove_path_if_in_use_stmt,
     path_transform::PathTransform,
     search::{FileReference, SearchScope},
+    source_change::SourceChangeBuilder,
     syntax_helpers::{insert_whitespace_into_node::insert_ws_into, node_ext::expr_as_name_ref},
     RootDatabase,
 };
@@ -100,18 +101,7 @@ pub(crate) fn inline_into_callers(acc: &mut Assists, ctx: &AssistContext<'_>) ->
                 builder.edit_file(file_id);
                 let count = refs.len();
                 // The collects are required as we are otherwise iterating while mutating 🙅‍♀️🙅‍♂️
-                let (name_refs, name_refs_use): (Vec<_>, Vec<_>) = refs
-                    .into_iter()
-                    .filter_map(|file_ref| match file_ref.name {
-                        ast::NameLike::NameRef(name_ref) => Some(name_ref),
-                        _ => None,
-                    })
-                    .partition_map(|name_ref| {
-                        match name_ref.syntax().ancestors().find_map(ast::UseTree::cast) {
-                            Some(use_tree) => Either::Right(builder.make_mut(use_tree)),
-                            None => Either::Left(name_ref),
-                        }
-                    });
+                let (name_refs, name_refs_use) = split_refs_and_uses(builder, refs, Some);
                 let call_infos: Vec<_> = name_refs
                     .into_iter()
                     .filter_map(CallInfo::from_name_ref)
@@ -130,11 +120,7 @@ pub(crate) fn inline_into_callers(acc: &mut Assists, ctx: &AssistContext<'_>) ->
                     .count();
                 if replaced + name_refs_use.len() == count {
                     // we replaced all usages in this file, so we can remove the imports
-                    name_refs_use.into_iter().for_each(|use_tree| {
-                        if let Some(path) = use_tree.path() {
-                            remove_path_if_in_use_stmt(&path);
-                        }
-                    })
+                    name_refs_use.iter().for_each(remove_path_if_in_use_stmt);
                 } else {
                     remove_def = false;
                 }
@@ -153,6 +139,23 @@ pub(crate) fn inline_into_callers(acc: &mut Assists, ctx: &AssistContext<'_>) ->
     )
 }
 
+pub(super) fn split_refs_and_uses<T: ast::AstNode>(
+    builder: &mut SourceChangeBuilder,
+    iter: impl IntoIterator<Item = FileReference>,
+    mut map_ref: impl FnMut(ast::NameRef) -> Option<T>,
+) -> (Vec<T>, Vec<ast::Path>) {
+    iter.into_iter()
+        .filter_map(|file_ref| match file_ref.name {
+            ast::NameLike::NameRef(name_ref) => Some(name_ref),
+            _ => None,
+        })
+        .filter_map(|name_ref| match name_ref.syntax().ancestors().find_map(ast::UseTree::cast) {
+            Some(use_tree) => builder.make_mut(use_tree).path().map(Either::Right),
+            None => map_ref(name_ref).map(Either::Left),
+        })
+        .partition_map(|either| either)
+}
+
 // Assist: inline_call
 //
 // Inlines a function or method body creating a `let` statement per parameter unless the parameter
diff --git a/crates/ide-assists/src/handlers/inline_type_alias.rs b/crates/ide-assists/src/handlers/inline_type_alias.rs
index ee560c6193c..4afe890c783 100644
--- a/crates/ide-assists/src/handlers/inline_type_alias.rs
+++ b/crates/ide-assists/src/handlers/inline_type_alias.rs
@@ -3,7 +3,9 @@
 // - Remove unused aliases if there are no longer any users, see inline_call.rs.
 
 use hir::{HasSource, PathResolution};
-use ide_db::{defs::Definition, search::FileReference};
+use ide_db::{
+    defs::Definition, imports::insert_use::remove_path_if_in_use_stmt, search::FileReference,
+};
 use itertools::Itertools;
 use std::collections::HashMap;
 use syntax::{
@@ -16,6 +18,8 @@ use crate::{
     AssistId, AssistKind,
 };
 
+use super::inline_call::split_refs_and_uses;
+
 // Assist: inline_type_alias_uses
 //
 // Inline a type alias into all of its uses where possible.
@@ -31,7 +35,7 @@ use crate::{
 // ```
 // ->
 // ```
-// 
+//
 // fn id(x: i32) -> i32 {
 //     x
 // };
@@ -62,15 +66,10 @@ pub(crate) fn inline_type_alias_uses(acc: &mut Assists, ctx: &AssistContext<'_>)
             let mut inline_refs_for_file = |file_id, refs: Vec<FileReference>| {
                 builder.edit_file(file_id);
 
-                let path_types: Vec<ast::PathType> = refs
-                    .into_iter()
-                    .filter_map(|file_ref| match file_ref.name {
-                        ast::NameLike::NameRef(path_type) => {
-                            path_type.syntax().ancestors().nth(3).and_then(ast::PathType::cast)
-                        }
-                        _ => None,
-                    })
-                    .collect();
+                let (path_types, path_type_uses) =
+                    split_refs_and_uses(builder, refs, |path_type| {
+                        path_type.syntax().ancestors().nth(3).and_then(ast::PathType::cast)
+                    });
 
                 for (target, replacement) in path_types.into_iter().filter_map(|path_type| {
                     let replacement = inline(&ast_alias, &path_type)?.to_text(&concrete_type);
@@ -79,6 +78,10 @@ pub(crate) fn inline_type_alias_uses(acc: &mut Assists, ctx: &AssistContext<'_>)
                 }) {
                     builder.replace(target, replacement);
                 }
+                if !path_type_uses.is_empty() {
+                    builder.edit_file(file_id);
+                    path_type_uses.iter().for_each(remove_path_if_in_use_stmt);
+                }
             };
 
             for (file_id, refs) in usages.into_iter() {
@@ -993,7 +996,12 @@ fn foo() {
 }
 "#,
                 r#"
-use super::I;
+//- /lib.rs
+mod foo;
+
+
+//- /foo.rs
+
 fn foo() {
     let _: i32 = 0;
 }