Skip to content

Commit e59acd5

Browse files
committed
Add ide-assist: extract_to_default_generic
Extracts selected type to default generic parameter. ```rust struct Foo(u32, $0String$0); ``` -> ```rust struct Foo<T$0 = String>(u32, T); ```
1 parent 529d3b9 commit e59acd5

File tree

5 files changed

+278
-0
lines changed

5 files changed

+278
-0
lines changed
Lines changed: 245 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,245 @@
1+
use ast::Name;
2+
use either::Either::{self, Left, Right};
3+
use ide_db::{source_change::SourceChangeBuilder, syntax_helpers::suggest_name::NameGenerator};
4+
use syntax::{
5+
ast::{self, AstNode, HasGenericParams, HasName, make},
6+
syntax_editor::{Position, SyntaxEditor},
7+
};
8+
9+
use crate::{AssistContext, AssistId, Assists};
10+
11+
// Assist: extract_to_default_generic
12+
//
13+
// Extracts selected type to default generic parameter.
14+
//
15+
// ```
16+
// struct Foo(u32, $0String$0);
17+
// ```
18+
// ->
19+
// ```
20+
// struct Foo<T$0 = String>(u32, T);
21+
// ```
22+
pub(crate) fn extract_to_default_generic(acc: &mut Assists, ctx: &AssistContext<'_>) -> Option<()> {
23+
if ctx.has_empty_selection() {
24+
return None;
25+
}
26+
27+
let ty: Either<ast::Type, ast::ConstArg> = ctx.find_node_at_range()?;
28+
let adt: Either<Either<ast::Adt, ast::TypeAlias>, ast::Fn> =
29+
ty.syntax().ancestors().find_map(AstNode::cast)?;
30+
31+
extract_to_default_generic_impl(acc, ctx, adt, ty)
32+
}
33+
34+
fn extract_to_default_generic_impl(
35+
acc: &mut Assists,
36+
ctx: &AssistContext<'_>,
37+
adt: Either<impl HasName + HasGenericParams, ast::Fn>,
38+
ty: Either<ast::Type, ast::ConstArg>,
39+
) -> Option<()> {
40+
let name = adt.name()?;
41+
let has_default = adt.is_left();
42+
43+
let label = if has_default {
44+
"Extract type as default generic parameter"
45+
} else {
46+
"Extract type as generic parameter"
47+
};
48+
let target = ty.syntax().text_range();
49+
acc.add(AssistId::refactor_extract("extract_to_default_generic"), label, target, |edit| {
50+
let mut editor = edit.make_editor(adt.syntax());
51+
let generic_list = get_or_create_generic_param_list(&name, &adt, &mut editor, edit);
52+
53+
let generic_name = generic_name(&generic_list, ty.is_right());
54+
55+
editor.replace(ty.syntax(), generic_name.syntax());
56+
57+
match ty {
58+
Left(ty) => {
59+
let param = if has_default {
60+
make::type_default_param(generic_name, None, ty)
61+
} else {
62+
make::type_param(generic_name, None)
63+
}
64+
.clone_for_update();
65+
66+
generic_list.add_generic_param(param.into());
67+
}
68+
Right(n) => {
69+
let param = if has_default {
70+
make::const_default_param(generic_name, const_ty(ctx, &n), n)
71+
} else {
72+
make::const_param(generic_name, const_ty(ctx, &n))
73+
}
74+
.clone_for_update();
75+
76+
generic_list.add_generic_param(param.into());
77+
}
78+
}
79+
80+
if let Some(cap) = ctx.config.snippet_cap
81+
&& let Some(last) = generic_list.generic_params().last()
82+
{
83+
if let ast::GenericParam::TypeParam(param) = &last
84+
&& let Some(name) = param.name()
85+
{
86+
let annotation = edit.make_tabstop_after(cap);
87+
editor.add_annotation(name.syntax(), annotation);
88+
} else if let ast::GenericParam::ConstParam(param) = &last
89+
&& let Some(ast::Type::InferType(ty)) = param.ty()
90+
{
91+
let annotation = edit.make_placeholder_snippet(cap);
92+
editor.add_annotation(ty.syntax(), annotation);
93+
}
94+
}
95+
96+
edit.add_file_edits(ctx.vfs_file_id(), editor);
97+
})
98+
}
99+
100+
fn array_index_type(n: &ast::ConstArg) -> Option<ast::Type> {
101+
let kind = n.syntax().parent()?.kind();
102+
103+
if ast::ArrayType::can_cast(kind) || ast::ArrayExpr::can_cast(kind) {
104+
Some(make::ty("usize"))
105+
} else {
106+
None
107+
}
108+
}
109+
110+
fn generic_name(generic_list: &ast::GenericParamList, is_const_param: bool) -> Name {
111+
let exist_names = generic_list
112+
.generic_params()
113+
.filter_map(|it| match it {
114+
ast::GenericParam::ConstParam(const_param) => const_param.name(),
115+
ast::GenericParam::TypeParam(type_param) => type_param.name(),
116+
ast::GenericParam::LifetimeParam(_) => None,
117+
})
118+
.map(|name| name.to_string())
119+
.collect::<Vec<_>>();
120+
121+
let mut name_gen = NameGenerator::new_with_names(exist_names.iter().map(|name| name.as_str()));
122+
123+
make::name(&if is_const_param {
124+
name_gen.suggest_name("N")
125+
} else {
126+
name_gen.suggest_name("T")
127+
})
128+
.clone_for_update()
129+
}
130+
131+
fn const_ty(ctx: &AssistContext<'_>, n: &ast::ConstArg) -> ast::Type {
132+
if let Some(expr) = n.expr()
133+
&& let Some(ty_info) = ctx.sema.type_of_expr(&expr)
134+
&& let Some(builtin) = ty_info.adjusted().as_builtin()
135+
{
136+
make::ty(builtin.name().as_str())
137+
} else if let Some(array_index_ty) = array_index_type(n) {
138+
array_index_ty
139+
} else {
140+
make::ty_placeholder()
141+
}
142+
}
143+
144+
fn get_or_create_generic_param_list(
145+
name: &ast::Name,
146+
adt: &impl HasGenericParams,
147+
editor: &mut SyntaxEditor,
148+
edit: &mut SourceChangeBuilder,
149+
) -> ast::GenericParamList {
150+
if let Some(list) = adt.generic_param_list() {
151+
edit.make_mut(list)
152+
} else {
153+
let generic = make::generic_param_list([]).clone_for_update();
154+
editor.insert(Position::after(name.syntax()), generic.syntax());
155+
generic
156+
}
157+
}
158+
159+
#[cfg(test)]
160+
mod tests {
161+
use super::*;
162+
163+
use crate::tests::check_assist;
164+
165+
#[test]
166+
fn test_extract_to_default_generic() {
167+
check_assist(
168+
extract_to_default_generic,
169+
r#"type X = ($0i32$0, i64);"#,
170+
r#"type X<T$0 = i32> = (T, i64);"#,
171+
);
172+
173+
check_assist(
174+
extract_to_default_generic,
175+
r#"type X<T> = ($0i32$0, T);"#,
176+
r#"type X<T, T1$0 = i32> = (T1, T);"#,
177+
);
178+
}
179+
180+
#[test]
181+
fn test_extract_to_default_generic_on_adt() {
182+
check_assist(
183+
extract_to_default_generic,
184+
r#"struct Foo($0i32$0);"#,
185+
r#"struct Foo<T$0 = i32>(T);"#,
186+
);
187+
188+
check_assist(
189+
extract_to_default_generic,
190+
r#"struct Foo<T>(T, $0i32$0);"#,
191+
r#"struct Foo<T, T1$0 = i32>(T, T1);"#,
192+
);
193+
194+
check_assist(
195+
extract_to_default_generic,
196+
r#"enum Foo { A($0i32$0), B, C(i64) };"#,
197+
r#"enum Foo<T$0 = i32> { A(T), B, C(i64) };"#,
198+
);
199+
}
200+
201+
#[test]
202+
fn test_extract_to_generic_on_fn() {
203+
check_assist(
204+
extract_to_default_generic,
205+
r#"fn foo(x: $0i32$0) {}"#,
206+
r#"fn foo<T$0>(x: T) {}"#,
207+
);
208+
209+
check_assist(
210+
extract_to_default_generic,
211+
r#"fn foo(x: [i32; $02$0]) {}"#,
212+
r#"fn foo<const N: usize>(x: [i32; N]) {}"#,
213+
);
214+
}
215+
216+
#[test]
217+
fn test_extract_to_default_generic_const() {
218+
check_assist(
219+
extract_to_default_generic,
220+
r#"type A = [i32; $08$0];"#,
221+
r#"type A<const N: usize = 8> = [i32; N];"#,
222+
);
223+
224+
check_assist(
225+
extract_to_default_generic,
226+
r#"type A<T> = [T; $08$0];"#,
227+
r#"type A<T, const N: usize = 8> = [T; N];"#,
228+
);
229+
}
230+
231+
#[test]
232+
fn test_extract_to_default_generic_const_non_array() {
233+
check_assist(
234+
extract_to_default_generic,
235+
r#"
236+
struct Foo<const N: usize>([(); N]);
237+
type A = Foo<$08$0>;
238+
"#,
239+
r#"
240+
struct Foo<const N: usize>([(); N]);
241+
type A<const N: ${0:_} = 8> = Foo<N>;
242+
"#,
243+
);
244+
}
245+
}

