5 files changed, 85 insertions, 37 deletions

diff --git a/crates/libeditor/src/code_actions.rs b/crates/libeditor/src/code_actions.rs
index 6e2c73f29..88c348436 100644
--- a/crates/libeditor/src/code_actions.rs
+++ b/crates/libeditor/src/code_actions.rs
@@ -10,7 +10,17 @@ use libsyntax2::{
     },
 };
 
-pub fn flip_comma<'a>(file: &'a File, offset: TextUnit) -> Option<impl FnOnce() -> Edit + 'a> {
+pub struct ActionResult {
+    pub edit: Edit,
+    pub cursor_position: CursorPosition,
+}
+
+pub enum CursorPosition {
+    Same,
+    Offset(TextUnit),
+}
+
+pub fn flip_comma<'a>(file: &'a File, offset: TextUnit) -> Option<impl FnOnce() -> ActionResult + 'a> {
     let syntax = file.syntax();
     let syntax = syntax.as_ref();
 
@@ -21,18 +31,27 @@ pub fn flip_comma<'a>(file: &'a File, offset: TextUnit) -> Option<impl FnOnce()
         let mut edit = EditBuilder::new();
         edit.replace(left.range(), right.text());
         edit.replace(right.range(), left.text());
-        edit.finish()
+        ActionResult {
+            edit: edit.finish(),
+            cursor_position: CursorPosition::Same,
+        }
     })
 }
 
-pub fn add_derive<'a>(file: &'a File, offset: TextUnit) -> Option<impl FnOnce() -> Edit + 'a> {
+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)?;
     Some(move || {
         let mut edit = EditBuilder::new();
-        edit.insert(nominal.syntax().range().start(), "#[derive()]\n".to_string());
-        edit.finish()
+        let node_start = nominal.syntax().range().start();
+        edit.insert(node_start, "#[derive()]\n".to_string());
+        ActionResult {
+            edit: edit.finish(),
+            cursor_position: CursorPosition::Offset(
+                node_start + TextUnit::of_str("#[derive(")
+            ),
+        }
     })
 }
 
diff --git a/crates/libeditor/src/edit.rs b/crates/libeditor/src/edit.rs
index 15a2a904f..3edd0809d 100644
--- a/crates/libeditor/src/edit.rs
+++ b/crates/libeditor/src/edit.rs
@@ -67,6 +67,21 @@ impl Edit {
         assert_eq!(buf.len(), total_len);
         buf
     }
