Skip to content

Commit 0f1673c

Browse files
alibektasalibektas
committed
v3
1 parent 6f74604 commit 0f1673c

File tree

1 file changed

+88
-94
lines changed

1 file changed

+88
-94
lines changed

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

+88-94
Original file line numberDiff line numberDiff line change
@@ -103,49 +103,50 @@ pub(crate) fn extract_function(acc: &mut Assists, ctx: &AssistContext<'_>) -> Op
103103

104104
let scope = ImportScope::find_insert_use_container(&node, &ctx.sema)?;
105105

106-
let outliving_locals: Vec<_> = ret_values.collect();
107-
if stdx::never!(!outliving_locals.is_empty() && !ret_ty.is_unit()) {
108-
// We should not have variables that outlive body if we have expression block
109-
return None;
110-
}
111-
112-
let params = body.extracted_function_params(ctx, &container_info, locals_used.iter().copied());
113-
114-
let name = make_function_name(&semantics_scope);
115-
116-
let has_impl_wrapper =
117-
insert_after.ancestors().any(|a| a.kind() == SyntaxKind::IMPL && a != insert_after);
118-
119-
let fun = Function {
120-
name,
121-
self_param,
122-
params,
123-
control_flow,
124-
ret_ty,
125-
body,
126-
outliving_locals,
127-
contains_tail_expr,
128-
mods: container_info,
129-
};
130-
131-
let new_indent = IndentLevel::from_node(&insert_after);
132-
let old_indent = fun.body.indent_level();
133-
134-
let fn_def = match fun.self_param_adt(ctx) {
135-
Some(adt) if anchor == Anchor::Method && !has_impl_wrapper => {
136-
let fn_def = format_function(ctx, module, &fun, old_indent, new_indent + 1)?;
137-
generate_impl_text(&adt, &fn_def).replace("{\n\n", "{")
138-
}
139-
_ => format_function(ctx, module, &fun, old_indent, new_indent)?,
140-
};
141-
142106
acc.add(
143107
AssistId("extract_function", crate::AssistKind::RefactorExtract),
144108
"Extract into function",
145109
target_range,
146110
move |builder| {
111+
let outliving_locals: Vec<_> = ret_values.collect();
112+
if stdx::never!(!outliving_locals.is_empty() && !ret_ty.is_unit()) {
113+
// We should not have variables that outlive body if we have expression block
114+
return;
115+
}
116+
117+
let params =
118+
body.extracted_function_params(ctx, &container_info, locals_used.iter().copied());
119+
120+
let name = make_function_name(&semantics_scope);
121+
122+
let fun = Function {
123+
name,
124+
self_param,
125+
params,
126+
control_flow,
127+
ret_ty,
128+
body,
129+
outliving_locals,
130+
contains_tail_expr,
131+
mods: container_info,
132+
};
133+
134+
let new_indent = IndentLevel::from_node(&insert_after);
135+
let old_indent = fun.body.indent_level();
136+
147137
builder.replace(target_range, make_call(ctx, &fun, old_indent));
148138

139+
let has_impl_wrapper =
140+
insert_after.ancestors().any(|a| a.kind() == SyntaxKind::IMPL && a != insert_after);
141+
142+
let fn_def = match fun.self_param_adt(ctx) {
143+
Some(adt) if anchor == Anchor::Method && !has_impl_wrapper => {
144+
let fn_def = format_function(ctx, module, &fun, old_indent, new_indent + 1);
145+
generate_impl_text(&adt, &fn_def).replace("{\n\n", "{")
146+
}
147+
_ => format_function(ctx, module, &fun, old_indent, new_indent),
148+
};
149+
149150
if fn_def.contains("ControlFlow") {
150151
let scope = match scope {
151152
ImportScope::File(it) => ImportScope::File(builder.make_mut(it)),
@@ -530,7 +531,7 @@ impl FunctionBody {
530531

531532
fn extracted_from_trait_impl(&self) -> bool {
532533
match self.node().ancestors().find_map(ast::Impl::cast) {
533-
Some(c) => return c.trait_().is_some(),
534+
Some(c) => c.trait_().is_some(),
534535
None => false,
535536
}
536537
}
@@ -1047,23 +1048,17 @@ impl GenericParent {
10471048
fn generic_parents(parent: &SyntaxNode) -> Vec<GenericParent> {
10481049
let mut list = Vec::new();
10491050
if let Some(parent_item) = parent.ancestors().find_map(ast::Item::cast) {
1050-
match parent_item {
1051-
ast::Item::Fn(ref fn_) => {
1052-
if let Some(parent_parent) = parent_item
1053-
.syntax()
1054-
.parent()
1055-
.and_then(|it| it.parent())
1056-
.and_then(ast::Item::cast)
1057-
{
1058-
match parent_parent {
1059-
ast::Item::Impl(impl_) => list.push(GenericParent::Impl(impl_)),
1060-
ast::Item::Trait(trait_) => list.push(GenericParent::Trait(trait_)),
1061-
_ => (),
1062-
}
1051+
if let ast::Item::Fn(ref fn_) = parent_item {
1052+
if let Some(parent_parent) =
1053+
parent_item.syntax().parent().and_then(|it| it.parent()).and_then(ast::Item::cast)
1054+
{
1055+
match parent_parent {
1056+
ast::Item::Impl(impl_) => list.push(GenericParent::Impl(impl_)),
1057+
ast::Item::Trait(trait_) => list.push(GenericParent::Trait(trait_)),
1058+
_ => (),
10631059
}
1064-
list.push(GenericParent::Fn(fn_.clone()));
10651060
}
1066-
_ => (),
1061+
list.push(GenericParent::Fn(fn_.clone()));
10671062
}
10681063
}
10691064
list
@@ -1500,13 +1495,13 @@ fn format_function(
15001495
fun: &Function,
15011496
old_indent: IndentLevel,
15021497
new_indent: IndentLevel,
1503-
) -> Option<String> {
1498+
) -> String {
15041499
let mut fn_def = String::new();
15051500

15061501
let fun_name = &fun.name;
15071502
let params = fun.make_param_list(ctx, module);
15081503
let ret_ty = fun.make_ret_ty(ctx, module);
1509-
let body = make_body(ctx, old_indent, new_indent, fun)?;
1504+
let body = make_body(ctx, old_indent, new_indent, fun);
15101505
let const_kw = if fun.mods.is_const { "const " } else { "" };
15111506
let async_kw = if fun.control_flow.is_async { "async " } else { "" };
15121507
let unsafe_kw = if fun.control_flow.is_unsafe { "unsafe " } else { "" };
@@ -1534,7 +1529,7 @@ fn format_function(
15341529

15351530
format_to!(fn_def, " {body}");
15361531

1537-
Some(fn_def)
1532+
fn_def
15381533
}
15391534

15401535
fn make_generic_params_and_where_clause(
@@ -1720,15 +1715,14 @@ fn make_body(
17201715
old_indent: IndentLevel,
17211716
new_indent: IndentLevel,
17221717
fun: &Function,
1723-
) -> Option<ast::BlockExpr> {
1718+
) -> ast::BlockExpr {
17241719
let ret_ty = fun.return_type(ctx);
17251720
let handler = FlowHandler::from_ret_ty(fun, &ret_ty);
17261721

17271722
let block = match &fun.body {
17281723
FunctionBody::Expr(expr) => {
1729-
let expr = rewrite_body_segment(ctx, &fun.params, &handler, expr.syntax())
1730-
.expect("rewrite_body_segment failed.");
1731-
let expr = ast::Expr::cast(expr).expect("error while casting body segment to an expr.");
1724+
let expr = rewrite_body_segment(ctx, &fun.params, &handler, expr.syntax());
1725+
let expr = ast::Expr::cast(expr).expect("Body segment should be an expr");
17321726
match expr {
17331727
ast::Expr::BlockExpr(block) => {
17341728
// If the extracted expression is itself a block, there is no need to wrap it inside another block.
@@ -1749,16 +1743,12 @@ fn make_body(
17491743
.children_with_tokens()
17501744
.filter(|it| text_range.contains_range(it.text_range()))
17511745
.map(|it| match &it {
1752-
syntax::NodeOrToken::Node(n) => {
1753-
let bs = rewrite_body_segment(ctx, &fun.params, &handler, n);
1754-
match bs {
1755-
Some(n) => Some(syntax::NodeOrToken::Node(n)),
1756-
None => None,
1757-
}
1758-
}
1759-
_ => Some(it),
1746+
syntax::NodeOrToken::Node(n) => syntax::NodeOrToken::Node(
1747+
rewrite_body_segment(ctx, &fun.params, &handler, n),
1748+
),
1749+
_ => it,
17601750
})
1761-
.collect::<Option<_>>()?;
1751+
.collect();
17621752

17631753
let mut tail_expr = match &elements.last() {
17641754
Some(syntax::NodeOrToken::Node(node)) if ast::Expr::can_cast(node.kind()) => {
@@ -1842,7 +1832,7 @@ fn make_body(
18421832
}),
18431833
};
18441834

1845-
Some(block.indent(new_indent))
1835+
block.indent(new_indent)
18461836
}
18471837

18481838
fn map_tail_expr(block: ast::BlockExpr, f: impl FnOnce(ast::Expr) -> ast::Expr) -> ast::BlockExpr {
@@ -1872,9 +1862,8 @@ fn with_tail_expr(block: ast::BlockExpr, tail_expr: ast::Expr) -> ast::BlockExpr
18721862

18731863
if let Some(stmt_list) = block.stmt_list() {
18741864
stmt_list.syntax().children_with_tokens().for_each(|node_or_token| {
1875-
match &node_or_token {
1876-
syntax::NodeOrToken::Token(_) => elements.push(node_or_token),
1877-
_ => (),
1865+
if let syntax::NodeOrToken::Token(_) = &node_or_token {
1866+
elements.push(node_or_token)
18781867
};
18791868
});
18801869
}
@@ -1900,18 +1889,14 @@ fn rewrite_body_segment(
19001889
params: &[Param],
19011890
handler: &FlowHandler,
19021891
syntax: &SyntaxNode,
1903-
) -> Option<SyntaxNode> {
1904-
let syntax = fix_param_usages(ctx, params, syntax)?;
1892+
) -> SyntaxNode {
1893+
let syntax = fix_param_usages(ctx, params, syntax);
19051894
update_external_control_flow(handler, &syntax);
1906-
Some(syntax)
1895+
syntax
19071896
}
19081897

19091898
/// change all usages to account for added `&`/`&mut` for some params
1910-
fn fix_param_usages(
1911-
ctx: &AssistContext<'_>,
1912-
params: &[Param],
1913-
syntax: &SyntaxNode,
1914-
) -> Option<SyntaxNode> {
1899+
fn fix_param_usages(ctx: &AssistContext<'_>, params: &[Param], syntax: &SyntaxNode) -> SyntaxNode {
19151900
let mut usages_for_param: Vec<(&Param, Vec<ast::Expr>)> = Vec::new();
19161901

19171902
let tm = TreeMutator::new(syntax);
@@ -1935,26 +1920,35 @@ fn fix_param_usages(
19351920

19361921
for (param, usages) in usages_for_param {
19371922
for usage in usages {
1938-
let expr = usage.syntax().ancestors().skip(1).find_map(ast::Expr::cast);
1939-
if let Some(ast::Expr::MethodCallExpr(_) | ast::Expr::FieldExpr(_)) = expr {
1940-
continue;
1941-
}
1942-
1943-
if let Some(ast::Expr::RefExpr(node)) = expr {
1944-
if (param.kind() == ParamKind::MutRef && node.mut_token().is_some())
1945-
|| (param.kind() == ParamKind::SharedRef && node.mut_token().is_none())
1923+
match usage.syntax().ancestors().skip(1).find_map(ast::Expr::cast) {
1924+
Some(ast::Expr::MethodCallExpr(_) | ast::Expr::FieldExpr(_)) => {
1925+
// do nothing
1926+
}
1927+
Some(ast::Expr::RefExpr(node))
1928+
if param.kind() == ParamKind::MutRef && node.mut_token().is_some() =>
1929+
{
1930+
ted::replace(
1931+
node.syntax(),
1932+
node.expr().expect("RefExpr::expr() cannot be None").syntax(),
1933+
);
1934+
}
1935+
Some(ast::Expr::RefExpr(node))
1936+
if param.kind() == ParamKind::SharedRef && node.mut_token().is_none() =>
19461937
{
1947-
ted::replace(node.syntax(), node.expr()?.syntax());
1948-
continue;
1938+
ted::replace(
1939+
node.syntax(),
1940+
node.expr().expect("RefExpr::expr() cannot be None").syntax(),
1941+
);
1942+
}
1943+
Some(_) | None => {
1944+
let p = &make::expr_prefix(T![*], usage.clone()).clone_for_update();
1945+
ted::replace(usage.syntax(), p.syntax())
19491946
}
19501947
}
1951-
1952-
let p = &make::expr_prefix(T![*], usage.clone()).clone_for_update();
1953-
ted::replace(usage.syntax(), p.syntax())
19541948
}
19551949
}
19561950

1957-
Some(res)
1951+
res
19581952
}
19591953

19601954
fn update_external_control_flow(handler: &FlowHandler, syntax: &SyntaxNode) {

0 commit comments

Comments
 (0)