11 files changed, 109 insertions, 83 deletions

diff --git a/crates/rust-analyzer/src/config.rs b/crates/rust-analyzer/src/config.rs
index d75c48597..0e4412ade 100644
--- a/crates/rust-analyzer/src/config.rs
+++ b/crates/rust-analyzer/src/config.rs
@@ -102,6 +102,7 @@ pub struct ClientCapsConfig {
     pub hierarchical_symbols: bool,
     pub code_action_literals: bool,
     pub work_done_progress: bool,
+    pub code_action_group: bool,
 }
 
 impl Default for Config {
@@ -294,9 +295,13 @@ impl Config {
 
         self.assist.allow_snippets(false);
         if let Some(experimental) = &caps.experimental {
-            let enable =
+            let snippet_text_edit =
                 experimental.get("snippetTextEdit").and_then(|it| it.as_bool()) == Some(true);
-            self.assist.allow_snippets(enable);
+            self.assist.allow_snippets(snippet_text_edit);
+
+            let code_action_group =
+                experimental.get("codeActionGroup").and_then(|it| it.as_bool()) == Some(true);
+            self.client_caps.code_action_group = code_action_group
         }
     }
 }
diff --git a/crates/rust-analyzer/src/diagnostics/snapshots/rust_analyzer__diagnostics__to_proto__tests__snap_multi_line_fix.snap b/crates/rust-analyzer/src/diagnostics/snapshots/rust_analyzer__diagnostics__to_proto__tests__snap_multi_line_fix.snap
index 96466b5c9..c40cfdcdc 100644
--- a/crates/rust-analyzer/src/diagnostics/snapshots/rust_analyzer__diagnostics__to_proto__tests__snap_multi_line_fix.snap
+++ b/crates/rust-analyzer/src/diagnostics/snapshots/rust_analyzer__diagnostics__to_proto__tests__snap_multi_line_fix.snap
@@ -65,6 +65,7 @@ expression: diag
         fixes: [
             CodeAction {
                 title: "return the expression directly",
+                group: None,
                 kind: Some(
                     "quickfix",
                 ),
diff --git a/crates/rust-analyzer/src/diagnostics/snapshots/rust_analyzer__diagnostics__to_proto__tests__snap_rustc_unused_variable.snap b/crates/rust-analyzer/src/diagnostics/snapshots/rust_analyzer__diagnostics__to_proto__tests__snap_rustc_unused_variable.snap
index 8f962277f..6dd3fcb2e 100644
--- a/crates/rust-analyzer/src/diagnostics/snapshots/rust_analyzer__diagnostics__to_proto__tests__snap_rustc_unused_variable.snap
+++ b/crates/rust-analyzer/src/diagnostics/snapshots/rust_analyzer__diagnostics__to_proto__tests__snap_rustc_unused_variable.snap
@@ -50,6 +50,7 @@ expression: diag
         fixes: [
             CodeAction {
                 title: "consider prefixing with an underscore",
+                group: None,
                 kind: Some(
                     "quickfix",
                 ),
diff --git a/crates/rust-analyzer/src/diagnostics/to_proto.rs b/crates/rust-analyzer/src/diagnostics/to_proto.rs
index afea59525..a500d670a 100644
--- a/crates/rust-analyzer/src/diagnostics/to_proto.rs
+++ b/crates/rust-analyzer/src/diagnostics/to_proto.rs
@@ -145,6 +145,7 @@ fn map_rust_child_diagnostic(
     } else {
         MappedRustChildDiagnostic::SuggestedFix(lsp_ext::CodeAction {
             title: rd.message.clone(),
+            group: None,
             kind: Some("quickfix".to_string()),
             edit: Some(lsp_ext::SnippetWorkspaceEdit {
                 // FIXME: there's no good reason to use edit_map here....
diff --git a/crates/rust-analyzer/src/lsp_ext.rs b/crates/rust-analyzer/src/lsp_ext.rs
index 0fd60caf4..c25d90a50 100644
--- a/crates/rust-analyzer/src/lsp_ext.rs
+++ b/crates/rust-analyzer/src/lsp_ext.rs
@@ -133,14 +133,6 @@ pub struct Runnable {
     pub cwd: Option<PathBuf>,
 }
 
-#[derive(Deserialize, Serialize, Debug)]
-#[serde(rename_all = "camelCase")]
-pub struct SourceChange {
-    pub label: String,
-    pub workspace_edit: SnippetWorkspaceEdit,
-    pub cursor_position: Option<lsp_types::TextDocumentPositionParams>,
-}
-
 pub enum InlayHints {}
 
 impl Request for InlayHints {
@@ -196,6 +188,8 @@ impl Request for CodeActionRequest {
 pub struct CodeAction {
     pub title: String,
     #[serde(skip_serializing_if = "Option::is_none")]
+    pub group: Option<String>,
+    #[serde(skip_serializing_if = "Option::is_none")]
     pub kind: Option<String>,
     #[serde(skip_serializing_if = "Option::is_none")]
     pub command: Option<lsp_types::Command>,
diff --git a/crates/rust-analyzer/src/main_loop/handlers.rs b/crates/rust-analyzer/src/main_loop/handlers.rs
index 25e660bd5..89144f743 100644
--- a/crates/rust-analyzer/src/main_loop/handlers.rs
+++ b/crates/rust-analyzer/src/main_loop/handlers.rs
@@ -18,7 +18,7 @@ use lsp_types::{
     SemanticTokensResult, SymbolInformation, TextDocumentIdentifier, Url, WorkspaceEdit,
 };
 use ra_ide::{
-    Assist, FileId, FilePosition, FileRange, Query, RangeInfo, Runnable, RunnableKind, SearchScope,
+    FileId, FilePosition, FileRange, Query, RangeInfo, Runnable, RunnableKind, SearchScope,
     TextEdit,
 };
 use ra_prof::profile;
@@ -720,6 +720,7 @@ pub fn handle_code_action(
     let file_id = from_proto::file_id(&world, &params.text_document.uri)?;
     let line_index = world.analysis().file_line_index(file_id)?;
     let range = from_proto::text_range(&line_index, params.range);
+    let frange = FileRange { file_id, range };
 
     let diagnostics = world.analysis().diagnostics(file_id)?;
     let mut res: Vec<lsp_ext::CodeAction> = Vec::new();
@@ -733,7 +734,8 @@ pub fn handle_code_action(
     for source_edit in fixes_from_diagnostics {
         let title = source_edit.label.clone();
         let edit = to_proto::snippet_workspace_edit(&world, source_edit)?;
-        let action = lsp_ext::CodeAction { title, kind: None, edit: Some(edit), command: None };
+        let action =
+            lsp_ext::CodeAction { title, group: None, kind: None, edit: Some(edit), command: None };
         res.push(action);
     }
 
@@ -745,53 +747,9 @@ pub fn handle_code_action(
         res.push(fix.action.clone());
     }
 
-    let mut grouped_assists: FxHashMap<String, (usize, Vec<Assist>)> = FxHashMap::default();
-    for assist in
-        world.analysis().assists(&world.config.assist, FileRange { file_id, range })?.into_iter()
-    {
-        match &assist.group_label {
-            Some(label) => grouped_assists
-                .entry(label.to_owned())
-                .or_insert_with(|| {
-                    let idx = res.len();
-                    let dummy = lsp_ext::CodeAction {
-                        title: String::new(),
-                        kind: None,
-                        command: None,
-                        edit: None,
-                    };
-                    res.push(dummy);
-                    (idx, Vec::new())
-                })
-                .1
-                .push(assist),
-            None => {
-                res.push(to_proto::code_action(&world, assist)?.into());
-            }
-        }
-    }
-
-    for (group_label, (idx, assists)) in grouped_assists {
-        if assists.len() == 1 {
-            res[idx] = to_proto::code_action(&world, assists.into_iter().next().unwrap())?.into();
-        } else {
-            let title = group_label;
-
-            let mut arguments = Vec::with_capacity(assists.len());
-            for assist in assists {
-                let source_change = to_proto::source_change(&world, assist.source_change)?;
-                arguments.push(to_value(source_change)?);
-            }
-
-            let command = Some(Command {
-                title: title.clone(),
-                command: "rust-analyzer.selectAndApplySourceChange".to_string(),
-                arguments: Some(vec![serde_json::Value::Array(arguments)]),
-            });
-            res[idx] = lsp_ext::CodeAction { title, kind: None, edit: None, command };
-        }
+    for assist in world.analysis().assists(&world.config.assist, frange)?.into_iter() {
+        res.push(to_proto::code_action(&world, assist)?.into());
     }
-
     Ok(Some(res))
 }
 
diff --git a/crates/rust-analyzer/src/to_proto.rs b/crates/rust-analyzer/src/to_proto.rs
index f6f4bb134..461944ada 100644
--- a/crates/rust-analyzer/src/to_proto.rs
+++ b/crates/rust-analyzer/src/to_proto.rs
@@ -478,15 +478,6 @@ pub(crate) fn resource_op(
     Ok(res)
 }
 
-pub(crate) fn source_change(
-    world: &WorldSnapshot,
-    source_change: SourceChange,
-) -> Result<lsp_ext::SourceChange> {
-    let label = source_change.label.clone();
-    let workspace_edit = self::snippet_workspace_edit(world, source_change)?;
-    Ok(lsp_ext::SourceChange { label, workspace_edit, cursor_position: None })
-}
-
 pub(crate) fn snippet_workspace_edit(
     world: &WorldSnapshot,
     source_change: SourceChange,
@@ -606,6 +597,7 @@ fn main() <fold>{
 pub(crate) fn code_action(world: &WorldSnapshot, assist: Assist) -> Result<lsp_ext::CodeAction> {
     let res = lsp_ext::CodeAction {
         title: assist.label,
+        group: if world.config.client_caps.code_action_group { assist.group_label } else { None },
         kind: Some(String::new()),
         edit: Some(snippet_workspace_edit(world, assist.source_change)?),
         command: None,
diff --git a/docs/dev/lsp-extensions.md b/docs/dev/lsp-extensions.md
index 7c45aef4c..d90875f8b 100644
--- a/docs/dev/lsp-extensions.md
+++ b/docs/dev/lsp-extensions.md
@@ -5,7 +5,7 @@ It's a best effort document, when in doubt, consult the source (and send a PR wi
 We aim to upstream all non Rust-specific extensions to the protocol, but this is not a top priority.
 All capabilities are enabled via `experimental` field of `ClientCapabilities`.
 
-## `SnippetTextEdit`
+## Snippet `TextEdit`
 
 **Client Capability:** `{ "snippetTextEdit": boolean }`
 
@@ -36,7 +36,7 @@ At the moment, rust-analyzer guarantees that only a single edit will have `Inser
 * Where exactly are `SnippetTextEdit`s allowed (only in code actions at the moment)?
 * Can snippets span multiple files (so far, no)?
 
-## `joinLines`
+## Join Lines
 
 **Server Capability:** `{ "joinLines": boolean }`
 
@@ -119,3 +119,48 @@ SSR with query `foo($a:expr, $b:expr) ==>> ($a).foo($b)` will transform, eg `foo
 
 * Probably needs search without replace mode
 * Needs a way to limit the scope to certain files.
+
+## `CodeAction` Groups
+
+**Client Capability:** `{ "codeActionGroup": boolean }`
+
+If this capability is set, `CodeAction` returned from the server contain an additional field, `group`:
+
+```typescript
+interface CodeAction {
+        title: string;
+    group?: string;
+    ...
+}
+```
+
+All code-actions with the same `group` should be grouped under single (extendable) entry in lightbulb menu.
+The set of actions `[ { title: "foo" }, { group: "frobnicate", title: "bar" }, { group: "frobnicate", title: "baz" }]` should be rendered as
+
+```
+💡
+  +-------------+
+  | foo         |
+  +-------------+-----+
+  | frobnicate >| bar |
+  +-------------+-----+
+                | baz |
+                +-----+
+```
+
+Alternatively, selecting `frobnicate` could present a user with an additional menu to choose between `bar` and `baz`.
+
+### Example
+
+```rust
+fn main() {
+    let x: Entry/*cursor here*/ = todo!();
+}
+```
+
+Invoking code action at this position will yield two code actions for importing `Entry` from either `collections::HashMap` or `collection::BTreeMap`, grouped under a single "import" group.
+
+### Unresolved Questions
+
+* Is a fixed two-level structure enough?
+* Should we devise a general way to encode custom interaction protocols for GUI refactorings?
diff --git a/editors/code/src/client.ts b/editors/code/src/client.ts
index fac1a0be3..d64f9a3f9 100644
--- a/editors/code/src/client.ts
+++ b/editors/code/src/client.ts
@@ -41,10 +41,12 @@ export function createClient(serverPath: string, cwd: string): lc.LanguageClient
                 return client.sendRequest(lc.CodeActionRequest.type, params, token).then((values) => {
                     if (values === null) return undefined;
                     const result: (vscode.CodeAction | vscode.Command)[] = [];
+                    const groups = new Map<string, { index: number; items: vscode.CodeAction[] }>();
                     for (const item of values) {
                         if (lc.CodeAction.is(item)) {
                             const action = client.protocol2CodeConverter.asCodeAction(item);
-                            if (isSnippetEdit(item)) {
+                            const group = actionGroup(item);
+                            if (isSnippetEdit(item) || group) {
                                 action.command = {
                                     command: "rust-analyzer.applySnippetWorkspaceEdit",
                                     title: "",
@@ -52,12 +54,38 @@ export function createClient(serverPath: string, cwd: string): lc.LanguageClient
                                 };
                                 action.edit = undefined;
                             }
-                            result.push(action);
+
+                            if (group) {
+                                let entry = groups.get(group);
+                                if (!entry) {
+                                    entry = { index: result.length, items: [] };
+                                    groups.set(group, entry);
+                                    result.push(action);
+                                }
+                                entry.items.push(action);
+                            } else {
+                                result.push(action);
+                            }
                         } else {
                             const command = client.protocol2CodeConverter.asCommand(item);
                             result.push(command);
                         }
                     }
+                    for (const [group, { index, items }] of groups) {
+                        if (items.length === 1) {
+                            result[index] = items[0];
+                        } else {
+                            const action = new vscode.CodeAction(group);
+                            action.command = {
+                                command: "rust-analyzer.applyActionGroup",
+                                title: "",
+                                arguments: [items.map((item) => {
+                                    return { label: item.title, edit: item.command!!.arguments!![0] };
+                                })],
+                            };
+                            result[index] = action;
+                        }
+                    }
                     return result;
                 },
                     (_error) => undefined
@@ -81,15 +109,16 @@ export function createClient(serverPath: string, cwd: string): lc.LanguageClient
     // implementations are still in the "proposed" category for 3.16.
     client.registerFeature(new CallHierarchyFeature(client));
     client.registerFeature(new SemanticTokensFeature(client));
-    client.registerFeature(new SnippetTextEditFeature());
+    client.registerFeature(new ExperimentalFeatures());
 
     return client;
 }
 
-class SnippetTextEditFeature implements lc.StaticFeature {
+class ExperimentalFeatures implements lc.StaticFeature {
     fillClientCapabilities(capabilities: lc.ClientCapabilities): void {
         const caps: any = capabilities.experimental ?? {};
         caps.snippetTextEdit = true;
+        caps.codeActionGroup = true;
         capabilities.experimental = caps;
     }
     initialize(_capabilities: lc.ServerCapabilities<any>, _documentSelector: lc.DocumentSelector | undefined): void {
@@ -107,3 +136,7 @@ function isSnippetEdit(action: lc.CodeAction): boolean {
     }
     return false;
 }
+
+function actionGroup(action: lc.CodeAction): string | undefined {
+    return (action as any).group;
+}
diff --git a/editors/code/src/commands/index.ts b/editors/code/src/commands/index.ts
index e5ed77e32..abb53a248 100644
--- a/editors/code/src/commands/index.ts
+++ b/editors/code/src/commands/index.ts
@@ -41,15 +41,11 @@ export function applySourceChange(ctx: Ctx): Cmd {
     };
 }
 
-export function selectAndApplySourceChange(ctx: Ctx): Cmd {
-    return async (changes: ra.SourceChange[]) => {
-        if (changes.length === 1) {
-            await sourceChange.applySourceChange(ctx, changes[0]);
-        } else if (changes.length > 0) {
-            const selectedChange = await vscode.window.showQuickPick(changes);
-            if (!selectedChange) return;
-            await sourceChange.applySourceChange(ctx, selectedChange);
-        }
+export function applyActionGroup(_ctx: Ctx): Cmd {
+    return async (actions: { label: string; edit: vscode.WorkspaceEdit }[]) => {
+        const selectedAction = await vscode.window.showQuickPick(actions);
+        if (!selectedAction) return;
+        await applySnippetWorkspaceEdit(selectedAction.edit);
     };
 }
 
diff --git a/editors/code/src/main.ts b/editors/code/src/main.ts
index 8b0a9d870..4d4513869 100644
--- a/editors/code/src/main.ts
+++ b/editors/code/src/main.ts
@@ -92,7 +92,7 @@ export async function activate(context: vscode.ExtensionContext) {
     ctx.registerCommand('showReferences', commands.showReferences);
     ctx.registerCommand('applySourceChange', commands.applySourceChange);
     ctx.registerCommand('applySnippetWorkspaceEdit', commands.applySnippetWorkspaceEditCommand);
-    ctx.registerCommand('selectAndApplySourceChange', commands.selectAndApplySourceChange);
+    ctx.registerCommand('applyActionGroup', commands.applyActionGroup);
 
     ctx.pushCleanup(activateTaskProvider(workspaceFolder));