Skip to content
This repository was archived by the owner on May 28, 2025. It is now read-only.

Commit 1e1761e

Browse files
DropDemBitsDropDemBits
committed
Migrate extract_variable to mutable ast
1 parent 457b966 commit 1e1761e

File tree

1 file changed

+130
-93
lines changed

1 file changed

+130
-93
lines changed

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

Lines changed: 130 additions & 93 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,8 @@
11
use hir::TypeInfo;
2-
use stdx::format_to;
32
use syntax::{
4-
ast::{self, AstNode},
5-
NodeOrToken,
6-
SyntaxKind::{
7-
BLOCK_EXPR, BREAK_EXPR, CLOSURE_EXPR, COMMENT, LOOP_EXPR, MATCH_ARM, MATCH_GUARD,
8-
PATH_EXPR, RETURN_EXPR,
9-
},
3+
ast::{self, edit::IndentLevel, edit_in_place::Indent, make, AstNode, HasName},
4+
ted, NodeOrToken,
5+
SyntaxKind::{BLOCK_EXPR, BREAK_EXPR, COMMENT, LOOP_EXPR, MATCH_GUARD, PATH_EXPR, RETURN_EXPR},
106
SyntaxNode,
117
};
128

@@ -66,98 +62,140 @@ pub(crate) fn extract_variable(acc: &mut Assists, ctx: &AssistContext<'_>) -> Op
6662
.as_ref()
6763
.map_or(false, |it| matches!(it, ast::Expr::FieldExpr(_) | ast::Expr::MethodCallExpr(_)));
6864

69-
let reference_modifier = match ty.filter(|_| needs_adjust) {
70-
Some(receiver_type) if receiver_type.is_mutable_reference() => "&mut ",
71-
Some(receiver_type) if receiver_type.is_reference() => "&",
72-
_ => "",
73-
};
74-
75-
let var_modifier = match parent {
76-
Some(ast::Expr::RefExpr(expr)) if expr.mut_token().is_some() => "mut ",
77-
_ => "",
78-
};
79-
8065
let anchor = Anchor::from(&to_extract)?;
81-
let indent = anchor.syntax().prev_sibling_or_token()?.as_token()?.clone();
8266
let target = to_extract.syntax().text_range();
8367
acc.add(
8468
AssistId("extract_variable", AssistKind::RefactorExtract),
8569
"Extract into variable",
8670
target,
8771
move |edit| {
88-
let field_shorthand =
89-
match to_extract.syntax().parent().and_then(ast::RecordExprField::cast) {
90-
Some(field) => field.name_ref(),
91-
None => None,
92-
};
93-
94-
let mut buf = String::new();
95-
96-
let var_name = match &field_shorthand {
97-
Some(it) => it.to_string(),
98-
None => suggest_name::for_variable(&to_extract, &ctx.sema),
72+
let field_shorthand = to_extract
73+
.syntax()
74+
.parent()
75+
.and_then(ast::RecordExprField::cast)
76+
.filter(|field| field.name_ref().is_some());
77+
78+
let (var_name, expr_replace) = match field_shorthand {
79+
Some(field) => (field.to_string(), field.syntax().clone()),
80+
None => (
81+
suggest_name::for_variable(&to_extract, &ctx.sema),
82+
to_extract.syntax().clone(),
83+
),
9984
};
100-
let expr_range = match &field_shorthand {
101-
Some(it) => it.syntax().text_range().cover(to_extract.syntax().text_range()),
102-
None => to_extract.syntax().text_range(),
85+
86+
let ident_pat = match parent {
87+
Some(ast::Expr::RefExpr(expr)) if expr.mut_token().is_some() => {
88+
make::ident_pat(false, true, make::name(&var_name))
89+
}
90+
_ => make::ident_pat(false, false, make::name(&var_name)),
10391
};
10492

105-
match anchor {
106-
Anchor::Before(_) | Anchor::Replace(_) => {
107-
format_to!(buf, "let {var_modifier}{var_name} = {reference_modifier}")
93+
let to_extract = match ty.as_ref().filter(|_| needs_adjust) {
94+
Some(receiver_type) if receiver_type.is_mutable_reference() => {
95+
make::expr_ref(to_extract, true)
10896
}
109-
Anchor::WrapInBlock(_) => {
110-
format_to!(buf, "{{ let {var_name} = {reference_modifier}")
97+
Some(receiver_type) if receiver_type.is_reference() => {
98+
make::expr_ref(to_extract, false)
11199
}
100+
_ => to_extract,
112101
};
113-
format_to!(buf, "{to_extract}");
114102

115-
if let Anchor::Replace(stmt) = anchor {
116-
cov_mark::hit!(test_extract_var_expr_stmt);
117-
if stmt.semicolon_token().is_none() {
118-
buf.push(';');
119-
}
120-
match ctx.config.snippet_cap {
121-
Some(cap) => {
122-
let snip = buf.replace(
123-
&format!("let {var_modifier}{var_name}"),
124-
&format!("let {var_modifier}$0{var_name}"),
125-
);
126-
edit.replace_snippet(cap, expr_range, snip)
103+
let expr_replace = edit.make_syntax_mut(expr_replace);
104+
let let_stmt =
105+
make::let_stmt(ident_pat.into(), None, Some(to_extract)).clone_for_update();
106+
let name_expr = make::expr_path(make::ext::ident_path(&var_name)).clone_for_update();
107+
108+
match anchor {
109+
Anchor::Before(place) => {
110+
let prev_ws = place.prev_sibling_or_token().and_then(|it| it.into_token());
111+
let indent_to = IndentLevel::from_node(&place);
112+
let insert_place = edit.make_syntax_mut(place);
113+
114+
// Adjust ws to insert depending on if this is all inline or on separate lines
115+
let trailing_ws = if prev_ws.is_some_and(|it| it.text().starts_with("\n")) {
116+
format!("\n{indent_to}")
117+
} else {
118+
format!(" ")
119+
};
120+
121+
ted::insert_all_raw(
122+
ted::Position::before(insert_place),
123+
vec![
124+
let_stmt.syntax().clone().into(),
125+
make::tokens::whitespace(&trailing_ws).into(),
126+
],
127+
);
128+
129+
ted::replace(expr_replace, name_expr.syntax());
130+
131+
if let Some(cap) = ctx.config.snippet_cap {
132+
if let Some(ast::Pat::IdentPat(ident_pat)) = let_stmt.pat() {
133+
if let Some(name) = ident_pat.name() {
134+
edit.add_tabstop_before(cap, name);
135+
}
136+
}
127137
}
128-
None => edit.replace(expr_range, buf),
129138
}
130-
return;
131-
}
139+
Anchor::Replace(stmt) => {
140+
cov_mark::hit!(test_extract_var_expr_stmt);
132141

133-
buf.push(';');
134-
135-
// We want to maintain the indent level,
136-
// but we do not want to duplicate possible
137-
// extra newlines in the indent block
138-
let text = indent.text();
139-
if text.starts_with('\n') {
140-
buf.push('\n');
141-
buf.push_str(text.trim_start_matches('\n'));
142-
} else {
143-
buf.push_str(text);
144-
}
142+
let stmt_replace = edit.make_mut(stmt);
143+
ted::replace(stmt_replace.syntax(), let_stmt.syntax());
145144

146-
edit.replace(expr_range, var_name.clone());
147-
let offset = anchor.syntax().text_range().start();
148-
match ctx.config.snippet_cap {
149-
Some(cap) => {
150-
let snip = buf.replace(
151-
&format!("let {var_modifier}{var_name}"),
152-
&format!("let {var_modifier}$0{var_name}"),
153-
);
154-
edit.insert_snippet(cap, offset, snip)
145+
if let Some(cap) = ctx.config.snippet_cap {
146+
if let Some(ast::Pat::IdentPat(ident_pat)) = let_stmt.pat() {
147+
if let Some(name) = ident_pat.name() {
148+
edit.add_tabstop_before(cap, name);
149+
}
150+
}
151+
}
155152
}
156-
None => edit.insert(offset, buf),
157-
}
153+
Anchor::WrapInBlock(to_wrap) => {
154+
let indent_to = to_wrap.indent_level();
155+
156+
let block = if to_wrap.syntax() == &expr_replace {
157+
// Since `expr_replace` is the same that needs to be wrapped in a block,
158+
// we can just directly replace it with a block
159+
let block =
160+
make::block_expr([let_stmt.into()], Some(name_expr)).clone_for_update();
161+
ted::replace(expr_replace, block.syntax());
162+
163+
block
164+
} else {
165+
// `expr_replace` is a descendant of `to_wrap`, so both steps need to be
166+
// handled seperately, otherwise we wrap the wrong expression
167+
let to_wrap = edit.make_mut(to_wrap);
168+
169+
// Replace the target expr first so that we don't need to find where
170+
// `expr_replace` is in the wrapped `to_wrap`
171+
ted::replace(expr_replace, name_expr.syntax());
172+
173+
// Wrap `to_wrap` in a block
174+
let block = make::block_expr([let_stmt.into()], Some(to_wrap.clone()))
175+
.clone_for_update();
176+
ted::replace(to_wrap.syntax(), block.syntax());
177+
178+
block
179+
};
180+
181+
if let Some(cap) = ctx.config.snippet_cap {
182+
// Adding a tabstop to `name` requires finding the let stmt again, since
183+
// the existing `let_stmt` is not actually added to the tree
184+
let pat = block.statements().find_map(|stmt| {
185+
let ast::Stmt::LetStmt(let_stmt) = stmt else { return None };
186+
let_stmt.pat()
187+
});
188+
189+
if let Some(ast::Pat::IdentPat(ident_pat)) = pat {
190+
if let Some(name) = ident_pat.name() {
191+
edit.add_tabstop_before(cap, name);
192+
}
193+
}
194+
}
158195

159-
if let Anchor::WrapInBlock(_) = anchor {
160-
edit.insert(anchor.syntax().text_range().end(), " }");
196+
// fixup indentation of block
197+
block.indent(indent_to);
198+
}
161199
}
162200
},
163201
)
@@ -181,7 +219,7 @@ fn valid_target_expr(node: SyntaxNode) -> Option<ast::Expr> {
181219
enum Anchor {
182220
Before(SyntaxNode),
183221
Replace(ast::ExprStmt),
184-
WrapInBlock(SyntaxNode),
222+
WrapInBlock(ast::Expr),
185223
}
186224

187225
impl Anchor {
@@ -204,16 +242,16 @@ impl Anchor {
204242
}
205243

206244
if let Some(parent) = node.parent() {
207-
if parent.kind() == CLOSURE_EXPR {
245+
if let Some(parent) = ast::ClosureExpr::cast(parent.clone()) {
208246
cov_mark::hit!(test_extract_var_in_closure_no_block);
209-
return Some(Anchor::WrapInBlock(node));
247+
return parent.body().map(Anchor::WrapInBlock);
210248
}
211-
if parent.kind() == MATCH_ARM {
249+
if let Some(parent) = ast::MatchArm::cast(parent) {
212250
if node.kind() == MATCH_GUARD {
213251
cov_mark::hit!(test_extract_var_in_match_guard);
214252
} else {
215253
cov_mark::hit!(test_extract_var_in_match_arm_no_block);
216-
return Some(Anchor::WrapInBlock(node));
254+
return parent.expr().map(Anchor::WrapInBlock);
217255
}
218256
}
219257
}
@@ -229,13 +267,6 @@ impl Anchor {
229267
None
230268
})
231269
}
232-
233-
fn syntax(&self) -> &SyntaxNode {
234-
match self {
235-
Anchor::Before(it) | Anchor::WrapInBlock(it) => it,
236-
Anchor::Replace(stmt) => stmt.syntax(),
237-
}
238-
}
239270
}
240271

241272
#[cfg(test)]
@@ -502,7 +533,10 @@ fn main() {
502533
fn main() {
503534
let x = true;
504535
let tuple = match x {
505-
true => { let $0var_name = 2 + 2; (var_name, true) }
536+
true => {
537+
let $0var_name = 2 + 2;
538+
(var_name, true)
539+
}
506540
_ => (0, false)
507541
};
508542
}
@@ -579,7 +613,10 @@ fn main() {
579613
"#,
580614
r#"
581615
fn main() {
582-
let lambda = |x: u32| { let $0var_name = x * 2; var_name };
616+
let lambda = |x: u32| {
617+
let $0var_name = x * 2;
618+
var_name
619+
};
583620
}
584621
"#,
585622
);

0 commit comments

Comments
 (0)