1use std::{collections::HashMap as StdHashMap, ops::Range, sync::Arc};
2
3use collections::{HashMap, HashSet};
4use futures::future::join_all;
5use gpui::{MouseButton, SharedString, Task, WeakEntity};
6use itertools::Itertools;
7use language::BufferId;
8use multi_buffer::{Anchor, MultiBufferRow, MultiBufferSnapshot, ToPoint as _};
9use project::{CodeAction, LspAction, TaskSourceKind, lsp_store::lsp_ext_command};
10use task::TaskContext;
11use text::Point;
12
13use ui::{Context, Window, div, prelude::*};
14
15use crate::{
16 Editor, LSP_REQUEST_DEBOUNCE_TIMEOUT, MultibufferSelectionMode, SelectionEffects,
17 actions::ToggleCodeLens,
18 display_map::{BlockPlacement, BlockProperties, BlockStyle, CustomBlockId},
19};
20
21#[derive(Clone, Debug)]
22struct CodeLensLine {
23 position: Anchor,
24 items: Vec<CodeLensItem>,
25}
26
27#[derive(Clone, Debug)]
28struct CodeLensItem {
29 title: SharedString,
30 action: CodeAction,
31}
32
33#[derive(Default)]
34pub(super) struct CodeLensState {
35 pub(super) block_ids: HashMap<BufferId, Vec<CustomBlockId>>,
36}
37
38impl CodeLensState {
39 fn all_block_ids(&self) -> HashSet<CustomBlockId> {
40 self.block_ids.values().flatten().copied().collect()
41 }
42}
43
44fn group_lenses_by_row(
45 lenses: Vec<(Anchor, CodeLensItem)>,
46 snapshot: &MultiBufferSnapshot,
47) -> Vec<CodeLensLine> {
48 let mut grouped: HashMap<MultiBufferRow, (Anchor, Vec<CodeLensItem>)> = HashMap::default();
49
50 for (position, item) in lenses {
51 let row = position.to_point(snapshot).row;
52 grouped
53 .entry(MultiBufferRow(row))
54 .or_insert_with(|| (position, Vec::new()))
55 .1
56 .push(item);
57 }
58
59 let mut result: Vec<CodeLensLine> = grouped
60 .into_iter()
61 .map(|(_, (position, items))| CodeLensLine { position, items })
62 .collect();
63
64 result.sort_by_key(|lens| lens.position.to_point(snapshot).row);
65 result
66}
67
68fn render_code_lens_line(
69 line_number: usize,
70 lens: CodeLensLine,
71 editor: WeakEntity<Editor>,
72) -> impl Fn(&mut crate::display_map::BlockContext) -> gpui::AnyElement {
73 move |cx| {
74 let mut children: Vec<gpui::AnyElement> = Vec::new();
75 let text_style = &cx.editor_style.text;
76 let font = text_style.font();
77 let font_size = text_style.font_size.to_pixels(cx.window.rem_size()) * 0.9;
78
79 for (i, item) in lens.items.iter().enumerate() {
80 if i > 0 {
81 children.push(
82 div()
83 .font(font.clone())
84 .text_size(font_size)
85 .text_color(cx.app.theme().colors().text_muted)
86 .child(" | ")
87 .into_any_element(),
88 );
89 }
90
91 let title = item.title.clone();
92 let action = item.action.clone();
93 let editor_handle = editor.clone();
94 let position = lens.position;
95 let id = (line_number as u64) << 32 | (i as u64);
96
97 children.push(
98 div()
99 .id(ElementId::Integer(id))
100 .font(font.clone())
101 .text_size(font_size)
102 .text_color(cx.app.theme().colors().text_muted)
103 .cursor_pointer()
104 .hover(|style| style.text_color(cx.app.theme().colors().text))
105 .child(title.clone())
106 .on_mouse_down(MouseButton::Left, |_, _, cx| {
107 cx.stop_propagation();
108 })
109 .on_mouse_down(MouseButton::Right, |_, _, cx| {
110 cx.stop_propagation();
111 })
112 .on_click({
113 move |_event, window, cx| {
114 if let Some(editor) = editor_handle.upgrade() {
115 editor.update(cx, |editor, cx| {
116 editor.change_selections(
117 SelectionEffects::default(),
118 window,
119 cx,
120 |s| {
121 s.select_anchor_ranges([position..position]);
122 },
123 );
124
125 let action = action.clone();
126 if let Some(workspace) = editor.workspace() {
127 if try_handle_client_command(
128 &action, editor, &workspace, window, cx,
129 ) {
130 return;
131 }
132
133 let project = workspace.read(cx).project().clone();
134 let buffer = editor.buffer().clone();
135 if let Some(excerpt_buffer) = buffer.read(cx).as_singleton()
136 {
137 project
138 .update(cx, |project, cx| {
139 project.apply_code_action(
140 excerpt_buffer.clone(),
141 action,
142 true,
143 cx,
144 )
145 })
146 .detach_and_log_err(cx);
147 }
148 }
149 });
150 }
151 }
152 })
153 .into_any_element(),
154 );
155 }
156
157 div()
158 .pl(cx.margins.gutter.full_width())
159 .h_full()
160 .flex()
161 .flex_row()
162 .items_end()
163 .children(children)
164 .into_any_element()
165 }
166}
167
168fn try_handle_client_command(
169 action: &CodeAction,
170 editor: &mut Editor,
171 workspace: &gpui::Entity<workspace::Workspace>,
172 window: &mut Window,
173 cx: &mut Context<Editor>,
174) -> bool {
175 let command = match &action.lsp_action {
176 LspAction::CodeLens(lens) => lens.command.as_ref(),
177 _ => None,
178 };
179 let Some(command) = command else {
180 return false;
181 };
182 let arguments = command.arguments.as_deref().unwrap_or_default();
183
184 match command.command.as_str() {
185 "rust-analyzer.runSingle" => {
186 try_schedule_runnable(arguments, action, editor, workspace, window, cx)
187 }
188 "rust-analyzer.showReferences" => {
189 try_show_references(arguments, action, editor, workspace, window, cx)
190 }
191 _ => false,
192 }
193}
194
195fn try_schedule_runnable(
196 arguments: &[serde_json::Value],
197 action: &CodeAction,
198 editor: &Editor,
199 workspace: &gpui::Entity<workspace::Workspace>,
200 window: &mut Window,
201 cx: &mut Context<Editor>,
202) -> bool {
203 let Some(first_arg) = arguments.first() else {
204 return false;
205 };
206 let Ok(runnable) = serde_json::from_value::<lsp_ext_command::Runnable>(first_arg.clone())
207 else {
208 return false;
209 };
210
211 let task_template = lsp_ext_command::runnable_to_task_template(runnable.label, runnable.args);
212 let task_context = TaskContext {
213 cwd: task_template.cwd.as_ref().map(std::path::PathBuf::from),
214 ..TaskContext::default()
215 };
216 let language_name = editor
217 .buffer()
218 .read(cx)
219 .as_singleton()
220 .and_then(|buffer| buffer.read(cx).language())
221 .map(|language| language.name());
222 let task_source_kind = match language_name {
223 Some(language_name) => TaskSourceKind::Lsp {
224 server: action.server_id,
225 language_name: SharedString::from(language_name),
226 },
227 None => TaskSourceKind::AbsPath {
228 id_base: "code-lens".into(),
229 abs_path: task_template
230 .cwd
231 .as_ref()
232 .map(std::path::PathBuf::from)
233 .unwrap_or_default(),
234 },
235 };
236
237 workspace.update(cx, |workspace, cx| {
238 workspace.schedule_task(
239 task_source_kind,
240 &task_template,
241 &task_context,
242 false,
243 window,
244 cx,
245 );
246 });
247 true
248}
249
250fn try_show_references(
251 arguments: &[serde_json::Value],
252 action: &CodeAction,
253 _editor: &mut Editor,
254 workspace: &gpui::Entity<workspace::Workspace>,
255 window: &mut Window,
256 cx: &mut Context<Editor>,
257) -> bool {
258 if arguments.len() < 3 {
259 return false;
260 }
261 let Ok(locations) = serde_json::from_value::<Vec<lsp::Location>>(arguments[2].clone()) else {
262 return false;
263 };
264 if locations.is_empty() {
265 return false;
266 }
267
268 let server_id = action.server_id;
269 let project = workspace.read(cx).project().clone();
270 let workspace = workspace.clone();
271
272 cx.spawn_in(window, async move |_editor, cx| {
273 let mut buffer_locations: StdHashMap<gpui::Entity<language::Buffer>, Vec<Range<Point>>> =
274 StdHashMap::default();
275
276 for location in &locations {
277 let open_task = cx.update(|_, cx| {
278 project.update(cx, |project, cx| {
279 let uri: lsp::Uri = location.uri.clone();
280 project.open_local_buffer_via_lsp(uri, server_id, cx)
281 })
282 })?;
283 let buffer = open_task.await?;
284
285 let range = range_from_lsp(location.range);
286 buffer_locations.entry(buffer).or_default().push(range);
287 }
288
289 workspace.update_in(cx, |workspace, window, cx| {
290 Editor::open_locations_in_multibuffer(
291 workspace,
292 buffer_locations,
293 "References".to_owned(),
294 false,
295 true,
296 MultibufferSelectionMode::First,
297 window,
298 cx,
299 );
300 })?;
301 anyhow::Ok(())
302 })
303 .detach_and_log_err(cx);
304
305 true
306}
307
308fn range_from_lsp(range: lsp::Range) -> Range<Point> {
309 let start = Point::new(range.start.line, range.start.character);
310 let end = Point::new(range.end.line, range.end.character);
311 start..end
312}
313
314impl Editor {
315 pub(super) fn refresh_code_lenses(
316 &mut self,
317 for_buffer: Option<BufferId>,
318 _window: &Window,
319 cx: &mut Context<Self>,
320 ) {
321 if !self.lsp_data_enabled() || self.code_lens.is_none() {
322 return;
323 }
324 let Some(project) = self.project.clone() else {
325 return;
326 };
327
328 let buffers_to_query = self
329 .visible_buffers(cx)
330 .into_iter()
331 .filter(|buffer| self.is_lsp_relevant(buffer.read(cx).file(), cx))
332 .chain(for_buffer.and_then(|buffer_id| self.buffer.read(cx).buffer(buffer_id)))
333 .filter(|editor_buffer| {
334 let editor_buffer_id = editor_buffer.read(cx).remote_id();
335 for_buffer.is_none_or(|buffer_id| buffer_id == editor_buffer_id)
336 && self.registered_buffers.contains_key(&editor_buffer_id)
337 })
338 .unique_by(|buffer| buffer.read(cx).remote_id())
339 .collect::<Vec<_>>();
340
341 if buffers_to_query.is_empty() {
342 return;
343 }
344
345 let project = project.downgrade();
346 self.refresh_code_lens_task = cx.spawn(async move |editor, cx| {
347 cx.background_executor()
348 .timer(LSP_REQUEST_DEBOUNCE_TIMEOUT)
349 .await;
350
351 let Some(tasks) = project
352 .update(cx, |project, cx| {
353 project.lsp_store().update(cx, |lsp_store, cx| {
354 buffers_to_query
355 .into_iter()
356 .map(|buffer| {
357 let buffer_id = buffer.read(cx).remote_id();
358 let task = lsp_store.code_lens_actions(&buffer, cx);
359 async move { (buffer_id, task.await) }
360 })
361 .collect::<Vec<_>>()
362 })
363 })
364 .ok()
365 else {
366 return;
367 };
368
369 let results = join_all(tasks).await;
370 if results.is_empty() {
371 return;
372 }
373
374 let Ok(multi_buffer_snapshot) =
375 editor.update(cx, |editor, cx| editor.buffer().read(cx).snapshot(cx))
376 else {
377 return;
378 };
379
380 let mut new_lenses_per_buffer: HashMap<BufferId, Vec<CodeLensLine>> =
381 HashMap::default();
382
383 for (buffer_id, result) in results {
384 match result {
385 Ok(Some(actions)) => {
386 let individual_lenses: Vec<(Anchor, CodeLensItem)> = actions
387 .into_iter()
388 .filter_map(|action| {
389 let title = match &action.lsp_action {
390 project::LspAction::CodeLens(lens) => lens
391 .command
392 .as_ref()
393 .map(|cmd| SharedString::from(&cmd.title)),
394 _ => None,
395 }?;
396 let position =
397 multi_buffer_snapshot.anchor_in_excerpt(action.range.start)?;
398 Some((position, CodeLensItem { title, action }))
399 })
400 .collect();
401
402 let grouped =
403 group_lenses_by_row(individual_lenses, &multi_buffer_snapshot);
404 new_lenses_per_buffer.insert(buffer_id, grouped);
405 }
406 Ok(None) => {}
407 Err(e) => {
408 log::error!("Failed to fetch code lenses for buffer {buffer_id:?}: {e:#}");
409 }
410 }
411 }
412
413 editor
414 .update(cx, |editor, cx| {
415 let code_lens = editor.code_lens.get_or_insert_with(CodeLensState::default);
416
417 let mut blocks_to_remove: HashSet<CustomBlockId> = HashSet::default();
418 for (buffer_id, _) in &new_lenses_per_buffer {
419 if let Some(old_ids) = code_lens.block_ids.remove(buffer_id) {
420 blocks_to_remove.extend(old_ids);
421 }
422 }
423
424 if !blocks_to_remove.is_empty() {
425 editor.remove_blocks(blocks_to_remove, None, cx);
426 }
427
428 let editor_handle = cx.entity().downgrade();
429
430 let mut all_new_blocks: Vec<(BufferId, Vec<BlockProperties<Anchor>>)> =
431 Vec::new();
432 for (buffer_id, lense_lines) in new_lenses_per_buffer {
433 if lense_lines.is_empty() {
434 continue;
435 }
436 let blocks: Vec<BlockProperties<Anchor>> = lense_lines
437 .into_iter()
438 .enumerate()
439 .map(|(line_number, lens_line)| {
440 let position = lens_line.position;
441 let render_fn = render_code_lens_line(
442 line_number,
443 lens_line,
444 editor_handle.clone(),
445 );
446 BlockProperties {
447 placement: BlockPlacement::Above(position),
448 height: Some(1),
449 style: BlockStyle::Flex,
450 render: Arc::new(render_fn),
451 priority: 0,
452 }
453 })
454 .collect();
455 all_new_blocks.push((buffer_id, blocks));
456 }
457
458 for (buffer_id, blocks) in all_new_blocks {
459 let block_ids = editor.insert_blocks(blocks, None, cx);
460 editor
461 .code_lens
462 .get_or_insert_with(CodeLensState::default)
463 .block_ids
464 .insert(buffer_id, block_ids);
465 }
466
467 cx.notify();
468 })
469 .ok();
470 });
471 }
472
473 pub fn supports_code_lens(&self, cx: &ui::App) -> bool {
474 let Some(project) = self.project.as_ref() else {
475 return false;
476 };
477 let lsp_store = project.read(cx).lsp_store().read(cx);
478 lsp_store
479 .lsp_server_capabilities
480 .values()
481 .any(|caps| caps.code_lens_provider.is_some())
482 }
483
484 pub fn code_lens_enabled(&self) -> bool {
485 self.code_lens.is_some()
486 }
487
488 pub fn toggle_code_lens_action(
489 &mut self,
490 _: &ToggleCodeLens,
491 window: &mut Window,
492 cx: &mut Context<Self>,
493 ) {
494 let currently_enabled = self.code_lens.is_some();
495 self.toggle_code_lens(!currently_enabled, window, cx);
496 }
497
498 pub(super) fn toggle_code_lens(
499 &mut self,
500 enabled: bool,
501 window: &mut Window,
502 cx: &mut Context<Self>,
503 ) {
504 if enabled {
505 self.code_lens.get_or_insert_with(CodeLensState::default);
506 self.refresh_code_lenses(None, window, cx);
507 } else {
508 self.clear_code_lenses(cx);
509 }
510 }
511
512 pub(super) fn clear_code_lenses(&mut self, cx: &mut Context<Self>) {
513 if let Some(code_lens) = self.code_lens.take() {
514 let all_blocks = code_lens.all_block_ids();
515 if !all_blocks.is_empty() {
516 self.remove_blocks(all_blocks, None, cx);
517 }
518 cx.notify();
519 }
520 self.refresh_code_lens_task = Task::ready(());
521 }
522}
523
524#[cfg(test)]
525mod tests {
526 use futures::StreamExt;
527 use gpui::TestAppContext;
528
529 use settings::CodeLens;
530
531 use crate::{
532 editor_tests::{init_test, update_test_editor_settings},
533 test::editor_lsp_test_context::EditorLspTestContext,
534 };
535
536 #[gpui::test]
537 async fn test_code_lens_blocks(cx: &mut TestAppContext) {
538 init_test(cx, |_| {});
539 update_test_editor_settings(cx, &|settings| {
540 settings.code_lens = Some(CodeLens::On);
541 });
542
543 let mut cx = EditorLspTestContext::new_typescript(
544 lsp::ServerCapabilities {
545 code_lens_provider: Some(lsp::CodeLensOptions {
546 resolve_provider: None,
547 }),
548 execute_command_provider: Some(lsp::ExecuteCommandOptions {
549 commands: vec!["lens_cmd".to_string()],
550 ..lsp::ExecuteCommandOptions::default()
551 }),
552 ..lsp::ServerCapabilities::default()
553 },
554 cx,
555 )
556 .await;
557
558 let mut code_lens_request =
559 cx.set_request_handler::<lsp::request::CodeLensRequest, _, _>(move |_, _, _| async {
560 Ok(Some(vec![
561 lsp::CodeLens {
562 range: lsp::Range::new(lsp::Position::new(0, 0), lsp::Position::new(0, 19)),
563 command: Some(lsp::Command {
564 title: "2 references".to_owned(),
565 command: "lens_cmd".to_owned(),
566 arguments: None,
567 }),
568 data: None,
569 },
570 lsp::CodeLens {
571 range: lsp::Range::new(lsp::Position::new(1, 0), lsp::Position::new(1, 19)),
572 command: Some(lsp::Command {
573 title: "0 references".to_owned(),
574 command: "lens_cmd".to_owned(),
575 arguments: None,
576 }),
577 data: None,
578 },
579 ]))
580 });
581
582 cx.set_state("ˇfunction hello() {}\nfunction world() {}");
583
584 assert!(
585 code_lens_request.next().await.is_some(),
586 "should have received a code lens request"
587 );
588 cx.run_until_parked();
589
590 cx.editor.read_with(&cx.cx.cx, |editor, _cx| {
591 assert_eq!(
592 editor.code_lens_enabled(),
593 true,
594 "code lens should be enabled"
595 );
596 let total_blocks: usize = editor
597 .code_lens
598 .as_ref()
599 .map(|s| s.block_ids.values().map(|v| v.len()).sum())
600 .unwrap_or(0);
601 assert_eq!(total_blocks, 2, "Should have inserted two code lens blocks");
602 });
603 }
604
605 #[gpui::test]
606 async fn test_code_lens_disabled_by_default(cx: &mut TestAppContext) {
607 init_test(cx, |_| {});
608
609 let mut cx = EditorLspTestContext::new_typescript(
610 lsp::ServerCapabilities {
611 code_lens_provider: Some(lsp::CodeLensOptions {
612 resolve_provider: None,
613 }),
614 execute_command_provider: Some(lsp::ExecuteCommandOptions {
615 commands: vec!["lens_cmd".to_string()],
616 ..lsp::ExecuteCommandOptions::default()
617 }),
618 ..lsp::ServerCapabilities::default()
619 },
620 cx,
621 )
622 .await;
623
624 cx.lsp
625 .set_request_handler::<lsp::request::CodeLensRequest, _, _>(|_, _| async move {
626 panic!("Should not request code lenses when disabled");
627 });
628
629 cx.set_state("ˇfunction hello() {}");
630 cx.run_until_parked();
631
632 cx.editor.read_with(&cx.cx.cx, |editor, _cx| {
633 assert_eq!(
634 editor.code_lens_enabled(),
635 false,
636 "code lens should not be enabled when setting is off"
637 );
638 });
639 }
640
641 #[gpui::test]
642 async fn test_code_lens_toggling(cx: &mut TestAppContext) {
643 init_test(cx, |_| {});
644 update_test_editor_settings(cx, &|settings| {
645 settings.code_lens = Some(CodeLens::On);
646 });
647
648 let mut cx = EditorLspTestContext::new_typescript(
649 lsp::ServerCapabilities {
650 code_lens_provider: Some(lsp::CodeLensOptions {
651 resolve_provider: None,
652 }),
653 execute_command_provider: Some(lsp::ExecuteCommandOptions {
654 commands: vec!["lens_cmd".to_string()],
655 ..lsp::ExecuteCommandOptions::default()
656 }),
657 ..lsp::ServerCapabilities::default()
658 },
659 cx,
660 )
661 .await;
662
663 let mut code_lens_request =
664 cx.set_request_handler::<lsp::request::CodeLensRequest, _, _>(move |_, _, _| async {
665 Ok(Some(vec![lsp::CodeLens {
666 range: lsp::Range::new(lsp::Position::new(0, 0), lsp::Position::new(0, 19)),
667 command: Some(lsp::Command {
668 title: "1 reference".to_owned(),
669 command: "lens_cmd".to_owned(),
670 arguments: None,
671 }),
672 data: None,
673 }]))
674 });
675
676 cx.set_state("ˇfunction hello() {}");
677
678 assert!(
679 code_lens_request.next().await.is_some(),
680 "should have received a code lens request"
681 );
682 cx.run_until_parked();
683
684 cx.editor.read_with(&cx.cx.cx, |editor, _cx| {
685 assert_eq!(
686 editor.code_lens_enabled(),
687 true,
688 "code lens should be enabled"
689 );
690 let total_blocks: usize = editor
691 .code_lens
692 .as_ref()
693 .map(|s| s.block_ids.values().map(|v| v.len()).sum())
694 .unwrap_or(0);
695 assert_eq!(total_blocks, 1, "Should have one code lens block");
696 });
697
698 cx.update_editor(|editor, _window, cx| {
699 editor.clear_code_lenses(cx);
700 });
701
702 cx.editor.read_with(&cx.cx.cx, |editor, _cx| {
703 assert_eq!(
704 editor.code_lens_enabled(),
705 false,
706 "code lens should be disabled after clearing"
707 );
708 });
709 }
710
711 #[gpui::test]
712 async fn test_code_lens_resolve(cx: &mut TestAppContext) {
713 init_test(cx, |_| {});
714 update_test_editor_settings(cx, &|settings| {
715 settings.code_lens = Some(CodeLens::On);
716 });
717
718 let mut cx = EditorLspTestContext::new_typescript(
719 lsp::ServerCapabilities {
720 code_lens_provider: Some(lsp::CodeLensOptions {
721 resolve_provider: Some(true),
722 }),
723 ..lsp::ServerCapabilities::default()
724 },
725 cx,
726 )
727 .await;
728
729 let mut code_lens_request =
730 cx.set_request_handler::<lsp::request::CodeLensRequest, _, _>(move |_, _, _| async {
731 Ok(Some(vec![
732 lsp::CodeLens {
733 range: lsp::Range::new(lsp::Position::new(0, 0), lsp::Position::new(0, 19)),
734 command: None,
735 data: Some(serde_json::json!({"id": "lens_1"})),
736 },
737 lsp::CodeLens {
738 range: lsp::Range::new(lsp::Position::new(1, 0), lsp::Position::new(1, 19)),
739 command: None,
740 data: Some(serde_json::json!({"id": "lens_2"})),
741 },
742 ]))
743 });
744
745 cx.lsp
746 .set_request_handler::<lsp::request::CodeLensResolve, _, _>(|lens, _| async move {
747 let id = lens
748 .data
749 .as_ref()
750 .and_then(|d| d.get("id"))
751 .and_then(|v| v.as_str())
752 .unwrap_or("unknown");
753 let title = match id {
754 "lens_1" => "3 references",
755 "lens_2" => "1 implementation",
756 _ => "unknown",
757 };
758 Ok(lsp::CodeLens {
759 command: Some(lsp::Command {
760 title: title.to_owned(),
761 command: format!("resolved_{id}"),
762 arguments: None,
763 }),
764 ..lens
765 })
766 });
767
768 cx.set_state("ˇfunction hello() {}\nfunction world() {}");
769
770 assert!(
771 code_lens_request.next().await.is_some(),
772 "should have received a code lens request"
773 );
774 cx.run_until_parked();
775
776 cx.editor.read_with(&cx.cx.cx, |editor, _cx| {
777 let total_blocks: usize = editor
778 .code_lens
779 .as_ref()
780 .map(|s| s.block_ids.values().map(|v| v.len()).sum())
781 .unwrap_or(0);
782 assert_eq!(
783 total_blocks, 2,
784 "Unresolved lenses should have been resolved and displayed"
785 );
786 });
787 }
788}