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 `{:?}` impl using `{}`", deref_type_to_generate, 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 `{:?}` impl using `{}`", deref_type_to_generate, field.syntax()),
 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 = {0};
 136
 137     fn deref(&self) -> &Self::Target {{
 138         &self.{1}
 139     }}"#,
 140             field_type_syntax, field_name
 141         ),
 142         DerefType::DerefMut => format!(
 143             r#"    fn deref_mut(&mut self) -> &mut Self::Target {{
 144         &mut self.{}
 145     }}"#,
 146             field_name
 147         ),
 148     };
 149     let strukt_adt = ast::Adt::Struct(strukt);
 150     let deref_impl = generate_trait_impl_text(&strukt_adt, &trait_path.to_string(), &impl_code);
 151     edit.insert(start_offset, deref_impl);
 152 }
 153
 154 fn existing_deref_impl(
 155     sema: &hir::Semantics<'_, RootDatabase>,
 156     strukt: &ast::Struct,
 157 ) -> Option<DerefType> {
 158     let strukt = sema.to_def(strukt)?;
 159     let krate = strukt.module(sema.db).krate();
 160
 161     let deref_trait = FamousDefs(sema, krate).core_ops_Deref()?;
 162     let deref_mut_trait = FamousDefs(sema, krate).core_ops_DerefMut()?;
 163     let strukt_type = strukt.ty(sema.db);
 164
 165     if strukt_type.impls_trait(sema.db, deref_trait, &[]) {
 166         if strukt_type.impls_trait(sema.db, deref_mut_trait, &[]) {
 167             Some(DerefType::DerefMut)
 168         } else {
 169             Some(DerefType::Deref)
 170         }
 171     } else {
 172         None
 173     }
 174 }
 175
 176 #[derive(Debug)]
 177 enum DerefType {
 178     Deref,
 179     DerefMut,
 180 }
 181
 182 impl DerefType {
 183     fn to_trait(
 184         &self,
 185         sema: &hir::Semantics<'_, RootDatabase>,
 186         krate: hir::Crate,
 187     ) -> Option<hir::Trait> {
 188         match self {
 189             DerefType::Deref => FamousDefs(sema, krate).core_ops_Deref(),
 190             DerefType::DerefMut => FamousDefs(sema, krate).core_ops_DerefMut(),
 191         }
 192     }
 193 }
 194
 195 #[cfg(test)]
 196 mod tests {
 197     use crate::tests::{check_assist, check_assist_not_applicable};
 198
 199     use super::*;
 200
 201     #[test]
 202     fn test_generate_record_deref() {
 203         check_assist(
 204             generate_deref,
 205             r#"
 206 //- minicore: deref
 207 struct A { }
 208 struct B { $0a: A }"#,
 209             r#"
 210 struct A { }
 211 struct B { a: A }
 212
 213 impl core::ops::Deref for B {
 214     type Target = A;
 215
 216     fn deref(&self) -> &Self::Target {
 217         &self.a
 218     }
 219 }"#,
 220         );
 221     }
 222
 223     #[test]
 224     fn test_generate_record_deref_short_path() {
 225         check_assist(
 226             generate_deref,
 227             r#"
 228 //- minicore: deref
 229 use core::ops::Deref;
 230 struct A { }
 231 struct B { $0a: A }"#,
 232             r#"
 233 use core::ops::Deref;
 234 struct A { }
 235 struct B { a: A }
 236
 237 impl Deref for B {
 238     type Target = A;
 239
 240     fn deref(&self) -> &Self::Target {
 241         &self.a
 242     }
 243 }"#,
 244         );
 245     }
 246
 247     #[test]
 248     fn test_generate_field_deref_idx_0() {
 249         check_assist(
 250             generate_deref,
 251             r#"
 252 //- minicore: deref
 253 struct A { }
 254 struct B($0A);"#,
 255             r#"
 256 struct A { }
 257 struct B(A);
 258
 259 impl core::ops::Deref for B {
 260     type Target = A;
 261
 262     fn deref(&self) -> &Self::Target {
 263         &self.0
 264     }
 265 }"#,
 266         );
 267     }
 268     #[test]
 269     fn test_generate_field_deref_idx_1() {
 270         check_assist(
 271             generate_deref,
 272             r#"
 273 //- minicore: deref
 274 struct A { }
 275 struct B(u8, $0A);"#,
 276             r#"
 277 struct A { }
 278 struct B(u8, A);
 279
 280 impl core::ops::Deref for B {
 281     type Target = A;
 282
 283     fn deref(&self) -> &Self::Target {
 284         &self.1
 285     }
 286 }"#,
 287         );
 288     }
 289
 290     #[test]
 291     fn test_generates_derefmut_when_deref_present() {
 292         check_assist(
 293             generate_deref,
 294             r#"
 295 //- minicore: deref, deref_mut
 296 struct B { $0a: u8 }
 297
 298 impl core::ops::Deref for B {}
 299 "#,
 300             r#"
 301 struct B { a: u8 }
 302
 303 impl core::ops::DerefMut for B {
 304     fn deref_mut(&mut self) -> &mut Self::Target {
 305         &mut self.a
 306     }
 307 }
 308
 309 impl core::ops::Deref for B {}
 310 "#,
 311         );
 312     }
 313
 314     #[test]
 315     fn test_generate_record_deref_not_applicable_if_already_impl() {
 316         cov_mark::check!(test_add_record_deref_impl_already_exists);
 317         check_assist_not_applicable(
 318             generate_deref,
 319             r#"
 320 //- minicore: deref, deref_mut
 321 struct A { }
 322 struct B { $0a: A }
 323
 324 impl core::ops::Deref for B {}
 325 impl core::ops::DerefMut for B {}
 326 "#,
 327         )
 328     }
 329
 330     #[test]
 331     fn test_generate_field_deref_not_applicable_if_already_impl() {
 332         cov_mark::check!(test_add_field_deref_impl_already_exists);
 333         check_assist_not_applicable(
 334             generate_deref,
 335             r#"
 336 //- minicore: deref, deref_mut
 337 struct A { }
 338 struct B($0A)
 339
 340 impl core::ops::Deref for B {}
 341 impl core::ops::DerefMut for B {}
 342 "#,
 343         )
 344     }
 345 }