|
| 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 = 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<ast::Adt, Either<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: impl HasName + HasGenericParams, |
| 38 | + ty: Either<ast::Type, ast::ConstArg>, |
| 39 | +) -> Option<()> { |
| 40 | + let name = adt.name()?; |
| 41 | + |
| 42 | + let target = ty.syntax().text_range(); |
| 43 | + acc.add( |
| 44 | + AssistId::refactor_extract("extract_to_default_generic"), |
| 45 | + "Extract type as default generic parameter", |
| 46 | + target, |
| 47 | + |edit| { |
| 48 | + let mut editor = edit.make_editor(adt.syntax()); |
| 49 | + let generic_list = get_or_create_generic_param_list(&name, &adt, &mut editor, edit); |
| 50 | + |
| 51 | + let generic_name = generic_name(&generic_list, ty.is_right()); |
| 52 | + |
| 53 | + editor.replace(ty.syntax(), generic_name.syntax()); |
| 54 | + |
| 55 | + match ty { |
| 56 | + Left(ty) => { |
| 57 | + let param = make::type_default_param(generic_name, None, ty).clone_for_update(); |
| 58 | + generic_list.add_generic_param(param.into()); |
| 59 | + } |
| 60 | + Right(n) => { |
| 61 | + let param = make::const_default_param(generic_name, const_ty(ctx, &n), n) |
| 62 | + .clone_for_update(); |
| 63 | + generic_list.add_generic_param(param.into()); |
| 64 | + |
| 65 | + if let Some(ast::GenericParam::ConstParam(param)) = |
| 66 | + generic_list.generic_params().last() |
| 67 | + && let Some(ast::Type::InferType(ty)) = param.ty() |
| 68 | + && let Some(cap) = ctx.config.snippet_cap |
| 69 | + { |
| 70 | + let annotation = edit.make_placeholder_snippet(cap); |
| 71 | + editor.add_annotation(ty.syntax(), annotation); |
| 72 | + } |
| 73 | + } |
| 74 | + } |
| 75 | + |
| 76 | + edit.add_file_edits(ctx.vfs_file_id(), editor); |
| 77 | + }, |
| 78 | + ) |
| 79 | +} |
| 80 | + |
| 81 | +fn array_index_type(n: &ast::ConstArg) -> Option<ast::Type> { |
| 82 | + let kind = n.syntax().parent()?.kind(); |
| 83 | + |
| 84 | + if ast::ArrayType::can_cast(kind) || ast::ArrayExpr::can_cast(kind) { |
| 85 | + Some(make::ty("usize")) |
| 86 | + } else { |
| 87 | + None |
| 88 | + } |
| 89 | +} |
| 90 | + |
| 91 | +fn generic_name(generic_list: &ast::GenericParamList, is_const_param: bool) -> Name { |
| 92 | + let exist_names = generic_list |
| 93 | + .generic_params() |
| 94 | + .filter_map(|it| match it { |
| 95 | + ast::GenericParam::ConstParam(const_param) => const_param.name(), |
| 96 | + ast::GenericParam::TypeParam(type_param) => type_param.name(), |
| 97 | + ast::GenericParam::LifetimeParam(_) => None, |
| 98 | + }) |
| 99 | + .map(|name| name.to_string()) |
| 100 | + .collect::<Vec<_>>(); |
| 101 | + |
| 102 | + let mut name_gen = NameGenerator::new_with_names(exist_names.iter().map(|name| name.as_str())); |
| 103 | + |
| 104 | + make::name(&if is_const_param { |
| 105 | + name_gen.suggest_name("N") |
| 106 | + } else { |
| 107 | + name_gen.suggest_name("T") |
| 108 | + }) |
| 109 | + .clone_for_update() |
| 110 | +} |
| 111 | + |
| 112 | +fn const_ty(ctx: &AssistContext<'_>, n: &ast::ConstArg) -> ast::Type { |
| 113 | + if let Some(expr) = n.expr() |
| 114 | + && let Some(ty_info) = ctx.sema.type_of_expr(&expr) |
| 115 | + && let Some(builtin) = ty_info.adjusted().as_builtin() |
| 116 | + { |
| 117 | + make::ty(builtin.name().as_str()) |
| 118 | + } else if let Some(array_index_ty) = array_index_type(n) { |
| 119 | + array_index_ty |
| 120 | + } else { |
| 121 | + make::ty_placeholder() |
| 122 | + } |
| 123 | +} |
| 124 | + |
| 125 | +fn get_or_create_generic_param_list( |
| 126 | + name: &ast::Name, |
| 127 | + adt: &impl HasGenericParams, |
| 128 | + editor: &mut SyntaxEditor, |
| 129 | + edit: &mut SourceChangeBuilder, |
| 130 | +) -> ast::GenericParamList { |
| 131 | + if let Some(list) = adt.generic_param_list() { |
| 132 | + edit.make_mut(list) |
| 133 | + } else { |
| 134 | + let generic = make::generic_param_list([]).clone_for_update(); |
| 135 | + editor.insert(Position::after(name.syntax()), generic.syntax()); |
| 136 | + generic |
| 137 | + } |
| 138 | +} |
| 139 | + |
| 140 | +#[cfg(test)] |
| 141 | +mod tests { |
| 142 | + use super::*; |
| 143 | + |
| 144 | + use crate::tests::check_assist; |
| 145 | + |
| 146 | + #[test] |
| 147 | + fn test_extract_to_default_generic() { |
| 148 | + check_assist( |
| 149 | + extract_to_default_generic, |
| 150 | + r#"type X = ($0i32$0, i64);"#, |
| 151 | + r#"type X<T = i32> = (T, i64);"#, |
| 152 | + ); |
| 153 | + |
| 154 | + check_assist( |
| 155 | + extract_to_default_generic, |
| 156 | + r#"type X<T> = ($0i32$0, T);"#, |
| 157 | + r#"type X<T, T1 = i32> = (T1, T);"#, |
| 158 | + ); |
| 159 | + } |
| 160 | + |
| 161 | + #[test] |
| 162 | + fn test_extract_to_default_generic_on_adt() { |
| 163 | + check_assist( |
| 164 | + extract_to_default_generic, |
| 165 | + r#"struct Foo($0i32$0);"#, |
| 166 | + r#"struct Foo<T = i32>(T);"#, |
| 167 | + ); |
| 168 | + |
| 169 | + check_assist( |
| 170 | + extract_to_default_generic, |
| 171 | + r#"struct Foo<T>(T, $0i32$0);"#, |
| 172 | + r#"struct Foo<T, T1 = i32>(T, T1);"#, |
| 173 | + ); |
| 174 | + |
| 175 | + check_assist( |
| 176 | + extract_to_default_generic, |
| 177 | + r#"enum Foo { A($0i32$0), B, C(i64) };"#, |
| 178 | + r#"enum Foo<T = i32> { A(T), B, C(i64) };"#, |
| 179 | + ); |
| 180 | + } |
| 181 | + |
| 182 | + #[test] |
| 183 | + fn test_extract_to_default_generic_on_fn() { |
| 184 | + check_assist( |
| 185 | + extract_to_default_generic, |
| 186 | + r#"fn foo(x: $0i32$0) {}"#, |
| 187 | + r#"fn foo<T = i32>(x: T) {}"#, |
| 188 | + ); |
| 189 | + } |
| 190 | + |
| 191 | + #[test] |
| 192 | + fn test_extract_to_default_generic_const() { |
| 193 | + check_assist( |
| 194 | + extract_to_default_generic, |
| 195 | + r#"type A = [i32; $08$0];"#, |
| 196 | + r#"type A<const N: usize = 8> = [i32; N];"#, |
| 197 | + ); |
| 198 | + |
| 199 | + check_assist( |
| 200 | + extract_to_default_generic, |
| 201 | + r#"type A<T> = [T; $08$0];"#, |
| 202 | + r#"type A<T, const N: usize = 8> = [T; N];"#, |
| 203 | + ); |
| 204 | + } |
| 205 | + |
| 206 | + #[test] |
| 207 | + fn test_extract_to_default_generic_const_non_array() { |
| 208 | + check_assist( |
| 209 | + extract_to_default_generic, |
| 210 | + r#" |
| 211 | + struct Foo<const N: usize>([(); N]); |
| 212 | + type A = Foo<$08$0>; |
| 213 | + "#, |
| 214 | + r#" |
| 215 | + struct Foo<const N: usize>([(); N]); |
| 216 | + type A<const N: ${0:_} = 8> = Foo<N>; |
| 217 | + "#, |
| 218 | + ); |
| 219 | + } |
| 220 | +} |
0 commit comments