fix: simplification on extract_module

roife · roife · commit 5b2809f329a1 · 2024-03-15T14:24:16.000+08:00
diff --git a/crates/ide-assists/src/handlers/extract_module.rs b/crates/ide-assists/src/handlers/extract_module.rs
@@ -1,5 +1,6 @@
 use std::iter;
 
+use either::Either;
 use hir::{HasSource, HirFileIdExt, ModuleSource};
 use ide_db::{
     assists::{AssistId, AssistKind},
@@ -10,15 +11,14 @@ use ide_db::{
 };
 use itertools::Itertools;
 use smallvec::SmallVec;
-use stdx::format_to;
 use syntax::{
     algo::find_node_at_range,
     ast::{
         self,
         edit::{AstNodeEdit, IndentLevel},
         make, HasVisibility,
     },
-    match_ast, ted, AstNode, SourceFile,
+    match_ast, ted, AstNode,
     SyntaxKind::{self, WHITESPACE},
     SyntaxNode, TextRange,
 };
@@ -114,63 +114,23 @@ pub(crate) fn extract_module(acc: &mut Assists, ctx: &AssistContext<'_>) -> Opti
             let import_paths_to_be_removed = module.resolve_imports(curr_parent_module, ctx);
             module.change_visibility(record_fields);
 
-            let mut body_items: Vec<String> = Vec::new();
-            let mut items_to_be_processed: Vec<ast::Item> = module.body_items.clone();
+            let module_def = generate_module_def(&impl_parent, &mut module, old_item_indent);
 
-            let new_item_indent = if impl_parent.is_some() {
-                old_item_indent + 2
-            } else {
-                items_to_be_processed = [module.use_items.clone(), items_to_be_processed].concat();
-                old_item_indent + 1
-            };
-
-            for item in items_to_be_processed {
-                let item = item.indent(IndentLevel(1));
-                let mut indented_item = String::new();
-                format_to!(indented_item, "{new_item_indent}{item}");
-                body_items.push(indented_item);
-            }
-
-            let mut body = body_items.join("\n\n");
-
-            if let Some(impl_) = &impl_parent {
-                if let Some(self_ty) = impl_.self_ty() {
-                    let impl_indent = old_item_indent + 1;
-                    let mut impl_body_def = String::new();
-                    format_to!(
-                        impl_body_def,
-                        "{impl_indent}impl {self_ty} {{\n{body}\n{impl_indent}}}",
-                    );
-                    body = impl_body_def;
-
-                    // Add the import for enum/struct corresponding to given impl block
-                    module.make_use_stmt_of_node_with_super(self_ty.syntax());
-                    for item in module.use_items {
-                        let item_indent = old_item_indent + 1;
-                        body = format!("{item_indent}{item}\n\n{body}");
-                    }
-                }
-            }
-
-            let mut module_def = String::new();
-            let module_name = module.name;
-            format_to!(module_def, "mod {module_name} {{\n{body}\n{old_item_indent}}}");
-
-            let mut usages_to_be_updated_for_curr_file = vec![];
-            for usages_to_be_updated_for_file in usages_to_be_processed {
-                if usages_to_be_updated_for_file.0 == ctx.file_id() {
-                    usages_to_be_updated_for_curr_file = usages_to_be_updated_for_file.1;
+            let mut usages_to_be_processed_for_cur_file = vec![];
+            for (file_id, usages) in usages_to_be_processed {
+                if file_id == ctx.file_id() {
+                    usages_to_be_processed_for_cur_file = usages;
                     continue;
                 }
-                builder.edit_file(usages_to_be_updated_for_file.0);
-                for usage_to_be_processed in usages_to_be_updated_for_file.1 {
-                    builder.replace(usage_to_be_processed.0, usage_to_be_processed.1)
+                builder.edit_file(file_id);
+                for (text_range, usage) in usages {
+                    builder.replace(text_range, usage)
                 }
             }
 
             builder.edit_file(ctx.file_id());
-            for usage_to_be_processed in usages_to_be_updated_for_curr_file {
-                builder.replace(usage_to_be_processed.0, usage_to_be_processed.1)
+            for (text_range, usage) in usages_to_be_processed_for_cur_file {
+                builder.replace(text_range, usage);
             }
 
             if let Some(impl_) = impl_parent {
@@ -205,6 +165,37 @@ pub(crate) fn extract_module(acc: &mut Assists, ctx: &AssistContext<'_>) -> Opti
     )
 }
 
+fn generate_module_def(
+    parent_impl: &Option<ast::Impl>,
+    module: &mut Module,
+    old_indent: IndentLevel,
+) -> String {
+    let (items_to_be_processed, new_item_indent) = if parent_impl.is_some() {
+        (Either::Left(module.body_items.iter()), old_indent + 2)
+    } else {
+        (Either::Right(module.use_items.iter().chain(module.body_items.iter())), old_indent + 1)
+    };
+
+    let mut body = items_to_be_processed
+        .map(|item| item.indent(IndentLevel(1)))
+        .map(|item| format!("{new_item_indent}{item}"))
+        .join("\n\n");
+
+    if let Some(self_ty) = parent_impl.as_ref().and_then(|imp| imp.self_ty()) {
+        let impl_indent = old_indent + 1;
+        body = format!("{impl_indent}impl {self_ty} {{\n{body}\n{impl_indent}}}");
+
+        // Add the import for enum/struct corresponding to given impl block
+        module.make_use_stmt_of_node_with_super(self_ty.syntax());
+        for item in module.use_items.iter() {
+            body = format!("{impl_indent}{item}\n\n{body}");
+        }
+    }
+
+    let module_name = module.name;
+    format!("mod {module_name} {{\n{body}\n{old_indent}}}")
+}
+
 #[derive(Debug)]
 struct Module {
     text_range: TextRange,
@@ -240,6 +231,7 @@ impl Module {
         //Here impl is not included as each item inside impl will be tied to the parent of
         //implementing block(a struct, enum, etc), if the parent is in selected module, it will
         //get updated by ADT section given below or if it is not, then we dont need to do any operation
+
         for item in &self.body_items {
             match_ast! {
                 match (item.syntax()) {
@@ -320,48 +312,38 @@ impl Module {
         node_def: Definition,
         refs_in_files: &mut FxHashMap<FileId, Vec<(TextRange, String)>>,
     ) {
-        for (file_id, references) in node_def.usages(&ctx.sema).all() {
+        let mod_name = self.name;
+        let out_of_sel = |node: &SyntaxNode| !self.text_range.contains_range(node.text_range());
+        for (file_id, refs) in node_def.usages(&ctx.sema).all() {
             let source_file = ctx.sema.parse(file_id);
-            let usages_in_file = references
-                .into_iter()
-                .filter_map(|usage| self.get_usage_to_be_processed(&source_file, usage));
-            refs_in_files.entry(file_id).or_default().extend(usages_in_file);
-        }
-    }
-
-    fn get_usage_to_be_processed(
-        &self,
-        source_file: &SourceFile,
-        FileReference { range, name, .. }: FileReference,
-    ) -> Option<(TextRange, String)> {
-        let path: ast::Path = find_node_at_range(source_file.syntax(), range)?;
-
-        for desc in path.syntax().descendants() {
-            if desc.to_string() == name.syntax().to_string()
-                && !self.text_range.contains_range(desc.text_range())
-            {
-                if let Some(name_ref) = ast::NameRef::cast(desc) {
-                    let mod_name = self.name;
-                    return Some((
-                        name_ref.syntax().text_range(),
-                        format!("{mod_name}::{name_ref}"),
-                    ));
+            let usages = refs.into_iter().filter_map(|FileReference { range, name, .. }| {
+                let path: ast::Path = find_node_at_range(source_file.syntax(), range)?;
+                let name = name.syntax().to_string();
+
+                for desc in path.syntax().descendants() {
+                    if desc.to_string() == name && out_of_sel(&desc) {
+                        if let Some(name_ref) = ast::NameRef::cast(desc) {
+                            let new_ref = format!("{mod_name}::{name_ref}");
+                            return Some((name_ref.syntax().text_range(), new_ref));
+                        }
+                    }
                 }
-            }
-        }
 
-        None
+                None
+            });
+            refs_in_files.entry(file_id).or_default().extend(usages);
+        }
     }
 
     fn change_visibility(&mut self, record_fields: Vec<SyntaxNode>) {
         let (mut replacements, record_field_parents, impls) =
             get_replacements_for_visibility_change(&mut self.body_items, false);
 
-        let mut impl_items: Vec<ast::Item> = impls
+        let mut impl_items = impls
             .into_iter()
             .flat_map(|impl_| impl_.syntax().descendants())
             .filter_map(ast::Item::cast)
-            .collect();
+            .collect_vec();
 
         let (mut impl_item_replacements, _, _) =
             get_replacements_for_visibility_change(&mut impl_items, true);
@@ -444,19 +426,19 @@ impl Module {
         let file = ctx.sema.parse(file_id);
 
         // track uses which does not exists in `Use`
-        let mut exists_inside_sel = false;
-        let mut exists_outside_sel = false;
+        let mut uses_exist_in_sel = false;
+        let mut uses_exist_out_sel = false;
         'outside: for (_, refs) in usage_res.iter() {
             for x in refs
                 .iter()
                 .filter(|x| find_node_at_range::<ast::Use>(file.syntax(), x.range).is_none())
                 .filter_map(|x| find_node_at_range::<ast::Path>(file.syntax(), x.range))
             {
                 let in_selectin = selection_range.contains_range(x.syntax().text_range());
-                exists_inside_sel |= in_selectin;
-                exists_outside_sel |= !in_selectin;
+                uses_exist_in_sel |= in_selectin;
+                uses_exist_out_sel |= !in_selectin;
 
-                if exists_inside_sel && exists_outside_sel {
+                if uses_exist_in_sel && uses_exist_out_sel {
                     break 'outside;
                 }
             }
@@ -475,7 +457,7 @@ impl Module {
             !selection_range.contains_range(use_stmt.syntax().text_range())
         });
 
-        let mut use_tree_str_opt: Option<Vec<ast::Path>> = None;
+        let mut use_tree_paths: Option<Vec<ast::Path>> = None;
         //Exists inside and outside selection
         // - Use stmt for item is present -> get the use_tree_str and reconstruct the path in new
         // module
@@ -489,13 +471,13 @@ impl Module {
         //get the use_tree_str, reconstruct the use stmt in new module
 
         let mut import_path_to_be_removed: Option<TextRange> = None;
-        if exists_inside_sel && exists_outside_sel {
+        if uses_exist_in_sel && uses_exist_out_sel {
             //Changes to be made only inside new module
 
             //If use_stmt exists, find the use_tree_str, reconstruct it inside new module
             //If not, insert a use stmt with super and the given nameref
             match self.process_use_stmt_for_import_resolve(use_stmt, use_node) {
-                Some((use_tree_str, _)) => use_tree_str_opt = Some(use_tree_str),
+                Some((use_tree_str, _)) => use_tree_paths = Some(use_tree_str),
                 None if def_in_mod && def_out_sel => {
                     //Considered only after use_stmt is not present
                     //def_in_mod && def_out_sel | exists_outside_sel(exists_inside_sel =
@@ -509,7 +491,7 @@ impl Module {
                 }
                 None => {}
             }
-        } else if exists_inside_sel && !exists_outside_sel {
+        } else if uses_exist_in_sel && !uses_exist_out_sel {
             //Changes to be made inside new module, and remove import from outside
 
             if let Some((mut use_tree_str, text_range_opt)) =
@@ -531,43 +513,42 @@ impl Module {
                     }
                 }
 
-                use_tree_str_opt = Some(use_tree_str);
+                use_tree_paths = Some(use_tree_str);
             } else if def_in_mod && def_out_sel {
                 self.make_use_stmt_of_node_with_super(use_node);
             }
         }
 
-        if let Some(use_tree_str) = use_tree_str_opt {
-            let mut use_tree_str = use_tree_str;
-            use_tree_str.reverse();
+        if let Some(mut use_tree_paths) = use_tree_paths {
+            use_tree_paths.reverse();
 
-            if exists_outside_sel || !exists_inside_sel || !def_in_mod || !def_out_sel {
-                if let Some(first_path_in_use_tree) = use_tree_str.first() {
+            if uses_exist_out_sel || !uses_exist_in_sel || !def_in_mod || !def_out_sel {
+                if let Some(first_path_in_use_tree) = use_tree_paths.first() {
                     if first_path_in_use_tree.to_string().contains("super") {
-                        use_tree_str.insert(0, make::ext::ident_path("super"));
+                        use_tree_paths.insert(0, make::ext::ident_path("super"));
                     }
                 }
             }
 
-            let use_ =
-                make::use_(None, make::use_tree(make::join_paths(use_tree_str), None, None, false));
-            let item = ast::Item::from(use_);
-
-            let is_item = match def {
-                Definition::Macro(_) => true,
-                Definition::Module(_) => true,
-                Definition::Function(_) => true,
-                Definition::Adt(_) => true,
-                Definition::Const(_) => true,
-                Definition::Static(_) => true,
-                Definition::Trait(_) => true,
-                Definition::TraitAlias(_) => true,
-                Definition::TypeAlias(_) => true,
-                _ => false,
-            };
+            let is_item = matches!(
+                def,
+                Definition::Macro(_)
+                    | Definition::Module(_)
+                    | Definition::Function(_)
+                    | Definition::Adt(_)
+                    | Definition::Const(_)
+                    | Definition::Static(_)
+                    | Definition::Trait(_)
+                    | Definition::TraitAlias(_)
+                    | Definition::TypeAlias(_)
+            );
 
             if (def_out_sel || !is_item) && use_stmt_not_in_sel {
-                self.use_items.insert(0, item.clone());
+                let use_ = make::use_(
+                    None,
+                    make::use_tree(make::join_paths(use_tree_paths), None, None, false),
+                );
+                self.use_items.insert(0, ast::Item::from(use_));
             }
         }
 
