crates/ide-assists/src/handlers/generate_deref.rs

use std::fmt::Display;

use hir::{ModPath, ModuleDef};
use ide_db::{RootDatabase, famous_defs::FamousDefs};
use syntax::{
    AstNode, Edition, SyntaxNode,
    ast::{self, HasName},
};

use crate::{
    AssistId,
    assist_context::{AssistContext, Assists, SourceChangeBuilder},
    utils::generate_trait_impl_text_intransitive,
};

// Assist: generate_deref
//
// Generate `Deref` impl using the given struct field.
//
// ```
// # //- minicore: deref, deref_mut
// struct A;
// struct B {
//    $0a: A
// }
// ```
// ->
// ```
// struct A;
// struct B {
//    a: A
// }
//
// impl core::ops::Deref for B {
//     type Target = A;
//
//     fn deref(&self) -> &Self::Target {
//         &self.a
//     }
// }
// ```
pub(crate) fn generate_deref(acc: &mut Assists, ctx: &AssistContext<'_>) -> Option<()> {
    generate_record_deref(acc, ctx).or_else(|| generate_tuple_deref(acc, ctx))
}

fn generate_record_deref(acc: &mut Assists, ctx: &AssistContext<'_>) -> Option<()> {
    let strukt = ctx.find_node_at_offset::<ast::Struct>()?;
    let field = ctx.find_node_at_offset::<ast::RecordField>()?;

    let deref_type_to_generate = match existing_deref_impl(&ctx.sema, &strukt) {
        None => DerefType::Deref,
        Some(DerefType::Deref) => DerefType::DerefMut,
        Some(DerefType::DerefMut) => {
            cov_mark::hit!(test_add_record_deref_impl_already_exists);
            return None;
        }
    };

    let module = ctx.sema.to_def(&strukt)?.module(ctx.db());
    let trait_ = deref_type_to_generate.to_trait(&ctx.sema, module.krate())?;
    let trait_path =
        module.find_path(ctx.db(), ModuleDef::Trait(trait_), ctx.config.import_path_config())?;

    let field_type = field.ty()?;
    let field_name = field.name()?;
    let target = field.syntax().text_range();
    acc.add(
        AssistId::generate("generate_deref"),
        format!("Generate `{deref_type_to_generate:?}` impl using `{field_name}`"),
        target,
        |edit| {
            generate_edit(
                ctx.db(),
                edit,
                strukt,
                field_type.syntax(),
                field_name.syntax(),
                deref_type_to_generate,
                trait_path,
                module.krate().edition(ctx.db()),
            )
        },
    )
}

fn generate_tuple_deref(acc: &mut Assists, ctx: &AssistContext<'_>) -> Option<()> {
    let strukt = ctx.find_node_at_offset::<ast::Struct>()?;
    let field = ctx.find_node_at_offset::<ast::TupleField>()?;
    let field_list = ctx.find_node_at_offset::<ast::TupleFieldList>()?;
    let field_list_index = field_list.syntax().children().position(|s| &s == field.syntax())?;

    let deref_type_to_generate = match existing_deref_impl(&ctx.sema, &strukt) {
        None => DerefType::Deref,
        Some(DerefType::Deref) => DerefType::DerefMut,
        Some(DerefType::DerefMut) => {
            cov_mark::hit!(test_add_field_deref_impl_already_exists);
            return None;
        }
    };

    let module = ctx.sema.to_def(&strukt)?.module(ctx.db());
    let trait_ = deref_type_to_generate.to_trait(&ctx.sema, module.krate())?;
    let trait_path =
        module.find_path(ctx.db(), ModuleDef::Trait(trait_), ctx.config.import_path_config())?;

    let field_type = field.ty()?;
    let target = field.syntax().text_range();
    acc.add(
        AssistId::generate("generate_deref"),
        format!("Generate `{deref_type_to_generate:?}` impl using `{field}`"),
        target,
        |edit| {
            generate_edit(
                ctx.db(),
                edit,
                strukt,
                field_type.syntax(),
                field_list_index,
                deref_type_to_generate,
                trait_path,
                module.krate().edition(ctx.db()),
            )
        },
    )
}

fn generate_edit(
    db: &RootDatabase,
    edit: &mut SourceChangeBuilder,
    strukt: ast::Struct,
    field_type_syntax: &SyntaxNode,
    field_name: impl Display,
    deref_type: DerefType,
    trait_path: ModPath,
    edition: Edition,
) {
    let start_offset = strukt.syntax().text_range().end();
    let impl_code = match deref_type {
        DerefType::Deref => format!(
            r#"    type Target = {field_type_syntax};

    fn deref(&self) -> &Self::Target {{
        &self.{field_name}
    }}"#,
        ),
        DerefType::DerefMut => format!(
            r#"    fn deref_mut(&mut self) -> &mut Self::Target {{
        &mut self.{field_name}
    }}"#,
        ),
    };
    let strukt_adt = ast::Adt::Struct(strukt);
    let deref_impl = generate_trait_impl_text_intransitive(
        &strukt_adt,
        &trait_path.display(db, edition).to_string(),
        &impl_code,
    );
    edit.insert(start_offset, deref_impl);
}

