Skip to content

Commit 875c1a3

Browse files
committed
Add AsRef and Borrow for generate_mut_trait_impl
- AsRef -> AsMut - Borrow -> BorrowMut Example ==================== ```rust //- minicore: as_ref struct Foo(i32); impl<T> core::convert::AsRef$0<i32> for Foo { fn as_ref(&self) -> &i32 { &self.0 } } ``` -> ```rust struct Foo(i32); $0impl<T> core::convert::AsMut<i32> for Foo { fn as_mut(&mut self) -> &mut i32 { &self.0 } } impl<T> core::convert::AsRef<i32> for Foo { fn as_ref(&self) -> &i32 { &self.0 } } ```
1 parent cf969d2 commit 875c1a3

File tree

2 files changed

+100
-18
lines changed

2 files changed

+100
-18
lines changed

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

Lines changed: 96 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -48,36 +48,42 @@ pub(crate) fn generate_mut_trait_impl(acc: &mut Assists, ctx: &AssistContext<'_>
4848
let impl_def = ctx.find_node_at_offset::<ast::Impl>()?.clone_for_update();
4949
let indent = impl_def.indent_level();
5050

51+
let (apply_trait, new_apply_trait) = impl_def
52+
.syntax()
53+
.descendants()
54+
.filter_map(ast::NameRef::cast)
55+
.find_map(process_trait_name)?;
56+
5157
let trait_ = impl_def.trait_()?;
5258
if let ast::Type::PathType(trait_path) = trait_ {
5359
let trait_type = ctx.sema.resolve_trait(&trait_path.path()?)?;
5460
let scope = ctx.sema.scope(trait_path.syntax())?;
55-
if trait_type != FamousDefs(&ctx.sema, scope.krate()).core_convert_Index()? {
61+
let famous_defs = FamousDefs(&ctx.sema, scope.krate());
62+
if trait_type != get_famous(&apply_trait.text(), famous_defs)? {
5663
return None;
5764
}
5865
}
5966

6067
// Index -> IndexMut
61-
let index_trait = impl_def
62-
.syntax()
63-
.descendants()
64-
.filter_map(ast::NameRef::cast)
65-
.find(|it| it.text() == "Index")?;
6668
ted::replace(
67-
index_trait.syntax(),
68-
make::path_segment(make::name_ref("IndexMut")).clone_for_update().syntax(),
69+
apply_trait.syntax(),
70+
make::path_segment(make::name_ref(new_apply_trait)).clone_for_update().syntax(),
6971
);
7072

7173
// index -> index_mut
72-
let trait_method_name = impl_def
74+
let (trait_method_name, new_trait_method_name) = impl_def
7375
.syntax()
7476
.descendants()
7577
.filter_map(ast::Name::cast)
76-
.find(|it| it.text() == "index")?;
77-
ted::replace(trait_method_name.syntax(), make::name("index_mut").clone_for_update().syntax());
78+
.find_map(process_method_name)?;
79+
ted::replace(
80+
trait_method_name.syntax(),
81+
make::name(new_trait_method_name).clone_for_update().syntax(),
82+
);
7883

79-
let type_alias = impl_def.syntax().descendants().find_map(ast::TypeAlias::cast)?;
80-
ted::remove(type_alias.syntax());
84+
if let Some(type_alias) = impl_def.syntax().descendants().find_map(ast::TypeAlias::cast) {
85+
ted::remove(type_alias.syntax());
86+
}
8187

8288
// &self -> &mut self
8389
let mut_self_param = make::mut_self_param();
@@ -87,10 +93,8 @@ pub(crate) fn generate_mut_trait_impl(acc: &mut Assists, ctx: &AssistContext<'_>
8793

8894
// &Self::Output -> &mut Self::Output
8995
let ret_type = impl_def.syntax().descendants().find_map(ast::RetType::cast)?;
90-
ted::replace(
91-
ret_type.syntax(),
92-
make::ret_type(make::ty("&mut Self::Output")).clone_for_update().syntax(),
93-
);
96+
let new_ret_type = process_ret_type(&ret_type)?;
97+
ted::replace(ret_type.syntax(), make::ret_type(new_ret_type).clone_for_update().syntax());
9498

9599
let fn_ = impl_def.assoc_item_list()?.assoc_items().find_map(|it| match it {
96100
ast::AssocItem::Fn(f) => Some(f),
@@ -104,14 +108,51 @@ pub(crate) fn generate_mut_trait_impl(acc: &mut Assists, ctx: &AssistContext<'_>
104108
let target = impl_def.syntax().text_range();
105109
acc.add(
106110
AssistId::generate("generate_mut_trait_impl"),
107-
"Generate `IndexMut` impl from this `Index` trait",
111+
format!("Generate `{new_apply_trait}` impl from this `{apply_trait}` trait"),
108112
target,
109113
|edit| {
110114
edit.insert(target.start(), format!("$0{impl_def}\n\n{indent}"));
111115
},
112116
)
113117
}
114118

119+
fn get_famous(apply_trait: &str, famous: FamousDefs<'_, '_>) -> Option<hir::Trait> {
120+
match apply_trait {
121+
"Index" => famous.core_convert_Index(),
122+
"AsRef" => famous.core_convert_AsRef(),
123+
"Borrow" => famous.core_borrow_Borrow(),
124+
_ => None,
125+
}
126+
}
127+
128+
fn process_trait_name(name: ast::NameRef) -> Option<(ast::NameRef, &'static str)> {
129+
let new_name = match &*name.text() {
130+
"Index" => "IndexMut",
131+
"AsRef" => "AsMut",
132+
"Borrow" => "BorrowMut",
133+
_ => return None,
134+
};
135+
Some((name, new_name))
136+
}
137+
138+
fn process_method_name(name: ast::Name) -> Option<(ast::Name, &'static str)> {
139+
let new_name = match &*name.text() {
140+
"index" => "index_mut",
141+
"as_ref" => "as_mut",
142+
"borrow" => "borrow_mut",
143+
_ => return None,
144+
};
145+
Some((name, new_name))
146+
}
147+
148+
fn process_ret_type(ref_ty: &ast::RetType) -> Option<ast::Type> {
149+
let ty = ref_ty.ty()?;
150+
let ast::Type::RefType(ref_type) = ty else {
151+
return None;
152+
};
153+
Some(make::ty_ref(ref_type.ty()?, true))
154+
}
155+
115156
#[cfg(test)]
116157
mod tests {
117158
use crate::tests::{check_assist, check_assist_not_applicable};
@@ -186,6 +227,35 @@ impl<T> core::ops::Index<Axis> for [T; 3] where T: Copy {
186227
var_name
187228
}
188229
}
230+
"#,
231+
);
232+
233+
check_assist(
234+
generate_mut_trait_impl,
235+
r#"
236+
//- minicore: as_ref
237+
struct Foo(i32);
238+
239+
impl<T> core::convert::AsRef$0<i32> for Foo {
240+
fn as_ref(&self) -> &i32 {
241+
&self.0
242+
}
243+
}
244+
"#,
245+
r#"
246+
struct Foo(i32);
247+
248+
$0impl<T> core::convert::AsMut<i32> for Foo {
249+
fn as_mut(&mut self) -> &mut i32 {
250+
&self.0
251+
}
252+
}
253+
254+
impl<T> core::convert::AsRef<i32> for Foo {
255+
fn as_ref(&self) -> &i32 {
256+
&self.0
257+
}
258+
}
189259
"#,
190260
);
191261
}
@@ -285,6 +355,14 @@ mod foo {
285355
pub trait Index<Idx: ?Sized> {}
286356
287357
impl<T> Index$0<i32> for [T; 3] {}
358+
"#,
359+
);
360+
check_assist_not_applicable(
361+
generate_mut_trait_impl,
362+
r#"
363+
pub trait AsRef<T: ?Sized> {}
364+
365+
impl<T> AsRef$0<i32> for [T; 3] {}
288366
"#,
289367
);
290368
}

crates/ide-db/src/famous_defs.rs

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -70,6 +70,10 @@ impl FamousDefs<'_, '_> {
7070
self.find_trait("core:ops:Index")
7171
}
7272

73+
pub fn core_borrow_Borrow(&self) -> Option<Trait> {
74+
self.find_trait("core:borrow:Borrow")
75+
}
76+
7377
pub fn core_option_Option(&self) -> Option<Enum> {
7478
self.find_enum("core:option:Option")
7579
}

0 commit comments

Comments
 (0)