Auto merge of #18196 - DropDemBits:sed-syntax-factory, r=Veykril

internal: Add `SyntaxFactory` to ease generating nodes with syntax mappings

Part of [#​15710](https://github.com/rust-lang/rust-analyzer/issues/15710)

Instead of requiring passing a `&mut SyntaxEditor` to every make constructor to generate mappings, we instead wrap that logic in `SyntaxFactory`, and afterwards add all the mappings to the `SyntaxEditor`.

Includes an example of using `SyntaxEditor` & `SyntaxFactory` in the `extract_variable` assist.

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

@@ -1,8 +1,12 @@
 use hir::TypeInfo;
 use ide_db::syntax_helpers::suggest_name;
 use syntax::{
-    ast::{self, edit::IndentLevel, edit_in_place::Indent, make, AstNode, HasName},
-    ted, NodeOrToken,
+    ast::{
+        self, edit::IndentLevel, edit_in_place::Indent, make, syntax_factory::SyntaxFactory,
+        AstNode,
+    },
+    syntax_editor::Position,
+    NodeOrToken,
     SyntaxKind::{BLOCK_EXPR, BREAK_EXPR, COMMENT, LOOP_EXPR, MATCH_GUARD, PATH_EXPR, RETURN_EXPR},
     SyntaxNode, T,
 };
@@ -105,39 +109,46 @@ pub(crate) fn extract_variable(acc: &mut Assists, ctx: &AssistContext<'_>) -> Op
                 ),
             };
 
+            let make = SyntaxFactory::new();
+            let mut editor = edit.make_editor(&expr_replace);
+
+            let pat_name = make.name(&var_name);
+            let name_expr = make.expr_path(make::ext::ident_path(&var_name));
+
+            if let Some(cap) = ctx.config.snippet_cap {
+                let tabstop = edit.make_tabstop_before(cap);
+                editor.add_annotation(pat_name.syntax().clone(), tabstop);
+            }
+
             let ident_pat = match parent {
                 Some(ast::Expr::RefExpr(expr)) if expr.mut_token().is_some() => {
-                    make::ident_pat(false, true, make::name(&var_name))
+                    make.ident_pat(false, true, pat_name)
                 }
                 _ if needs_adjust
                     && !needs_ref
                     && ty.as_ref().is_some_and(|ty| ty.is_mutable_reference()) =>
                 {
-                    make::ident_pat(false, true, make::name(&var_name))
+                    make.ident_pat(false, true, pat_name)
                 }
-                _ => make::ident_pat(false, false, make::name(&var_name)),
+                _ => make.ident_pat(false, false, pat_name),
             };
 
             let to_extract_no_ref = match ty.as_ref().filter(|_| needs_ref) {
                 Some(receiver_type) if receiver_type.is_mutable_reference() => {
-                    make::expr_ref(to_extract_no_ref, true)
+                    make.expr_ref(to_extract_no_ref, true)
                 }
                 Some(receiver_type) if receiver_type.is_reference() => {
-                    make::expr_ref(to_extract_no_ref, false)
+                    make.expr_ref(to_extract_no_ref, false)
                 }
                 _ => to_extract_no_ref,
             };
 
