tenant.openagents/omega

1use crate::{
2    HighlightId, HighlightMap, LanguageConfig, LanguageConfigOverride, LanguageName,
3    LanguageQueries, language_config::BracketPairConfig,
4};
5use anyhow::{Context as _, Result};
6use collections::HashMap;
7use gpui_shared_string::SharedString;
8use parking_lot::Mutex;
9use std::sync::atomic::{AtomicUsize, Ordering::SeqCst};
10use tree_sitter::Query;
11
12pub static NEXT_GRAMMAR_ID: AtomicUsize = AtomicUsize::new(0);
13
14#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Clone, Copy)]
15pub struct GrammarId(pub usize);
16
17impl GrammarId {
18    pub fn new() -> Self {
19        Self(NEXT_GRAMMAR_ID.fetch_add(1, SeqCst))
20    }
21}
22
23impl Default for GrammarId {
24    fn default() -> Self {
25        Self::new()
26    }
27}
28
29pub struct Grammar {
30    id: GrammarId,
31    pub ts_language: tree_sitter::Language,
32    pub error_query: Option<Query>,
33    pub highlights_config: Option<HighlightsConfig>,
34    pub brackets_config: Option<BracketsConfig>,
35    pub redactions_config: Option<RedactionConfig>,
36    pub runnable_config: Option<RunnableConfig>,
37    pub indents_config: Option<IndentConfig>,
38    pub outline_config: Option<OutlineConfig>,
39    pub text_object_config: Option<TextObjectConfig>,
40    pub injection_config: Option<InjectionConfig>,
41    pub override_config: Option<OverrideConfig>,
42    pub debug_variables_config: Option<DebugVariablesConfig>,
43    pub highlight_map: Mutex<HighlightMap>,
44}
45
46pub struct HighlightsConfig {
47    pub query: Query,
48    pub identifier_capture_indices: Vec<u32>,
49}
50
51pub struct IndentConfig {
52    pub query: Query,
53    pub indent_capture_ix: u32,
54    pub start_capture_ix: Option<u32>,
55    pub end_capture_ix: Option<u32>,
56    pub outdent_capture_ix: Option<u32>,
57    pub suffixed_start_captures: HashMap<u32, SharedString>,
58}
59
60pub struct OutlineConfig {
61    pub query: Query,
62    pub item_capture_ix: u32,
63    pub name_capture_ix: u32,
64    pub context_capture_ix: Option<u32>,
65    pub extra_context_capture_ix: Option<u32>,
66    pub open_capture_ix: Option<u32>,
67    pub close_capture_ix: Option<u32>,
68    pub annotation_capture_ix: Option<u32>,
69}
70
71#[derive(Debug, Clone, Copy, PartialEq)]
72pub enum DebuggerTextObject {
73    Variable,
74    Scope,
75}
76
77impl DebuggerTextObject {
78    pub fn from_capture_name(name: &str) -> Option<DebuggerTextObject> {
79        match name {
80            "debug-variable" => Some(DebuggerTextObject::Variable),
81            "debug-scope" => Some(DebuggerTextObject::Scope),
82            _ => None,
83        }
84    }
85}
86
87#[derive(Debug, Clone, Copy, PartialEq)]
88pub enum TextObject {
89    InsideFunction,
90    AroundFunction,
91    InsideClass,
92    AroundClass,
93    InsideComment,
94    AroundComment,
95}
96
97impl TextObject {
98    pub fn from_capture_name(name: &str) -> Option<TextObject> {
99        match name {
100            "function.inside" => Some(TextObject::InsideFunction),
101            "function.around" => Some(TextObject::AroundFunction),
102            "class.inside" => Some(TextObject::InsideClass),
103            "class.around" => Some(TextObject::AroundClass),
104            "comment.inside" => Some(TextObject::InsideComment),
105            "comment.around" => Some(TextObject::AroundComment),
106            _ => None,
107        }
108    }
109
110    pub fn around(&self) -> Option<Self> {
111        match self {
112            TextObject::InsideFunction => Some(TextObject::AroundFunction),
113            TextObject::InsideClass => Some(TextObject::AroundClass),
114            TextObject::InsideComment => Some(TextObject::AroundComment),
115            _ => None,
116        }
117    }
118}
119
120pub struct TextObjectConfig {
121    pub query: Query,
122    pub text_objects_by_capture_ix: Vec<(u32, TextObject)>,
123}
124
125pub struct InjectionConfig {
126    pub query: Query,
127    pub content_capture_ix: u32,
128    pub language_capture_ix: Option<u32>,
129    pub patterns: Vec<InjectionPatternConfig>,
130}
131
132pub struct RedactionConfig {
133    pub query: Query,
134    pub redaction_capture_ix: u32,
135}
136
137#[derive(Clone, Debug, PartialEq)]
138pub enum RunnableCapture {
139    /// `@run` — marks the node that identifies the runnable (e.g. the test
140    /// function name, the subtest string literal).
141    Run,
142    /// `@run_item` — marks one candidate runnable inside a larger grammar
143    /// match. Used by the per-match dispatch pipeline so a single tree-sitter
144    /// match can produce multiple runnables (e.g. Go table-test rows).
145    RunItem,
146    /// Any other named capture, exposed by name to language resolvers and used
147    /// to populate `RunnableRange::extra_captures`.
148    Named(SharedString),
149}
150
151pub struct RunnableConfig {
152    pub query: Query,
153    /// A mapping from capture index to capture kind
154    pub extra_captures: Vec<RunnableCapture>,
155    /// `true` if the query uses `@run_item`, the marker for matches that
156    /// emit multiple runnables.
157    pub supports_grouped_runnables: bool,
158}
159
160pub struct OverrideConfig {
161    pub query: Query,
162    pub values: HashMap<u32, OverrideEntry>,
163}
164
165#[derive(Debug)]
166pub struct OverrideEntry {
167    pub name: String,
168    pub range_is_inclusive: bool,
169    pub value: LanguageConfigOverride,
170}
171
172#[derive(Default, Clone)]
173pub struct InjectionPatternConfig {
174    pub language: Option<Box<str>>,
175    pub combined: bool,
176}
177
178#[derive(Debug)]
179pub struct BracketsConfig {
180    pub query: Query,
181    pub open_capture_ix: u32,
182    pub close_capture_ix: u32,
183    pub patterns: Vec<BracketsPatternConfig>,
184}
185
186#[derive(Clone, Debug, Default)]
187pub struct BracketsPatternConfig {
188    pub newline_only: bool,
189    pub rainbow_exclude: bool,
190}
191
192pub struct DebugVariablesConfig {
193    pub query: Query,
194    pub objects_by_capture_ix: Vec<(u32, DebuggerTextObject)>,
195}
196
197enum Capture<'a> {
198    Required(&'static str, &'a mut u32),
199    Optional(&'static str, &'a mut Option<u32>),
200}
201
202fn populate_capture_indices(
203    query: &Query,
204    language_name: &LanguageName,
205    query_type: &str,
206    expected_prefixes: &[&str],
207    captures: &mut [Capture<'_>],
208) -> bool {
209    let mut found_required_indices = Vec::new();
210    'outer: for (ix, name) in query.capture_names().iter().enumerate() {
211        for (required_ix, capture) in captures.iter_mut().enumerate() {
212            match capture {
213                Capture::Required(capture_name, index) if capture_name == name => {
214                    **index = ix as u32;
215                    found_required_indices.push(required_ix);
216                    continue 'outer;
217                }
218                Capture::Optional(capture_name, index) if capture_name == name => {
219                    **index = Some(ix as u32);
220                    continue 'outer;
221                }
222                _ => {}
223            }
224        }
225        if !name.starts_with("_")
226            && !expected_prefixes
227                .iter()
228                .any(|&prefix| name.starts_with(prefix))
229        {
230            log::warn!(
231                "unrecognized capture name '{}' in {} {} TreeSitter query \
232                (suppress this warning by prefixing with '_')",
233                name,
234                language_name,
235                query_type
236            );
237        }
238    }
239    let mut missing_required_captures = Vec::new();
240    for (capture_ix, capture) in captures.iter().enumerate() {
241        if let Capture::Required(capture_name, _) = capture
242            && !found_required_indices.contains(&capture_ix)
243        {
244            missing_required_captures.push(*capture_name);
245        }
246    }
247    let success = missing_required_captures.is_empty();
248    if !success {
249        log::error!(
250            "missing required capture(s) in {} {} TreeSitter query: {}",
251            language_name,
252            query_type,
253            missing_required_captures.join(", ")
254        );
255    }
256    success
257}
258
259impl Grammar {
260    pub fn new(ts_language: tree_sitter::Language) -> Self {
261        Self {
262            id: GrammarId::new(),
263            highlights_config: None,
264            brackets_config: None,
265            outline_config: None,
266            text_object_config: None,
267            indents_config: None,
268            injection_config: None,
269            override_config: None,
270            redactions_config: None,
271            runnable_config: None,
272            error_query: Query::new(&ts_language, "(ERROR) @error").ok(),
273            debug_variables_config: None,
274            ts_language,
275            highlight_map: Default::default(),
276        }
277    }
278
279    pub fn id(&self) -> GrammarId {
280        self.id
281    }
282
283    pub fn highlight_map(&self) -> HighlightMap {
284        self.highlight_map.lock().clone()
285    }
286
287    pub fn highlight_id_for_name(&self, name: &str) -> Option<HighlightId> {
288        self.highlights_config
289            .as_ref()?
290            .query
291            .capture_index_for_name(name)
292            .and_then(|capture_id| self.highlight_map.lock().get(capture_id))
293    }
294
295    pub fn debug_variables_config(&self) -> Option<&DebugVariablesConfig> {
296        self.debug_variables_config.as_ref()
297    }
298
299    /// Load all queries from `LanguageQueries` into this grammar, mutating the
300    /// associated `LanguageConfig` (the override query clears
301    /// `brackets.disabled_scopes_by_bracket_ix`).
302    pub fn with_queries(
303        mut self,
304        queries: LanguageQueries,
305        config: &mut LanguageConfig,
306    ) -> Result<Self> {
307        let name = &config.name;
308        if let Some(query) = queries.highlights {
309            self = self
310                .with_highlights_query(query.as_ref())
311                .context("Error loading highlights query")?;
312        }
313        if let Some(query) = queries.brackets {
314            self = self
315                .with_brackets_query(query.as_ref(), name)
316                .context("Error loading brackets query")?;
317        }
318        if let Some(query) = queries.indents {
319            self = self
320                .with_indents_query(query.as_ref(), name)
321                .context("Error loading indents query")?;
322        }
323        if let Some(query) = queries.outline {
324            self = self
325                .with_outline_query(query.as_ref(), name)
326                .context("Error loading outline query")?;
327        }
328        if let Some(query) = queries.injections {
329            self = self
330                .with_injection_query(query.as_ref(), name)
331                .context("Error loading injection query")?;
332        }
333        if let Some(query) = queries.overrides {
334            self = self
335                .with_override_query(
336                    query.as_ref(),
337                    name,
338                    &config.overrides,
339                    &mut config.brackets,
340                    &config.scope_opt_in_language_servers,
341                )
342                .context("Error loading override query")?;
343        }
344        if let Some(query) = queries.redactions {
345            self = self
346                .with_redaction_query(query.as_ref(), name)
347                .context("Error loading redaction query")?;
348        }
349        if let Some(query) = queries.runnables {
350            self = self
351                .with_runnable_query(query.as_ref())
352                .context("Error loading runnables query")?;
353        }
354        if let Some(query) = queries.text_objects {
355            self = self
356                .with_text_object_query(query.as_ref(), name)
357                .context("Error loading textobject query")?;
358        }
359        if let Some(query) = queries.debugger {
360            self = self
361                .with_debug_variables_query(query.as_ref(), name)
362                .context("Error loading debug variables query")?;
363        }
364        Ok(self)
365    }
366
367    pub fn with_highlights_query(mut self, source: &str) -> Result<Self> {
368        let query = Query::new(&self.ts_language, source)?;
369
370        let mut identifier_capture_indices = Vec::new();
371        for name in [
372            "variable",
373            "constant",
374            "constructor",
375            "function",
376            "function.method",
377            "function.method.call",
378            "function.special",
379            "property",
380            "type",
381            "type.interface",
382        ] {
383            identifier_capture_indices.extend(query.capture_index_for_name(name));
384        }
385
386        self.highlights_config = Some(HighlightsConfig {
387            query,
388            identifier_capture_indices,
389        });
390
391        Ok(self)
392    }
393
394    pub fn with_runnable_query(mut self, source: &str) -> Result<Self> {
395        let query = Query::new(&self.ts_language, source)?;
396        let extra_captures: Vec<_> = query
397            .capture_names()
398            .iter()
399            .map(|&name| match name {
400                "run" => RunnableCapture::Run,
401                "run_item" => RunnableCapture::RunItem,
402                name => RunnableCapture::Named(name.to_string().into()),
403            })
404            .collect();
405        let supports_grouped_runnables = extra_captures
406            .iter()
407            .any(|capture| matches!(capture, RunnableCapture::RunItem));
408
409        self.runnable_config = Some(RunnableConfig {
410            extra_captures,
411            query,
412            supports_grouped_runnables,
413        });
414
415        Ok(self)
416    }
417
418    pub fn with_outline_query(
419        mut self,
420        source: &str,
421        language_name: &LanguageName,
422    ) -> Result<Self> {
423        let query = Query::new(&self.ts_language, source)?;
424        let mut item_capture_ix = 0;
425        let mut name_capture_ix = 0;
426        let mut context_capture_ix = None;
427        let mut extra_context_capture_ix = None;
428        let mut open_capture_ix = None;
429        let mut close_capture_ix = None;
430        let mut annotation_capture_ix = None;
431        if populate_capture_indices(
432            &query,
433            language_name,
434            "outline",
435            &[],
436            &mut [
437                Capture::Required("item", &mut item_capture_ix),
438                Capture::Required("name", &mut name_capture_ix),
439                Capture::Optional("context", &mut context_capture_ix),
440                Capture::Optional("context.extra", &mut extra_context_capture_ix),
441                Capture::Optional("open", &mut open_capture_ix),
442                Capture::Optional("close", &mut close_capture_ix),
443                Capture::Optional("annotation", &mut annotation_capture_ix),
444            ],
445        ) {
446            self.outline_config = Some(OutlineConfig {
447                query,
448                item_capture_ix,
449                name_capture_ix,
450                context_capture_ix,
451                extra_context_capture_ix,
452                open_capture_ix,
453                close_capture_ix,
454                annotation_capture_ix,
455            });
456        }
457        Ok(self)
458    }
459
460    pub fn with_text_object_query(
461        mut self,
462        source: &str,
463        language_name: &LanguageName,
464    ) -> Result<Self> {
465        let query = Query::new(&self.ts_language, source)?;
466
467        let mut text_objects_by_capture_ix = Vec::new();
468        for (ix, name) in query.capture_names().iter().enumerate() {
469            if let Some(text_object) = TextObject::from_capture_name(name) {
470                text_objects_by_capture_ix.push((ix as u32, text_object));
471            } else {
472                log::warn!(
473                    "unrecognized capture name '{}' in {} textobjects TreeSitter query",
474                    name,
475                    language_name,
476                );
477            }
478        }
479
480        self.text_object_config = Some(TextObjectConfig {
481            query,
482            text_objects_by_capture_ix,
483        });
484        Ok(self)
485    }
486
487    pub fn with_debug_variables_query(
488        mut self,
489        source: &str,
490        language_name: &LanguageName,
491    ) -> Result<Self> {
492        let query = Query::new(&self.ts_language, source)?;
493
494        let mut objects_by_capture_ix = Vec::new();
495        for (ix, name) in query.capture_names().iter().enumerate() {
496            if let Some(text_object) = DebuggerTextObject::from_capture_name(name) {
497                objects_by_capture_ix.push((ix as u32, text_object));
498            } else {
499                log::warn!(
500                    "unrecognized capture name '{}' in {} debugger TreeSitter query",
501                    name,
502                    language_name,
503                );
504            }
505        }
506
507        self.debug_variables_config = Some(DebugVariablesConfig {
508            query,
509            objects_by_capture_ix,
510        });
511        Ok(self)
512    }
513
514    pub fn with_brackets_query(
515        mut self,
516        source: &str,
517        language_name: &LanguageName,
518    ) -> Result<Self> {
519        let query = Query::new(&self.ts_language, source)?;
520        let mut open_capture_ix = 0;
521        let mut close_capture_ix = 0;
522        if populate_capture_indices(
523            &query,
524            language_name,
525            "brackets",
526            &[],
527            &mut [
528                Capture::Required("open", &mut open_capture_ix),
529                Capture::Required("close", &mut close_capture_ix),
530            ],
531        ) {
532            let patterns = (0..query.pattern_count())
533                .map(|ix| {
534                    let mut config = BracketsPatternConfig::default();
535                    for setting in query.property_settings(ix) {
536                        let setting_key = setting.key.as_ref();
537                        if setting_key == "newline.only" {
538                            config.newline_only = true
539                        }
540                        if setting_key == "rainbow.exclude" {
541                            config.rainbow_exclude = true
542                        }
543                    }
544                    config
545                })
546                .collect();
547            self.brackets_config = Some(BracketsConfig {
548                query,
549                open_capture_ix,
550                close_capture_ix,
551                patterns,
552            });
553        }
554        Ok(self)
555    }
556
557    pub fn with_indents_query(
558        mut self,
559        source: &str,
560        language_name: &LanguageName,
561    ) -> Result<Self> {
562        let query = Query::new(&self.ts_language, source)?;
563        let mut indent_capture_ix = 0;
564        let mut start_capture_ix = None;
565        let mut end_capture_ix = None;
566        let mut outdent_capture_ix = None;
567        if populate_capture_indices(
568            &query,
569            language_name,
570            "indents",
571            &["start."],
572            &mut [
573                Capture::Required("indent", &mut indent_capture_ix),
574                Capture::Optional("start", &mut start_capture_ix),
575                Capture::Optional("end", &mut end_capture_ix),
576                Capture::Optional("outdent", &mut outdent_capture_ix),
577            ],
578        ) {
579            let mut suffixed_start_captures = HashMap::default();
580            for (ix, name) in query.capture_names().iter().enumerate() {
581                if let Some(suffix) = name.strip_prefix("start.") {
582                    suffixed_start_captures.insert(ix as u32, suffix.to_owned().into());
583                }
584            }
585
586            self.indents_config = Some(IndentConfig {
587                query,
588                indent_capture_ix,
589                start_capture_ix,
590                end_capture_ix,
591                outdent_capture_ix,
592                suffixed_start_captures,
593            });
594        }
595        Ok(self)
596    }
597
598    pub fn with_injection_query(
599        mut self,
600        source: &str,
601        language_name: &LanguageName,
602    ) -> Result<Self> {
603        let query = Query::new(&self.ts_language, source)?;
604        let mut language_capture_ix = None;
605        let mut injection_language_capture_ix = None;
606        let mut content_capture_ix = None;
607        let mut injection_content_capture_ix = None;
608        if populate_capture_indices(
609            &query,
610            language_name,
611            "injections",
612            &[],
613            &mut [
614                Capture::Optional("language", &mut language_capture_ix),
615                Capture::Optional("injection.language", &mut injection_language_capture_ix),
616                Capture::Optional("content", &mut content_capture_ix),
617                Capture::Optional("injection.content", &mut injection_content_capture_ix),
618            ],
619        ) {
620            language_capture_ix = match (language_capture_ix, injection_language_capture_ix) {
621                (None, Some(ix)) => Some(ix),
622                (Some(_), Some(_)) => {
623                    anyhow::bail!("both language and injection.language captures are present");
624                }
625                _ => language_capture_ix,
626            };
627            content_capture_ix = match (content_capture_ix, injection_content_capture_ix) {
628                (None, Some(ix)) => Some(ix),
629                (Some(_), Some(_)) => {
630                    anyhow::bail!("both content and injection.content captures are present")
631                }
632                _ => content_capture_ix,
633            };
634            let patterns = (0..query.pattern_count())
635                .map(|ix| {
636                    let mut config = InjectionPatternConfig::default();
637                    for setting in query.property_settings(ix) {
638                        match setting.key.as_ref() {
639                            "language" | "injection.language" => {
640                                config.language.clone_from(&setting.value);
641                            }
642                            "combined" | "injection.combined" => {
643                                config.combined = true;
644                            }
645                            _ => {}
646                        }
647                    }
648                    config
649                })
650                .collect();
651            if let Some(content_capture_ix) = content_capture_ix {
652                self.injection_config = Some(InjectionConfig {
653                    query,
654                    language_capture_ix,
655                    content_capture_ix,
656                    patterns,
657                });
658            } else {
659                log::error!(
660                    "missing required capture in injections {} TreeSitter query: \
661                    content or injection.content",
662                    language_name,
663                );
664            }
665        }
666        Ok(self)
667    }
668
669    pub fn with_override_query(
670        mut self,
671        source: &str,
672        language_name: &LanguageName,
673        overrides: &HashMap<String, LanguageConfigOverride>,
674        brackets: &mut BracketPairConfig,
675        scope_opt_in_language_servers: &[SharedString],
676    ) -> Result<Self> {
677        let query = Query::new(&self.ts_language, source)?;
678
679        let mut override_configs_by_id = HashMap::default();
680        for (ix, mut name) in query.capture_names().iter().copied().enumerate() {
681            let mut range_is_inclusive = false;
682            if name.starts_with('_') {
683                continue;
684            }
685            if let Some(prefix) = name.strip_suffix(".inclusive") {
686                name = prefix;
687                range_is_inclusive = true;
688            }
689
690            let value = overrides.get(name).cloned().unwrap_or_default();
691            for server_name in &value.opt_into_language_servers {
692                if !scope_opt_in_language_servers.contains(server_name) {
693                    gpui_util::debug_panic!(
694                        "Server {server_name:?} has been opted-in by scope {name:?} but has not been marked as an opt-in server"
695                    );
696                }
697            }
698
699            override_configs_by_id.insert(
700                ix as u32,
701                OverrideEntry {
702                    name: name.to_string(),
703                    range_is_inclusive,
704                    value,
705                },
706            );
707        }
708
709        let referenced_override_names = overrides
710            .keys()
711            .chain(brackets.disabled_scopes_by_bracket_ix.iter().flatten());
712
713        for referenced_name in referenced_override_names {
714            if !override_configs_by_id
715                .values()
716                .any(|entry| entry.name == *referenced_name)
717            {
718                anyhow::bail!(
719                    "language {:?} has overrides in config not in query: {referenced_name:?}",
720                    language_name
721                );
722            }
723        }
724
725        for entry in override_configs_by_id.values_mut() {
726            entry.value.disabled_bracket_ixs = brackets
727                .disabled_scopes_by_bracket_ix
728                .iter()
729                .enumerate()
730                .filter_map(|(ix, disabled_scope_names)| {
731                    if disabled_scope_names.contains(&entry.name) {
732                        Some(ix as u16)
733                    } else {
734                        None
735                    }
736                })
737                .collect();
738        }
739
740        brackets.disabled_scopes_by_bracket_ix.clear();
741
742        self.override_config = Some(OverrideConfig {
743            query,
744            values: override_configs_by_id,
745        });
746        Ok(self)
747    }
748
749    pub fn with_redaction_query(
750        mut self,
751        source: &str,
752        language_name: &LanguageName,
753    ) -> Result<Self> {
754        let query = Query::new(&self.ts_language, source)?;
755        let mut redaction_capture_ix = 0;
756        if populate_capture_indices(
757            &query,
758            language_name,
759            "redactions",
760            &[],
761            &mut [Capture::Required("redact", &mut redaction_capture_ix)],
762        ) {
763            self.redactions_config = Some(RedactionConfig {
764                query,
765                redaction_capture_ix,
766            });
767        }
768        Ok(self)
769    }
770}
771