Skip to content

Commit b89fef5

Browse files
Merge #7562
7562: add `generate_enum_match` assist r=matklad a=yoshuawuyts This adds a `generate_enum_match` assist, which generates `is_` variants for enums (e.g. `Option::{is_none,is_some}` in std). This is my first attempt at contributing to Rust-Analyzer, so I'm not sure if I've gotten everything right. Thanks! ## Example **Input** ```rust pub(crate) enum Variant { Undefined, Minor, // cursor here Major, } ``` **Output** ```rust pub(crate) enum Variant { Undefined, Minor, Major, } impl Variant { pub(crate) fn is_minor(&self) -> bool { matches!(self, Self::Minor) } } ``` ## Future Directions I made this as a stepping stone for some of the more involved refactors (e.g. #5944). I'm not sure yet how to create, use, and test `window.showQuickPick`-based asssists in RA. But once that's possible, it'd probably be nice to be able to generate match methods in bulk through the quickpick UI rather than one-by-one: ``` [x] Select enum members to generate methods for. (3 selected) [ OK ] --------------------------------------------------------------------------- [x] Undefined [x] Minor [x] Major ``` Co-authored-by: Yoshua Wuyts <[email protected]> Co-authored-by: Yoshua Wuyts <[email protected]>
2 parents 5009958 + dfd7513 commit b89fef5

File tree

5 files changed

+317
-64
lines changed

5 files changed

