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, InlayId, InlaySplice, RangeToAnchorExt,
17 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 cx.spawn(async move |editor, cx| {
197 let all_colors = join_all(all_colors_task).await;
198 if all_colors.is_empty() {
199 return;
200 }
201 let Ok((multi_buffer_snapshot, editor_excerpts)) = editor.update(cx, |editor, cx| {
202 let multi_buffer_snapshot = editor.buffer().read(cx).snapshot(cx);
203 let editor_excerpts = multi_buffer_snapshot.excerpts().fold(
204 HashMap::default(),
205 |mut acc, (excerpt_id, buffer_snapshot, excerpt_range)| {
206 let excerpt_data = acc
207 .entry(buffer_snapshot.remote_id())
208 .or_insert_with(Vec::new);
209 let excerpt_point_range =
210 excerpt_range.context.to_point_utf16(buffer_snapshot);
211 excerpt_data.push((
212 excerpt_id,
213 buffer_snapshot.clone(),
214 excerpt_point_range,
215 ));
216 acc
217 },
218 );
219 (multi_buffer_snapshot, editor_excerpts)
220 }) else {
221 return;
222 };
223
224 let mut new_editor_colors = HashMap::default();
225 for (buffer_id, colors) in all_colors {
226 let Some(excerpts) = editor_excerpts.get(&buffer_id) else {
227 continue;
228 };
229 match colors {
230 Ok(colors) => {
231 if colors.colors.is_empty() {
232 let new_entry =
233 new_editor_colors.entry(buffer_id).or_insert_with(|| {
234 (Vec::<(Range<Anchor>, DocumentColor)>::new(), None)
235 });
236 new_entry.0.clear();
237 new_entry.1 = colors.cache_version;
238 } else {
239 for color in colors.colors {
240 let color_start = point_from_lsp(color.lsp_range.start);
241 let color_end = point_from_lsp(color.lsp_range.end);
242
243 for (excerpt_id, buffer_snapshot, excerpt_range) in excerpts {
244 if !excerpt_range.contains(&color_start.0)
245 || !excerpt_range.contains(&color_end.0)
246 {
247 continue;
248 }
249 let Some(color_start_anchor) = multi_buffer_snapshot
250 .anchor_in_excerpt(
251 *excerpt_id,
252 buffer_snapshot.anchor_before(
253 buffer_snapshot
254 .clip_point_utf16(color_start, Bias::Left),
255 ),
256 )
257 else {
258 continue;
259 };
260 let Some(color_end_anchor) = multi_buffer_snapshot
261 .anchor_in_excerpt(
262 *excerpt_id,
263 buffer_snapshot.anchor_after(
264 buffer_snapshot
265 .clip_point_utf16(color_end, Bias::Right),
266 ),
267 )
268 else {
269 continue;
270 };
271
272 let new_entry =
273 new_editor_colors.entry(buffer_id).or_insert_with(|| {
274 (Vec::<(Range<Anchor>, DocumentColor)>::new(), None)
275 });
276 new_entry.1 = colors.cache_version;
277 let new_buffer_colors = &mut new_entry.0;
278
279 let (Ok(i) | Err(i)) =
280 new_buffer_colors.binary_search_by(|(probe, _)| {
281 probe
282 .start
283 .cmp(&color_start_anchor, &multi_buffer_snapshot)
284 .then_with(|| {
285 probe.end.cmp(
286 &color_end_anchor,
287 &multi_buffer_snapshot,
288 )
289 })
290 });
291 new_buffer_colors
292 .insert(i, (color_start_anchor..color_end_anchor, color));
293 break;
294 }
295 }
296 }
297 }
298 Err(e) => log::error!("Failed to retrieve document colors: {e}"),
299 }
300 }
301
302 editor
303 .update(cx, |editor, cx| {
304 let mut colors_splice = InlaySplice::default();
305 let Some(colors) = &mut editor.colors else {
306 return;
307 };
308 let mut updated = false;
309 for (buffer_id, (new_buffer_colors, new_cache_version)) in new_editor_colors {
310 let mut new_buffer_color_inlays =
311 Vec::with_capacity(new_buffer_colors.len());
312 let mut existing_buffer_colors = colors
313 .buffer_colors
314 .entry(buffer_id)
315 .or_default()
316 .colors
317 .iter()
318 .peekable();
319 for (new_range, new_color) in new_buffer_colors {
320 let rgba_color = Rgba {
321 r: new_color.color.red,
322 g: new_color.color.green,
323 b: new_color.color.blue,
324 a: new_color.color.alpha,
325 };
326
327 loop {
328 match existing_buffer_colors.peek() {
329 Some((existing_range, existing_color, existing_inlay_id)) => {
330 match existing_range
331 .start
332 .cmp(&new_range.start, &multi_buffer_snapshot)
333 .then_with(|| {
334 existing_range
335 .end
336 .cmp(&new_range.end, &multi_buffer_snapshot)
337 }) {
338 cmp::Ordering::Less => {
339 colors_splice.to_remove.push(*existing_inlay_id);
340 existing_buffer_colors.next();
341 continue;
342 }
343 cmp::Ordering::Equal => {
344 if existing_color == &new_color {
345 new_buffer_color_inlays.push((
346 new_range,
347 new_color,
348 *existing_inlay_id,
349 ));
350 } else {
351 colors_splice
352 .to_remove
353 .push(*existing_inlay_id);
354
355 let inlay = Inlay::color(
356 post_inc(&mut editor.next_color_inlay_id),
357 new_range.start,
358 rgba_color,
359 );
360 let inlay_id = inlay.id;
361 colors_splice.to_insert.push(inlay);
362 new_buffer_color_inlays
363 .push((new_range, new_color, inlay_id));
364 }
365 existing_buffer_colors.next();
366 break;
367 }
368 cmp::Ordering::Greater => {
369 let inlay = Inlay::color(
370 post_inc(&mut editor.next_color_inlay_id),
371 new_range.start,
372 rgba_color,
373 );
374 let inlay_id = inlay.id;
375 colors_splice.to_insert.push(inlay);
376 new_buffer_color_inlays
377 .push((new_range, new_color, inlay_id));
378 break;
379 }
380 }
381 }
382 None => {
383 let inlay = Inlay::color(
384 post_inc(&mut editor.next_color_inlay_id),
385 new_range.start,
386 rgba_color,
387 );
388 let inlay_id = inlay.id;
389 colors_splice.to_insert.push(inlay);
390 new_buffer_color_inlays
391 .push((new_range, new_color, inlay_id));
392 break;
393 }
394 }
395 }
396 }
397
398 if existing_buffer_colors.peek().is_some() {
399 colors_splice
400 .to_remove
401 .extend(existing_buffer_colors.map(|(_, _, id)| *id));
402 }
403 updated |= colors.set_colors(
404 buffer_id,
405 new_buffer_color_inlays,
406 new_cache_version,
407 );
408 }
409
410 if colors.render_mode == DocumentColorsRenderMode::Inlay
411 && (!colors_splice.to_insert.is_empty()
412 || !colors_splice.to_remove.is_empty())
413 {
414 editor.splice_inlays(&colors_splice.to_remove, colors_splice.to_insert, cx);
415 updated = true;
416 }
417
418 if updated {
419 cx.notify();
420 }
421 })
422 .ok();
423 })
424 .detach();
425 }
426}