1use std::{cmp, ops::Range};
2
3use collections::HashMap;
4use futures::future::join_all;
5use gpui::{Hsla, Rgba, Task};
6use itertools::Itertools;
7use language::point_from_lsp;
8use multi_buffer::Anchor;
9use project::DocumentColor;
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_for_visible_range(
147 &mut self,
148 buffer_id: Option<BufferId>,
149 _: &Window,
150 cx: &mut Context<Self>,
151 ) {
152 if self.ignore_lsp_data() {
153 return;
154 }
155 let Some(project) = self.project.clone() else {
156 return;
157 };
158 if self
159 .colors
160 .as_ref()
161 .is_none_or(|colors| colors.render_mode == DocumentColorsRenderMode::None)
162 {
163 return;
164 }
165
166 let visible_buffers = self
167 .visible_excerpts(None, cx)
168 .into_values()
169 .map(|(buffer, ..)| buffer)
170 .filter(|editor_buffer| {
171 let editor_buffer_id = editor_buffer.read(cx).remote_id();
172 buffer_id.is_none_or(|buffer_id| buffer_id == editor_buffer_id)
173 && self.registered_buffers.contains_key(&editor_buffer_id)
174 })
175 .unique_by(|buffer| buffer.read(cx).remote_id())
176 .collect::<Vec<_>>();
177
178 let all_colors_task = project.read(cx).lsp_store().update(cx, |lsp_store, cx| {
179 visible_buffers
180 .into_iter()
181 .filter_map(|buffer| {
182 let buffer_id = buffer.read(cx).remote_id();
183 let known_cache_version = self.colors.as_ref().and_then(|colors| {
184 Some(colors.buffer_colors.get(&buffer_id)?.cache_version_used)
185 });
186 let colors_task = lsp_store.document_colors(known_cache_version, buffer, cx)?;
187 Some(async move { (buffer_id, colors_task.await) })
188 })
189 .collect::<Vec<_>>()
190 });
191
192 if all_colors_task.is_empty() {
193 self.refresh_colors_task = Task::ready(());
194 return;
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 start = buffer_snapshot.anchor_before(
255 buffer_snapshot.clip_point_utf16(color_start, Bias::Left),
256 );
257 let end = buffer_snapshot.anchor_after(
258 buffer_snapshot.clip_point_utf16(color_end, Bias::Right),
259 );
260 let Some(range) = multi_buffer_snapshot
261 .anchor_range_in_excerpt(*excerpt_id, start..end)
262 else {
263 continue;
264 };
265
266 let new_entry =
267 new_editor_colors.entry(buffer_id).or_insert_with(|| {
268 (Vec::<(Range<Anchor>, DocumentColor)>::new(), None)
269 });
270 new_entry.1 = colors.cache_version;
271 let new_buffer_colors = &mut new_entry.0;
272
273 let (Ok(i) | Err(i)) =
274 new_buffer_colors.binary_search_by(|(probe, _)| {
275 probe
276 .start
277 .cmp(&range.start, &multi_buffer_snapshot)
278 .then_with(|| {
279 probe
280 .end
281 .cmp(&range.end, &multi_buffer_snapshot)
282 })
283 });
284 new_buffer_colors.insert(i, (range, color));
285 break;
286 }
287 }
288 }
289 }
290 Err(e) => log::error!("Failed to retrieve document colors: {e}"),
291 }
292 }
293
294 editor
295 .update(cx, |editor, cx| {
296 let mut colors_splice = InlaySplice::default();
297 let Some(colors) = &mut editor.colors else {
298 return;
299 };
300 let mut updated = false;
301 for (buffer_id, (new_buffer_colors, new_cache_version)) in new_editor_colors {
302 let mut new_buffer_color_inlays =
303 Vec::with_capacity(new_buffer_colors.len());
304 let mut existing_buffer_colors = colors
305 .buffer_colors
306 .entry(buffer_id)
307 .or_default()
308 .colors
309 .iter()
310 .peekable();
311 for (new_range, new_color) in new_buffer_colors {
312 let rgba_color = Rgba {
313 r: new_color.color.red,
314 g: new_color.color.green,
315 b: new_color.color.blue,
316 a: new_color.color.alpha,
317 };
318
319 loop {
320 match existing_buffer_colors.peek() {
321 Some((existing_range, existing_color, existing_inlay_id)) => {
322 match existing_range
323 .start
324 .cmp(&new_range.start, &multi_buffer_snapshot)
325 .then_with(|| {
326 existing_range
327 .end
328 .cmp(&new_range.end, &multi_buffer_snapshot)
329 }) {
330 cmp::Ordering::Less => {
331 colors_splice.to_remove.push(*existing_inlay_id);
332 existing_buffer_colors.next();
333 continue;
334 }
335 cmp::Ordering::Equal => {
336 if existing_color == &new_color {
337 new_buffer_color_inlays.push((
338 new_range,
339 new_color,
340 *existing_inlay_id,
341 ));
342 } else {
343 colors_splice
344 .to_remove
345 .push(*existing_inlay_id);
346
347 let inlay = Inlay::color(
348 post_inc(&mut editor.next_color_inlay_id),
349 new_range.start,
350 rgba_color,
351 );
352 let inlay_id = inlay.id;
353 colors_splice.to_insert.push(inlay);
354 new_buffer_color_inlays
355 .push((new_range, new_color, inlay_id));
356 }
357 existing_buffer_colors.next();
358 break;
359 }
360 cmp::Ordering::Greater => {
361 let inlay = Inlay::color(
362 post_inc(&mut editor.next_color_inlay_id),
363 new_range.start,
364 rgba_color,
365 );
366 let inlay_id = inlay.id;
367 colors_splice.to_insert.push(inlay);
368 new_buffer_color_inlays
369 .push((new_range, new_color, inlay_id));
370 break;
371 }
372 }
373 }
374 None => {
375 let inlay = Inlay::color(
376 post_inc(&mut editor.next_color_inlay_id),
377 new_range.start,
378 rgba_color,
379 );
380 let inlay_id = inlay.id;
381 colors_splice.to_insert.push(inlay);
382 new_buffer_color_inlays
383 .push((new_range, new_color, inlay_id));
384 break;
385 }
386 }
387 }
388 }
389
390 if existing_buffer_colors.peek().is_some() {
391 colors_splice
392 .to_remove
393 .extend(existing_buffer_colors.map(|(_, _, id)| *id));
394 }
395 updated |= colors.set_colors(
396 buffer_id,
397 new_buffer_color_inlays,
398 new_cache_version,
399 );
400 }
401
402 if colors.render_mode == DocumentColorsRenderMode::Inlay
403 && (!colors_splice.to_insert.is_empty()
404 || !colors_splice.to_remove.is_empty())
405 {
406 editor.splice_inlays(&colors_splice.to_remove, colors_splice.to_insert, cx);
407 updated = true;
408 }
409
410 if updated {
411 cx.notify();
412 }
413 })
414 .ok();
415 });
416 }
417}