+317
-64
lines changed
Lines changed: 213 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,213 @@
1+
use stdx::{format_to, to_lower_snake_case};
2+
use syntax::ast::{self, AstNode, NameOwner};
3+
use syntax::{ast::VisibilityOwner, T};
4+
use test_utils::mark;
5+
6+
use crate::{utils::find_struct_impl, AssistContext, AssistId, AssistKind, Assists};
7+
8+
// Assist: generate_enum_match_method
9+
//
10+
// Generate an `is_` method for an enum variant.
11+
//
12+
// ```
13+
// enum Version {
14+
// Undefined,
15+
// Minor$0,
16+
// Major,
17+
// }
18+
// ```
19+
// ->
20+
// ```
21+
// enum Version {
22+
// Undefined,
23+
// Minor,
24+
// Major,
25+
// }
26+
//
27+
// impl Version {
28+
// fn is_minor(&self) -> bool {
29+
// matches!(self, Self::Minor)
30+
// }
31+
// }
32+
// ```
33+
pub(crate) fn generate_enum_match_method(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
34+
let variant = ctx.find_node_at_offset::<ast::Variant>()?;
35+
let variant_name = variant.name()?;
36+
let parent_enum = variant.parent_enum();
37+
if !matches!(variant.kind(), ast::StructKind::Unit) {
38+
mark::hit!(test_gen_enum_match_on_non_unit_variant_not_implemented);
39+
return None;
40+
}
41+
42+
let fn_name = to_lower_snake_case(&variant_name.to_string());
43+
44+
// Return early if we've found an existing new fn
45+
let impl_def = find_struct_impl(
46+
&ctx,
47+
&ast::AdtDef::Enum(parent_enum.clone()),
48+
format!("is_{}", fn_name).as_str(),
49+
)?;
50+
51+
let target = variant.syntax().text_range();
52+
acc.add(
53+
AssistId("generate_enum_match_method", AssistKind::Generate),
54+
"Generate an `is_` method for an enum variant",
55+
target,
56+
|builder| {
57+
let mut buf = String::with_capacity(512);
58+
59+
if impl_def.is_some() {
60+
buf.push('\n');
61+
}
62+
63+
let vis = parent_enum.visibility().map_or(String::new(), |v| format!("{} ", v));
64+
65+
format_to!(
66+
buf,
67+
" {}fn is_{}(&self) -> bool {{
68+
matches!(self, Self::{})
69+
}}",
70+
vis,
71+
fn_name,
72+
variant_name
73+
);
74+
75+
let start_offset = impl_def
76+
.and_then(|impl_def| {
77+
buf.push('\n');
78+
let start = impl_def
79+
.syntax()
80+
.descendants_with_tokens()
81+
.find(|t| t.kind() == T!['{'])?
82+
.text_range()
83+
.end();
84+
85+
Some(start)
86+
})
87+
.unwrap_or_else(|| {
88+
buf = generate_impl_text(&parent_enum, &buf);
89+
parent_enum.syntax().text_range().end()
90+
});
91+
92+
builder.insert(start_offset, buf);
93+
},
94+
)
95+
}
96+
97+
// Generates the surrounding `impl Type { <code> }` including type and lifetime
98+
// parameters
99+
fn generate_impl_text(strukt: &ast::Enum, code: &str) -> String {
100+
let mut buf = String::with_capacity(code.len());
101+
buf.push_str("\n\nimpl ");
102+
buf.push_str(strukt.name().unwrap().text());
103+
format_to!(buf, " {{\n{}\n}}", code);
104+
buf
105+
}
106+
107+
#[cfg(test)]
108+
mod tests {
109+
use test_utils::mark;
110+
111+
use crate::tests::{check_assist, check_assist_not_applicable};
112+
113+
use super::*;
114+
115+
fn check_not_applicable(ra_fixture: &str) {
116+
check_assist_not_applicable(generate_enum_match_method, ra_fixture)
117+
}
118+
119+
#[test]
120+
fn test_generate_enum_match_from_variant() {
121+
check_assist(
122+
generate_enum_match_method,
123+
r#"
124+
enum Variant {
125+
Undefined,
126+
Minor$0,
127+
Major,
128+
}"#,
129+
r#"enum Variant {
130+
Undefined,
131+
Minor,
132+
Major,
133+
}
134+
135+
impl Variant {
136+
fn is_minor(&self) -> bool {
137+
matches!(self, Self::Minor)
138+
}
139+
}"#,
140+
);
141+
}
142+
143+
#[test]
144+
fn test_generate_enum_match_already_implemented() {
145+
check_not_applicable(
146+
r#"
147+
enum Variant {
148+
Undefined,
149+
Minor$0,
150+
Major,
151+
}
152+
153+
impl Variant {
154+
fn is_minor(&self) -> bool {
155+
matches!(self, Self::Minor)
156+
}
157+
}"#,
158+
);
159+
}
160+
161+
#[test]
162+
fn test_add_from_impl_no_element() {
163+
mark::check!(test_gen_enum_match_on_non_unit_variant_not_implemented);
164+
check_not_applicable(
165+
r#"
166+
enum Variant {
167+
Undefined,
168+
Minor(u32)$0,
169+
Major,
170+
}"#,
171+
);
172+
}
173+
174+
#[test]
175+
fn test_generate_enum_match_from_variant_with_one_variant() {
176+
check_assist(
177+
generate_enum_match_method,
178+
r#"enum Variant { Undefi$0ned }"#,
179+
r#"
180+
enum Variant { Undefined }
181+
182+
impl Variant {
183+
fn is_undefined(&self) -> bool {
184+
matches!(self, Self::Undefined)
185+
}
186+
}"#,
187+
);
188+
}
189+
190+
#[test]
191+
fn test_generate_enum_match_from_variant_with_visibility_marker() {
192+
check_assist(
193+
generate_enum_match_method,
194+
r#"
195+
pub(crate) enum Variant {
196+
Undefined,
197+
Minor$0,
198+
Major,
199+
}"#,
200+
r#"pub(crate) enum Variant {
201+
Undefined,
202+
Minor,
203+
Major,
204+
}
205+
206+
impl Variant {
207+
pub(crate) fn is_minor(&self) -> bool {
208+
matches!(self, Self::Minor)
209+
}
210+
}"#,
211+
);
212+
}
213+
}

crates/assists/src/handlers/generate_new.rs

Lines changed: 2 additions & 62 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,11 @@
1-
use hir::Adt;
21
use itertools::Itertools;
32
use stdx::format_to;
43
use syntax::{
54
ast::{self, AstNode, GenericParamsOwner, NameOwner, StructKind, VisibilityOwner},
65
SmolStr, T,
76
};
87

9-
use crate::{AssistContext, AssistId, AssistKind, Assists};
8+
use crate::{utils::find_struct_impl, AssistContext, AssistId, AssistKind, Assists};
109

1110
// Assist: generate_new
1211
//
@@ -38,7 +37,7 @@ pub(crate) fn generate_new(acc: &mut Assists, ctx: &AssistContext) -> Option<()>
3837
};
3938

4039
// Return early if we've found an existing new fn
41-
let impl_def = find_struct_impl(&ctx, &strukt)?;
40+
let impl_def = find_struct_impl(&ctx, &ast::AdtDef::Struct(strukt.clone()), "new")?;
4241

