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 autosave: true,
606 },
607 project.clone(),
608 window,
609 cx,
610 )
611 });
612 save.await.unwrap();
613
614 color_request_handle.next().await.unwrap();
615 cx.run_until_parked();
616 assert_eq!(
617 2,
618 requests_made.load(atomic::Ordering::Acquire),
619 "Should query for colors once per save (deduplicated) and once per formatting after save"
620 );
621
622 drop(editor);
623 let close = workspace.update_in(cx, |workspace, window, cx| {
624 workspace.active_pane().update(cx, |pane, cx| {
625 pane.close_active_item(&CloseActiveItem::default(), window, cx)
626 })
627 });
628 close.await.unwrap();
629 let close = workspace.update_in(cx, |workspace, window, cx| {
630 workspace.active_pane().update(cx, |pane, cx| {
631 pane.close_active_item(&CloseActiveItem::default(), window, cx)
632 })
633 });
634 close.await.unwrap();
635 assert_eq!(
636 2,
637 requests_made.load(atomic::Ordering::Acquire),
638 "After saving and closing all editors, no extra requests should be made"
639 );
640 workspace.update_in(cx, |workspace, _, cx| {
641 assert!(
642 workspace.active_item(cx).is_none(),
643 "Should close all editors"
644 )
645 });
646
647 workspace.update_in(cx, |workspace, window, cx| {
648 workspace.active_pane().update(cx, |pane, cx| {
649 pane.navigate_backward(&workspace::GoBack, window, cx);
650 })
651 });
652 cx.executor().advance_clock(LSP_REQUEST_DEBOUNCE_TIMEOUT);
653 cx.run_until_parked();
654 let editor = workspace.update_in(cx, |workspace, _, cx| {
655 workspace
656 .active_item(cx)
657 .expect("Should have reopened the editor again after navigating back")
658 .downcast::<Editor>()
659 .expect("Should be an editor")
660 });
661
662 assert_eq!(
663 2,
664 requests_made.load(atomic::Ordering::Acquire),
665 "Cache should be reused on buffer close and reopen"
666 );
667 editor.update(cx, |editor, cx| {
668 assert_eq!(
669 vec![expected_color],
670 extract_color_inlays(editor, cx),
671 "Should have an initial inlay"
672 );
673 });
674
675 drop(color_request_handle);
676 let closure_requests_made = Arc::clone(&requests_made);
677 let mut empty_color_request_handle = fake_language_server
678 .set_request_handler::<lsp::request::DocumentColor, _, _>(move |params, _| {
679 let requests_made = Arc::clone(&closure_requests_made);
680 async move {
681 assert_eq!(
682 params.text_document.uri,
683 lsp::Uri::from_file_path(path!("/a/first.rs")).unwrap()
684 );
685 requests_made.fetch_add(1, atomic::Ordering::Release);
686 Ok(Vec::new())
687 }
688 });
689 let save = editor.update_in(cx, |editor, window, cx| {
690 editor.move_to_end(&MoveToEnd, window, cx);
691 editor.handle_input("dirty_again", window, cx);
692 editor.save(
693 SaveOptions {
694 format: false,
695 autosave: true,
696 },
697 project.clone(),
698 window,
699 cx,
700 )
701 });
702 save.await.unwrap();
703
704 cx.executor().advance_clock(LSP_REQUEST_DEBOUNCE_TIMEOUT);
705 empty_color_request_handle.next().await.unwrap();
706 cx.run_until_parked();
707 assert_eq!(
708 3,
709 requests_made.load(atomic::Ordering::Acquire),
710 "Should query for colors once per save only, as formatting was not requested"
711 );
712 editor.update(cx, |editor, cx| {
713 assert_eq!(
714 Vec::<Rgba>::new(),
715 extract_color_inlays(editor, cx),
716 "Should clear all colors when the server returns an empty response"
717 );
718 });
719 }
720}