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