-            let expr_replace = edit.make_syntax_mut(expr_replace);
-            let let_stmt =
-                make::let_stmt(ident_pat.into(), None, Some(to_extract_no_ref)).clone_for_update();
-            let name_expr = make::expr_path(make::ext::ident_path(&var_name)).clone_for_update();
+            let let_stmt = make.let_stmt(ident_pat.into(), None, Some(to_extract_no_ref));
 
             match anchor {
                 Anchor::Before(place) => {
                     let prev_ws = place.prev_sibling_or_token().and_then(|it| it.into_token());
                     let indent_to = IndentLevel::from_node(&place);
-                    let insert_place = edit.make_syntax_mut(place);
 
                     // Adjust ws to insert depending on if this is all inline or on separate lines
                     let trailing_ws = if prev_ws.is_some_and(|it| it.text().starts_with('\n')) {
@@ -146,37 +157,20 @@ pub(crate) fn extract_variable(acc: &mut Assists, ctx: &AssistContext<'_>) -> Op
                         " ".to_owned()
                     };
 
-                    ted::insert_all_raw(
-                        ted::Position::before(insert_place),
+                    editor.insert_all(
+                        Position::before(place),
                         vec![
                             let_stmt.syntax().clone().into(),
                             make::tokens::whitespace(&trailing_ws).into(),
                         ],
                     );
 
-                    ted::replace(expr_replace, name_expr.syntax());
-
-                    if let Some(cap) = ctx.config.snippet_cap {
-                        if let Some(ast::Pat::IdentPat(ident_pat)) = let_stmt.pat() {
-                            if let Some(name) = ident_pat.name() {
-                                edit.add_tabstop_before(cap, name);
-                            }
-                        }
-                    }
+                    editor.replace(expr_replace, name_expr.syntax());
                 }
                 Anchor::Replace(stmt) => {
                     cov_mark::hit!(test_extract_var_expr_stmt);
 
-                    let stmt_replace = edit.make_mut(stmt);
-                    ted::replace(stmt_replace.syntax(), let_stmt.syntax());
-
-                    if let Some(cap) = ctx.config.snippet_cap {
-                        if let Some(ast::Pat::IdentPat(ident_pat)) = let_stmt.pat() {
-                            if let Some(name) = ident_pat.name() {
-                                edit.add_tabstop_before(cap, name);
-                            }
-                        }
-                    }
+                    editor.replace(stmt.syntax(), let_stmt.syntax());
                 }
                 Anchor::WrapInBlock(to_wrap) => {
                     let indent_to = to_wrap.indent_level();
@@ -184,47 +178,22 @@ pub(crate) fn extract_variable(acc: &mut Assists, ctx: &AssistContext<'_>) -> Op
                     let block = if to_wrap.syntax() == &expr_replace {
                         // Since `expr_replace` is the same that needs to be wrapped in a block,
                         // we can just directly replace it with a block
-                        let block =
-                            make::block_expr([let_stmt.into()], Some(name_expr)).clone_for_update();
-                        ted::replace(expr_replace, block.syntax());
-
-                        block
+                        make.block_expr([let_stmt.into()], Some(name_expr))
                     } else {
-                        // `expr_replace` is a descendant of `to_wrap`, so both steps need to be
-                        // handled separately, otherwise we wrap the wrong expression
-                        let to_wrap = edit.make_mut(to_wrap);
-
-                        // Replace the target expr first so that we don't need to find where
-                        // `expr_replace` is in the wrapped `to_wrap`
-                        ted::replace(expr_replace, name_expr.syntax());
-
-                        // Wrap `to_wrap` in a block
-                        let block = make::block_expr([let_stmt.into()], Some(to_wrap.clone()))
-                            .clone_for_update();
-                        ted::replace(to_wrap.syntax(), block.syntax());
-
-                        block
+                        // `expr_replace` is a descendant of `to_wrap`, so we just replace it with `name_expr`.
+                        editor.replace(expr_replace, name_expr.syntax());
+                        make.block_expr([let_stmt.into()], Some(to_wrap.clone()))
                     };
 
-                    if let Some(cap) = ctx.config.snippet_cap {
-                        // Adding a tabstop to `name` requires finding the let stmt again, since
-                        // the existing `let_stmt` is not actually added to the tree
-                        let pat = block.statements().find_map(|stmt| {
-                            let ast::Stmt::LetStmt(let_stmt) = stmt else { return None };
-                            let_stmt.pat()
-                        });
-
-                        if let Some(ast::Pat::IdentPat(ident_pat)) = pat {
-                            if let Some(name) = ident_pat.name() {
-                                edit.add_tabstop_before(cap, name);
-                            }
-                        }
-                    }
+                    editor.replace(to_wrap.syntax(), block.syntax());
 
                     // fixup indentation of block
                     block.indent(indent_to);
                 }
             }
+
+            editor.add_mappings(make.finish_with_mappings());
+            edit.add_file_edits(ctx.file_id(), editor);
             edit.rename();
         },
     )

crates/syntax/src/ast.rs

@@ -8,6 +8,7 @@ pub mod make;
 mod node_ext;
 mod operators;
 pub mod prec;
+pub mod syntax_factory;
 mod token_ext;
 mod traits;
 

crates/syntax/src/ast/syntax_factory.rs

@@ -0,0 +1,45 @@
+//! Builds upon [`crate::ast::make`] constructors to create ast fragments with
+//! optional syntax mappings.
+//!
+//! Instead of forcing make constructors to perform syntax mapping, we instead
+//! let [`SyntaxFactory`] handle constructing the mappings. Care must be taken
+//! to remember to feed the syntax mappings into a [`SyntaxEditor`](crate::syntax_editor::SyntaxEditor),
+//! if applicable.
+
+mod constructors;
+
+use std::cell::{RefCell, RefMut};
+
+use crate::syntax_editor::SyntaxMapping;
+
+pub struct SyntaxFactory {
+    // Stored in a refcell so that the factory methods can be &self
+    mappings: Option<RefCell<SyntaxMapping>>,
+}
+
+impl SyntaxFactory {
+    /// Creates a new [`SyntaxFactory`], generating mappings between input nodes and generated nodes.
+    pub fn new() -> Self {
+        Self { mappings: Some(RefCell::new(SyntaxMapping::new())) }
+    }
+
+    /// Creates a [`SyntaxFactory`] without generating mappings.
+    pub fn without_mappings() -> Self {
+        Self { mappings: None }
+    }
+
+    /// Gets all of the tracked syntax mappings, if any.
+    pub fn finish_with_mappings(self) -> SyntaxMapping {
+        self.mappings.unwrap_or_default().into_inner()
+    }
+
+    fn mappings(&self) -> Option<RefMut<'_, SyntaxMapping>> {
+        self.mappings.as_ref().map(|it| it.borrow_mut())
+    }
+}
+
+impl Default for SyntaxFactory {
+    fn default() -> Self {
+        Self::without_mappings()
+    }
+}