fn existing_deref_impl(
    sema: &hir::Semantics<'_, RootDatabase>,
    strukt: &ast::Struct,
) -> Option<DerefType> {
    let strukt = sema.to_def(strukt)?;
    let krate = strukt.module(sema.db).krate();

    let deref_trait = FamousDefs(sema, krate).core_ops_Deref()?;
    let deref_mut_trait = FamousDefs(sema, krate).core_ops_DerefMut()?;
    let strukt_type = strukt.ty(sema.db);

    if strukt_type.impls_trait(sema.db, deref_trait, &[]) {
        if strukt_type.impls_trait(sema.db, deref_mut_trait, &[]) {
            Some(DerefType::DerefMut)
        } else {
            Some(DerefType::Deref)
        }
    } else {
        None
    }
}

#[derive(Debug)]
enum DerefType {
    Deref,
    DerefMut,
}

impl DerefType {
    fn to_trait(
        &self,
        sema: &hir::Semantics<'_, RootDatabase>,
        krate: hir::Crate,
    ) -> Option<hir::Trait> {
        match self {
            DerefType::Deref => FamousDefs(sema, krate).core_ops_Deref(),
            DerefType::DerefMut => FamousDefs(sema, krate).core_ops_DerefMut(),
        }
    }
}

#[cfg(test)]
mod tests {
    use crate::tests::{check_assist, check_assist_not_applicable};

    use super::*;

    #[test]
    fn test_generate_record_deref() {
        check_assist(
            generate_deref,
            r#"
//- minicore: deref
struct A { }
struct B { $0a: A }"#,
            r#"
struct A { }
struct B { a: A }

impl core::ops::Deref for B {
    type Target = A;

    fn deref(&self) -> &Self::Target {
        &self.a
    }
}"#,
        );
    }

    #[test]
    fn test_generate_record_deref_with_generic() {
        check_assist(
            generate_deref,
            r#"
//- minicore: deref
struct A<T>($0T);
"#,
            r#"
struct A<T>(T);

impl<T> core::ops::Deref for A<T> {
    type Target = T;

    fn deref(&self) -> &Self::Target {
        &self.0
    }
}
"#,
        );
    }

    #[test]
    fn test_generate_record_deref_short_path() {
        check_assist(
            generate_deref,
            r#"
//- minicore: deref
use core::ops::Deref;
struct A { }
struct B { $0a: A }"#,
            r#"
use core::ops::Deref;
struct A { }
struct B { a: A }

impl Deref for B {
    type Target = A;

    fn deref(&self) -> &Self::Target {
        &self.a
    }
}"#,
        );
    }

    #[test]
    fn test_generate_field_deref_idx_0() {
        check_assist(
            generate_deref,
            r#"
//- minicore: deref
struct A { }
struct B($0A);"#,
            r#"
struct A { }
struct B(A);

impl core::ops::Deref for B {
    type Target = A;

    fn deref(&self) -> &Self::Target {
        &self.0
    }
}"#,
        );
    }
    #[test]
    fn test_generate_field_deref_idx_1() {
        check_assist(
            generate_deref,
            r#"
//- minicore: deref
struct A { }
struct B(u8, $0A);"#,
            r#"
struct A { }
struct B(u8, A);

impl core::ops::Deref for B {
    type Target = A;

    fn deref(&self) -> &Self::Target {
        &self.1
    }
}"#,
        );
    }

    #[test]
    fn test_generates_derefmut_when_deref_present() {
        check_assist(
            generate_deref,
            r#"
//- minicore: deref, deref_mut
struct B { $0a: u8 }

impl core::ops::Deref for B {}
"#,
            r#"
struct B { a: u8 }

impl core::ops::DerefMut for B {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.a
    }
}

impl core::ops::Deref for B {}
"#,
        );
    }

    #[test]
    fn test_generate_record_deref_not_applicable_if_already_impl() {
        cov_mark::check!(test_add_record_deref_impl_already_exists);
        check_assist_not_applicable(
            generate_deref,
            r#"
//- minicore: deref, deref_mut
struct A { }
struct B { $0a: A }

impl core::ops::Deref for B {}
impl core::ops::DerefMut for B {}
"#,
        )
    }

    #[test]
    fn test_generate_field_deref_not_applicable_if_already_impl() {
        cov_mark::check!(test_add_field_deref_impl_already_exists);
        check_assist_not_applicable(
            generate_deref,
            r#"
//- minicore: deref, deref_mut
struct A { }
struct B($0A)

impl core::ops::Deref for B {}
impl core::ops::DerefMut for B {}
"#,
        )
    }
}