In extract_module.rs, generate ast::Module instead of String

Author: Hmikihiro

crates/ide-assists/src/handlers/extract_module.rs

@@ -1,6 +1,5 @@
 use std::iter;
 
-use either::Either;
 use hir::{HasSource, ModuleSource};
 use ide_db::{
     FileId, FxHashMap, FxHashSet,
@@ -120,7 +119,8 @@ 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 module_def = generate_module_def(&impl_parent, &mut module, old_item_indent);
+            let module_def =
+                generate_module_def(&impl_parent, &mut module, old_item_indent).to_string();
 
             let mut usages_to_be_processed_for_cur_file = vec![];
             for (file_id, usages) in usages_to_be_processed {
@@ -175,31 +175,36 @@ 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)
+) -> ast::Module {
+    let items = if let Some(self_ty) = parent_impl.as_ref().and_then(|imp| imp.self_ty()) {
+        let assoc_items = module
+            .body_items
+            .iter()
+            .filter_map(|item| match item {
+                ast::Item::Const(it) => Some(ast::AssocItem::Const(it.clone())),
+                ast::Item::Fn(it) => Some(ast::AssocItem::Fn(it.clone())),
+                ast::Item::MacroCall(it) => Some(ast::AssocItem::MacroCall(it.clone())),
+                ast::Item::TypeAlias(it) => Some(ast::AssocItem::TypeAlias(it.clone())),
+                _ => None,
+            })
+            .map(|it| it.indent(IndentLevel(1)))
+            .collect_vec();
+        let assoc_item_list = make::assoc_item_list(Some(assoc_items));
+        let impl_ = make::impl_(None, None, self_ty.clone(), None, Some(assoc_item_list));
+        // Add the import for enum/struct corresponding to given impl block
+        module.make_use_stmt_of_node_with_super(self_ty.syntax());
+        let mut module_body_items = module.use_items.clone();
+        module_body_items.push(ast::Item::Impl(impl_));
+        module_body_items
     } else {
-        (Either::Right(module.use_items.iter().chain(module.body_items.iter())), old_indent + 1)
+        [module.use_items.clone(), module.body_items.clone()].concat()
     };
 
-    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}}}");
+    let items = items.into_iter().map(|it| it.reset_indent().indent(IndentLevel(1))).collect_vec();
+    let module_body = make::item_list(Some(items));
 
-        // 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}}}")
+    let module_name = make::name(module.name);
+    make::mod_(module_name, Some(module_body)).indent(old_indent)
 }
 
 #[derive(Debug)]
@@ -1422,10 +1427,10 @@ $0fn foo(x: B) {}$0
             struct B {}
 
 mod modname {
-    use super::B;
-
     use super::A;
 
+    use super::B;
+
     impl A {
         pub(crate) fn foo(x: B) {}
     }

crates/syntax/src/ast/make.rs

@@ -229,6 +229,24 @@ pub fn ty_fn_ptr<I: Iterator<Item = Param>>(
     }
 }
 
+pub fn item_list(body: Option<Vec<ast::Item>>) -> ast::ItemList {
+    let is_break_braces = body.is_some();
+    let body_newline = if is_break_braces { "\n".to_owned() } else { String::new() };
+    let body_indent = if is_break_braces { "    ".to_owned() } else { String::new() };
+
+    let body = match body {
+        Some(bd) => bd.iter().map(|elem| elem.to_string()).join("\n\n    "),
+        None => String::new(),
+    };
+    ast_from_text(&format!("mod C {{{body_newline}{body_indent}{body}{body_newline}}}"))
+}
+
+pub fn mod_(name: ast::Name, body: Option<ast::ItemList>) -> ast::Module {
+    let name = name.to_string();
+    let body = body.map_or(";".to_owned(), |body| format!(" {body}"));
+    ast_from_text(&format!("mod {name}{body}"))
+}
+
 pub fn assoc_item_list(body: Option<Vec<ast::AssocItem>>) -> ast::AssocItemList {
     let is_break_braces = body.is_some();
     let body_newline = if is_break_braces { "\n".to_owned() } else { String::new() };