1use std::{cmp, ops::Range};
2
3use collections::HashMap;
4use futures::future::join_all;
5use gpui::{Hsla, Rgba};
6use itertools::Itertools;
7use language::point_from_lsp;
8use multi_buffer::Anchor;
9use project::{DocumentColor, InlayId};
10use settings::Settings as _;
11use text::{Bias, BufferId};
12use ui::{App, Context, Window};
13use util::post_inc;
14
15use crate::{
16 DisplayPoint, Editor, EditorSettings, EditorSnapshot, InlaySplice,
17 LSP_REQUEST_DEBOUNCE_TIMEOUT, RangeToAnchorExt, editor_settings::DocumentColorsRenderMode,
18 inlays::Inlay,
19};
20
21#[derive(Debug)]
22pub(super) struct LspColorData {
23 buffer_colors: HashMap<BufferId, BufferColors>,
24 render_mode: DocumentColorsRenderMode,
25}
26
27#[derive(Debug, Default)]
28struct BufferColors {
29 colors: Vec<(Range<Anchor>, DocumentColor, InlayId)>,
30 inlay_colors: HashMap<InlayId, usize>,
31}
32
33impl LspColorData {
34 pub fn new(cx: &App) -> Self {
35 Self {
36 buffer_colors: HashMap::default(),
37 render_mode: EditorSettings::get_global(cx).lsp_document_colors,
38 }
39 }
40
41 pub fn render_mode_updated(
42 &mut self,
43 new_render_mode: DocumentColorsRenderMode,
44 ) -> Option<InlaySplice> {
45 if self.render_mode == new_render_mode {
46 return None;
47 }
48 self.render_mode = new_render_mode;
49 match new_render_mode {
50 DocumentColorsRenderMode::Inlay => Some(InlaySplice {
51 to_remove: Vec::new(),
52 to_insert: self
53 .buffer_colors
54 .iter()
55 .flat_map(|(_, buffer_colors)| buffer_colors.colors.iter())
56 .map(|(range, color, id)| {
57 Inlay::color(
58 id.id(),
59 range.start,
60 Rgba {
61 r: color.color.red,
62 g: color.color.green,
63 b: color.color.blue,
64 a: color.color.alpha,
65 },
66 )
67 })
68 .collect(),
69 }),
70 DocumentColorsRenderMode::None => Some(InlaySplice {
71 to_remove: self
72 .buffer_colors
73 .drain()
74 .flat_map(|(_, buffer_colors)| buffer_colors.inlay_colors)
75 .map(|(id, _)| id)
76 .collect(),
77 to_insert: Vec::new(),
78 }),
79 DocumentColorsRenderMode::Border | DocumentColorsRenderMode::Background => {
80 Some(InlaySplice {
81 to_remove: self
82 .buffer_colors
83 .iter_mut()
84 .flat_map(|(_, buffer_colors)| buffer_colors.inlay_colors.drain())
85 .map(|(id, _)| id)
86 .collect(),
87 to_insert: Vec::new(),
88 })
89 }
90 }
91 }
92
93 fn set_colors(
94 &mut self,
95 buffer_id: BufferId,
96 colors: Vec<(Range<Anchor>, DocumentColor, InlayId)>,
97 ) -> bool {
98 let buffer_colors = self.buffer_colors.entry(buffer_id).or_default();
99 if buffer_colors.colors == colors {
100 return false;
101 }
102
103 buffer_colors.inlay_colors = colors
104 .iter()
105 .enumerate()
106 .map(|(i, (_, _, id))| (*id, i))
107 .collect();
108 buffer_colors.colors = colors;
109 true
110 }
111
112 pub fn editor_display_highlights(
113 &self,
114 snapshot: &EditorSnapshot,
115 ) -> (DocumentColorsRenderMode, Vec<(Range<DisplayPoint>, Hsla)>) {
116 let render_mode = self.render_mode;
117 let highlights = if render_mode == DocumentColorsRenderMode::None
118 || render_mode == DocumentColorsRenderMode::Inlay
119 {
120 Vec::new()
121 } else {
122 self.buffer_colors
123 .iter()
124 .flat_map(|(_, buffer_colors)| &buffer_colors.colors)
125 .map(|(range, color, _)| {
126 let display_range = range.clone().to_display_points(snapshot);
127 let color = Hsla::from(Rgba {
128 r: color.color.red,
129 g: color.color.green,
130 b: color.color.blue,
131 a: color.color.alpha,
132 });
133 (display_range, color)
134 })
135 .collect()
136 };
137 (render_mode, highlights)
138 }
139}
140
141impl Editor {
142 pub(super) fn refresh_document_colors(
143 &mut self,
144 buffer_id: Option<BufferId>,
145 _: &Window,
146 cx: &mut Context<Self>,
147 ) {
148 if !self.lsp_data_enabled() {
149 return;
150 }
151 let Some(project) = self.project.as_ref() else {
152 return;
153 };
154 if self
155 .colors
156 .as_ref()
157 .is_none_or(|colors| colors.render_mode == DocumentColorsRenderMode::None)
158 {
159 return;
160 }
161
162 let buffers_to_query = self
163 .visible_buffers(cx)
164 .into_iter()
165 .filter(|buffer| self.is_lsp_relevant(buffer.read(cx).file(), cx))
166 .chain(buffer_id.and_then(|buffer_id| self.buffer.read(cx).buffer(buffer_id)))
167 .filter(|editor_buffer| {
168 let editor_buffer_id = editor_buffer.read(cx).remote_id();
169 buffer_id.is_none_or(|buffer_id| buffer_id == editor_buffer_id)
170 && self.registered_buffers.contains_key(&editor_buffer_id)
171 })
172 .unique_by(|buffer| buffer.read(cx).remote_id())
173 .collect::<Vec<_>>();
174
175 let project = project.downgrade();
176 self.refresh_colors_task = cx.spawn(async move |editor, cx| {
177 cx.background_executor()
178 .timer(LSP_REQUEST_DEBOUNCE_TIMEOUT)
179 .await;
180
181 let Some(all_colors_task) = project
182 .update(cx, |project, cx| {
183 project.lsp_store().update(cx, |lsp_store, cx| {
184 buffers_to_query
185 .into_iter()
186 .filter_map(|buffer| {
187 let buffer_snapshot = buffer.read(cx).snapshot();
188 let colors_task = lsp_store.document_colors(buffer, cx)?;
189 Some(async move { (buffer_snapshot, colors_task.await) })
190 })
191 .collect::<Vec<_>>()
192 })
193 })
194 .ok()
195 else {
196 return;
197 };
198
199 let all_colors = join_all(all_colors_task).await;
200 if all_colors.is_empty() {
201 return;
202 }
203 let Some(multi_buffer_snapshot) = editor
204 .update(cx, |editor, cx| editor.buffer.read(cx).snapshot(cx))
205 .ok()
206 else {
207 return;
208 };
209
210 let mut new_editor_colors: HashMap<BufferId, Vec<(Range<Anchor>, DocumentColor)>> =
211 HashMap::default();
212 for (buffer_snapshot, colors) in all_colors {
213 match colors {
214 Ok(colors) => {
215 if colors.colors.is_empty() {
216 new_editor_colors
217 .entry(buffer_snapshot.remote_id())
218 .or_insert_with(Vec::new)
219 .clear();
220 } else {
221 for color in colors.colors {
222 let color_start = point_from_lsp(color.lsp_range.start);
223 let color_end = point_from_lsp(color.lsp_range.end);
224
225 let Some(range) = multi_buffer_snapshot
226 .buffer_anchor_range_to_anchor_range(
227 buffer_snapshot.anchor_range_outside(
228 buffer_snapshot
229 .clip_point_utf16(color_start, Bias::Left)
230 ..buffer_snapshot
231 .clip_point_utf16(color_end, Bias::Right),
232 ),
233 )
234 else {
235 continue;
236 };
237
238 let new_buffer_colors = new_editor_colors
239 .entry(buffer_snapshot.remote_id())
240 .or_insert_with(Vec::new);
241
242 let (Ok(i) | Err(i)) =
243 new_buffer_colors.binary_search_by(|(probe, _)| {
244 probe
245 .start
246 .cmp(&range.start, &multi_buffer_snapshot)
247 .then_with(|| {
248 probe.end.cmp(&range.end, &multi_buffer_snapshot)
249 })
250 });
251 new_buffer_colors.insert(i, (range, color));
252 }
253 }
254 }
255 Err(e) => log::error!("Failed to retrieve document colors: {e}"),
256 }
257 }
258
259 editor
260 .update(cx, |editor, cx| {
261 let mut colors_splice = InlaySplice::default();
262 let Some(colors) = &mut editor.colors else {
263 return;
264 };
265 let mut updated = false;
266 for (buffer_id, new_buffer_colors) in new_editor_colors {
267 let mut new_buffer_color_inlays =
268 Vec::with_capacity(new_buffer_colors.len());
269 let mut existing_buffer_colors = colors
270 .buffer_colors
271 .entry(buffer_id)
272 .or_default()
273 .colors
274 .iter()
275 .peekable();
276 for (new_range, new_color) in new_buffer_colors {
277 let rgba_color = Rgba {
278 r: new_color.color.red,
279 g: new_color.color.green,
280 b: new_color.color.blue,
281 a: new_color.color.alpha,
282 };
283
284 loop {
285 match existing_buffer_colors.peek() {
286 Some((existing_range, existing_color, existing_inlay_id)) => {
287 match existing_range
288 .start
289 .cmp(&new_range.start, &multi_buffer_snapshot)
290 .then_with(|| {
291 existing_range
292 .end
293 .cmp(&new_range.end, &multi_buffer_snapshot)
294 }) {
295 cmp::Ordering::Less => {
296 colors_splice.to_remove.push(*existing_inlay_id);
297 existing_buffer_colors.next();
298 continue;
299 }
300 cmp::Ordering::Equal => {
301 if existing_color == &new_color {
302 new_buffer_color_inlays.push((
303 new_range,
304 new_color,
305 *existing_inlay_id,
306 ));
307 } else {
308 colors_splice
309 .to_remove
310 .push(*existing_inlay_id);
311
312 let inlay = Inlay::color(
313 post_inc(&mut editor.next_color_inlay_id),
314 new_range.start,
315 rgba_color,
316 );
317 let inlay_id = inlay.id;
318 colors_splice.to_insert.push(inlay);
319 new_buffer_color_inlays
320 .push((new_range, new_color, inlay_id));
321 }
322 existing_buffer_colors.next();
323 break;
324 }
325 cmp::Ordering::Greater => {
326 let inlay = Inlay::color(
327 post_inc(&mut editor.next_color_inlay_id),
328 new_range.start,
329 rgba_color,
330 );
331 let inlay_id = inlay.id;
332 colors_splice.to_insert.push(inlay);
333 new_buffer_color_inlays
334 .push((new_range, new_color, inlay_id));
335 break;
336 }
337 }
338 }
339 None => {
340 let inlay = Inlay::color(
341 post_inc(&mut editor.next_color_inlay_id),
342 new_range.start,
343 rgba_color,
344 );
345 let inlay_id = inlay.id;
346 colors_splice.to_insert.push(inlay);
347 new_buffer_color_inlays
348 .push((new_range, new_color, inlay_id));
349 break;
350 }
351 }
352 }
353 }
354
355 if existing_buffer_colors.peek().is_some() {
356 colors_splice
357 .to_remove
358 .extend(existing_buffer_colors.map(|(_, _, id)| *id));
359 }
360 updated |= colors.set_colors(buffer_id, new_buffer_color_inlays);
361 }
362
363 if colors.render_mode == DocumentColorsRenderMode::Inlay
364 && !colors_splice.is_empty()
365 {
366 editor.splice_inlays(&colors_splice.to_remove, colors_splice.to_insert, cx);
367 updated = true;
368 }
369
370 if updated {
371 cx.notify();
372 }
373 })
374 .ok();
375 });
376 }
377}
378
379#[cfg(test)]
380mod tests {
381 use std::{
382 path::PathBuf,
383 sync::{
384 Arc,
385 atomic::{self, AtomicUsize},
386 },
387 time::Duration,
388 };
389
390 use futures::StreamExt;
391 use gpui::{Rgba, TestAppContext};
392 use language::FakeLspAdapter;
393 use languages::rust_lang;
394 use project::{FakeFs, Project};
395 use serde_json::json;
396 use util::{path, rel_path::rel_path};
397 use workspace::{
398 CloseActiveItem, MoveItemToPaneInDirection, MultiWorkspace, OpenOptions,
399 item::{Item as _, SaveOptions},
400 };
401
402 use crate::{
403 Editor, LSP_REQUEST_DEBOUNCE_TIMEOUT, actions::MoveToEnd, editor_tests::init_test,
404 };
405
406 fn extract_color_inlays(editor: &Editor, cx: &gpui::App) -> Vec<Rgba> {
407 editor
408 .all_inlays(cx)
409 .into_iter()
410 .filter_map(|inlay| inlay.get_color())
411 .map(Rgba::from)
412 .collect()
413 }
414
415 #[gpui::test(iterations = 10)]
416 async fn test_document_colors(cx: &mut TestAppContext) {
417 let expected_color = Rgba {
418 r: 0.33,
419 g: 0.33,
420 b: 0.33,
421 a: 0.33,
422 };
423
424 init_test(cx, |_| {});
425
426 let fs = FakeFs::new(cx.executor());
427 fs.insert_tree(
428 path!("/a"),
429 json!({
430 "first.rs": "fn main() { let a = 5; }",
431 }),
432 )
433 .await;
434
435 let project = Project::test(fs, [path!("/a").as_ref()], cx).await;
436 let (multi_workspace, cx) =
437 cx.add_window_view(|window, cx| MultiWorkspace::test_new(project.clone(), window, cx));
438 let workspace = multi_workspace.read_with(cx, |mw, _| mw.workspace().clone());
439
440 let language_registry = project.read_with(cx, |project, _| project.languages().clone());
441 language_registry.add(rust_lang());
442 let mut fake_servers = language_registry.register_fake_lsp(
443 "Rust",
444 FakeLspAdapter {
445 capabilities: lsp::ServerCapabilities {
446 color_provider: Some(lsp::ColorProviderCapability::Simple(true)),
447 ..lsp::ServerCapabilities::default()
448 },
449 name: "rust-analyzer",
450 ..FakeLspAdapter::default()
451 },
452 );
453 let mut fake_servers_without_capabilities = language_registry.register_fake_lsp(
454 "Rust",
455 FakeLspAdapter {
456 capabilities: lsp::ServerCapabilities {
457 color_provider: Some(lsp::ColorProviderCapability::Simple(false)),
458 ..lsp::ServerCapabilities::default()
459 },
460 name: "not-rust-analyzer",
461 ..FakeLspAdapter::default()
462 },
463 );
464
465 let editor = workspace
466 .update_in(cx, |workspace, window, cx| {
467 workspace.open_abs_path(
468 PathBuf::from(path!("/a/first.rs")),
469 OpenOptions::default(),
470 window,
471 cx,
472 )
473 })
474 .await
475 .unwrap()
476 .downcast::<Editor>()
477 .unwrap();
478 let fake_language_server = fake_servers.next().await.unwrap();
479 let fake_language_server_without_capabilities =
480 fake_servers_without_capabilities.next().await.unwrap();
481 let requests_made = Arc::new(AtomicUsize::new(0));
482 let closure_requests_made = Arc::clone(&requests_made);
483 let mut color_request_handle = fake_language_server
484 .set_request_handler::<lsp::request::DocumentColor, _, _>(move |params, _| {
485 let requests_made = Arc::clone(&closure_requests_made);
486 async move {
487 assert_eq!(
488 params.text_document.uri,
489 lsp::Uri::from_file_path(path!("/a/first.rs")).unwrap()
490 );
491 requests_made.fetch_add(1, atomic::Ordering::Release);
492 Ok(vec![
493 lsp::ColorInformation {
494 range: lsp::Range {
495 start: lsp::Position {
496 line: 0,
497 character: 0,
498 },
499 end: lsp::Position {
500 line: 0,
501 character: 1,
502 },
503 },
504 color: lsp::Color {
505 red: 0.33,
506 green: 0.33,
507 blue: 0.33,
508 alpha: 0.33,
509 },
510 },
511 lsp::ColorInformation {
512 range: lsp::Range {
513 start: lsp::Position {
514 line: 0,
515 character: 0,
516 },
517 end: lsp::Position {
518 line: 0,
519 character: 1,
520 },
521 },
522 color: lsp::Color {
523 red: 0.33,
524 green: 0.33,
525 blue: 0.33,
526 alpha: 0.33,
527 },
528 },
529 ])
530 }
531 });
532
533 let _handle = fake_language_server_without_capabilities
534 .set_request_handler::<lsp::request::DocumentColor, _, _>(move |_, _| async move {
535 panic!("Should not be called");
536 });
537 cx.executor().advance_clock(LSP_REQUEST_DEBOUNCE_TIMEOUT);
538 color_request_handle.next().await.unwrap();
539 cx.run_until_parked();
540 assert_eq!(
541 1,
542 requests_made.load(atomic::Ordering::Acquire),
543 "Should query for colors once per editor open"
544 );
545 editor.update_in(cx, |editor, _, cx| {
546 assert_eq!(
547 vec![expected_color],
548 extract_color_inlays(editor, cx),
549 "Should have an initial inlay"
550 );
551 });
552
553 // opening another file in a split should not influence the LSP query counter
554 workspace.update_in(cx, |workspace, window, cx| {
555 assert_eq!(
556 workspace.panes().len(),
557 1,
558 "Should have one pane with one editor"
559 );
560 workspace.move_item_to_pane_in_direction(
561 &MoveItemToPaneInDirection {
562 direction: workspace::SplitDirection::Right,
563 focus: false,
564 clone: true,
565 },
566 window,
567 cx,
568 );
569 });
570 cx.run_until_parked();
571 workspace.update_in(cx, |workspace, _, cx| {
572 let panes = workspace.panes();
573 assert_eq!(panes.len(), 2, "Should have two panes after splitting");
574 for pane in panes {
575 let editor = pane
576 .read(cx)
577 .active_item()
578 .and_then(|item| item.downcast::<Editor>())
579 .expect("Should have opened an editor in each split");
580 let editor_file = editor
581 .read(cx)
582 .buffer()
583 .read(cx)
584 .as_singleton()
585 .expect("test deals with singleton buffers")
586 .read(cx)
587 .file()
588 .expect("test buffese should have a file")
589 .path();
590 assert_eq!(
591 editor_file.as_ref(),
592 rel_path("first.rs"),
593 "Both editors should be opened for the same file"
594 )
595 }
596 });
597
598 cx.executor().advance_clock(Duration::from_millis(500));
599 let save = editor.update_in(cx, |editor, window, cx| {
600 editor.move_to_end(&MoveToEnd, window, cx);
601 editor.handle_input("dirty", window, cx);
602 editor.save(
603 SaveOptions {
604 format: true,
605 force_format: false,
606 autosave: true,
607 },
608 project.clone(),
609 window,
610 cx,
611 )
612 });
613 save.await.unwrap();
614
615 color_request_handle.next().await.unwrap();
616 cx.run_until_parked();
617 assert_eq!(
618 2,
619 requests_made.load(atomic::Ordering::Acquire),
620 "Should query for colors once per save (deduplicated) and once per formatting after save"
621 );
622
623 drop(editor);
624 let close = workspace.update_in(cx, |workspace, window, cx| {
625 workspace.active_pane().update(cx, |pane, cx| {
626 pane.close_active_item(&CloseActiveItem::default(), window, cx)
627 })
628 });
629 close.await.unwrap();
630 let close = workspace.update_in(cx, |workspace, window, cx| {
631 workspace.active_pane().update(cx, |pane, cx| {
632 pane.close_active_item(&CloseActiveItem::default(), window, cx)
633 })
634 });
635 close.await.unwrap();
636 assert_eq!(
637 2,
638 requests_made.load(atomic::Ordering::Acquire),
639 "After saving and closing all editors, no extra requests should be made"
640 );
641 workspace.update_in(cx, |workspace, _, cx| {
642 assert!(
643 workspace.active_item(cx).is_none(),
644 "Should close all editors"
645 )
646 });
647
648 workspace.update_in(cx, |workspace, window, cx| {
649 workspace.active_pane().update(cx, |pane, cx| {
650 pane.navigate_backward(&workspace::GoBack, window, cx);
651 })
652 });
653 cx.executor().advance_clock(LSP_REQUEST_DEBOUNCE_TIMEOUT);
654 cx.run_until_parked();
655 let editor = workspace.update_in(cx, |workspace, _, cx| {
656 workspace
657 .active_item(cx)
658 .expect("Should have reopened the editor again after navigating back")
659 .downcast::<Editor>()
660 .expect("Should be an editor")
661 });
662
663 assert_eq!(
664 2,
665 requests_made.load(atomic::Ordering::Acquire),
666 "Cache should be reused on buffer close and reopen"
667 );
668 editor.update(cx, |editor, cx| {
669 assert_eq!(
670 vec![expected_color],
671 extract_color_inlays(editor, cx),
672 "Should have an initial inlay"
673 );
674 });
675
676 drop(color_request_handle);
677 let closure_requests_made = Arc::clone(&requests_made);
678 let mut empty_color_request_handle = fake_language_server
679 .set_request_handler::<lsp::request::DocumentColor, _, _>(move |params, _| {
680 let requests_made = Arc::clone(&closure_requests_made);
681 async move {
682 assert_eq!(
683 params.text_document.uri,
684 lsp::Uri::from_file_path(path!("/a/first.rs")).unwrap()
685 );
686 requests_made.fetch_add(1, atomic::Ordering::Release);
687 Ok(Vec::new())
688 }
689 });
690 let save = editor.update_in(cx, |editor, window, cx| {
691 editor.move_to_end(&MoveToEnd, window, cx);
692 editor.handle_input("dirty_again", window, cx);
693 editor.save(
694 SaveOptions {
695 format: false,
696 force_format: false,
697 autosave: true,
698 },
699 project.clone(),
700 window,
701 cx,
702 )
703 });
704 save.await.unwrap();
705
706 cx.executor().advance_clock(LSP_REQUEST_DEBOUNCE_TIMEOUT);
707 empty_color_request_handle.next().await.unwrap();
708 cx.run_until_parked();
709 assert_eq!(
710 3,
711 requests_made.load(atomic::Ordering::Acquire),
712 "Should query for colors once per save only, as formatting was not requested"
713 );
714 editor.update(cx, |editor, cx| {
715 assert_eq!(
716 Vec::<Rgba>::new(),
717 extract_color_inlays(editor, cx),
718 "Should clear all colors when the server returns an empty response"
719 );
720 });
721 }
722}