5 files changed, 51 insertions, 20 deletions

diff --git a/crates/libeditor/src/code_actions.rs b/crates/libeditor/src/code_actions.rs
index 4987964d2..6df64be12 100644
--- a/crates/libeditor/src/code_actions.rs
+++ b/crates/libeditor/src/code_actions.rs
@@ -1,6 +1,6 @@
 use {TextUnit, File, EditBuilder, Edit};
 use libsyntax2::{
-    ast::{self, AstNode},
+    ast::{self, AstNode, AttrsOwner},
     SyntaxKind::COMMA,
     SyntaxNodeRef,
     SyntaxRoot,
@@ -39,18 +39,28 @@ pub fn flip_comma<'a>(file: &'a File, offset: TextUnit) -> Option<impl FnOnce()
 }
 
 pub fn add_derive<'a>(file: &'a File, offset: TextUnit) -> Option<impl FnOnce() -> ActionResult + 'a> {
-    let syntax = file.syntax();
-    let syntax = syntax.as_ref();
-    let nominal = find_node::<ast::NominalDef<_>>(syntax, offset)?;
+    let nominal = find_node::<ast::NominalDef<_>>(file.syntax_ref(), offset)?;
     Some(move || {
+        let derive_attr = nominal
+            .attrs()
+            .filter_map(|x| x.as_call())
+            .filter(|(name, _arg)| name == "derive")
+            .map(|(_name, arg)| arg)
+            .next();
         let mut edit = EditBuilder::new();
-        let node_start = nominal.syntax().range().start();
-        edit.insert(node_start, "#[derive()]\n".to_string());
+        let offset = match derive_attr {
+            None => {
+                let node_start = nominal.syntax().range().start();
+                edit.insert(node_start, "#[derive()]\n".to_string());
+                node_start + TextUnit::of_str("#[derive(")
+            }
+            Some(tt) => {
+                tt.syntax().range().end() - TextUnit::of_char(')')
+            }
+        };
         ActionResult {
             edit: edit.finish(),
-            cursor_position: CursorPosition::Offset(
-                node_start + TextUnit::of_str("#[derive(")
-            ),
+            cursor_position: CursorPosition::Offset(offset),
         }
     })
 }
diff --git a/crates/libeditor/tests/test.rs b/crates/libeditor/tests/test.rs
index d5df9d0cc..97919d347 100644
--- a/crates/libeditor/tests/test.rs
+++ b/crates/libeditor/tests/test.rs
@@ -116,7 +116,12 @@ fn test_add_derive() {
         "struct Foo { a: i32, <|>}",
         "#[derive(<|>)]\nstruct Foo { a: i32, }",
         |file, off| add_derive(file, off).map(|f| f()),
-    )
+    );
+    check_action(
+        "#[derive(Clone)]\nstruct Foo { a: i32<|>, }",
+        "#[derive(Clone<|>)]\nstruct Foo { a: i32, }",
+        |file, off| add_derive(file, off).map(|f| f()),
+    );
 }
 
 #[test]
diff --git a/crates/libsyntax2/src/ast/generated.rs b/crates/libsyntax2/src/ast/generated.rs
index c575e15df..be3f73f7e 100644
--- a/crates/libsyntax2/src/ast/generated.rs
+++ b/crates/libsyntax2/src/ast/generated.rs
@@ -339,6 +339,7 @@ impl<R: TreeRoot> AstNode<R> for NominalDef<R> {
     }
 }
 
+impl<R: TreeRoot> ast::AttrsOwner<R> for NominalDef<R> {}
 impl<R: TreeRoot> NominalDef<R> {}
 
 // ParenType
diff --git a/crates/libsyntax2/src/ast/mod.rs b/crates/libsyntax2/src/ast/mod.rs
index fe8f91d15..d53b12ab8 100644
--- a/crates/libsyntax2/src/ast/mod.rs
+++ b/crates/libsyntax2/src/ast/mod.rs
@@ -52,19 +52,31 @@ impl<R: TreeRoot> File<R> {
 impl<R: TreeRoot> FnDef<R> {
     pub fn has_atom_attr(&self, atom: &str) -> bool {
         self.attrs()
-            .filter_map(|x| x.value())
-            .filter_map(|x| as_atom(x))
+            .filter_map(|x| x.as_atom())
             .any(|x| x == atom)
     }
 }
 
-fn as_atom<R: TreeRoot>(tt: TokenTree<R>) -> Option<SmolStr> {
-    let syntax = tt.syntax_ref();
-    let (_bra, attr, _ket) = syntax.children().collect_tuple()?;
-    if attr.kind() == IDENT {
-        Some(attr.leaf_text().unwrap())
-    } else {
-        None
+impl<R: TreeRoot> Attr<R> {
+    pub fn as_atom(&self) -> Option<SmolStr> {
+        let tt = self.value()?;
+        let (_bra, attr, _ket) = tt.syntax().children().collect_tuple()?;
+        if attr.kind() == IDENT {
+            Some(attr.leaf_text().unwrap())
+        } else {
+            None
+        }
+    }
+
+    pub fn as_call(&self) -> Option<(SmolStr, TokenTree<R>)> {
+        let tt = self.value()?;
+        let (_bra, attr, args, _ket) = tt.syntax().children().collect_tuple()?;
+        let args = TokenTree::cast(args)?;
+        if attr.kind() == IDENT {
+            Some((attr.leaf_text().unwrap(), args))
+        } else {
+            None
+        }
     }
 }
 
diff --git a/crates/libsyntax2/src/grammar.ron b/crates/libsyntax2/src/grammar.ron
index 4e523da9a..abeffb2c3 100644
--- a/crates/libsyntax2/src/grammar.ron
+++ b/crates/libsyntax2/src/grammar.ron
@@ -272,6 +272,9 @@ Grammar(
             "DynTraitType",
         ]),
 
-        "NominalDef": ( enum: ["StructDef", "EnumDef"]),
+        "NominalDef": (
+            enum: ["StructDef", "EnumDef"],
+            traits: [ "AttrsOwner" ],
+        ),
     },
 )