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, OffsetRangeExt as _};
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.mode().is_full() {
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_excerpts(true, cx)
164 .into_values()
165 .map(|(buffer, ..)| buffer)
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_id = buffer.read(cx).remote_id();
188 let colors_task = lsp_store.document_colors(buffer, cx)?;
189 Some(async move { (buffer_id, 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 Ok((multi_buffer_snapshot, editor_excerpts)) = editor.update(cx, |editor, cx| {
204 let multi_buffer_snapshot = editor.buffer().read(cx).snapshot(cx);
205 let editor_excerpts = multi_buffer_snapshot.excerpts().fold(
206 HashMap::default(),
207 |mut acc, (excerpt_id, buffer_snapshot, excerpt_range)| {
208 let excerpt_data = acc
209 .entry(buffer_snapshot.remote_id())
210 .or_insert_with(Vec::new);
211 let excerpt_point_range =
212 excerpt_range.context.to_point_utf16(buffer_snapshot);
213 excerpt_data.push((
214 excerpt_id,
215 buffer_snapshot.clone(),
216 excerpt_point_range,
217 ));
218 acc
219 },
220 );
221 (multi_buffer_snapshot, editor_excerpts)
222 }) else {
223 return;
224 };
225
226 let mut new_editor_colors: HashMap<BufferId, Vec<(Range<Anchor>, DocumentColor)>> =
227 HashMap::default();
228 for (buffer_id, colors) in all_colors {
229 let Some(excerpts) = editor_excerpts.get(&buffer_id) else {
230 continue;
231 };
232 match colors {
233 Ok(colors) => {
234 if colors.colors.is_empty() {
235 new_editor_colors
236 .entry(buffer_id)
237 .or_insert_with(Vec::new)
238 .clear();
239 } else {
240 for color in colors.colors {
241 let color_start = point_from_lsp(color.lsp_range.start);
242 let color_end = point_from_lsp(color.lsp_range.end);
243
244 for (excerpt_id, buffer_snapshot, excerpt_range) in excerpts {
245 if !excerpt_range.contains(&color_start.0)
246 || !excerpt_range.contains(&color_end.0)
247 {
248 continue;
249 }
250 let start = buffer_snapshot.anchor_before(
251 buffer_snapshot.clip_point_utf16(color_start, Bias::Left),
252 );
253 let end = buffer_snapshot.anchor_after(
254 buffer_snapshot.clip_point_utf16(color_end, Bias::Right),
255 );
256 let Some(range) = multi_buffer_snapshot
257 .anchor_range_in_excerpt(*excerpt_id, start..end)
258 else {
259 continue;
260 };
261
262 let new_buffer_colors =
263 new_editor_colors.entry(buffer_id).or_insert_with(Vec::new);
264
265 let (Ok(i) | Err(i)) =
266 new_buffer_colors.binary_search_by(|(probe, _)| {
267 probe
268 .start
269 .cmp(&range.start, &multi_buffer_snapshot)
270 .then_with(|| {
271 probe
272 .end
273 .cmp(&range.end, &multi_buffer_snapshot)
274 })
275 });
276 new_buffer_colors.insert(i, (range, color));
277 break;
278 }
279 }
280 }
281 }
282 Err(e) => log::error!("Failed to retrieve document colors: {e}"),
283 }
284 }
285
286 editor
287 .update(cx, |editor, cx| {
288 let mut colors_splice = InlaySplice::default();
289 let Some(colors) = &mut editor.colors else {
290 return;
291 };
292 let mut updated = false;
293 for (buffer_id, new_buffer_colors) in new_editor_colors {
294 let mut new_buffer_color_inlays =
295 Vec::with_capacity(new_buffer_colors.len());
296 let mut existing_buffer_colors = colors
297 .buffer_colors
298 .entry(buffer_id)
299 .or_default()
300 .colors
301 .iter()
302 .peekable();
303 for (new_range, new_color) in new_buffer_colors {
304 let rgba_color = Rgba {
305 r: new_color.color.red,
306 g: new_color.color.green,
307 b: new_color.color.blue,
308 a: new_color.color.alpha,
309 };
310
311 loop {
312 match existing_buffer_colors.peek() {
313 Some((existing_range, existing_color, existing_inlay_id)) => {
314 match existing_range
315 .start
316 .cmp(&new_range.start, &multi_buffer_snapshot)
317 .then_with(|| {
318 existing_range
319 .end
320 .cmp(&new_range.end, &multi_buffer_snapshot)
321 }) {
322 cmp::Ordering::Less => {
323 colors_splice.to_remove.push(*existing_inlay_id);
324 existing_buffer_colors.next();
325 continue;
326 }
327 cmp::Ordering::Equal => {
328 if existing_color == &new_color {
329 new_buffer_color_inlays.push((
330 new_range,
331 new_color,
332 *existing_inlay_id,
333 ));
334 } else {
335 colors_splice
336 .to_remove
337 .push(*existing_inlay_id);
338
339 let inlay = Inlay::color(
340 post_inc(&mut editor.next_color_inlay_id),
341 new_range.start,
342 rgba_color,
343 );
344 let inlay_id = inlay.id;
345 colors_splice.to_insert.push(inlay);
346 new_buffer_color_inlays
347 .push((new_range, new_color, inlay_id));
348 }
349 existing_buffer_colors.next();
350 break;
351 }
352 cmp::Ordering::Greater => {
353 let inlay = Inlay::color(
354 post_inc(&mut editor.next_color_inlay_id),
355 new_range.start,
356 rgba_color,
357 );
358 let inlay_id = inlay.id;
359 colors_splice.to_insert.push(inlay);
360 new_buffer_color_inlays
361 .push((new_range, new_color, inlay_id));
362 break;
363 }
364 }
365 }
366 None => {
367 let inlay = Inlay::color(
368 post_inc(&mut editor.next_color_inlay_id),
369 new_range.start,
370 rgba_color,
371 );
372 let inlay_id = inlay.id;
373 colors_splice.to_insert.push(inlay);
374 new_buffer_color_inlays
375 .push((new_range, new_color, inlay_id));
376 break;
377 }
378 }
379 }
380 }
381
382 if existing_buffer_colors.peek().is_some() {
383 colors_splice
384 .to_remove
385 .extend(existing_buffer_colors.map(|(_, _, id)| *id));
386 }
387 updated |= colors.set_colors(buffer_id, new_buffer_color_inlays);
388 }
389
390 if colors.render_mode == DocumentColorsRenderMode::Inlay
391 && !colors_splice.is_empty()
392 {
393 editor.splice_inlays(&colors_splice.to_remove, colors_splice.to_insert, cx);
394 updated = true;
395 }
396
397 if updated {
398 cx.notify();
399 }
400 })
401 .ok();
402 });
403 }
404}
405
406#[cfg(test)]
407mod tests {
408 use std::{
409 path::PathBuf,
410 sync::{
411 Arc,
412 atomic::{self, AtomicUsize},
413 },
414 time::Duration,
415 };
416
417 use futures::StreamExt;
418 use gpui::{Rgba, TestAppContext};
419 use language::FakeLspAdapter;
420 use languages::rust_lang;
421 use project::{FakeFs, Project};
422 use serde_json::json;
423 use util::{path, rel_path::rel_path};
424 use workspace::{
425 CloseActiveItem, MoveItemToPaneInDirection, MultiWorkspace, OpenOptions,
426 item::{Item as _, SaveOptions},
427 };
428
429 use crate::{
430 Editor, LSP_REQUEST_DEBOUNCE_TIMEOUT, actions::MoveToEnd, editor_tests::init_test,
431 };
432
433 fn extract_color_inlays(editor: &Editor, cx: &gpui::App) -> Vec<Rgba> {
434 editor
435 .all_inlays(cx)
436 .into_iter()
437 .filter_map(|inlay| inlay.get_color())
438 .map(Rgba::from)
439 .collect()
440 }
441
442 #[gpui::test(iterations = 10)]
443 async fn test_document_colors(cx: &mut TestAppContext) {
444 let expected_color = Rgba {
445 r: 0.33,
446 g: 0.33,
447 b: 0.33,
448 a: 0.33,
449 };
450
451 init_test(cx, |_| {});
452
453 let fs = FakeFs::new(cx.executor());
454 fs.insert_tree(
455 path!("/a"),
456 json!({
457 "first.rs": "fn main() { let a = 5; }",
458 }),
459 )
460 .await;
461
462 let project = Project::test(fs, [path!("/a").as_ref()], cx).await;
463 let (multi_workspace, cx) =
464 cx.add_window_view(|window, cx| MultiWorkspace::test_new(project.clone(), window, cx));
465 let workspace = multi_workspace.read_with(cx, |mw, _| mw.workspace().clone());
466
467 let language_registry = project.read_with(cx, |project, _| project.languages().clone());
468 language_registry.add(rust_lang());
469 let mut fake_servers = language_registry.register_fake_lsp(
470 "Rust",
471 FakeLspAdapter {
472 capabilities: lsp::ServerCapabilities {
473 color_provider: Some(lsp::ColorProviderCapability::Simple(true)),
474 ..lsp::ServerCapabilities::default()
475 },
476 name: "rust-analyzer",
477 ..FakeLspAdapter::default()
478 },
479 );
480 let mut fake_servers_without_capabilities = language_registry.register_fake_lsp(
481 "Rust",
482 FakeLspAdapter {
483 capabilities: lsp::ServerCapabilities {
484 color_provider: Some(lsp::ColorProviderCapability::Simple(false)),
485 ..lsp::ServerCapabilities::default()
486 },
487 name: "not-rust-analyzer",
488 ..FakeLspAdapter::default()
489 },
490 );
491
492 let editor = workspace
493 .update_in(cx, |workspace, window, cx| {
494 workspace.open_abs_path(
495 PathBuf::from(path!("/a/first.rs")),
496 OpenOptions::default(),
497 window,
498 cx,
499 )
500 })
501 .await
502 .unwrap()
503 .downcast::<Editor>()
504 .unwrap();
505 let fake_language_server = fake_servers.next().await.unwrap();
506 let fake_language_server_without_capabilities =
507 fake_servers_without_capabilities.next().await.unwrap();
508 let requests_made = Arc::new(AtomicUsize::new(0));
509 let closure_requests_made = Arc::clone(&requests_made);
510 let mut color_request_handle = fake_language_server
511 .set_request_handler::<lsp::request::DocumentColor, _, _>(move |params, _| {
512 let requests_made = Arc::clone(&closure_requests_made);
513 async move {
514 assert_eq!(
515 params.text_document.uri,
516 lsp::Uri::from_file_path(path!("/a/first.rs")).unwrap()
517 );
518 requests_made.fetch_add(1, atomic::Ordering::Release);
519 Ok(vec![
520 lsp::ColorInformation {
521 range: lsp::Range {
522 start: lsp::Position {
523 line: 0,
524 character: 0,
525 },
526 end: lsp::Position {
527 line: 0,
528 character: 1,
529 },
530 },
531 color: lsp::Color {
532 red: 0.33,
533 green: 0.33,
534 blue: 0.33,
535 alpha: 0.33,
536 },
537 },
538 lsp::ColorInformation {
539 range: lsp::Range {
540 start: lsp::Position {
541 line: 0,
542 character: 0,
543 },
544 end: lsp::Position {
545 line: 0,
546 character: 1,
547 },
548 },
549 color: lsp::Color {
550 red: 0.33,
551 green: 0.33,
552 blue: 0.33,
553 alpha: 0.33,
554 },
555 },
556 ])
557 }
558 });
559
560 let _handle = fake_language_server_without_capabilities
561 .set_request_handler::<lsp::request::DocumentColor, _, _>(move |_, _| async move {
562 panic!("Should not be called");
563 });
564 cx.executor().advance_clock(LSP_REQUEST_DEBOUNCE_TIMEOUT);
565 color_request_handle.next().await.unwrap();
566 cx.run_until_parked();
567 assert_eq!(
568 1,
569 requests_made.load(atomic::Ordering::Acquire),
570 "Should query for colors once per editor open"
571 );
572 editor.update_in(cx, |editor, _, cx| {
573 assert_eq!(
574 vec![expected_color],
575 extract_color_inlays(editor, cx),
576 "Should have an initial inlay"
577 );
578 });
579
580 // opening another file in a split should not influence the LSP query counter
581 workspace.update_in(cx, |workspace, window, cx| {
582 assert_eq!(
583 workspace.panes().len(),
584 1,
585 "Should have one pane with one editor"
586 );
587 workspace.move_item_to_pane_in_direction(
588 &MoveItemToPaneInDirection {
589 direction: workspace::SplitDirection::Right,
590 focus: false,
591 clone: true,
592 },
593 window,
594 cx,
595 );
596 });
597 cx.run_until_parked();
598 workspace.update_in(cx, |workspace, _, cx| {
599 let panes = workspace.panes();
600 assert_eq!(panes.len(), 2, "Should have two panes after splitting");
601 for pane in panes {
602 let editor = pane
603 .read(cx)
604 .active_item()
605 .and_then(|item| item.downcast::<Editor>())
606 .expect("Should have opened an editor in each split");
607 let editor_file = editor
608 .read(cx)
609 .buffer()
610 .read(cx)
611 .as_singleton()
612 .expect("test deals with singleton buffers")
613 .read(cx)
614 .file()
615 .expect("test buffese should have a file")
616 .path();
617 assert_eq!(
618 editor_file.as_ref(),
619 rel_path("first.rs"),
620 "Both editors should be opened for the same file"
621 )
622 }
623 });
624
625 cx.executor().advance_clock(Duration::from_millis(500));
626 let save = editor.update_in(cx, |editor, window, cx| {
627 editor.move_to_end(&MoveToEnd, window, cx);
628 editor.handle_input("dirty", window, cx);
629 editor.save(
630 SaveOptions {
631 format: true,
632 autosave: true,
633 },
634 project.clone(),
635 window,
636 cx,
637 )
638 });
639 save.await.unwrap();
640
641 color_request_handle.next().await.unwrap();
642 cx.run_until_parked();
643 assert_eq!(
644 2,
645 requests_made.load(atomic::Ordering::Acquire),
646 "Should query for colors once per save (deduplicated) and once per formatting after save"
647 );
648
649 drop(editor);
650 let close = workspace.update_in(cx, |workspace, window, cx| {
651 workspace.active_pane().update(cx, |pane, cx| {
652 pane.close_active_item(&CloseActiveItem::default(), window, cx)
653 })
654 });
655 close.await.unwrap();
656 let close = workspace.update_in(cx, |workspace, window, cx| {
657 workspace.active_pane().update(cx, |pane, cx| {
658 pane.close_active_item(&CloseActiveItem::default(), window, cx)
659 })
660 });
661 close.await.unwrap();
662 assert_eq!(
663 2,
664 requests_made.load(atomic::Ordering::Acquire),
665 "After saving and closing all editors, no extra requests should be made"
666 );
667 workspace.update_in(cx, |workspace, _, cx| {
668 assert!(
669 workspace.active_item(cx).is_none(),
670 "Should close all editors"
671 )
672 });
673
674 workspace.update_in(cx, |workspace, window, cx| {
675 workspace.active_pane().update(cx, |pane, cx| {
676 pane.navigate_backward(&workspace::GoBack, window, cx);
677 })
678 });
679 cx.executor().advance_clock(LSP_REQUEST_DEBOUNCE_TIMEOUT);
680 cx.run_until_parked();
681 let editor = workspace.update_in(cx, |workspace, _, cx| {
682 workspace
683 .active_item(cx)
684 .expect("Should have reopened the editor again after navigating back")
685 .downcast::<Editor>()
686 .expect("Should be an editor")
687 });
688
689 assert_eq!(
690 2,
691 requests_made.load(atomic::Ordering::Acquire),
692 "Cache should be reused on buffer close and reopen"
693 );
694 editor.update(cx, |editor, cx| {
695 assert_eq!(
696 vec![expected_color],
697 extract_color_inlays(editor, cx),
698 "Should have an initial inlay"
699 );
700 });
701
702 drop(color_request_handle);
703 let closure_requests_made = Arc::clone(&requests_made);
704 let mut empty_color_request_handle = fake_language_server
705 .set_request_handler::<lsp::request::DocumentColor, _, _>(move |params, _| {
706 let requests_made = Arc::clone(&closure_requests_made);
707 async move {
708 assert_eq!(
709 params.text_document.uri,
710 lsp::Uri::from_file_path(path!("/a/first.rs")).unwrap()
711 );
712 requests_made.fetch_add(1, atomic::Ordering::Release);
713 Ok(Vec::new())
714 }
715 });
716 let save = editor.update_in(cx, |editor, window, cx| {
717 editor.move_to_end(&MoveToEnd, window, cx);
718 editor.handle_input("dirty_again", window, cx);
719 editor.save(
720 SaveOptions {
721 format: false,
722 autosave: true,
723 },
724 project.clone(),
725 window,
726 cx,
727 )
728 });
729 save.await.unwrap();
730
731 cx.executor().advance_clock(LSP_REQUEST_DEBOUNCE_TIMEOUT);
732 empty_color_request_handle.next().await.unwrap();
733 cx.run_until_parked();
734 assert_eq!(
735 3,
736 requests_made.load(atomic::Ordering::Acquire),
737 "Should query for colors once per save only, as formatting was not requested"
738 );
739 editor.update(cx, |editor, cx| {
740 assert_eq!(
741 Vec::<Rgba>::new(),
742 extract_color_inlays(editor, cx),
743 "Should clear all colors when the server returns an empty response"
744 );
745 });
746 }
747}