|
| 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 | +} |
0 commit comments