rust-analyzer/crates/ra_assists/src/add_missing_impl_members.rs

use std::fmt::Write;

use crate::{Assist, AssistId, AssistCtx};

use hir::db::HirDatabase;
use ra_syntax::{SmolStr, SyntaxKind, TextRange, TextUnit, TreeArc};
use ra_syntax::ast::{self, AstNode, AstToken, FnDef, ImplItem, ImplItemKind, NameOwner};
use ra_db::FilePosition;
use ra_fmt::{leading_indent, reindent};

use itertools::Itertools;

enum AddMissingImplMembersMode {
    DefaultMethodsOnly,
    NoDefaultMethods,
}

pub(crate) fn add_missing_impl_members(ctx: AssistCtx<impl HirDatabase>) -> Option<Assist> {
    add_missing_impl_members_inner(
        ctx,
        AddMissingImplMembersMode::NoDefaultMethods,
        "add_impl_missing_members",
        "add missing impl members",
    )
}

pub(crate) fn add_missing_default_members(ctx: AssistCtx<impl HirDatabase>) -> Option<Assist> {
    add_missing_impl_members_inner(
        ctx,
        AddMissingImplMembersMode::DefaultMethodsOnly,
        "add_impl_default_members",
        "add impl default members",
    )
}

fn add_missing_impl_members_inner(
    mut ctx: AssistCtx<impl HirDatabase>,
    mode: AddMissingImplMembersMode,
    assist_id: &'static str,
    label: &'static str,
) -> Option<Assist> {
    let impl_node = ctx.node_at_offset::<ast::ImplBlock>()?;
    let impl_item_list = impl_node.item_list()?;

    let trait_def = {
        let file_id = ctx.frange.file_id;
        let position = FilePosition { file_id, offset: impl_node.syntax().range().start() };
        let analyser = hir::SourceAnalyser::new(ctx.db, position.file_id, impl_node.syntax());

        resolve_target_trait_def(ctx.db, &analyser, impl_node)?
    };

    let missing_fns: Vec<_> = {
        let fn_def_opt = |kind| if let ImplItemKind::FnDef(def) = kind { Some(def) } else { None };
        let def_name = |def| -> Option<&SmolStr> { FnDef::name(def).map(ast::Name::text) };

        let trait_items =
            trait_def.syntax().descendants().find_map(ast::ItemList::cast)?.impl_items();
        let impl_items = impl_item_list.impl_items();

        let trait_fns = trait_items.map(ImplItem::kind).filter_map(fn_def_opt).collect::<Vec<_>>();
        let impl_fns = impl_items.map(ImplItem::kind).filter_map(fn_def_opt).collect::<Vec<_>>();

        trait_fns
            .into_iter()
            .filter(|t| def_name(t).is_some())
            .filter(|t| match mode {
                AddMissingImplMembersMode::DefaultMethodsOnly => t.body().is_some(),
                AddMissingImplMembersMode::NoDefaultMethods => t.body().is_none(),
            })
            .filter(|t| impl_fns.iter().all(|i| def_name(i) != def_name(t)))
            .collect()
    };
    if missing_fns.is_empty() {
        return None;
    }

    ctx.add_action(AssistId(assist_id), label, |edit| {
        let (parent_indent, indent) = {
            // FIXME: Find a way to get the indent already used in the file.
            // Now, we copy the indent of first item or indent with 4 spaces relative to impl block
            const DEFAULT_INDENT: &str = "    ";
            let first_item = impl_item_list.impl_items().next();
            let first_item_indent =
                first_item.and_then(|i| leading_indent(i.syntax())).map(ToOwned::to_owned);
            let impl_block_indent = leading_indent(impl_node.syntax()).unwrap_or_default();

            (
                impl_block_indent.to_owned(),
                first_item_indent.unwrap_or_else(|| impl_block_indent.to_owned() + DEFAULT_INDENT),
            )
        };

        let changed_range = {
            let children = impl_item_list.syntax().children_with_tokens();
            let last_whitespace =
                children.filter_map(|it| ast::Whitespace::cast(it.as_token()?)).last();

            last_whitespace.map(|w| w.syntax().range()).unwrap_or_else(|| {
                let in_brackets = impl_item_list.syntax().range().end() - TextUnit::of_str("}");
                TextRange::from_to(in_brackets, in_brackets)
            })
        };

        let func_bodies = format!("\n{}", missing_fns.into_iter().map(build_func_body).join("\n"));
        let trailing_whitespace = format!("\n{}", parent_indent);
        let func_bodies = reindent(&func_bodies, &indent) + &trailing_whitespace;

        let replaced_text_range = TextUnit::of_str(&func_bodies);

        edit.replace(changed_range, func_bodies);
        edit.set_cursor(
            changed_range.start() + replaced_text_range - TextUnit::of_str(&trailing_whitespace),
        );
    });

    ctx.build()
}

