Skip to content

Commit fed32aa

Browse files
ChayimFriedman2ChayimFriedman2
authored
Merge pull request #21487 from A4-Tacks/extract-var-in-macro
feat: Support extract variable in macro call
2 parents bb0e103 + 144cce8 commit fed32aa

1 file changed

Lines changed: 255 additions & 23 deletions

File tree

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

Lines changed: 255 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -4,14 +4,14 @@ use ide_db::{
44
syntax_helpers::{LexedStr, suggest_name},
55
};
66
use syntax::{
7-
NodeOrToken, SyntaxKind, SyntaxNode, T,
8-
algo::ancestors_at_offset,
7+
Direction, NodeOrToken, SyntaxKind, SyntaxNode, SyntaxToken, T, TextRange,
8+
algo::{ancestors_at_offset, skip_trivia_token},
99
ast::{
1010
self, AstNode,
1111
edit::{AstNodeEdit, IndentLevel},
1212
syntax_factory::SyntaxFactory,
1313
},
14-
syntax_editor::Position,
14+
syntax_editor::{Element, Position},
1515
};
1616

1717
use crate::{AssistContext, AssistId, Assists, utils::is_body_const};
@@ -92,27 +92,54 @@ pub(crate) fn extract_variable(acc: &mut Assists, ctx: &AssistContext<'_>) -> Op
9292
let node = node.ancestors().take_while(|anc| anc.text_range() == node.text_range()).last()?;
9393
let range = node.text_range();
9494

95-
let to_extract = node
96-
.descendants()
97-
.take_while(|it| range.contains_range(it.text_range()))
98-
.find_map(valid_target_expr(ctx))?;
95+
let (to_replace, analysis) = if node.kind() == SyntaxKind::TOKEN_TREE {
96+
let (first, last) = extract_token_range_of(&node, ctx.selection_trimmed())?;
9997

100-
let ty = ctx.sema.type_of_expr(&to_extract).map(TypeInfo::adjusted);
98+
let first_descend = ctx.sema.descend_into_macros_single_exact(first.clone());
99+
let last_descend = ctx.sema.descend_into_macros_single_exact(last.clone());
100+
let range = first_descend.text_range().cover(last_descend.text_range());
101+
102+
if first_descend.parent_ancestors().last() != last_descend.parent_ancestors().last() {
103+
return None;
104+
}
105+
106+
let expr = first_descend
107+
.parent_ancestors()
108+
.skip_while(|it| !it.text_range().contains_range(range))
109+
.find_map(valid_target_expr(ctx))?;
110+
let original_range = ctx.sema.original_range(expr.syntax());
111+
let (first, last) = extract_token_range_of(&node, original_range.range)?;
112+
let to_extract = first.syntax_element()..=last.syntax_element();
113+
(to_extract, expr)
114+
} else {
115+
let expr = node
116+
.descendants()
117+
.take_while(|it| range.contains_range(it.text_range()))
118+
.find_map(valid_target_expr(ctx))?;
119+
let to_extract = expr.syntax().syntax_element();
120+
(to_extract.clone()..=to_extract, expr)
121+
};
122+
let place = match to_replace.start() {
123+
NodeOrToken::Node(node) => node.clone(),
124+
NodeOrToken::Token(t) => t.parent()?,
125+
};
126+
127+
let ty = ctx.sema.type_of_expr(&analysis).map(TypeInfo::adjusted);
101128
if matches!(&ty, Some(ty_info) if ty_info.is_unit()) {
102129
return None;
103130
}
104131

105-
let parent = to_extract.syntax().parent().and_then(ast::Expr::cast);
132+
let parent = analysis.syntax().parent().and_then(ast::Expr::cast);
106133
// Any expression that autoderefs may need adjustment.
107134
let mut needs_adjust = parent.as_ref().is_some_and(|it| match it {
108135
ast::Expr::FieldExpr(_)
109136
| ast::Expr::MethodCallExpr(_)
110137
| ast::Expr::CallExpr(_)
111138
| ast::Expr::AwaitExpr(_) => true,
112-
ast::Expr::IndexExpr(index) if index.base().as_ref() == Some(&to_extract) => true,
139+
ast::Expr::IndexExpr(index) if index.base().as_ref() == Some(&analysis) => true,
113140
_ => false,
114141
});
115-
let mut to_extract_no_ref = peel_parens(to_extract.clone());
142+
let mut to_extract_no_ref = peel_parens(analysis.clone());
116143
let needs_ref = needs_adjust
117144
&& match &to_extract_no_ref {
118145
ast::Expr::FieldExpr(_)
@@ -127,14 +154,14 @@ pub(crate) fn extract_variable(acc: &mut Assists, ctx: &AssistContext<'_>) -> Op
127154
}
128155
_ => false,
129156
};
130-
let module = ctx.sema.scope(to_extract.syntax())?.module();
131-
let target = to_extract.syntax().text_range();
157+
let module = ctx.sema.scope(analysis.syntax())?.module();
158+
let target = to_replace.start().text_range().cover(to_replace.end().text_range());
132159
let needs_mut = match &parent {
133160
Some(ast::Expr::RefExpr(expr)) => expr.mut_token().is_some(),
134161
_ => needs_adjust && !needs_ref && ty.as_ref().is_some_and(|ty| ty.is_mutable_reference()),
135162
};
136163
for kind in ExtractionKind::ALL {
137-
let Some(anchor) = Anchor::from(&to_extract, kind) else {
164+
let Some(anchor) = Anchor::from(&place, kind) else {
138165
continue;
139166
};
140167

@@ -169,10 +196,18 @@ pub(crate) fn extract_variable(acc: &mut Assists, ctx: &AssistContext<'_>) -> Op
169196
kind.label(),
170197
target,
171198
|edit| {
172-
let (var_name, expr_replace) = kind.get_name_and_expr(ctx, &to_extract);
199+
let (var_name, expr_replace) = kind.get_name_and_expr(ctx, &analysis);
200+
201+
let to_replace =
202+
if expr_replace.ancestors().last() == to_replace.start().ancestors().last() {
203+
let element = expr_replace.clone().syntax_element();
204+
element.clone()..=element
205+
} else {
206+
to_replace.clone()
207+
};
173208

174209
let make = SyntaxFactory::with_mappings();
175-
let mut editor = edit.make_editor(&expr_replace);
210+
let mut editor = edit.make_editor(&place);
176211

177212
let pat_name = make.name(&var_name);
178213
let name_expr = make.expr_path(make.ident_path(&var_name));
@@ -236,7 +271,7 @@ pub(crate) fn extract_variable(acc: &mut Assists, ctx: &AssistContext<'_>) -> Op
236271
],
237272
);
238273

239-
editor.replace(expr_replace, name_expr.syntax());
274+
editor.replace_all(to_replace, vec![name_expr.syntax().syntax_element()]);
240275
}
241276
Anchor::Replace(stmt) => {
242277
cov_mark::hit!(test_extract_var_expr_stmt);
@@ -252,7 +287,8 @@ pub(crate) fn extract_variable(acc: &mut Assists, ctx: &AssistContext<'_>) -> Op
252287
make.block_expr([new_stmt], Some(name_expr))
253288
} else {
254289
// `expr_replace` is a descendant of `to_wrap`, so we just replace it with `name_expr`.
255-
editor.replace(expr_replace, name_expr.syntax());
290+
editor
291+
.replace_all(to_replace, vec![name_expr.syntax().syntax_element()]);
256292
make.block_expr([new_stmt], Some(to_wrap.clone()))
257293
}
258294
// fixup indentation of block
@@ -272,6 +308,23 @@ pub(crate) fn extract_variable(acc: &mut Assists, ctx: &AssistContext<'_>) -> Op
272308
Some(())
273309
}
274310

311+
fn extract_token_range_of(
312+
node: &SyntaxNode,
313+
range: TextRange,
314+
) -> Option<(SyntaxToken, SyntaxToken)> {
315+
let first = node.token_at_offset(range.start()).right_biased()?;
316+
let last = node.token_at_offset(range.end()).left_biased()?;
317+
318+
let first = skip_trivia_token(first, Direction::Next)?;
319+
let last = skip_trivia_token(last, Direction::Next)?;
320+
321+
if first.text_range().ordering(last.text_range()).is_gt() {
322+
return None;
323+
}
324+
325+
Some((first, last))
326+
}
327+
275328
fn peel_parens(mut expr: ast::Expr) -> ast::Expr {
276329
while let ast::Expr::ParenExpr(parens) = &expr {
277330
let Some(expr_inside) = parens.expr() else { break };
@@ -401,9 +454,8 @@ enum Anchor {
401454
}
402455

403456
impl Anchor {
404-
fn from(to_extract: &ast::Expr, kind: &ExtractionKind) -> Option<Anchor> {
405-
let result = to_extract
406-
.syntax()
457+
fn from(place: &SyntaxNode, kind: &ExtractionKind) -> Option<Anchor> {
458+
let result = place
407459
.ancestors()
408460
.take_while(|it| !ast::Item::can_cast(it.kind()) || ast::MacroCall::can_cast(it.kind()))
409461
.find_map(|node| {
@@ -435,7 +487,7 @@ impl Anchor {
435487

436488
if let Some(stmt) = ast::Stmt::cast(node.clone()) {
437489
if let ast::Stmt::ExprStmt(stmt) = stmt
438-
&& stmt.expr().as_ref() == Some(to_extract)
490+
&& stmt.expr().is_some_and(|it| it.syntax() == place)
439491
{
440492
return Some(Anchor::Replace(stmt));
441493
}
@@ -446,7 +498,7 @@ impl Anchor {
446498

447499
match kind {
448500
ExtractionKind::Constant | ExtractionKind::Static if result.is_none() => {
449-
to_extract.syntax().ancestors().find_map(|node| {
501+
place.ancestors().find_map(|node| {
450502
let item = ast::Item::cast(node.clone())?;
451503
let parent = item.syntax().parent()?;
452504
match parent.kind() {
@@ -2771,6 +2823,186 @@ fn main() {
27712823
let t2 = t;
27722824
let x = s;
27732825
}
2826+
"#,
2827+
"Extract into variable",
2828+
);
2829+
}
2830+
2831+
#[test]
2832+
fn extract_variable_in_token_tree() {
2833+
// FIXME: Keep the original trivia instead of extracting macro expanded?
2834+
check_assist_by_label(
2835+
extract_variable,
2836+
r#"
2837+
macro_rules! foo {
2838+
(= $($t:tt)*) => {
2839+
$($t)*
2840+
};
2841+
}
2842+
2843+
fn main() {
2844+
let x = foo!(= $02 + 3$0 + 4);
2845+
}
2846+
"#,
2847+
r#"
2848+
macro_rules! foo {
2849+
(= $($t:tt)*) => {
2850+
$($t)*
2851+
};
2852+
}
2853+
2854+
fn main() {
2855+
let $0var_name = 2+3;
2856+
let x = foo!(= var_name + 4);
2857+
}
2858+
"#,
2859+
"Extract into variable",
2860+
);
2861+
2862+
check_assist_by_label(
2863+
extract_variable,
2864+
r#"
2865+
macro_rules! foo {
2866+
(= $($t:tt)*) => {
2867+
$($t)*
2868+
};
2869+
}
2870+
2871+
fn main() {
2872+
let x = foo!(= $02 +$0 3 + 4);
2873+
}
2874+
"#,
2875+
r#"
2876+
macro_rules! foo {
2877+
(= $($t:tt)*) => {
2878+
$($t)*
2879+
};
2880+
}
2881+
2882+
fn main() {
2883+
let $0var_name = 2+3;
2884+
let x = foo!(= var_name + 4);
2885+
}
2886+
"#,
2887+
"Extract into variable",
2888+
);
2889+
2890+
check_assist_by_label(
2891+
extract_variable,
2892+
r#"
2893+
macro_rules! foo {
2894+
(= $($t:tt)*) => {
2895+
$($t)*
2896+
};
2897+
}
2898+
2899+
fn main() {
2900+
let x = foo!(= $02 + 3 + 4$0);
2901+
}
2902+
"#,
2903+
r#"
2904+
macro_rules! foo {
2905+
(= $($t:tt)*) => {
2906+
$($t)*
2907+
};
2908+
}
2909+
2910+
fn main() {
2911+
let $0var_name = 2+3+4;
2912+
let x = foo!(= var_name);
2913+
}
2914+
"#,
2915+
"Extract into variable",
2916+
);
2917+
2918+
// FIXME: Extract to inside the macro instead of outside the macro
2919+
check_assist_by_label(
2920+
extract_variable,
2921+
r#"
2922+
macro_rules! foo {
2923+
(= $($t:tt)*) => {
2924+
$($t)*
2925+
};
2926+
}
2927+
2928+
fn main() {
2929+
let x = foo!(= {
2930+
$02 + 3 + 4$0
2931+
});
2932+
}
2933+
"#,
2934+
r#"
2935+
macro_rules! foo {
2936+
(= $($t:tt)*) => {
2937+
$($t)*
2938+
};
2939+
}
2940+
2941+
fn main() {
2942+
let $0var_name = 2+3+4;
2943+
let x = foo!(= {
2944+
var_name
2945+
});
2946+
}
2947+
"#,
2948+
"Extract into variable",
2949+
);
2950+
}
2951+
2952+
#[test]
2953+
fn extract_variable_in_token_tree_record_expr() {
2954+
check_assist_by_label(
2955+
extract_variable,
2956+
r#"
2957+
macro_rules! foo {
2958+
(= $($t:tt)*) => {
2959+
$($t)*
2960+
};
2961+
}
2962+
2963+
fn main() {
2964+
let x = foo!(= Foo { x: $02 + 3$0 });
2965+
}
2966+
"#,
2967+
r#"
2968+
macro_rules! foo {
2969+
(= $($t:tt)*) => {
2970+
$($t)*
2971+
};
2972+
}
2973+
2974+
fn main() {
2975+
let $0x = 2+3;
2976+
let x = foo!(= Foo { x: x });
2977+
}
2978+
"#,
2979+
"Extract into variable",
2980+
);
2981+
2982+
check_assist_by_label(
2983+
extract_variable,
2984+
r#"
2985+
macro_rules! foo {
2986+
(= $($t:tt)*) => {
2987+
$($t)*
2988+
};
2989+
}
2990+
2991+
fn main() {
2992+
let x = foo!(= Foo { x: $02 + 3$0 + 4 });
2993+
}
2994+
"#,
2995+
r#"
2996+
macro_rules! foo {
2997+
(= $($t:tt)*) => {
2998+
$($t)*
2999+
};
3000+
}
3001+
3002+
fn main() {
3003+
let $0var_name = 2+3;
3004+
let x = foo!(= Foo { x: var_name + 4 });
3005+
}
27743006
"#,
27753007
"Extract into variable",
27763008
);

0 commit comments

Comments
 (0)