crates/syntax/src/ast/syntax_factory/constructors.rs

@@ -0,0 +1,110 @@
+//! Wrappers over [`make`] constructors
+use itertools::Itertools;
+
+use crate::{
+    ast::{self, make, HasName},
+    syntax_editor::SyntaxMappingBuilder,
+    AstNode,
+};
+
+use super::SyntaxFactory;
+
+impl SyntaxFactory {
+    pub fn name(&self, name: &str) -> ast::Name {
+        make::name(name).clone_for_update()
+    }
+
+    pub fn ident_pat(&self, ref_: bool, mut_: bool, name: ast::Name) -> ast::IdentPat {
+        let ast = make::ident_pat(ref_, mut_, name.clone()).clone_for_update();
+
+        if let Some(mut mapping) = self.mappings() {
+            let mut builder = SyntaxMappingBuilder::new(ast.syntax().clone());
+            builder.map_node(name.syntax().clone(), ast.name().unwrap().syntax().clone());
+            builder.finish(&mut mapping);
+        }
+
+        ast
+    }
+
+    pub fn block_expr(
+        &self,
+        stmts: impl IntoIterator<Item = ast::Stmt>,
+        tail_expr: Option<ast::Expr>,
+    ) -> ast::BlockExpr {
+        let stmts = stmts.into_iter().collect_vec();
+        let input = stmts.iter().map(|it| it.syntax().clone()).collect_vec();
+
+        let ast = make::block_expr(stmts, tail_expr.clone()).clone_for_update();
+
+        if let Some((mut mapping, stmt_list)) = self.mappings().zip(ast.stmt_list()) {
+            let mut builder = SyntaxMappingBuilder::new(stmt_list.syntax().clone());
+
+            builder.map_children(
+                input.into_iter(),
+                stmt_list.statements().map(|it| it.syntax().clone()),
+            );
+
+            if let Some((input, output)) = tail_expr.zip(stmt_list.tail_expr()) {
+                builder.map_node(input.syntax().clone(), output.syntax().clone());
+            }
+
+            builder.finish(&mut mapping);
+        }
+
+        ast
+    }
+
+    pub fn expr_path(&self, path: ast::Path) -> ast::Expr {
+        let ast::Expr::PathExpr(ast) = make::expr_path(path.clone()).clone_for_update() else {
+            unreachable!()
+        };
+
+        if let Some(mut mapping) = self.mappings() {
+            let mut builder = SyntaxMappingBuilder::new(ast.syntax().clone());
+            builder.map_node(path.syntax().clone(), ast.path().unwrap().syntax().clone());
+            builder.finish(&mut mapping);
+        }
+
+        ast.into()
+    }
+
+    pub fn expr_ref(&self, expr: ast::Expr, exclusive: bool) -> ast::Expr {
+        let ast::Expr::RefExpr(ast) = make::expr_ref(expr.clone(), exclusive).clone_for_update()
+        else {
+            unreachable!()
+        };
+
+        if let Some(mut mapping) = self.mappings() {
+            let mut builder = SyntaxMappingBuilder::new(ast.syntax().clone());
+            builder.map_node(expr.syntax().clone(), ast.expr().unwrap().syntax().clone());
+            builder.finish(&mut mapping);
+        }
+
+        ast.into()
+    }
+
+    pub fn let_stmt(
+        &self,
+        pattern: ast::Pat,
+        ty: Option<ast::Type>,
+        initializer: Option<ast::Expr>,
+    ) -> ast::LetStmt {
+        let ast =
+            make::let_stmt(pattern.clone(), ty.clone(), initializer.clone()).clone_for_update();
+
+        if let Some(mut mapping) = self.mappings() {
+            let mut builder = SyntaxMappingBuilder::new(ast.syntax().clone());
+            builder.map_node(pattern.syntax().clone(), ast.pat().unwrap().syntax().clone());
+            if let Some(input) = ty {
+                builder.map_node(input.syntax().clone(), ast.ty().unwrap().syntax().clone());
+            }
+            if let Some(input) = initializer {
+                builder
+                    .map_node(input.syntax().clone(), ast.initializer().unwrap().syntax().clone());
+            }
+            builder.finish(&mut mapping);
+        }
+
+        ast
+    }
+}