4342
let target = strukt.syntax().text_range();
4443
acc.add(AssistId("generate_new", AssistKind::Generate), "Generate `new`", target, |builder| {
@@ -111,65 +110,6 @@ fn generate_impl_text(strukt: &ast::Struct, code: &str) -> String {
111110
buf
112111
}
113112

114-
// Uses a syntax-driven approach to find any impl blocks for the struct that
115-
// exist within the module/file
116-
//
117-
// Returns `None` if we've found an existing `new` fn
118-
//
119-
// FIXME: change the new fn checking to a more semantic approach when that's more
120-
// viable (e.g. we process proc macros, etc)
121-
fn find_struct_impl(ctx: &AssistContext, strukt: &ast::Struct) -> Option<Option<ast::Impl>> {
122-
let db = ctx.db();
123-
let module = strukt.syntax().ancestors().find(|node| {
124-
ast::Module::can_cast(node.kind()) || ast::SourceFile::can_cast(node.kind())
125-
})?;
126-
127-
let struct_def = ctx.sema.to_def(strukt)?;
128-
129-
let block = module.descendants().filter_map(ast::Impl::cast).find_map(|impl_blk| {
130-
let blk = ctx.sema.to_def(&impl_blk)?;
131-
132-
// FIXME: handle e.g. `struct S<T>; impl<U> S<U> {}`
133-
// (we currently use the wrong type parameter)
134-
// also we wouldn't want to use e.g. `impl S<u32>`
135-
let same_ty = match blk.target_ty(db).as_adt() {
136-
Some(def) => def == Adt::Struct(struct_def),
137-
None => false,
138-
};
139-
let not_trait_impl = blk.target_trait(db).is_none();
140-
141-
if !(same_ty && not_trait_impl) {
142-
None
143-
} else {
144-
Some(impl_blk)
145-
}
146-
});
147-
148-
if let Some(ref impl_blk) = block {
149-
if has_new_fn(impl_blk) {
150-
return None;
151-
}
152-
}
153-
154-
Some(block)
155-
}
156-
157-
fn has_new_fn(imp: &ast::Impl) -> bool {
158-
if let Some(il) = imp.assoc_item_list() {
159-
for item in il.assoc_items() {
160-
if let ast::AssocItem::Fn(f) = item {
161-
if let Some(name) = f.name() {
162-
if name.text().eq_ignore_ascii_case("new") {
163-
return true;
164-
}
165-
}
166-
}
167-
}
168-
}
169-
170-
false
171-
}
172-
173113
#[cfg(test)]
174114
mod tests {
175115
use crate::tests::{check_assist, check_assist_not_applicable, check_assist_target};

crates/assists/src/lib.rs

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -127,6 +127,7 @@ mod handlers {
127127
mod flip_trait_bound;
128128
mod generate_default_from_enum_variant;
129129
mod generate_derive;
130+
mod generate_enum_match_method;
130131
mod generate_from_impl_for_enum;
131132
mod generate_function;
132133
mod generate_impl;
@@ -185,6 +186,7 @@ mod handlers {
185186
flip_trait_bound::flip_trait_bound,
186187
generate_default_from_enum_variant::generate_default_from_enum_variant,
187188
generate_derive::generate_derive,
189+
generate_enum_match_method::generate_enum_match_method,
188190
generate_from_impl_for_enum::generate_from_impl_for_enum,
189191
generate_function::generate_function,
190192
generate_impl::generate_impl,

crates/assists/src/tests/generated.rs

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -459,6 +459,33 @@ struct Point {
459459
)
460460
}
461461

462+
#[test]
463+
fn doctest_generate_enum_match_method() {
464+
check_doc_test(
465+
"generate_enum_match_method",
466+
r#####"
467+
enum Version {
468+
Undefined,
469+
Minor$0,
470+
Major,
471+
}
472+
"#####,
473+
r#####"
474+
enum Version {
475+
Undefined,
476+
Minor,
477+
Major,
478+
}
479+
480+
impl Version {
481+
fn is_minor(&self) -> bool {
482+
matches!(self, Self::Minor)
483+
}
484+
}
485+
"#####,
486+
)
487+
}
488+
462489
#[test]
463490
fn doctest_generate_from_impl_for_enum() {
464491
check_doc_test(

0 commit comments

Comments
 (0)