crates/ide-assists/src/lib.rs

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -146,6 +146,7 @@ mod handlers {
146146
mod extract_function;
147147
mod extract_module;
148148
mod extract_struct_from_enum_variant;
149+
mod extract_to_default_generic;
149150
mod extract_type_alias;
150151
mod extract_variable;
151152
mod fix_visibility;
@@ -281,6 +282,7 @@ mod handlers {
281282
extract_expressions_from_format_string::extract_expressions_from_format_string,
282283
extract_struct_from_enum_variant::extract_struct_from_enum_variant,
283284
extract_type_alias::extract_type_alias,
285+
extract_to_default_generic::extract_to_default_generic,
284286
fix_visibility::fix_visibility,
285287
flip_binexpr::flip_binexpr,
286288
flip_comma::flip_comma,

crates/ide-assists/src/tests/generated.rs

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1187,6 +1187,19 @@ enum A { One(One) }
11871187
)
11881188
}
11891189

1190+
#[test]
1191+
fn doctest_extract_to_default_generic() {
1192+
check_doc_test(
1193+
"extract_to_default_generic",
1194+
r#####"
1195+
struct Foo(u32, $0String$0);
1196+
"#####,
1197+
r#####"
1198+
struct Foo<T$0 = String>(u32, T);
1199+
"#####,
1200+
)
1201+
}
1202+
11901203
#[test]
11911204
fn doctest_extract_type_alias() {
11921205
check_doc_test(

crates/syntax/src/ast/make.rs

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1085,10 +1085,27 @@ pub fn type_param(name: ast::Name, bounds: Option<ast::TypeBoundList>) -> ast::T
10851085
ast_from_text(&format!("fn f<{name}{bounds}>() {{ }}"))
10861086
}
10871087

1088+
pub fn type_default_param(
1089+
name: ast::Name,
1090+
bounds: Option<ast::TypeBoundList>,
1091+
default: ast::Type,
1092+
) -> ast::TypeParam {
1093+
let bounds = bounds.map_or_else(String::new, |it| format!(": {it}"));
1094+
ast_from_text(&format!("fn f<{name}{bounds} = {default}>() {{ }}"))
1095+
}
1096+
10881097
pub fn const_param(name: ast::Name, ty: ast::Type) -> ast::ConstParam {
10891098
ast_from_text(&format!("fn f<const {name}: {ty}>() {{ }}"))
10901099
}
10911100

1101+
pub fn const_default_param(
1102+
name: ast::Name,
1103+
ty: ast::Type,
1104+
default: ast::ConstArg,
1105+
) -> ast::ConstParam {
1106+
ast_from_text(&format!("fn f<const {name}: {ty} = {default}>() {{ }}"))
1107+
}
1108+
10921109
pub fn lifetime_param(lifetime: ast::Lifetime) -> ast::LifetimeParam {
10931110
ast_from_text(&format!("fn f<{lifetime}>() {{ }}"))
10941111
}

crates/syntax/src/ast/traits.rs

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -161,3 +161,4 @@ impl Iterator for AttrDocCommentIter {
161161
}
162162

163163
impl<A: HasName, B: HasName> HasName for Either<A, B> {}
164+
impl<A: HasGenericParams, B: HasGenericParams> HasGenericParams for Either<A, B> {}

0 commit comments

Comments
 (0)