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, lsp_store::LspFetchStrategy};
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, FETCH_COLORS_DEBOUNCE_TIMEOUT, InlayId,
17 InlaySplice, RangeToAnchorExt, display_map::Inlay, editor_settings::DocumentColorsRenderMode,
18};
19
20#[derive(Debug)]
21pub(super) struct LspColorData {
22 buffer_colors: HashMap<BufferId, BufferColors>,
23 render_mode: DocumentColorsRenderMode,
24}
25
26#[derive(Debug, Default)]
27struct BufferColors {
28 colors: Vec<(Range<Anchor>, DocumentColor, InlayId)>,
29 inlay_colors: HashMap<InlayId, usize>,
30 cache_version_used: 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 cache_version: Option<usize>,
98 ) -> bool {
99 let buffer_colors = self.buffer_colors.entry(buffer_id).or_default();
100 if let Some(cache_version) = cache_version {
101 buffer_colors.cache_version_used = cache_version;
102 }
103 if buffer_colors.colors == colors {
104 return false;
105 }
106
107 buffer_colors.inlay_colors = colors
108 .iter()
109 .enumerate()
110 .map(|(i, (_, _, id))| (*id, i))
111 .collect();
112 buffer_colors.colors = colors;
113 true
114 }
115
116 pub fn editor_display_highlights(
117 &self,
118 snapshot: &EditorSnapshot,
119 ) -> (DocumentColorsRenderMode, Vec<(Range<DisplayPoint>, Hsla)>) {
120 let render_mode = self.render_mode;
121 let highlights = if render_mode == DocumentColorsRenderMode::None
122 || render_mode == DocumentColorsRenderMode::Inlay
123 {
124 Vec::new()
125 } else {
126 self.buffer_colors
127 .iter()
128 .flat_map(|(_, buffer_colors)| &buffer_colors.colors)
129 .map(|(range, color, _)| {
130 let display_range = range.clone().to_display_points(snapshot);
131 let color = Hsla::from(Rgba {
132 r: color.color.red,
133 g: color.color.green,
134 b: color.color.blue,
135 a: color.color.alpha,
136 });
137 (display_range, color)
138 })
139 .collect()
140 };
141 (render_mode, highlights)
142 }
143}
144
145impl Editor {
146 pub(super) fn refresh_colors(
147 &mut self,
148 ignore_cache: bool,
149 buffer_id: Option<BufferId>,
150 _: &Window,
151 cx: &mut Context<Self>,
152 ) {
153 if !self.mode().is_full() {
154 return;
155 }
156 let Some(project) = self.project.clone() else {
157 return;
158 };
159 if self
160 .colors
161 .as_ref()
162 .is_none_or(|colors| colors.render_mode == DocumentColorsRenderMode::None)
163 {
164 return;
165 }
166
167 let visible_buffers = self
168 .visible_excerpts(None, cx)
169 .into_values()
170 .map(|(buffer, ..)| buffer)
171 .filter(|editor_buffer| {
172 buffer_id.is_none_or(|buffer_id| buffer_id == editor_buffer.read(cx).remote_id())
173 })
174 .unique_by(|buffer| buffer.read(cx).remote_id())
175 .collect::<Vec<_>>();
176
177 let all_colors_task = project.read(cx).lsp_store().update(cx, |lsp_store, cx| {
178 visible_buffers
179 .into_iter()
180 .filter_map(|buffer| {
181 let buffer_id = buffer.read(cx).remote_id();
182 let fetch_strategy = if ignore_cache {
183 LspFetchStrategy::IgnoreCache
184 } else {
185 LspFetchStrategy::UseCache {
186 known_cache_version: self.colors.as_ref().and_then(|colors| {
187 Some(colors.buffer_colors.get(&buffer_id)?.cache_version_used)
188 }),
189 }
190 };
191 let colors_task = lsp_store.document_colors(fetch_strategy, buffer, cx)?;
192 Some(async move { (buffer_id, colors_task.await) })
193 })
194 .collect::<Vec<_>>()
195 });
196
197 self.refresh_colors_task = cx.spawn(async move |editor, cx| {
198 cx.background_executor()
199 .timer(FETCH_COLORS_DEBOUNCE_TIMEOUT)
200 .await;
201
202 let all_colors = join_all(all_colors_task).await;
203 if all_colors.is_empty() {
204 return;
205 }
206 let Ok((multi_buffer_snapshot, editor_excerpts)) = editor.update(cx, |editor, cx| {
207 let multi_buffer_snapshot = editor.buffer().read(cx).snapshot(cx);
208 let editor_excerpts = multi_buffer_snapshot.excerpts().fold(
209 HashMap::default(),
210 |mut acc, (excerpt_id, buffer_snapshot, excerpt_range)| {
211 let excerpt_data = acc
212 .entry(buffer_snapshot.remote_id())
213 .or_insert_with(Vec::new);
214 let excerpt_point_range =
215 excerpt_range.context.to_point_utf16(buffer_snapshot);
216 excerpt_data.push((
217 excerpt_id,
218 buffer_snapshot.clone(),
219 excerpt_point_range,
220 ));
221 acc
222 },
223 );
224 (multi_buffer_snapshot, editor_excerpts)
225 }) else {
226 return;
227 };
228
229 let mut new_editor_colors = HashMap::default();
230 for (buffer_id, colors) in all_colors {
231 let Some(excerpts) = editor_excerpts.get(&buffer_id) else {
232 continue;
233 };
234 match colors {
235 Ok(colors) => {
236 if colors.colors.is_empty() {
237 let new_entry =
238 new_editor_colors.entry(buffer_id).or_insert_with(|| {
239 (Vec::<(Range<Anchor>, DocumentColor)>::new(), None)
240 });
241 new_entry.0.clear();
242 new_entry.1 = colors.cache_version;
243 } else {
244 for color in colors.colors {
245 let color_start = point_from_lsp(color.lsp_range.start);
246 let color_end = point_from_lsp(color.lsp_range.end);
247
248 for (excerpt_id, buffer_snapshot, excerpt_range) in excerpts {
249 if !excerpt_range.contains(&color_start.0)
250 || !excerpt_range.contains(&color_end.0)
251 {
252 continue;
253 }
254 let Some(color_start_anchor) = multi_buffer_snapshot
255 .anchor_in_excerpt(
256 *excerpt_id,
257 buffer_snapshot.anchor_before(
258 buffer_snapshot
259 .clip_point_utf16(color_start, Bias::Left),
260 ),
261 )
262 else {
263 continue;
264 };
265 let Some(color_end_anchor) = multi_buffer_snapshot
266 .anchor_in_excerpt(
267 *excerpt_id,
268 buffer_snapshot.anchor_after(
269 buffer_snapshot
270 .clip_point_utf16(color_end, Bias::Right),
271 ),
272 )
273 else {
274 continue;
275 };
276
277 let new_entry =
278 new_editor_colors.entry(buffer_id).or_insert_with(|| {
279 (Vec::<(Range<Anchor>, DocumentColor)>::new(), None)
280 });
281 new_entry.1 = colors.cache_version;
282 let new_buffer_colors = &mut new_entry.0;
283
284 let (Ok(i) | Err(i)) =
285 new_buffer_colors.binary_search_by(|(probe, _)| {
286 probe
287 .start
288 .cmp(&color_start_anchor, &multi_buffer_snapshot)
289 .then_with(|| {
290 probe.end.cmp(
291 &color_end_anchor,
292 &multi_buffer_snapshot,
293 )
294 })
295 });
296 new_buffer_colors
297 .insert(i, (color_start_anchor..color_end_anchor, color));
298 break;
299 }
300 }
301 }
302 }
303 Err(e) => log::error!("Failed to retrieve document colors: {e}"),
304 }
305 }
306
307 editor
308 .update(cx, |editor, cx| {
309 let mut colors_splice = InlaySplice::default();
310 let Some(colors) = &mut editor.colors else {
311 return;
312 };
313 let mut updated = false;
314 for (buffer_id, (new_buffer_colors, new_cache_version)) in new_editor_colors {
315 let mut new_buffer_color_inlays =
316 Vec::with_capacity(new_buffer_colors.len());
317 let mut existing_buffer_colors = colors
318 .buffer_colors
319 .entry(buffer_id)
320 .or_default()
321 .colors
322 .iter()
323 .peekable();
324 for (new_range, new_color) in new_buffer_colors {
325 let rgba_color = Rgba {
326 r: new_color.color.red,
327 g: new_color.color.green,
328 b: new_color.color.blue,
329 a: new_color.color.alpha,
330 };
331
332 loop {
333 match existing_buffer_colors.peek() {
334 Some((existing_range, existing_color, existing_inlay_id)) => {
335 match existing_range
336 .start
337 .cmp(&new_range.start, &multi_buffer_snapshot)
338 .then_with(|| {
339 existing_range
340 .end
341 .cmp(&new_range.end, &multi_buffer_snapshot)
342 }) {
343 cmp::Ordering::Less => {
344 colors_splice.to_remove.push(*existing_inlay_id);
345 existing_buffer_colors.next();
346 continue;
347 }
348 cmp::Ordering::Equal => {
349 if existing_color == &new_color {
350 new_buffer_color_inlays.push((
351 new_range,
352 new_color,
353 *existing_inlay_id,
354 ));
355 } else {
356 colors_splice
357 .to_remove
358 .push(*existing_inlay_id);
359
360 let inlay = Inlay::color(
361 post_inc(&mut editor.next_color_inlay_id),
362 new_range.start,
363 rgba_color,
364 );
365 let inlay_id = inlay.id;
366 colors_splice.to_insert.push(inlay);
367 new_buffer_color_inlays
368 .push((new_range, new_color, inlay_id));
369 }
370 existing_buffer_colors.next();
371 break;
372 }
373 cmp::Ordering::Greater => {
374 let inlay = Inlay::color(
375 post_inc(&mut editor.next_color_inlay_id),
376 new_range.start,
377 rgba_color,
378 );
379 let inlay_id = inlay.id;
380 colors_splice.to_insert.push(inlay);
381 new_buffer_color_inlays
382 .push((new_range, new_color, inlay_id));
383 break;
384 }
385 }
386 }
387 None => {
388 let inlay = Inlay::color(
389 post_inc(&mut editor.next_color_inlay_id),
390 new_range.start,
391 rgba_color,
392 );
393 let inlay_id = inlay.id;
394 colors_splice.to_insert.push(inlay);
395 new_buffer_color_inlays
396 .push((new_range, new_color, inlay_id));
397 break;
398 }
399 }
400 }
401 }
402
403 if existing_buffer_colors.peek().is_some() {
404 colors_splice
405 .to_remove
406 .extend(existing_buffer_colors.map(|(_, _, id)| *id));
407 }
408 updated |= colors.set_colors(
409 buffer_id,
410 new_buffer_color_inlays,
411 new_cache_version,
412 );
413 }
414
415 if colors.render_mode == DocumentColorsRenderMode::Inlay
416 && (!colors_splice.to_insert.is_empty()
417 || !colors_splice.to_remove.is_empty())
418 {
419 editor.splice_inlays(&colors_splice.to_remove, colors_splice.to_insert, cx);
420 updated = true;
421 }
422
423 if updated {
424 cx.notify();
425 }
426 })
427 .ok();
428 });
429 }
430}