+
+    pub fn apply_to_offset(&self, offset: TextUnit) -> Option<TextUnit> {
+        let mut res = offset;
+        for atom in self.atoms.iter() {
+            if atom.delete.start() >= offset {
+                break;
+            }
+            if offset < atom.delete.end() {
+                return None
+            }
+            res += TextUnit::of_str(&atom.insert);
+            res -= atom.delete.len();
+        }
+        Some(res)
+    }
 }
 
 impl AtomEdit {
diff --git a/crates/libeditor/src/lib.rs b/crates/libeditor/src/lib.rs
index 2c46ca45f..9e44f5d92 100644
--- a/crates/libeditor/src/lib.rs
+++ b/crates/libeditor/src/lib.rs
@@ -21,7 +21,7 @@ pub use self::{
     extend_selection::extend_selection,
     symbols::{StructureNode, file_structure, FileSymbol, file_symbols},
     edit::{EditBuilder, Edit, AtomEdit},
-    code_actions::{flip_comma, add_derive},
+    code_actions::{flip_comma, add_derive, ActionResult, CursorPosition},
 };
 
 #[derive(Debug)]
@@ -37,13 +37,6 @@ pub struct Diagnostic {
 }
 
 #[derive(Debug)]
-pub struct Symbol {
-    // pub parent: ???,
-    pub name: String,
-    pub range: TextRange,
-}
-
-#[derive(Debug)]
 pub struct Runnable {
     pub range: TextRange,
     pub kind: RunnableKind,
diff --git a/crates/libeditor/tests/test.rs b/crates/libeditor/tests/test.rs
index dd2c99c3f..7063425ce 100644
--- a/crates/libeditor/tests/test.rs
+++ b/crates/libeditor/tests/test.rs
@@ -6,10 +6,10 @@ extern crate assert_eq_text;
 
 use std::fmt;
 use itertools::Itertools;
-use libsyntax2::AstNode;
 use libeditor::{
-    File, TextUnit, TextRange,
-    highlight, runnables, extend_selection, file_structure, flip_comma,
+    File, TextUnit, TextRange, ActionResult, CursorPosition,
+    highlight, runnables, extend_selection, file_structure,
+    flip_comma, add_derive,
 };
 
 #[test]
@@ -103,13 +103,19 @@ impl fmt::Debug for E {}
 
 #[test]
 fn test_swap_comma() {
-    check_modification(
+    check_action(
         "fn foo(x: i32,<|> y: Result<(), ()>) {}",
-        "fn foo(y: Result<(), ()>, x: i32) {}",
-        &|file, offset| {
-            let edit = flip_comma(file, offset).unwrap()();
-            edit.apply(&file.syntax().text())
-        },
+        "fn foo(y: Result<(), ()>,<|> x: i32) {}",
+        |file, off| flip_comma(file, off).map(|f| f()),
+    )
+}
+
+#[test]
+fn test_add_derive() {
+    check_action(
+        "struct Foo { a: i32, <|>}",
+        "#[derive(<|>)]\nstruct Foo { a: i32, }",
+        |file, off| add_derive(file, off).map(|f| f()),
     )
 }
 
@@ -123,21 +129,36 @@ fn dbg_eq(expected: &str, actual: &impl fmt::Debug) {
     assert_eq!(expected, actual);
 }
 
-fn check_modification(
+fn check_action<F: Fn(&File, TextUnit) -> Option<ActionResult>>(
     before: &str,
     after: &str,
-    f: &impl Fn(&File, TextUnit) -> String,
+    f: F,
 ) {
+    let (before_cursor_pos, before) = extract_cursor(before);
+    let file = file(&before);
+    let result = f(&file, before_cursor_pos).expect("code action is not applicable");
+    let actual = result.edit.apply(&before);
+    let actual_cursor_pos: u32 = match result.cursor_position {
+        CursorPosition::Same => result.edit.apply_to_offset(before_cursor_pos).unwrap(),
+        CursorPosition::Offset(off) => off,
+    }.into();
+    let actual_cursor_pos = actual_cursor_pos as usize;
+    let mut actual_with_cursor = String::new();
+    actual_with_cursor.push_str(&actual[..actual_cursor_pos]);
+    actual_with_cursor.push_str("<|>");
+    actual_with_cursor.push_str(&actual[actual_cursor_pos..]);
+    assert_eq_text!(after, &actual_with_cursor);
+}
+
+fn extract_cursor(text: &str) -> (TextUnit, String) {
     let cursor = "<|>";
-    let cursor_pos = match before.find(cursor) {
-        None => panic!("before text should contain cursor marker"),
+    let cursor_pos = match text.find(cursor) {
+        None => panic!("text should contain cursor marker"),
         Some(pos) => pos,
     };
-    let mut text = String::with_capacity(before.len() - cursor.len());
-    text.push_str(&before[..cursor_pos]);
-    text.push_str(&before[cursor_pos + cursor.len()..]);
+    let mut new_text = String::with_capacity(text.len() - cursor.len());
+    new_text.push_str(&text[..cursor_pos]);
+    new_text.push_str(&text[cursor_pos + cursor.len()..]);
     let cursor_pos = TextUnit::from(cursor_pos as u32);
-    let file = file(&text);
-    let actual = f(&file, cursor_pos);
-    assert_eq_text!(after, &actual);
+    (cursor_pos, new_text)
 }
diff --git a/crates/server/src/main_loop/handlers.rs b/crates/server/src/main_loop/handlers.rs
index 637bf1b24..078abfbfa 100644
--- a/crates/server/src/main_loop/handlers.rs
+++ b/crates/server/src/main_loop/handlers.rs
@@ -187,12 +187,12 @@ pub fn handle_execute_command(
     let arg: ActionRequest = from_value(arg)?;
     let file_id = arg.text_document.try_conv_with(&path_map)?;
     let file = world.file_syntax(file_id)?;
-    let edit = match arg.id {
-        ActionId::FlipComma => libeditor::flip_comma(&file, arg.offset).map(|edit| edit()),
-        ActionId::AddDerive => libeditor::add_derive(&file, arg.offset).map(|edit| edit()),
+    let action_result = match arg.id {
+        ActionId::FlipComma => libeditor::flip_comma(&file, arg.offset).map(|f| f()),
+        ActionId::AddDerive => libeditor::add_derive(&file, arg.offset).map(|f| f()),
     };
-    let edit = match edit {
-        Some(edit) => edit,
+    let edit = match action_result {
+        Some(action_result) => action_result.edit,
         None => bail!("command not applicable"),
     };
     let line_index = world.file_line_index(file_id)?;