rust/crates/ide-assists/src/handlers/generate_mut_trait_impl.rs

use syntax::{
    ast::{self, make},
    ted, AstNode,
};

use crate::{AssistContext, AssistId, AssistKind, Assists};

// Generate proper `index_mut` method body refer to `index` method body may impossible due to the unpredicable case [#15581].
// Here just leave the `index_mut` method body be same as `index` method body, user can modify it manually to meet their need.

// Assist: generate_mut_trait_impl
//
// Adds a IndexMut impl from the `Index` trait.
//
// ```
// pub enum Axis { X = 0, Y = 1, Z = 2 }
//
// impl<T> Index$0<Axis> for [T; 3] {
//     type Output = T;
//
//     fn index(&self, index: Axis) -> &Self::Output {
//         &self[index as usize]
//     }
// }
// ```
// ->
// ```
// pub enum Axis { X = 0, Y = 1, Z = 2 }
//
// $0impl<T> IndexMut<Axis> for [T; 3] {
//     fn index_mut(&mut self, index: Axis) -> &mut Self::Output {
//         &self[index as usize]
//     }
// }
//
// impl<T> Index<Axis> for [T; 3] {
//     type Output = T;
//
//     fn index(&self, index: Axis) -> &Self::Output {
//         &self[index as usize]
//     }
// }
// ```
pub(crate) fn generate_mut_trait_impl(acc: &mut Assists, ctx: &AssistContext<'_>) -> Option<()> {
    let impl_def = ctx.find_node_at_offset::<ast::Impl>()?.clone_for_update();

    let trait_ = impl_def.trait_()?;
    if let ast::Type::PathType(trait_type) = trait_.clone() {
        let trait_name = trait_type.path()?.segment()?.name_ref()?.to_string();
        if trait_name != "Index" {
            return None;
        }
    }

    // Index -> IndexMut
    let index_trait = impl_def
        .syntax()
        .descendants()
        .filter_map(ast::NameRef::cast)
        .find(|it| it.text() == "Index")?;
    ted::replace(
        index_trait.syntax(),
        make::path_segment(make::name_ref("IndexMut")).clone_for_update().syntax(),
    );

    // index -> index_mut
    let trait_method_name = impl_def
        .syntax()
        .descendants()
        .filter_map(ast::Name::cast)
        .find(|it| it.text() == "index")?;
    ted::replace(trait_method_name.syntax(), make::name("index_mut").clone_for_update().syntax());

    let type_alias = impl_def.syntax().descendants().find_map(ast::TypeAlias::cast)?;
    ted::remove(type_alias.syntax());

    // &self -> &mut self
    let mut_self_param = make::mut_self_param();
    let self_param: ast::SelfParam =
        impl_def.syntax().descendants().find_map(ast::SelfParam::cast)?;
    ted::replace(self_param.syntax(), mut_self_param.clone_for_update().syntax());

    // &Self::Output -> &mut Self::Output
    let ret_type = impl_def.syntax().descendants().find_map(ast::RetType::cast)?;
    ted::replace(
        ret_type.syntax(),
        make::ret_type(make::ty("&mut Self::Output")).clone_for_update().syntax(),
    );

    let fn_ = impl_def.assoc_item_list()?.assoc_items().find_map(|it| match it {
        ast::AssocItem::Fn(f) => Some(f),
        _ => None,
    })?;

    let assoc_list = make::assoc_item_list().clone_for_update();
    assoc_list.add_item(syntax::ast::AssocItem::Fn(fn_));
    ted::replace(impl_def.assoc_item_list()?.syntax(), assoc_list.syntax());

    let target = impl_def.syntax().text_range();
    acc.add(
        AssistId("generate_mut_trait_impl", AssistKind::Generate),
        "Generate `IndexMut` impl from this `Index` trait",
        target,
        |edit| {
            edit.insert(target.start(), format!("$0{}\n\n", impl_def.to_string()));
        },
    )
}

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

    use super::*;

    #[test]
    fn test_generate_mut_trait_impl() {
        check_assist(
            generate_mut_trait_impl,
            r#"
pub enum Axis { X = 0, Y = 1, Z = 2 }

impl<T> Index$0<Axis> for [T; 3] {
    type Output = T;

    fn index(&self, index: Axis) -> &Self::Output {
        &self[index as usize]
    }
}
"#,
            r#"
pub enum Axis { X = 0, Y = 1, Z = 2 }

$0impl<T> IndexMut<Axis> for [T; 3] {
    fn index_mut(&mut self, index: Axis) -> &mut Self::Output {
        &self[index as usize]
    }
}

impl<T> Index<Axis> for [T; 3] {
    type Output = T;

    fn index(&self, index: Axis) -> &Self::Output {
        &self[index as usize]
    }
}
"#,
        );

        check_assist(
            generate_mut_trait_impl,
            r#"
pub enum Axis { X = 0, Y = 1, Z = 2 }

impl<T> Index$0<Axis> for [T; 3] where T: Copy {
    type Output = T;

    fn index(&self, index: Axis) -> &Self::Output {
        let var_name = &self[index as usize];
        var_name
    }
}
"#,
            r#"
pub enum Axis { X = 0, Y = 1, Z = 2 }

$0impl<T> IndexMut<Axis> for [T; 3] where T: Copy {
    fn index_mut(&mut self, index: Axis) -> &mut Self::Output {
        let var_name = &self[index as usize];
        var_name
    }
}

impl<T> Index<Axis> for [T; 3] where T: Copy {
    type Output = T;

    fn index(&self, index: Axis) -> &Self::Output {
        let var_name = &self[index as usize];
        var_name
    }
}
"#,
        );
    }

    #[test]
    fn test_generate_mut_trait_impl_not_applicable() {
        check_assist_not_applicable(
            generate_mut_trait_impl,
            r#"
impl<T> Add$0<i32> for [T; 3] {}
"#,
        );
    }
}