src/tools/rust-analyzer/crates/ide-assists/src/handlers/generate_deref.rs

   1 use std::fmt::Display;
   2
   3 use hir::{ModPath, ModuleDef};
   4 use ide_db::{famous_defs::FamousDefs, RootDatabase};
   5 use syntax::{
   6     ast::{self, HasName},
   7     AstNode, SyntaxNode,
   8 };
   9
  10 use crate::{
  11     assist_context::{AssistContext, Assists, SourceChangeBuilder},
  12     utils::generate_trait_impl_text,
  13     AssistId, AssistKind,
  14 };
  15
  16 // Assist: generate_deref
  17 //
  18 // Generate `Deref` impl using the given struct field.
  19 //
  20 // ```
  21 // # //- minicore: deref, deref_mut
  22 // struct A;
  23 // struct B {
  24 //    $0a: A
  25 // }
  26 // ```
  27 // ->
  28 // ```
  29 // struct A;
  30 // struct B {
  31 //    a: A
  32 // }
  33 //
  34 // impl core::ops::Deref for B {
  35 //     type Target = A;
  36 //
  37 //     fn deref(&self) -> &Self::Target {
  38 //         &self.a
  39 //     }
  40 // }
  41 // ```
  42 pub(crate) fn generate_deref(acc: &mut Assists, ctx: &AssistContext<'_>) -> Option<()> {
  43     generate_record_deref(acc, ctx).or_else(|| generate_tuple_deref(acc, ctx))
  44 }
  45
  46 fn generate_record_deref(acc: &mut Assists, ctx: &AssistContext<'_>) -> Option<()> {
  47     let strukt = ctx.find_node_at_offset::<ast::Struct>()?;
  48     let field = ctx.find_node_at_offset::<ast::RecordField>()?;
  49
  50     let deref_type_to_generate = match existing_deref_impl(&ctx.sema, &strukt) {
  51         None => DerefType::Deref,
  52         Some(DerefType::Deref) => DerefType::DerefMut,
  53         Some(DerefType::DerefMut) => {
  54             cov_mark::hit!(test_add_record_deref_impl_already_exists);
  55             return None;
  56         }
  57     };
  58
  59     let module = ctx.sema.to_def(&strukt)?.module(ctx.db());
  60     let trait_ = deref_type_to_generate.to_trait(&ctx.sema, module.krate())?;
  61     let trait_path =
  62         module.find_use_path(ctx.db(), ModuleDef::Trait(trait_), ctx.config.prefer_no_std)?;
  63
  64     let field_type = field.ty()?;
  65     let field_name = field.name()?;
  66     let target = field.syntax().text_range();
  67     acc.add(
  68         AssistId("generate_deref", AssistKind::Generate),
  69         format!("Generate `{deref_type_to_generate:?}` impl using `{field_name}`"),
  70         target,
  71         |edit| {
  72             generate_edit(
  73                 edit,
  74                 strukt,
  75                 field_type.syntax(),
  76                 field_name.syntax(),
  77                 deref_type_to_generate,
  78                 trait_path,
  79             )
  80         },
  81     )
  82 }
  83
  84 fn generate_tuple_deref(acc: &mut Assists, ctx: &AssistContext<'_>) -> Option<()> {
  85     let strukt = ctx.find_node_at_offset::<ast::Struct>()?;
  86     let field = ctx.find_node_at_offset::<ast::TupleField>()?;
  87     let field_list = ctx.find_node_at_offset::<ast::TupleFieldList>()?;
  88     let field_list_index =
  89         field_list.syntax().children().into_iter().position(|s| &s == field.syntax())?;
  90
  91     let deref_type_to_generate = match existing_deref_impl(&ctx.sema, &strukt) {
  92         None => DerefType::Deref,
  93         Some(DerefType::Deref) => DerefType::DerefMut,
  94         Some(DerefType::DerefMut) => {
  95             cov_mark::hit!(test_add_field_deref_impl_already_exists);
  96             return None;
  97         }
  98     };
  99
 100     let module = ctx.sema.to_def(&strukt)?.module(ctx.db());
 101     let trait_ = deref_type_to_generate.to_trait(&ctx.sema, module.krate())?;
 102     let trait_path =
 103         module.find_use_path(ctx.db(), ModuleDef::Trait(trait_), ctx.config.prefer_no_std)?;
 104
 105     let field_type = field.ty()?;
 106     let target = field.syntax().text_range();
 107     acc.add(
 108         AssistId("generate_deref", AssistKind::Generate),
 109         format!("Generate `{deref_type_to_generate:?}` impl using `{field}`"),
 110         target,
 111         |edit| {
 112             generate_edit(
 113                 edit,
 114                 strukt,
 115                 field_type.syntax(),
 116                 field_list_index,
 117                 deref_type_to_generate,
 118                 trait_path,
 119             )
 120         },
 121     )
 122 }
 123
 124 fn generate_edit(
 125     edit: &mut SourceChangeBuilder,
 126     strukt: ast::Struct,
 127     field_type_syntax: &SyntaxNode,
 128     field_name: impl Display,
 129     deref_type: DerefType,
 130     trait_path: ModPath,
 131 ) {
 132     let start_offset = strukt.syntax().text_range().end();
 133     let impl_code = match deref_type {
 134         DerefType::Deref => format!(
 135             r#"    type Target = {field_type_syntax};
 136
 137     fn deref(&self) -> &Self::Target {{
 138         &self.{field_name}
 139     }}"#,
 140         ),
 141         DerefType::DerefMut => format!(
 142             r#"    fn deref_mut(&mut self) -> &mut Self::Target {{
 143         &mut self.{field_name}
 144     }}"#,
 145         ),
 146     };
 147     let strukt_adt = ast::Adt::Struct(strukt);
 148     let deref_impl = generate_trait_impl_text(&strukt_adt, &trait_path.to_string(), &impl_code);
 149     edit.insert(start_offset, deref_impl);
 150 }
 151
 152 fn existing_deref_impl(
 153     sema: &hir::Semantics<'_, RootDatabase>,
 154     strukt: &ast::Struct,
 155 ) -> Option<DerefType> {
 156     let strukt = sema.to_def(strukt)?;
 157     let krate = strukt.module(sema.db).krate();
 158
 159     let deref_trait = FamousDefs(sema, krate).core_ops_Deref()?;
 160     let deref_mut_trait = FamousDefs(sema, krate).core_ops_DerefMut()?;
 161     let strukt_type = strukt.ty(sema.db);
 162
 163     if strukt_type.impls_trait(sema.db, deref_trait, &[]) {
 164         if strukt_type.impls_trait(sema.db, deref_mut_trait, &[]) {
 165             Some(DerefType::DerefMut)
 166         } else {
 167             Some(DerefType::Deref)
 168         }
 169     } else {
 170         None
 171     }
 172 }
 173
 174 #[derive(Debug)]
 175 enum DerefType {
 176     Deref,
 177     DerefMut,
 178 }
 179
 180 impl DerefType {
 181     fn to_trait(
 182         &self,
 183         sema: &hir::Semantics<'_, RootDatabase>,
 184         krate: hir::Crate,
 185     ) -> Option<hir::Trait> {
 186         match self {
 187             DerefType::Deref => FamousDefs(sema, krate).core_ops_Deref(),
 188             DerefType::DerefMut => FamousDefs(sema, krate).core_ops_DerefMut(),
 189         }
 190     }
 191 }
 192
 193 #[cfg(test)]
 194 mod tests {
 195     use crate::tests::{check_assist, check_assist_not_applicable};
 196
 197     use super::*;
 198
 199     #[test]
 200     fn test_generate_record_deref() {
 201         check_assist(
 202             generate_deref,
 203             r#"
 204 //- minicore: deref
 205 struct A { }
 206 struct B { $0a: A }"#,
 207             r#"
 208 struct A { }
 209 struct B { a: A }
 210
 211 impl core::ops::Deref for B {
 212     type Target = A;
 213
 214     fn deref(&self) -> &Self::Target {
 215         &self.a
 216     }
 217 }"#,
 218         );
 219     }
 220
 221     #[test]
 222     fn test_generate_record_deref_short_path() {
 223         check_assist(
 224             generate_deref,
 225             r#"
 226 //- minicore: deref
 227 use core::ops::Deref;
 228 struct A { }
 229 struct B { $0a: A }"#,
 230             r#"
 231 use core::ops::Deref;
 232 struct A { }
 233 struct B { a: A }
 234
 235 impl Deref for B {
 236     type Target = A;
 237
 238     fn deref(&self) -> &Self::Target {
 239         &self.a
 240     }
 241 }"#,
 242         );
 243     }
 244
 245     #[test]
 246     fn test_generate_field_deref_idx_0() {
 247         check_assist(
 248             generate_deref,
 249             r#"
 250 //- minicore: deref
 251 struct A { }
 252 struct B($0A);"#,
 253             r#"
 254 struct A { }
 255 struct B(A);
 256
 257 impl core::ops::Deref for B {
 258     type Target = A;
 259
 260     fn deref(&self) -> &Self::Target {
 261         &self.0
 262     }
 263 }"#,
 264         );
 265     }
 266     #[test]
 267     fn test_generate_field_deref_idx_1() {
 268         check_assist(
 269             generate_deref,
 270             r#"
 271 //- minicore: deref
 272 struct A { }
 273 struct B(u8, $0A);"#,
 274             r#"
 275 struct A { }
 276 struct B(u8, A);
 277
 278 impl core::ops::Deref for B {
 279     type Target = A;
 280
 281     fn deref(&self) -> &Self::Target {
 282         &self.1
 283     }
 284 }"#,
 285         );
 286     }
 287
 288     #[test]
 289     fn test_generates_derefmut_when_deref_present() {
 290         check_assist(
 291             generate_deref,
 292             r#"
 293 //- minicore: deref, deref_mut
 294 struct B { $0a: u8 }
 295
 296 impl core::ops::Deref for B {}
 297 "#,
 298             r#"
 299 struct B { a: u8 }
 300
 301 impl core::ops::DerefMut for B {
 302     fn deref_mut(&mut self) -> &mut Self::Target {
 303         &mut self.a
 304     }
 305 }
 306
 307 impl core::ops::Deref for B {}
 308 "#,
 309         );
 310     }
 311
 312     #[test]
 313     fn test_generate_record_deref_not_applicable_if_already_impl() {
 314         cov_mark::check!(test_add_record_deref_impl_already_exists);
 315         check_assist_not_applicable(
 316             generate_deref,
 317             r#"
 318 //- minicore: deref, deref_mut
 319 struct A { }
 320 struct B { $0a: A }
 321
 322 impl core::ops::Deref for B {}
 323 impl core::ops::DerefMut for B {}
 324 "#,
 325         )
 326     }
 327
 328     #[test]
 329     fn test_generate_field_deref_not_applicable_if_already_impl() {
 330         cov_mark::check!(test_add_field_deref_impl_already_exists);
 331         check_assist_not_applicable(
 332             generate_deref,
 333             r#"
 334 //- minicore: deref, deref_mut
 335 struct A { }
 336 struct B($0A)
 337
 338 impl core::ops::Deref for B {}
 339 impl core::ops::DerefMut for B {}
 340 "#,
 341         )
 342     }
 343 }