Skip to content

Commit e0e4209

Browse files
viorinaviorina
committed
extend add_impl_members to constants and types
1 parent f15f0d1 commit e0e4209

File tree

2 files changed

+87
-50
lines changed

2 files changed

+87
-50
lines changed

‎crates/ra_assists/src/add_missing_impl_members.rs

Lines changed: 68 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -2,9 +2,10 @@ use crate::{Assist, AssistId, AssistCtx, ast_editor::{AstEditor, AstBuilder}};
22

33
use hir::{HasSource, db::HirDatabase};
44
use ra_syntax::{SmolStr, TreeArc};
5-
use ra_syntax::ast::{self, AstNode, FnDef, ImplItem, ImplItemKind, NameOwner};
5+
use ra_syntax::ast::{self, AstNode, ImplItem, ImplItemKind, NameOwner};
66
use ra_db::FilePosition;
77

8+
#[derive(PartialEq)]
89
enum AddMissingImplMembersMode {
910
DefaultMethodsOnly,
1011
NoDefaultMethods,
@@ -45,39 +46,50 @@ fn add_missing_impl_members_inner(
4546
resolve_target_trait_def(ctx.db, &analyzer, impl_node)?
4647
};
4748

48-
let missing_fns: Vec<_> = {
49-
let fn_def_opt = |kind| if let ImplItemKind::FnDef(def) = kind { Some(def) } else { None };
50-
let def_name = |def| -> Option<&SmolStr> { FnDef::name(def).map(ast::Name::text) };
51-
52-
let trait_items =
53-
trait_def.syntax().descendants().find_map(ast::ItemList::cast)?.impl_items();
54-
let impl_items = impl_item_list.impl_items();
55-
56-
let trait_fns = trait_items.map(ImplItem::kind).filter_map(fn_def_opt);
57-
let impl_fns = impl_items.map(ImplItem::kind).filter_map(fn_def_opt).collect::<Vec<_>>();
58-
59-
trait_fns
60-
.filter(|t| def_name(t).is_some())
61-
.filter(|t| match mode {
62-
AddMissingImplMembersMode::DefaultMethodsOnly => t.body().is_some(),
63-
AddMissingImplMembersMode::NoDefaultMethods => t.body().is_none(),
64-
})
65-
.filter(|t| impl_fns.iter().all(|i| def_name(i) != def_name(t)))
66-
.collect()
49+
let def_name = |kind| -> Option<&SmolStr> {
50+
match kind {
51+
ImplItemKind::FnDef(def) => def.name(),
52+
ImplItemKind::TypeAliasDef(def) => def.name(),
53+
ImplItemKind::ConstDef(def) => def.name(),
54+
}
55+
.map(ast::Name::text)
6756
};
68-
if missing_fns.is_empty() {
57+
58+
let trait_items = trait_def.item_list()?.impl_items();
59+
let impl_items = impl_item_list.impl_items().collect::<Vec<_>>();
60+
61+
let missing_items: Vec<_> = trait_items
62+
.filter(|t| def_name(t.kind()).is_some())
63+
.filter(|t| match t.kind() {
64+
ImplItemKind::FnDef(def) => match mode {
65+
AddMissingImplMembersMode::DefaultMethodsOnly => def.body().is_some(),
66+
AddMissingImplMembersMode::NoDefaultMethods => def.body().is_none(),
67+
},
68+
_ => mode == AddMissingImplMembersMode::NoDefaultMethods,
69+
})
70+
.filter(|t| impl_items.iter().all(|i| def_name(i.kind()) != def_name(t.kind())))
71+
.collect();
72+
if missing_items.is_empty() {
6973
return None;
7074
}
7175

7276
ctx.add_action(AssistId(assist_id), label, |edit| {
7377
let n_existing_items = impl_item_list.impl_items().count();
74-
let fns = missing_fns.into_iter().map(add_body_and_strip_docstring).collect::<Vec<_>>();
75-
7678
let mut ast_editor = AstEditor::new(impl_item_list);
7779
if n_existing_items == 0 {
7880
ast_editor.make_multiline();
7981
}
80-
ast_editor.append_functions(fns.iter().map(|it| &**it));
82+
83+
for item in missing_items {
84+
let it = match item.kind() {
85+
ImplItemKind::FnDef(def) => {
86+
strip_docstring(ImplItem::cast(add_body(def).syntax()).unwrap())
87+
}
88+
_ => strip_docstring(item),
89+
};
90+
ast_editor.append_item(&it)
91+
}
92+
8193
let first_new_item = ast_editor.ast().impl_items().nth(n_existing_items).unwrap();
8294
let cursor_poisition = first_new_item.syntax().range().start();
8395
ast_editor.into_text_edit(edit.text_edit_builder());
@@ -88,14 +100,19 @@ fn add_missing_impl_members_inner(
88100
ctx.build()
89101
}
90102

91-
fn add_body_and_strip_docstring(fn_def: &ast::FnDef) -> TreeArc<ast::FnDef> {
103+
fn strip_docstring(item: &ast::ImplItem) -> TreeArc<ast::ImplItem> {
104+
let mut ast_editor = AstEditor::new(item);
105+
ast_editor.strip_attrs_and_docs();
106+
ast_editor.ast().to_owned()
107+
}
108+
109+
fn add_body(fn_def: &ast::FnDef) -> TreeArc<ast::FnDef> {
92110
let mut ast_editor = AstEditor::new(fn_def);
93111
if fn_def.body().is_none() {
94112
ast_editor.set_body(&AstBuilder::<ast::Block>::single_expr(
95113
&AstBuilder::<ast::Expr>::unimplemented(),
96114
));
97115
}
98-
ast_editor.strip_attrs_and_docs();
99116
ast_editor.ast().to_owned()
100117
}
101118

@@ -126,6 +143,10 @@ mod tests {
126143
add_missing_impl_members,
127144
"
128145
trait Foo {
146+
type Output;
147+
148+
const CONST: usize = 42;
149+
129150
fn foo(&self);
130151
fn bar(&self);
131152
fn baz(&self);
@@ -139,6 +160,10 @@ impl Foo for S {
139160
}",
140161
"
141162
trait Foo {
163+
type Output;
164+
165+
const CONST: usize = 42;
166+
142167
fn foo(&self);
143168
fn bar(&self);
144169
fn baz(&self);
@@ -148,7 +173,9 @@ struct S;
148173
149174
impl Foo for S {
150175
fn bar(&self) {}
151-
<|>fn foo(&self) { unimplemented!() }
176+
<|>type Output;
177+
const CONST: usize = 42;
178+
fn foo(&self) { unimplemented!() }
152179
fn baz(&self) { unimplemented!() }
153180
154181
}",
@@ -256,6 +283,8 @@ impl Foo for S { <|> }",
256283
#[doc(alias = "test alias")]
257284
trait Foo {
258285
/// doc string
286+
type Output;
287+
259288
#[must_use]
260289
fn foo(&self);
261290
}
@@ -265,12 +294,15 @@ impl Foo for S {}<|>"#,
265294
#[doc(alias = "test alias")]
266295
trait Foo {
267296
/// doc string
297+
type Output;
298+
268299
#[must_use]
269300
fn foo(&self);
270301
}
271302
struct S;
272303
impl Foo for S {
273-
<|>fn foo(&self) { unimplemented!() }
304+
<|>type Output;
305+
fn foo(&self) { unimplemented!() }
274306
}"#,
275307
)
276308
}
@@ -281,13 +313,21 @@ impl Foo for S {
281313
add_missing_default_members,
282314
"
283315
trait Foo {
316+
type Output;
317+
318+
const CONST: usize = 42;
319+
284320
fn valid(some: u32) -> bool { false }
285321
fn foo(some: u32) -> bool;
286322
}
287323
struct S;
288324
impl Foo for S { <|> }",
289325
"
290326
trait Foo {
327+
type Output;
328+
329+
const CONST: usize = 42;
330+
291331
fn valid(some: u32) -> bool { false }
292332
fn foo(some: u32) -> bool;
293333
}

‎crates/ra_assists/src/ast_editor.rs

Lines changed: 19 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -163,11 +163,7 @@ impl AstEditor<ast::ItemList> {
163163
self.do_make_multiline()
164164
}
165165

166-
pub fn append_functions<'a>(&mut self, fns: impl Iterator<Item = &'a ast::FnDef>) {
167-
fns.for_each(|it| self.append_function(it))
168-
}
169-
170-
pub fn append_function(&mut self, fn_def: &ast::FnDef) {
166+
pub fn append_item(&mut self, item: &ast::ImplItem) {
171167
let (indent, position) = match self.ast().impl_items().last() {
172168
Some(it) => (
173169
leading_indent(it.syntax()).unwrap_or("").to_string(),
@@ -182,8 +178,7 @@ impl AstEditor<ast::ItemList> {
182178
},
183179
};
184180
let ws = tokens::WsBuilder::new(&format!("\n{}", indent));
185-
let to_insert: ArrayVec<[SyntaxElement; 2]> =
186-
[ws.ws().into(), fn_def.syntax().into()].into();
181+
let to_insert: ArrayVec<[SyntaxElement; 2]> = [ws.ws().into(), item.syntax().into()].into();
187182
self.ast = self.insert_children(position, to_insert.into_iter());
188183
}
189184

@@ -192,6 +187,23 @@ impl AstEditor<ast::ItemList> {
192187
}
193188
}
194189

190+
impl AstEditor<ast::ImplItem> {
191+
pub fn strip_attrs_and_docs(&mut self) {
192+
while let Some(start) = self
193+
.ast()
194+
.syntax()
195+
.children_with_tokens()
196+
.find(|it| it.kind() == ATTR || it.kind() == COMMENT)
197+
{
198+
let end = match start.next_sibling_or_token() {
199+
Some(el) if el.kind() == WHITESPACE => el,
200+
Some(_) | None => start,
201+
};
202+
self.ast = self.replace_children(RangeInclusive::new(start, end), iter::empty());
203+
}
204+
}
205+
}
206+
195207
impl AstEditor<ast::FnDef> {
196208
pub fn set_body(&mut self, body: &ast::Block) {
197209
let mut to_insert: ArrayVec<[SyntaxElement; 2]> = ArrayVec::new();
@@ -210,21 +222,6 @@ impl AstEditor<ast::FnDef> {
210222
let replace_range = RangeInclusive::new(old_body_or_semi, old_body_or_semi);
211223
self.ast = self.replace_children(replace_range, to_insert.into_iter())
212224
}
213-
214-
pub fn strip_attrs_and_docs(&mut self) {
215-
while let Some(start) = self
216-
.ast()
217-
.syntax()
218-
.children_with_tokens()
219-
.find(|it| it.kind() == ATTR || it.kind() == COMMENT)
220-
{
221-
let end = match start.next_sibling_or_token() {
222-
Some(el) if el.kind() == WHITESPACE => el,
223-
Some(_) | None => start,
224-
};
225-
self.ast = self.replace_children(RangeInclusive::new(start, end), iter::empty());
226-
}
227-
}
228225
}
229226

230227
pub struct AstBuilder<N: AstNode> {

0 commit comments

Comments
 (0)