feat: add generate_mut_trait_impl assist

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

@@ -0,0 +1,195 @@
+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] {}
+"#,
+        );
+    }
+}

crates/ide-assists/src/lib.rs

@@ -160,6 +160,7 @@ mod handlers {
     mod generate_getter_or_setter;
     mod generate_impl;
     mod generate_is_empty_from_len;
+    mod generate_mut_trait_impl;
     mod generate_new;
     mod generate_delegate_methods;
     mod generate_trait_from_impl;
@@ -274,6 +275,7 @@ mod handlers {
             generate_function::generate_function,
             generate_impl::generate_impl,
             generate_impl::generate_trait_impl,
+            generate_mut_trait_impl::generate_mut_trait_impl,
             generate_is_empty_from_len::generate_is_empty_from_len,
             generate_new::generate_new,
             generate_trait_from_impl::generate_trait_from_impl,

crates/ide-assists/src/tests/generated.rs

@@ -1538,6 +1538,41 @@ impl MyStruct {
     )
 }
 
+#[test]
+fn doctest_generate_mut_trait_impl() {
+    check_doc_test(
+        "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]
+    }
+}
+"#####,
+    )
+}
+
 #[test]
 fn doctest_generate_new() {
     check_doc_test(

crates/syntax/src/ast/make.rs

@@ -852,6 +852,10 @@ pub fn self_param() -> ast::SelfParam {
     ast_from_text("fn f(&self) { }")
 }
 
+pub fn mut_self_param() -> ast::SelfParam {
+    ast_from_text("fn f(&mut self) { }")
+}
+
 pub fn ret_type(ty: ast::Type) -> ast::RetType {
     ast_from_text(&format!("fn f() -> {ty} {{ }}"))
 }