|
1 | 1 | use ide_db::famous_defs::FamousDefs; |
2 | 2 | use stdx::{format_to, to_lower_snake_case}; |
3 | | -use syntax::ast::{self, AstNode, HasName, HasVisibility}; |
| 3 | +use syntax::{ |
| 4 | + ast::{self, AstNode, HasName, HasVisibility}, |
| 5 | + SyntaxNode, TextRange, |
| 6 | +}; |
4 | 7 |
|
5 | 8 | use crate::{ |
6 | 9 | utils::{convert_reference_type, find_impl_block_end, find_struct_impl, generate_impl_text}, |
@@ -72,92 +75,229 @@ pub(crate) fn generate_getter_mut(acc: &mut Assists, ctx: &AssistContext<'_>) -> |
72 | 75 | generate_getter_impl(acc, ctx, true) |
73 | 76 | } |
74 | 77 |
|
| 78 | +#[derive(Clone, Debug)] |
| 79 | +struct RecordFieldInfo { |
| 80 | + field_name: syntax::ast::Name, |
| 81 | + field_ty: syntax::ast::Type, |
| 82 | + fn_name: String, |
| 83 | + target: TextRange, |
| 84 | +} |
| 85 | + |
| 86 | +struct GetterInfo { |
| 87 | + impl_def: Option<ast::Impl>, |
| 88 | + strukt: ast::Struct, |
| 89 | + mutable: bool, |
| 90 | +} |
| 91 | + |
75 | 92 | pub(crate) fn generate_getter_impl( |
76 | 93 | acc: &mut Assists, |
77 | 94 | ctx: &AssistContext<'_>, |
78 | 95 | mutable: bool, |
79 | 96 | ) -> Option<()> { |
80 | | - let strukt = ctx.find_node_at_offset::<ast::Struct>()?; |
81 | | - let field = ctx.find_node_at_offset::<ast::RecordField>()?; |
| 97 | + let (strukt, info_of_record_fields, fn_names) = if !ctx.has_empty_selection() { |
| 98 | + let node = ctx.covering_element(); |
82 | 99 |
|
83 | | - let field_name = field.name()?; |
84 | | - let field_ty = field.ty()?; |
| 100 | + let node = match node { |
| 101 | + syntax::NodeOrToken::Node(n) => n, |
| 102 | + syntax::NodeOrToken::Token(t) => t.parent()?, |
| 103 | + }; |
85 | 104 |
|
86 | | - // Return early if we've found an existing fn |
87 | | - let mut fn_name = to_lower_snake_case(&field_name.to_string()); |
88 | | - if mutable { |
89 | | - format_to!(fn_name, "_mut"); |
| 105 | + let mut parent_struct: Option<ast::Struct> = None; |
| 106 | + if let Some(struct_node) = node.ancestors().find(|it| ast::Struct::can_cast(it.kind())) { |
| 107 | + parent_struct = ast::Struct::cast(struct_node); |
| 108 | + } |
| 109 | + |
| 110 | + // If we don't have a parent struct, something is wrong, OPERATION ABORT 🚨 |
| 111 | + if parent_struct.is_none() { |
| 112 | + return None; |
| 113 | + } |
| 114 | + |
| 115 | + let (info_of_record_fields, field_names) = |
| 116 | + extract_and_parse_record_fields(&node, ctx.selection_trimmed(), mutable)?; |
| 117 | + |
| 118 | + (parent_struct.unwrap(), info_of_record_fields, field_names) |
| 119 | + } else { |
| 120 | + let strukt = ctx.find_node_at_offset::<ast::Struct>()?; |
| 121 | + let field = ctx.find_node_at_offset::<ast::RecordField>()?; |
| 122 | + |
| 123 | + let record_field_info = parse_record_field(field, mutable)?; |
| 124 | + |
| 125 | + let fn_name = record_field_info.fn_name.clone(); |
| 126 | + |
| 127 | + (strukt, vec![record_field_info], vec![fn_name]) |
| 128 | + }; |
| 129 | + |
| 130 | + // No record fields to do work on :( |
| 131 | + if info_of_record_fields.len() == 0 { |
| 132 | + return None; |
90 | 133 | } |
91 | | - let impl_def = find_struct_impl(ctx, &ast::Adt::Struct(strukt.clone()), fn_name.as_str())?; |
| 134 | + |
| 135 | + let impl_def = find_struct_impl(ctx, &ast::Adt::Struct(strukt.clone()), fn_names)?; |
92 | 136 |
|
93 | 137 | let (id, label) = if mutable { |
94 | 138 | ("generate_getter_mut", "Generate a mut getter method") |
95 | 139 | } else { |
96 | 140 | ("generate_getter", "Generate a getter method") |
97 | 141 | }; |
98 | | - let target = field.syntax().text_range(); |
| 142 | + |
| 143 | + let mut target: TextRange = info_of_record_fields[0].target; |
| 144 | + |
| 145 | + // Computing collective text range of all record fields in selected region |
| 146 | + info_of_record_fields |
| 147 | + .clone() |
| 148 | + .iter() |
| 149 | + .for_each(|record_field_info| target = target.cover(record_field_info.target)); |
| 150 | + |
| 151 | + let getter_info = GetterInfo { impl_def, strukt, mutable }; |
| 152 | + |
99 | 153 | acc.add_group( |
100 | 154 | &GroupLabel("Generate getter/setter".to_owned()), |
101 | 155 | AssistId(id, AssistKind::Generate), |
102 | 156 | label, |
103 | 157 | target, |
104 | 158 | |builder| { |
| 159 | + let record_fields_count = info_of_record_fields.len(); |
| 160 | + |
105 | 161 | let mut buf = String::with_capacity(512); |
106 | 162 |
|
107 | | - if impl_def.is_some() { |
108 | | - buf.push('\n'); |
| 163 | + // Check if an impl exists |
| 164 | + if let Some(ref impl_def) = getter_info.impl_def { |
| 165 | + // Check if impl is empty |
| 166 | + if let Some(assoc_item_list) = impl_def.assoc_item_list() { |
| 167 | + if assoc_item_list.assoc_items().next().is_some() { |
| 168 | + // If not then only insert a new line |
| 169 | + buf.push('\n'); |
| 170 | + } |
| 171 | + } |
109 | 172 | } |
110 | 173 |
|
111 | | - let vis = strukt.visibility().map_or(String::new(), |v| format!("{} ", v)); |
112 | | - let (ty, body) = if mutable { |
113 | | - (format!("&mut {}", field_ty), format!("&mut self.{}", field_name)) |
114 | | - } else { |
115 | | - (|| { |
116 | | - let krate = ctx.sema.scope(field_ty.syntax())?.krate(); |
117 | | - let famous_defs = &FamousDefs(&ctx.sema, krate); |
118 | | - ctx.sema |
119 | | - .resolve_type(&field_ty) |
120 | | - .and_then(|ty| convert_reference_type(ty, ctx.db(), famous_defs)) |
121 | | - .map(|conversion| { |
122 | | - cov_mark::hit!(convert_reference_type); |
123 | | - ( |
124 | | - conversion.convert_type(ctx.db()), |
125 | | - conversion.getter(field_name.to_string()), |
126 | | - ) |
127 | | - }) |
128 | | - })() |
129 | | - .unwrap_or_else(|| (format!("&{}", field_ty), format!("&self.{}", field_name))) |
130 | | - }; |
131 | | - |
132 | | - format_to!( |
133 | | - buf, |
134 | | - " {}fn {}(&{}self) -> {} {{ |
135 | | - {} |
136 | | - }}", |
137 | | - vis, |
138 | | - fn_name, |
139 | | - mutable.then(|| "mut ").unwrap_or_default(), |
140 | | - ty, |
141 | | - body, |
142 | | - ); |
143 | | - |
144 | | - let start_offset = impl_def |
145 | | - .and_then(|impl_def| find_impl_block_end(impl_def, &mut buf)) |
| 174 | + for i in 0..record_fields_count { |
| 175 | + let mut getter_buf = |
| 176 | + generate_getter_from_record_info(ctx, &getter_info, &info_of_record_fields[i]); |
| 177 | + |
| 178 | + if i == record_fields_count - 1 { |
| 179 | + getter_buf = getter_buf.replacen("fn ", "fn $0", 1); |
| 180 | + } |
| 181 | + |
| 182 | + if i == 0 { |
| 183 | + buf = buf + &getter_buf; |
| 184 | + } else { |
| 185 | + buf = buf + "\n" + &getter_buf; |
| 186 | + } |
| 187 | + |
| 188 | + if i < record_fields_count - 1 { |
| 189 | + buf = buf + "\n"; |
| 190 | + } |
| 191 | + } |
| 192 | + |
| 193 | + let start_offset = getter_info |
| 194 | + .impl_def |
| 195 | + .as_ref() |
| 196 | + .and_then(|impl_def| find_impl_block_end(impl_def.to_owned(), &mut buf)) |
146 | 197 | .unwrap_or_else(|| { |
147 | | - buf = generate_impl_text(&ast::Adt::Struct(strukt.clone()), &buf); |
148 | | - strukt.syntax().text_range().end() |
| 198 | + buf = generate_impl_text(&ast::Adt::Struct(getter_info.strukt.clone()), &buf); |
| 199 | + getter_info.strukt.syntax().text_range().end() |
149 | 200 | }); |
150 | 201 |
|
151 | 202 | match ctx.config.snippet_cap { |
152 | | - Some(cap) => { |
153 | | - builder.insert_snippet(cap, start_offset, buf.replacen("fn ", "fn $0", 1)) |
154 | | - } |
| 203 | + Some(cap) => builder.insert_snippet(cap, start_offset, buf), |
155 | 204 | None => builder.insert(start_offset, buf), |
156 | 205 | } |
157 | 206 | }, |
158 | 207 | ) |
159 | 208 | } |
160 | 209 |
|
| 210 | +fn generate_getter_from_record_info( |
| 211 | + ctx: &AssistContext<'_>, |
| 212 | + info: &GetterInfo, |
| 213 | + record_field_info: &RecordFieldInfo, |
| 214 | +) -> String { |
| 215 | + let mut buf = String::with_capacity(512); |
| 216 | + |
| 217 | + let vis = info.strukt.visibility().map_or(String::new(), |v| format!("{} ", v)); |
| 218 | + let (ty, body) = if info.mutable { |
| 219 | + ( |
| 220 | + format!("&mut {}", record_field_info.field_ty), |
| 221 | + format!("&mut self.{}", record_field_info.field_name), |
| 222 | + ) |
| 223 | + } else { |
| 224 | + (|| { |
| 225 | + let krate = ctx.sema.scope(record_field_info.field_ty.syntax())?.krate(); |
| 226 | + let famous_defs = &FamousDefs(&ctx.sema, krate); |
| 227 | + ctx.sema |
| 228 | + .resolve_type(&record_field_info.field_ty) |
| 229 | + .and_then(|ty| convert_reference_type(ty, ctx.db(), famous_defs)) |
| 230 | + .map(|conversion| { |
| 231 | + cov_mark::hit!(convert_reference_type); |
| 232 | + ( |
| 233 | + conversion.convert_type(ctx.db()), |
| 234 | + conversion.getter(record_field_info.field_name.to_string()), |
| 235 | + ) |
| 236 | + }) |
| 237 | + })() |
| 238 | + .unwrap_or_else(|| { |
| 239 | + ( |
| 240 | + format!("&{}", record_field_info.field_ty), |
| 241 | + format!("&self.{}", record_field_info.field_name), |
| 242 | + ) |
| 243 | + }) |
| 244 | + }; |
| 245 | + |
| 246 | + format_to!( |
| 247 | + buf, |
| 248 | + " {}fn {}(&{}self) -> {} {{ |
| 249 | + {} |
| 250 | + }}", |
| 251 | + vis, |
| 252 | + record_field_info.fn_name, |
| 253 | + info.mutable.then(|| "mut ").unwrap_or_default(), |
| 254 | + ty, |
| 255 | + body, |
| 256 | + ); |
| 257 | + |
| 258 | + buf |
| 259 | +} |
| 260 | + |
| 261 | +fn extract_and_parse_record_fields( |
| 262 | + node: &SyntaxNode, |
| 263 | + selection_range: TextRange, |
| 264 | + mutable: bool, |
| 265 | +) -> Option<(Vec<RecordFieldInfo>, Vec<String>)> { |
| 266 | + let mut field_names: Vec<String> = vec![]; |
| 267 | + let info_of_record_fields_in_selection = node |
| 268 | + .children() |
| 269 | + .filter_map(|node| { |
| 270 | + if selection_range.contains_range(node.text_range()) { |
| 271 | + let record_field_info = parse_record_field(ast::RecordField::cast(node)?, mutable)?; |
| 272 | + field_names.push(record_field_info.fn_name.clone()); |
| 273 | + return Some(record_field_info); |
| 274 | + } |
| 275 | + |
| 276 | + None |
| 277 | + }) |
| 278 | + .collect::<Vec<RecordFieldInfo>>(); |
| 279 | + |
| 280 | + if info_of_record_fields_in_selection.len() == 0 { |
| 281 | + return None; |
| 282 | + } |
| 283 | + |
| 284 | + Some((info_of_record_fields_in_selection, field_names)) |
| 285 | +} |
| 286 | + |
| 287 | +fn parse_record_field(record_field: ast::RecordField, mutable: bool) -> Option<RecordFieldInfo> { |
| 288 | + let field_name = record_field.name()?; |
| 289 | + let field_ty = record_field.ty()?; |
| 290 | + |
| 291 | + let mut fn_name = to_lower_snake_case(&field_name.to_string()); |
| 292 | + if mutable { |
| 293 | + format_to!(fn_name, "_mut"); |
| 294 | + } |
| 295 | + |
| 296 | + let target = record_field.syntax().text_range(); |
| 297 | + |
| 298 | + Some(RecordFieldInfo { field_name, field_ty, fn_name, target }) |
| 299 | +} |
| 300 | + |
161 | 301 | #[cfg(test)] |
162 | 302 | mod tests { |
163 | 303 | use crate::tests::{check_assist, check_assist_not_applicable}; |
@@ -489,4 +629,52 @@ impl Context { |
489 | 629 | "#, |
490 | 630 | ); |
491 | 631 | } |
| 632 | + |
| 633 | + #[test] |
| 634 | + fn test_generate_multiple_getters_from_selection() { |
| 635 | + check_assist( |
| 636 | + generate_getter, |
| 637 | + r#" |
| 638 | +struct Context { |
| 639 | + $0data: Data, |
| 640 | + count: usize,$0 |
| 641 | +} |
| 642 | + "#, |
| 643 | + r#" |
| 644 | +struct Context { |
| 645 | + data: Data, |
| 646 | + count: usize, |
| 647 | +} |
| 648 | +
|
| 649 | +impl Context { |
| 650 | + fn data(&self) -> &Data { |
| 651 | + &self.data |
| 652 | + } |
| 653 | +
|
| 654 | + fn $0count(&self) -> &usize { |
| 655 | + &self.count |
| 656 | + } |
| 657 | +} |
| 658 | + "#, |
| 659 | + ); |
| 660 | + } |
| 661 | + |
| 662 | + #[test] |
| 663 | + fn test_generate_multiple_getters_from_selection_one_already_exists() { |
| 664 | + check_assist_not_applicable( |
| 665 | + generate_getter, |
| 666 | + r#" |
| 667 | +struct Context { |
| 668 | + $0data: Data, |
| 669 | + count: usize,$0 |
| 670 | +} |
| 671 | +
|
| 672 | +impl Context { |
| 673 | + fn data(&self) -> &Data { |
| 674 | + &self.data |
| 675 | + } |
| 676 | +} |
| 677 | + "#, |
| 678 | + ); |
| 679 | + } |
492 | 680 | } |
0 commit comments