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