/// Given an `ast::ImplBlock`, resolves the target trait (the one being
/// implemented) to a `ast::TraitDef`.
fn resolve_target_trait_def(
    db: &impl HirDatabase,
    binder: &hir::SourceAnalyser,
    impl_block: &ast::ImplBlock,
) -> Option<TreeArc<ast::TraitDef>> {
    let ast_path =
        impl_block.target_trait().map(AstNode::syntax).and_then(ast::PathType::cast)?.path()?;

    match binder.resolve_path(db, &ast_path) {
        Some(hir::PathResolution::Def(hir::ModuleDef::Trait(def))) => Some(def.source(db).1),
        _ => None,
    }
}

fn build_func_body(def: &ast::FnDef) -> String {
    let mut buf = String::new();

    for child in def.syntax().children_with_tokens() {
        match (child.prev_sibling_or_token().map(|c| c.kind()), child.kind()) {
            (_, SyntaxKind::SEMI) => buf.push_str(" { unimplemented!() }"),
            (_, SyntaxKind::ATTR) | (_, SyntaxKind::COMMENT) => {}
            (Some(SyntaxKind::ATTR), SyntaxKind::WHITESPACE)
            | (Some(SyntaxKind::COMMENT), SyntaxKind::WHITESPACE) => {}
            _ => write!(buf, "{}", child).unwrap(),
        };
    }

    buf.trim_end().to_string()
}

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

    #[test]
    fn test_add_missing_impl_members() {
        check_assist(
            add_missing_impl_members,
            "
trait Foo {
    fn foo(&self);
    fn bar(&self);
    fn baz(&self);
}

struct S;

impl Foo for S {
    fn bar(&self) {}
    <|>
}",
            "
trait Foo {
    fn foo(&self);
    fn bar(&self);
    fn baz(&self);
}

struct S;

impl Foo for S {
    fn bar(&self) {}
    fn foo(&self) { unimplemented!() }
    fn baz(&self) { unimplemented!() }<|>
}",
        );
    }

    #[test]
    fn test_copied_overriden_members() {
        check_assist(
            add_missing_impl_members,
            "
trait Foo {
    fn foo(&self);
    fn bar(&self) -> bool { true }
    fn baz(&self) -> u32 { 42 }
}

struct S;

impl Foo for S {
    fn bar(&self) {}
    <|>
}",
            "
trait Foo {
    fn foo(&self);
    fn bar(&self) -> bool { true }
    fn baz(&self) -> u32 { 42 }
}

struct S;

impl Foo for S {
    fn bar(&self) {}
    fn foo(&self) { unimplemented!() }<|>
}",
        );
    }

    #[test]
    fn test_empty_impl_block() {
        check_assist(
            add_missing_impl_members,
            "
trait Foo { fn foo(&self); }
struct S;
impl Foo for S { <|> }",
            "
trait Foo { fn foo(&self); }
struct S;
impl Foo for S {
    fn foo(&self) { unimplemented!() }<|>
}",
        );
    }

    #[test]
    fn test_cursor_after_empty_impl_block() {
        check_assist(
            add_missing_impl_members,
            "
trait Foo { fn foo(&self); }
struct S;
impl Foo for S {}<|>",
            "
trait Foo { fn foo(&self); }
struct S;
impl Foo for S {
    fn foo(&self) { unimplemented!() }<|>
}",
        )
    }

    #[test]
    fn test_empty_trait() {
        check_assist_not_applicable(
            add_missing_impl_members,
            "
trait Foo;
struct S;
impl Foo for S { <|> }",
        )
    }

    #[test]
    fn test_ignore_unnamed_trait_members_and_default_methods() {
        check_assist_not_applicable(
            add_missing_impl_members,
            "
trait Foo {
    fn (arg: u32);
    fn valid(some: u32) -> bool { false }
}
struct S;
impl Foo for S { <|> }",
        )
    }

    #[test]
    fn test_indented_impl_block() {
        check_assist(
            add_missing_impl_members,
            "
trait Foo {
    fn valid(some: u32) -> bool;
}
struct S;

mod my_mod {
    impl crate::Foo for S { <|> }
}",
            "
trait Foo {
    fn valid(some: u32) -> bool;
}
struct S;

mod my_mod {
    impl crate::Foo for S {
        fn valid(some: u32) -> bool { unimplemented!() }<|>
    }
}",
        )
    }

    #[test]
    fn test_with_docstring_and_attrs() {
        check_assist(
            add_missing_impl_members,
            r#"
#[doc(alias = "test alias")]
trait Foo {
    /// doc string
    #[must_use]
    fn foo(&self);
}
struct S;
impl Foo for S {}<|>"#,
            r#"
#[doc(alias = "test alias")]
trait Foo {
    /// doc string
    #[must_use]
    fn foo(&self);
}
struct S;
impl Foo for S {
    fn foo(&self) { unimplemented!() }<|>
}"#,
        )
    }

    #[test]
    fn test_default_methods() {
        check_assist(
            add_missing_default_members,
            "
trait Foo {
    fn valid(some: u32) -> bool { false }
    fn foo(some: u32) -> bool;
}
struct S;
impl Foo for S { <|> }",
            "
trait Foo {
    fn valid(some: u32) -> bool { false }
    fn foo(some: u32) -> bool;
}
struct S;
impl Foo for S {
    fn valid(some: u32) -> bool { false }<|>
}",
        )
    }

}