crates/syntax/src/syntax_editor.rs

@@ -100,6 +100,10 @@ impl SyntaxEditor {
     pub fn finish(self) -> SyntaxEdit {
         edit_algo::apply_edits(self)
     }
+
+    pub fn add_mappings(&mut self, other: SyntaxMapping) {
+        self.mappings.merge(other);
+    }
 }
 
 /// Represents a completed [`SyntaxEditor`] operation.
@@ -319,85 +323,14 @@ fn is_ancestor_or_self_of_element(node: &SyntaxElement, ancestor: &SyntaxNode) -
 #[cfg(test)]
 mod tests {
     use expect_test::expect;
-    use itertools::Itertools;
 
     use crate::{
-        ast::{self, make, HasName},
+        ast::{self, make, syntax_factory::SyntaxFactory},
         AstNode,
     };
 
     use super::*;
 
-    fn make_ident_pat(
-        editor: Option<&mut SyntaxEditor>,
-        ref_: bool,
-        mut_: bool,
-        name: ast::Name,
-    ) -> ast::IdentPat {
-        let ast = make::ident_pat(ref_, mut_, name.clone()).clone_for_update();
-
-        if let Some(editor) = editor {
-            let mut mapping = SyntaxMappingBuilder::new(ast.syntax().clone());
-            mapping.map_node(name.syntax().clone(), ast.name().unwrap().syntax().clone());
-            mapping.finish(editor);
-        }
-
-        ast
-    }
-
-    fn make_let_stmt(
-        editor: Option<&mut SyntaxEditor>,
-        pattern: ast::Pat,
-        ty: Option<ast::Type>,
-        initializer: Option<ast::Expr>,
-    ) -> ast::LetStmt {
-        let ast =
-            make::let_stmt(pattern.clone(), ty.clone(), initializer.clone()).clone_for_update();
-
-        if let Some(editor) = editor {
-            let mut mapping = SyntaxMappingBuilder::new(ast.syntax().clone());
-            mapping.map_node(pattern.syntax().clone(), ast.pat().unwrap().syntax().clone());
-            if let Some(input) = ty {
-                mapping.map_node(input.syntax().clone(), ast.ty().unwrap().syntax().clone());
-            }
-            if let Some(input) = initializer {
-                mapping
-                    .map_node(input.syntax().clone(), ast.initializer().unwrap().syntax().clone());
-            }
-            mapping.finish(editor);
-        }
-
-        ast
-    }
-
-    fn make_block_expr(
-        editor: Option<&mut SyntaxEditor>,
-        stmts: impl IntoIterator<Item = ast::Stmt>,
-        tail_expr: Option<ast::Expr>,
-    ) -> ast::BlockExpr {
-        let stmts = stmts.into_iter().collect_vec();
-        let input = stmts.iter().map(|it| it.syntax().clone()).collect_vec();
-
-        let ast = make::block_expr(stmts, tail_expr.clone()).clone_for_update();
-
-        if let Some((editor, stmt_list)) = editor.zip(ast.stmt_list()) {
-            let mut mapping = SyntaxMappingBuilder::new(stmt_list.syntax().clone());
-
-            mapping.map_children(
-                input.into_iter(),
-                stmt_list.statements().map(|it| it.syntax().clone()),
-            );
-
-            if let Some((input, output)) = tail_expr.zip(stmt_list.tail_expr()) {
-                mapping.map_node(input.syntax().clone(), output.syntax().clone());
-            }
-
-            mapping.finish(editor);
-        }
-
-        ast
-    }
-
     #[test]
     fn basic_usage() {
         let root = make::match_arm(
@@ -417,6 +350,7 @@ mod tests {
         let to_replace = root.syntax().descendants().find_map(ast::BinExpr::cast).unwrap();
 
         let mut editor = SyntaxEditor::new(root.syntax().clone());
+        let make = SyntaxFactory::new();
 
         let name = make::name("var_name");
         let name_ref = make::name_ref("var_name").clone_for_update();
@@ -425,21 +359,20 @@ mod tests {
         editor.add_annotation(name.syntax(), placeholder_snippet);
         editor.add_annotation(name_ref.syntax(), placeholder_snippet);
 
-        let make_ident_pat = make_ident_pat(Some(&mut editor), false, false, name);
-        let make_let_stmt = make_let_stmt(
-            Some(&mut editor),
-            make_ident_pat.into(),
-            None,
-            Some(to_replace.clone().into()),
-        );
-        let new_block = make_block_expr(
-            Some(&mut editor),
-            [make_let_stmt.into()],
+        let new_block = make.block_expr(
+            [make
+                .let_stmt(
+                    make.ident_pat(false, false, name.clone()).into(),
+                    None,
+                    Some(to_replace.clone().into()),
+                )
+                .into()],
             Some(to_wrap.clone().into()),
         );
 
         editor.replace(to_replace.syntax(), name_ref.syntax());
         editor.replace(to_wrap.syntax(), new_block.syntax());
+        editor.add_mappings(make.finish_with_mappings());
 
         let edit = editor.finish();
 
@@ -473,11 +406,11 @@ mod tests {
         let second_let = root.syntax().descendants().find_map(ast::LetStmt::cast).unwrap();
 
         let mut editor = SyntaxEditor::new(root.syntax().clone());
+        let make = SyntaxFactory::without_mappings();
 
         editor.insert(
             Position::first_child_of(root.stmt_list().unwrap().syntax()),
-            make_let_stmt(
-                None,
+            make.let_stmt(
                 make::ext::simple_ident_pat(make::name("first")).into(),
                 None,
                 Some(make::expr_literal("1").into()),
@@ -487,8 +420,7 @@ mod tests {
 
         editor.insert(
             Position::after(second_let.syntax()),
-            make_let_stmt(
-                None,
+            make.let_stmt(
                 make::ext::simple_ident_pat(make::name("third")).into(),
                 None,
                 Some(make::expr_literal("3").into()),
@@ -528,19 +460,17 @@ mod tests {
         let second_let = root.syntax().descendants().find_map(ast::LetStmt::cast).unwrap();
 
         let mut editor = SyntaxEditor::new(root.syntax().clone());
+        let make = SyntaxFactory::new();
 
-        let new_block_expr =
-            make_block_expr(Some(&mut editor), [], Some(ast::Expr::BlockExpr(inner_block.clone())));
+        let new_block_expr = make.block_expr([], Some(ast::Expr::BlockExpr(inner_block.clone())));
 
-        let first_let = make_let_stmt(
-            Some(&mut editor),
+        let first_let = make.let_stmt(
             make::ext::simple_ident_pat(make::name("first")).into(),
             None,
             Some(make::expr_literal("1").into()),
         );
 
-        let third_let = make_let_stmt(
-            Some(&mut editor),
+        let third_let = make.let_stmt(
             make::ext::simple_ident_pat(make::name("third")).into(),
             None,
             Some(make::expr_literal("3").into()),
@@ -552,6 +482,7 @@ mod tests {
         );
         editor.insert(Position::after(second_let.syntax()), third_let.syntax());
         editor.replace(inner_block.syntax(), new_block_expr.syntax());
+        editor.add_mappings(make.finish_with_mappings());
 
         let edit = editor.finish();
 
@@ -581,12 +512,11 @@ mod tests {
         let inner_block = root.clone();
 
         let mut editor = SyntaxEditor::new(root.syntax().clone());
+        let make = SyntaxFactory::new();
 
-        let new_block_expr =
-            make_block_expr(Some(&mut editor), [], Some(ast::Expr::BlockExpr(inner_block.clone())));
+        let new_block_expr = make.block_expr([], Some(ast::Expr::BlockExpr(inner_block.clone())));
 
-        let first_let = make_let_stmt(
-            Some(&mut editor),
+        let first_let = make.let_stmt(
             make::ext::simple_ident_pat(make::name("first")).into(),
             None,
             Some(make::expr_literal("1").into()),
@@ -597,6 +527,7 @@ mod tests {
             first_let.syntax(),
         );
         editor.replace(inner_block.syntax(), new_block_expr.syntax());
+        editor.add_mappings(make.finish_with_mappings());
 
         let edit = editor.finish();
 

crates/syntax/src/syntax_editor/mapping.rs

@@ -7,8 +7,6 @@ use rustc_hash::FxHashMap;
 
 use crate::{SyntaxElement, SyntaxNode};
 
-use super::SyntaxEditor;
-
 #[derive(Debug, Default)]
 pub struct SyntaxMapping {
     // important information to keep track of:
@@ -209,7 +207,7 @@ impl SyntaxMapping {
         Some(output)
     }
 
-    fn add_mapping(&mut self, syntax_mapping: SyntaxMappingBuilder) {
+    pub fn add_mapping(&mut self, syntax_mapping: SyntaxMappingBuilder) {
         let SyntaxMappingBuilder { parent_node, node_mappings } = syntax_mapping;
 
         let parent_entry: u32 = self.entry_parents.len().try_into().unwrap();
@@ -257,8 +255,8 @@ impl SyntaxMappingBuilder {
         }
     }
 
-    pub fn finish(self, editor: &mut SyntaxEditor) {
-        editor.mappings.add_mapping(self);
+    pub fn finish(self, mappings: &mut SyntaxMapping) {
+        mappings.add_mapping(self);
     }
 }