tenant.openagents/omega

1use crate::{Diagnostic, range_to_lsp};
2use anyhow::Result;
3use collections::HashMap;
4use lsp::LanguageServerId;
5use serde::Serialize;
6use std::{
7    cmp::{Ordering, Reverse},
8    iter,
9    ops::Range,
10};
11use sum_tree::{self, Bias, SumTree};
12use text::{Anchor, FromAnchor, PointUtf16, ToOffset};
13
14/// A set of diagnostics associated with a given buffer, provided
15/// by a single language server.
16///
17/// The diagnostics are stored in a [`SumTree`], which allows this struct
18/// to be cheaply copied, and allows for efficient retrieval of the
19/// diagnostics that intersect a given range of the buffer.
20#[derive(Clone, Debug)]
21pub struct DiagnosticSet {
22    diagnostics: SumTree<DiagnosticEntry<Anchor>>,
23}
24
25/// A single diagnostic in a set. Generic over its range type, because
26/// the diagnostics are stored internally as [`Anchor`]s, but can be
27/// resolved to different coordinates types like [`usize`] byte offsets or
28/// [`Point`](gpui::Point)s.
29#[derive(Clone, Debug, PartialEq, Eq, Serialize)]
30pub struct DiagnosticEntry<T> {
31    /// The range of the buffer where the diagnostic applies.
32    pub range: Range<T>,
33    /// The information about the diagnostic.
34    pub diagnostic: Diagnostic,
35}
36
37/// A single diagnostic in a set. Generic over its range type, because
38/// the diagnostics are stored internally as [`Anchor`]s, but can be
39/// resolved to different coordinates types like [`usize`] byte offsets or
40/// [`Point`](gpui::Point)s.
41#[derive(Clone, Debug, PartialEq, Eq, Serialize)]
42pub struct DiagnosticEntryRef<'a, T> {
43    /// The range of the buffer where the diagnostic applies.
44    pub range: Range<T>,
45    /// The information about the diagnostic.
46    pub diagnostic: &'a Diagnostic,
47}
48
49impl<T: PartialEq> PartialEq<DiagnosticEntry<T>> for DiagnosticEntryRef<'_, T> {
50    fn eq(&self, other: &DiagnosticEntry<T>) -> bool {
51        self.range == other.range && *self.diagnostic == other.diagnostic
52    }
53}
54
55impl<T: PartialEq> PartialEq<DiagnosticEntryRef<'_, T>> for DiagnosticEntry<T> {
56    fn eq(&self, other: &DiagnosticEntryRef<'_, T>) -> bool {
57        self.range == other.range && self.diagnostic == *other.diagnostic
58    }
59}
60
61impl<T: Clone> DiagnosticEntryRef<'_, T> {
62    pub fn to_owned(&self) -> DiagnosticEntry<T> {
63        DiagnosticEntry {
64            range: self.range.clone(),
65            diagnostic: self.diagnostic.clone(),
66        }
67    }
68}
69
70impl<'a> DiagnosticEntryRef<'a, Anchor> {
71    /// Converts the [DiagnosticEntry] to a different buffer coordinate type.
72    pub fn resolve<O: FromAnchor>(
73        &self,
74        buffer: &text::BufferSnapshot,
75    ) -> DiagnosticEntryRef<'a, O> {
76        DiagnosticEntryRef {
77            range: O::from_anchor(&self.range.start, buffer)
78                ..O::from_anchor(&self.range.end, buffer),
79            diagnostic: &self.diagnostic,
80        }
81    }
82}
83
84/// A group of related diagnostics, ordered by their start position
85/// in the buffer.
86#[derive(Debug, Serialize)]
87pub struct DiagnosticGroup<'a, T> {
88    /// The diagnostics.
89    pub entries: Vec<DiagnosticEntryRef<'a, T>>,
90    /// The index into `entries` where the primary diagnostic is stored.
91    pub primary_ix: usize,
92}
93
94impl<'a> DiagnosticGroup<'a, Anchor> {
95    /// Converts the entries in this [`DiagnosticGroup`] to a different buffer coordinate type.
96    pub fn resolve<O: FromAnchor>(&self, buffer: &text::BufferSnapshot) -> DiagnosticGroup<'a, O> {
97        DiagnosticGroup {
98            entries: self
99                .entries
100                .iter()
101                .map(|entry| entry.resolve(buffer))
102                .collect(),
103            primary_ix: self.primary_ix,
104        }
105    }
106}
107
108#[derive(Clone, Debug)]
109pub struct Summary {
110    start: Anchor,
111    end: Anchor,
112    min_start: Anchor,
113    max_end: Anchor,
114    count: usize,
115}
116
117impl DiagnosticEntry<PointUtf16> {
118    /// Returns a raw LSP diagnostic used to provide diagnostic context to LSP
119    /// codeAction request
120    pub fn to_lsp_diagnostic_stub(&self) -> Result<lsp::Diagnostic> {
121        let range = range_to_lsp(self.range.clone())?;
122
123        Ok(lsp::Diagnostic {
124            range,
125            code: self.diagnostic.code.clone(),
126            severity: Some(self.diagnostic.severity),
127            source: self.diagnostic.source.clone(),
128            message: self.diagnostic.message.clone(),
129            data: self.diagnostic.data.clone(),
130            ..Default::default()
131        })
132    }
133}
134impl DiagnosticEntryRef<'_, PointUtf16> {
135    /// Returns a raw LSP diagnostic used to provide diagnostic context to LSP
136    /// codeAction request
137    pub fn to_lsp_diagnostic_stub(&self) -> Result<lsp::Diagnostic> {
138        let range = range_to_lsp(self.range.clone())?;
139
140        Ok(lsp::Diagnostic {
141            range,
142            code: self.diagnostic.code.clone(),
143            severity: Some(self.diagnostic.severity),
144            source: self.diagnostic.source.clone(),
145            message: self.diagnostic.message.clone(),
146            data: self.diagnostic.data.clone(),
147            ..Default::default()
148        })
149    }
150}
151
152impl DiagnosticSet {
153    /// Constructs a [DiagnosticSet] from a sequence of entries, ordered by
154    /// their position in the buffer.
155    pub fn from_sorted_entries<I>(iter: I, buffer: &text::BufferSnapshot) -> Self
156    where
157        I: IntoIterator<Item = DiagnosticEntry<Anchor>>,
158    {
159        Self {
160            diagnostics: SumTree::from_iter(iter, buffer),
161        }
162    }
163
164    /// Constructs a [DiagnosticSet] from a sequence of entries in an arbitrary order.
165    pub fn new<I>(iter: I, buffer: &text::BufferSnapshot) -> Self
166    where
167        I: IntoIterator<Item = DiagnosticEntry<PointUtf16>>,
168    {
169        let mut entries = iter.into_iter().collect::<Vec<_>>();
170        entries.sort_unstable_by_key(|entry| (entry.range.start, Reverse(entry.range.end)));
171        Self {
172            diagnostics: SumTree::from_iter(
173                entries.into_iter().map(|entry| DiagnosticEntry {
174                    range: buffer.anchor_before(entry.range.start)
175                        ..buffer.anchor_before(entry.range.end),
176                    diagnostic: entry.diagnostic,
177                }),
178                buffer,
179            ),
180        }
181    }
182
183    /// Returns the number of diagnostics in the set.
184    pub fn len(&self) -> usize {
185        self.diagnostics.summary().count
186    }
187    /// Returns true when there are no diagnostics in this diagnostic set
188    pub fn is_empty(&self) -> bool {
189        self.len() == 0
190    }
191
192    /// Returns an iterator over the diagnostic entries in the set.
193    pub fn iter(&self) -> impl Iterator<Item = &DiagnosticEntry<Anchor>> {
194        self.diagnostics.iter()
195    }
196
197    /// Returns an iterator over the diagnostic entries that intersect the
198    /// given range of the buffer.
199    pub fn range<'a, T, O>(
200        &'a self,
201        range: Range<T>,
202        buffer: &'a text::BufferSnapshot,
203        inclusive: bool,
204        reversed: bool,
205    ) -> impl 'a + Iterator<Item = DiagnosticEntryRef<'a, O>>
206    where
207        T: 'a + ToOffset,
208        O: FromAnchor,
209    {
210        let end_bias = if inclusive { Bias::Right } else { Bias::Left };
211        let range = buffer.anchor_before(range.start)..buffer.anchor_at(range.end, end_bias);
212        let mut cursor = self.diagnostics.filter::<_, ()>(buffer, {
213            move |summary: &Summary| {
214                let start_cmp = range.start.cmp(&summary.max_end, buffer);
215                let end_cmp = range.end.cmp(&summary.min_start, buffer);
216                if inclusive {
217                    start_cmp <= Ordering::Equal && end_cmp >= Ordering::Equal
218                } else {
219                    start_cmp == Ordering::Less && end_cmp == Ordering::Greater
220                }
221            }
222        });
223
224        if reversed {
225            cursor.prev();
226        } else {
227            cursor.next();
228        }
229        iter::from_fn({
230            move || {
231                if let Some(diagnostic) = cursor.item() {
232                    if reversed {
233                        cursor.prev();
234                    } else {
235                        cursor.next();
236                    }
237                    Some(diagnostic.resolve(buffer))
238                } else {
239                    None
240                }
241            }
242        })
243    }
244
245    /// Adds all of this set's diagnostic groups to the given output vector.
246    pub fn groups<'a>(
247        &'a self,
248        language_server_id: LanguageServerId,
249        output: &mut Vec<(LanguageServerId, DiagnosticGroup<'a, Anchor>)>,
250        buffer: &text::BufferSnapshot,
251    ) {
252        let mut groups = HashMap::default();
253        for entry in self.diagnostics.iter() {
254            groups
255                .entry(entry.diagnostic.group_id)
256                .or_insert(Vec::new())
257                .push(DiagnosticEntryRef {
258                    range: entry.range.clone(),
259                    diagnostic: &entry.diagnostic,
260                });
261        }
262
263        let start_ix = output.len();
264        output.extend(groups.into_values().filter_map(|mut entries| {
265            entries.sort_unstable_by(|a, b| a.range.start.cmp(&b.range.start, buffer));
266            entries
267                .iter()
268                .position(|entry| entry.diagnostic.is_primary)
269                .map(|primary_ix| {
270                    (
271                        language_server_id,
272                        DiagnosticGroup {
273                            entries,
274                            primary_ix,
275                        },
276                    )
277                })
278        }));
279        output[start_ix..].sort_unstable_by(|(id_a, group_a), (id_b, group_b)| {
280            group_a.entries[group_a.primary_ix]
281                .range
282                .start
283                .cmp(&group_b.entries[group_b.primary_ix].range.start, buffer)
284                .then_with(|| id_a.cmp(id_b))
285        });
286    }
287
288    /// Returns all of the diagnostics in a particular diagnostic group,
289    /// in order of their position in the buffer.
290    pub fn group<'a, O: FromAnchor>(
291        &'a self,
292        group_id: usize,
293        buffer: &'a text::BufferSnapshot,
294    ) -> impl 'a + Iterator<Item = DiagnosticEntryRef<'a, O>> {
295        self.iter()
296            .filter(move |entry| entry.diagnostic.group_id == group_id)
297            .map(|entry| entry.resolve(buffer))
298    }
299}
300
301impl sum_tree::Item for DiagnosticEntry<Anchor> {
302    type Summary = Summary;
303
304    fn summary(&self, _cx: &text::BufferSnapshot) -> Self::Summary {
305        Summary {
306            start: self.range.start,
307            end: self.range.end,
308            min_start: self.range.start,
309            max_end: self.range.end,
310            count: 1,
311        }
312    }
313}
314
315impl DiagnosticEntry<Anchor> {
316    /// Converts the [DiagnosticEntry] to a different buffer coordinate type.
317    pub fn resolve<'a, O: FromAnchor>(
318        &'a self,
319        buffer: &text::BufferSnapshot,
320    ) -> DiagnosticEntryRef<'a, O> {
321        DiagnosticEntryRef {
322            range: O::from_anchor(&self.range.start, buffer)
323                ..O::from_anchor(&self.range.end, buffer),
324            diagnostic: &self.diagnostic,
325        }
326    }
327}
328
329impl sum_tree::Summary for Summary {
330    type Context<'a> = &'a text::BufferSnapshot;
331
332    fn zero(buffer: &text::BufferSnapshot) -> Self {
333        Self {
334            start: Anchor::min_for_buffer(buffer.remote_id()),
335            end: Anchor::max_for_buffer(buffer.remote_id()),
336            min_start: Anchor::max_for_buffer(buffer.remote_id()),
337            max_end: Anchor::min_for_buffer(buffer.remote_id()),
338            count: 0,
339        }
340    }
341
342    fn add_summary(&mut self, other: &Self, buffer: Self::Context<'_>) {
343        if other.min_start.cmp(&self.min_start, buffer).is_lt() {
344            self.min_start = other.min_start;
345        }
346        if other.max_end.cmp(&self.max_end, buffer).is_gt() {
347            self.max_end = other.max_end;
348        }
349        self.start = other.start;
350        self.end = other.end;
351        self.count += other.count;
352    }
353}
354