tenant.openagents/omega

1use std::{cmp::Reverse, iter, ops::Range, sync::Arc};
2
3use collections::HashMap;
4use smallvec::SmallVec;
5use text::BufferId;
6use tree_sitter::QueryCapture;
7use util::RangeExt;
8
9use crate::{BufferSnapshot, Language, Runnable, RunnableCapture, RunnableConfig, RunnableTag};
10
11pub struct RunnableRange {
12    pub buffer_id: BufferId,
13    pub run_range: Range<usize>,
14    pub full_range: Range<usize>,
15    pub runnable: Runnable,
16    pub extra_captures: HashMap<String, String>,
17}
18
19#[derive(Clone, Debug)]
20pub struct RunnableMatchCapture {
21    range: Range<usize>,
22    capture: ResolverCapture,
23}
24
25/// The subset of `RunnableCapture` that's meaningful to resolvers.
26#[derive(Clone, Debug)]
27enum ResolverCapture {
28    Run,
29    Named(String),
30}
31
32impl RunnableMatchCapture {
33    pub fn range(&self) -> Range<usize> {
34        self.range.clone()
35    }
36
37    pub fn is_run(&self) -> bool {
38        matches!(self.capture, ResolverCapture::Run)
39    }
40
41    pub fn name(&self) -> Option<&str> {
42        match &self.capture {
43            ResolverCapture::Named(name) => Some(name.as_str()),
44            ResolverCapture::Run => None,
45        }
46    }
47}
48
49pub struct ResolvedRunnable {
50    pub run_range: Range<usize>,
51    pub extra_captures: SmallVec<[(String, String); 2]>,
52}
53
54pub trait RunnableResolver: Send + Sync {
55    fn resolve(
56        &self,
57        local_captures: &[RunnableMatchCapture],
58        shared_captures: &[RunnableMatchCapture],
59        buffer: &BufferSnapshot,
60    ) -> Option<ResolvedRunnable>;
61}
62
63pub(crate) fn runnable_ranges(
64    buffer: &BufferSnapshot,
65    offset_range: Range<usize>,
66) -> impl Iterator<Item = RunnableRange> + '_ {
67    let mut syntax_matches = buffer.matches(offset_range.clone(), |grammar| {
68        grammar.runnable_config.as_ref().map(|config| &config.query)
69    });
70
71    let runnable_configs = syntax_matches
72        .grammars()
73        .iter()
74        .map(|grammar| grammar.runnable_config.as_ref())
75        .collect::<Vec<_>>();
76
77    iter::from_fn(move || -> Option<SmallVec<[RunnableRange; 1]>> {
78        let mat = syntax_matches.peek()?;
79
80        let ranges = match runnable_configs[mat.grammar_index] {
81            Some(runnable_config) => {
82                let is_grouped = runnable_config.supports_grouped_runnables
83                    && mat.captures.iter().any(|capture| {
84                        matches!(
85                            runnable_config.extra_captures.get(capture.index as usize),
86                            Some(RunnableCapture::RunItem)
87                        )
88                    });
89                if is_grouped {
90                    runnable_ranges_from_grouped_matches(
91                        buffer,
92                        mat.captures,
93                        runnable_config,
94                        mat.pattern_index,
95                        mat.language.clone(),
96                        offset_range.clone(),
97                    )
98                } else {
99                    runnable_range_from_captures(
100                        buffer,
101                        mat.captures,
102                        runnable_config,
103                        mat.pattern_index,
104                        mat.language,
105                    )
106                    .into_iter()
107                    .collect()
108                }
109            }
110            None => SmallVec::new(),
111        };
112
113        syntax_matches.advance();
114        Some(ranges)
115    })
116    .flatten()
117}
118
119type RunnableMatchCaptures = SmallVec<[RunnableMatchCapture; 4]>;
120
121struct RunnableMatchGroup {
122    range: Range<usize>,
123    captures: RunnableMatchCaptures,
124}
125
126struct GroupedRunnableMatches {
127    groups: SmallVec<[RunnableMatchGroup; 1]>,
128    shared_captures: RunnableMatchCaptures,
129}
130
131fn runnable_tags_from_pattern(
132    query: &tree_sitter::Query,
133    pattern_index: usize,
134) -> SmallVec<[RunnableTag; 1]> {
135    query
136        .property_settings(pattern_index)
137        .iter()
138        .filter_map(|property| {
139            (*property.key == *"tag")
140                .then(|| {
141                    property
142                        .value
143                        .as_ref()
144                        .map(|value| RunnableTag(value.to_string().into()))
145                })
146                .flatten()
147        })
148        .collect()
149}
150
151/// `overlaps` rejects empty ranges, so handle a zero-width `offset_range` (cursor) separately.
152fn range_overlaps_or_contains(range: &Range<usize>, offset_range: &Range<usize>) -> bool {
153    if offset_range.is_empty() {
154        range.contains(&offset_range.start)
155    } else {
156        range.overlaps(offset_range)
157    }
158}
159
160fn group_runnable_matches(
161    captures: &[QueryCapture<'_>],
162    runnable_config: &RunnableConfig,
163    offset_range: Range<usize>,
164) -> GroupedRunnableMatches {
165    let mut sorted: SmallVec<[&QueryCapture<'_>; 16]> = captures.iter().collect();
166    sorted.sort_by_key(|capture| {
167        let range = capture.node.byte_range();
168        (range.start, Reverse(range.end))
169    });
170
171    let mut groups = SmallVec::new();
172    let mut shared_captures = SmallVec::new();
173    let mut current_group: Option<RunnableMatchGroup> = None;
174    let mut current_in_offset = false;
175
176    for capture in sorted {
177        let range = capture.node.byte_range();
178        let Some(kind) = runnable_config.extra_captures.get(capture.index as usize) else {
179            continue;
180        };
181
182        let resolver_capture = match kind {
183            RunnableCapture::RunItem => {
184                if let Some(group) = current_group.take()
185                    && current_in_offset
186                {
187                    groups.push(group);
188                }
189                current_in_offset = range_overlaps_or_contains(&range, &offset_range);
190                current_group = Some(RunnableMatchGroup {
191                    range,
192                    captures: SmallVec::new(),
193                });
194                continue;
195            }
196            RunnableCapture::Run => ResolverCapture::Run,
197            RunnableCapture::Named(name) => ResolverCapture::Named(name.to_string()),
198        };
199
200        let match_capture = RunnableMatchCapture {
201            range: range.clone(),
202            capture: resolver_capture,
203        };
204
205        match current_group.as_mut() {
206            Some(group) if group.range.contains_inclusive(&range) => {
207                if current_in_offset {
208                    group.captures.push(match_capture);
209                }
210            }
211            _ => {
212                if let Some(group) = current_group.take()
213                    && current_in_offset
214                {
215                    groups.push(group);
216                }
217                shared_captures.push(match_capture);
218            }
219        }
220    }
221    if let Some(group) = current_group.take()
222        && current_in_offset
223    {
224        groups.push(group);
225    }
226
227    GroupedRunnableMatches {
228        groups,
229        shared_captures,
230    }
231}
232
233fn runnable_ranges_from_grouped_matches(
234    buffer: &BufferSnapshot,
235    captures: &[QueryCapture<'_>],
236    runnable_config: &RunnableConfig,
237    pattern_index: usize,
238    language: Arc<Language>,
239    offset_range: Range<usize>,
240) -> SmallVec<[RunnableRange; 1]> {
241    let GroupedRunnableMatches {
242        groups,
243        shared_captures,
244    } = group_runnable_matches(captures, runnable_config, offset_range);
245
246    let shared_extras: SmallVec<[(String, String); 4]> = shared_captures
247        .iter()
248        .filter_map(|capture| {
249            capture.name().map(|name| {
250                (
251                    name.to_string(),
252                    buffer.text_for_range(capture.range()).collect::<String>(),
253                )
254            })
255        })
256        .collect();
257
258    let Some(resolver) = language
259        .context_provider()
260        .and_then(|provider| provider.runnable_resolver())
261    else {
262        return SmallVec::new();
263    };
264    let mut runnable_ranges = SmallVec::with_capacity(groups.len());
265
266    let tags = runnable_tags_from_pattern(&runnable_config.query, pattern_index);
267    let buffer_id = buffer.remote_id();
268    for group in groups {
269        let Some(ResolvedRunnable {
270            run_range,
271            extra_captures: local_extras,
272        }) = resolver.resolve(&group.captures, &shared_captures, buffer)
273        else {
274            continue;
275        };
276
277        let extra_captures = shared_extras.iter().cloned().chain(local_extras).collect();
278
279        runnable_ranges.push(RunnableRange {
280            run_range,
281            full_range: group.range,
282            runnable: Runnable {
283                tags: tags.clone(),
284                language: language.clone(),
285                buffer: buffer_id,
286            },
287            extra_captures,
288            buffer_id,
289        });
290    }
291
292    runnable_ranges
293}
294
295fn runnable_range_from_captures(
296    buffer: &BufferSnapshot,
297    captures: &[QueryCapture<'_>],
298    runnable_config: &RunnableConfig,
299    pattern_index: usize,
300    language: Arc<Language>,
301) -> Option<RunnableRange> {
302    let mut run_range = None;
303    let first_capture = captures.first()?;
304    let full_range =
305        captures
306            .iter()
307            .skip(1)
308            .fold(first_capture.node.byte_range(), |mut acc, next| {
309                let byte_range = next.node.byte_range();
310                if acc.start > byte_range.start {
311                    acc.start = byte_range.start;
312                }
313                if acc.end < byte_range.end {
314                    acc.end = byte_range.end;
315                }
316                acc
317            });
318
319    let extra_captures: SmallVec<[_; 1]> =
320        SmallVec::from_iter(captures.iter().filter_map(|capture| {
321            runnable_config
322                .extra_captures
323                .get(capture.index as usize)
324                .cloned()
325                .and_then(|tag_name| match tag_name {
326                    RunnableCapture::Named(name) => Some((capture.node.byte_range(), name)),
327                    RunnableCapture::Run => {
328                        let _ = run_range.insert(capture.node.byte_range());
329                        None
330                    }
331                    RunnableCapture::RunItem => None,
332                })
333        }));
334    let run_range = run_range?;
335    let tags = runnable_tags_from_pattern(&runnable_config.query, pattern_index);
336    let extra_captures = extra_captures
337        .into_iter()
338        .map(|(range, name)| {
339            (
340                name.to_string(),
341                buffer.text_for_range(range).collect::<String>(),
342            )
343        })
344        .collect();
345
346    // A runnable has one range, even when its query pattern contributes multiple tags.
347    let buffer_id = buffer.remote_id();
348    Some(RunnableRange {
349        run_range,
350        full_range,
351        runnable: Runnable {
352            tags,
353            language,
354            buffer: buffer_id,
355        },
356        extra_captures,
357        buffer_id,
358    })
359}
360
361#[cfg(test)]
362mod tests {
363    use super::*;
364    use crate::{
365        Buffer, ContextProvider, Language, LanguageConfig, LanguageMatcher, LanguageQueries,
366    };
367    use gpui::{AppContext as _, TestAppContext};
368    use indoc::indoc;
369    use std::{borrow::Cow, sync::Arc};
370
371    struct TestContextProvider {
372        resolver: Arc<dyn RunnableResolver>,
373    }
374
375    impl ContextProvider for TestContextProvider {
376        fn runnable_resolver(&self) -> Option<Arc<dyn RunnableResolver>> {
377            Some(self.resolver.clone())
378        }
379    }
380
381    fn make_language(
382        runnables_query: &'static str,
383        resolver: Option<Arc<dyn RunnableResolver>>,
384    ) -> Arc<Language> {
385        let language = Language::new(
386            LanguageConfig {
387                name: "Rust".into(),
388                matcher: (LanguageMatcher {
389                    path_suffixes: vec!["rs".to_string()],
390                    ..Default::default()
391                })
392                .into(),
393                ..Default::default()
394            },
395            Some(tree_sitter_rust::LANGUAGE.into()),
396        )
397        .with_queries(LanguageQueries {
398            runnables: Some(Cow::Borrowed(runnables_query)),
399            ..Default::default()
400        })
401        .expect("parse runnables query");
402        let context_provider = resolver
403            .map(|resolver| Arc::new(TestContextProvider { resolver }) as Arc<dyn ContextProvider>);
404        Arc::new(language.with_context_provider(context_provider))
405    }
406
407    fn collect_runnables(
408        cx: &mut TestAppContext,
409        source: &str,
410        runnables_query: &'static str,
411        resolver: Option<Arc<dyn RunnableResolver>>,
412    ) -> Vec<RunnableRange> {
413        collect_runnables_in(cx, source, runnables_query, resolver, None)
414    }
415
416    fn collect_runnables_in(
417        cx: &mut TestAppContext,
418        source: &str,
419        runnables_query: &'static str,
420        resolver: Option<Arc<dyn RunnableResolver>>,
421        offset_range: Option<Range<usize>>,
422    ) -> Vec<RunnableRange> {
423        let language = make_language(runnables_query, resolver);
424        let source_owned = source.to_string();
425        let buffer = cx
426            .new(|cx| Buffer::local(source_owned.clone(), cx).with_language(language.clone(), cx));
427        cx.executor().run_until_parked();
428        let range = offset_range.unwrap_or(0..source_owned.len());
429        buffer.update(cx, |buffer, _| {
430            buffer.snapshot().runnable_ranges(range).collect()
431        })
432    }
433
434    fn text_at(buffer: &BufferSnapshot, range: Range<usize>) -> String {
435        buffer.text_for_range(range).collect()
436    }
437
438    /// Picks the first `@run` capture, attaches no extras.
439    struct FirstRunResolver;
440
441    impl RunnableResolver for FirstRunResolver {
442        fn resolve(
443            &self,
444            local_captures: &[RunnableMatchCapture],
445            _shared_captures: &[RunnableMatchCapture],
446            _buffer: &BufferSnapshot,
447        ) -> Option<ResolvedRunnable> {
448            let run = local_captures.iter().find(|capture| capture.is_run())?;
449            Some(ResolvedRunnable {
450                run_range: run.range(),
451                extra_captures: SmallVec::new(),
452            })
453        }
454    }
455
456    /// Picks the first `@run` and surfaces every local named capture as an extra.
457    struct LocalExtrasResolver;
458
459    impl RunnableResolver for LocalExtrasResolver {
460        fn resolve(
461            &self,
462            local_captures: &[RunnableMatchCapture],
463            _shared_captures: &[RunnableMatchCapture],
464            buffer: &BufferSnapshot,
465        ) -> Option<ResolvedRunnable> {
466            let run = local_captures.iter().find(|capture| capture.is_run())?;
467            let extras = local_captures
468                .iter()
469                .filter_map(|capture| {
470                    capture
471                        .name()
472                        .map(|name| (name.to_string(), text_at(buffer, capture.range())))
473                })
474                .collect();
475            Some(ResolvedRunnable {
476                run_range: run.range(),
477                extra_captures: extras,
478            })
479        }
480    }
481
482    /// Skips groups whose `@run` text equals `skip_text`; otherwise picks the first `@run`.
483    struct SkipByTextResolver {
484        skip_text: &'static str,
485    }
486
487    impl RunnableResolver for SkipByTextResolver {
488        fn resolve(
489            &self,
490            local_captures: &[RunnableMatchCapture],
491            _shared_captures: &[RunnableMatchCapture],
492            buffer: &BufferSnapshot,
493        ) -> Option<ResolvedRunnable> {
494            let run = local_captures.iter().find(|capture| capture.is_run())?;
495            if text_at(buffer, run.range()) == self.skip_text {
496                return None;
497            }
498            Some(ResolvedRunnable {
499                run_range: run.range(),
500                extra_captures: SmallVec::new(),
501            })
502        }
503    }
504
505    /// Always emits `_outer = LOCAL` as a local extra (to exercise the shared/local merge).
506    struct OverrideSharedResolver;
507
508    impl RunnableResolver for OverrideSharedResolver {
509        fn resolve(
510            &self,
511            local_captures: &[RunnableMatchCapture],
512            _shared_captures: &[RunnableMatchCapture],
513            _buffer: &BufferSnapshot,
514        ) -> Option<ResolvedRunnable> {
515            let run = local_captures.iter().find(|capture| capture.is_run())?;
516            let mut extras: SmallVec<[(String, String); 2]> = SmallVec::new();
517            extras.push(("_outer".to_string(), "LOCAL".to_string()));
518            Some(ResolvedRunnable {
519                run_range: run.range(),
520                extra_captures: extras,
521            })
522        }
523    }
524
525    const GROUPED_QUERY: &str = indoc! {r#"
526        (function_item
527          name: (identifier) @_outer
528          body: (block
529            ((expression_statement
530               (call_expression
531                 function: (identifier) @run @_call)) @run_item)+))
532    "#};
533
534    const GROUPED_SOURCE: &str = indoc! {r#"
535        fn outer() {
536            alpha();
537            beta();
538            gamma();
539        }
540    "#};
541
542    #[gpui::test]
543    fn test_single_match_emits_one_runnable_per_match(cx: &mut TestAppContext) {
544        let query = indoc! {r#"
545            ((function_item
546               name: (identifier) @run
547               (#match? @run "^test_")) @_decl)
548        "#};
549        let source = indoc! {r#"
550            fn test_alpha() {}
551            fn helper() {}
552            fn test_beta() {}
553        "#};
554
555        let runnables = collect_runnables(cx, source, query, None);
556        let run_texts: Vec<String> = runnables
557            .iter()
558            .map(|range| source[range.run_range.clone()].to_string())
559            .collect();
560        assert_eq!(run_texts, vec!["test_alpha", "test_beta"]);
561
562        let decls: Vec<&str> = runnables
563            .iter()
564            .filter_map(|range| range.extra_captures.get("_decl").map(String::as_str))
565            .collect();
566        assert_eq!(decls, vec!["fn test_alpha() {}", "fn test_beta() {}"]);
567    }
568
569    #[gpui::test]
570    fn test_single_match_without_run_capture_skipped(cx: &mut TestAppContext) {
571        // Pattern with only a named capture and no `@run`: should silently produce nothing.
572        let query = indoc! {r#"
573            (function_item) @_decl
574        "#};
575        let source = indoc! {r#"
576            fn helper() {}
577            fn another() {}
578        "#};
579
580        let runnables = collect_runnables(cx, source, query, None);
581        assert!(
582            runnables.is_empty(),
583            "matches without @run should produce no runnables, got {}",
584            runnables.len()
585        );
586    }
587
588    #[gpui::test]
589    fn test_match_with_no_runnable_does_not_terminate_iteration(cx: &mut TestAppContext) {
590        // A syntax match yielding no runnable must not terminate the
591        // outer iterator before later matches that DO have `@run` are visited.
592        let query = indoc! {r#"
593            ((function_item
594               name: (identifier) @_helper
595               (#match? @_helper "^helper")) @_decl_no_run)
596
597            ((function_item
598               name: (identifier) @run
599               (#match? @run "^test_")) @_decl)
600        "#};
601        let source = indoc! {r#"
602            fn helper() {}
603            fn test_alpha() {}
604        "#};
605
606        let runnables = collect_runnables(cx, source, query, None);
607        let run_texts: Vec<String> = runnables
608            .iter()
609            .map(|range| source[range.run_range.clone()].to_string())
610            .collect();
611        assert_eq!(
612            run_texts,
613            vec!["test_alpha"],
614            "syntax matches that produce no runnable must not terminate iteration"
615        );
616    }
617
618    #[gpui::test]
619    fn test_grouped_match_without_resolver_emits_nothing(cx: &mut TestAppContext) {
620        // `@run_item` is present but no resolver is registered on the language.
621        let runnables = collect_runnables(cx, GROUPED_SOURCE, GROUPED_QUERY, None);
622        assert!(
623            runnables.is_empty(),
624            "grouped path with no resolver should emit nothing, got {}",
625            runnables.len()
626        );
627    }
628
629    #[gpui::test]
630    fn test_grouped_match_emits_one_runnable_per_run_item(cx: &mut TestAppContext) {
631        let resolver: Arc<dyn RunnableResolver> = Arc::new(FirstRunResolver);
632        let runnables = collect_runnables(cx, GROUPED_SOURCE, GROUPED_QUERY, Some(resolver));
633
634        let run_texts: Vec<String> = runnables
635            .iter()
636            .map(|range| GROUPED_SOURCE[range.run_range.clone()].to_string())
637            .collect();
638        assert_eq!(run_texts, vec!["alpha", "beta", "gamma"]);
639    }
640
641    #[gpui::test]
642    fn test_grouped_match_shared_captures_propagate(cx: &mut TestAppContext) {
643        let resolver: Arc<dyn RunnableResolver> = Arc::new(FirstRunResolver);
644        let runnables = collect_runnables(cx, GROUPED_SOURCE, GROUPED_QUERY, Some(resolver));
645
646        for range in &runnables {
647            assert_eq!(
648                range.extra_captures.get("_outer").map(String::as_str),
649                Some("outer"),
650                "every grouped runnable should inherit the shared `_outer` capture"
651            );
652        }
653        assert_eq!(runnables.len(), 3);
654    }
655
656    #[gpui::test]
657    fn test_grouped_match_local_extras_are_per_group(cx: &mut TestAppContext) {
658        let resolver: Arc<dyn RunnableResolver> = Arc::new(LocalExtrasResolver);
659        let runnables = collect_runnables(cx, GROUPED_SOURCE, GROUPED_QUERY, Some(resolver));
660
661        let calls: Vec<&str> = runnables
662            .iter()
663            .filter_map(|range| range.extra_captures.get("_call").map(String::as_str))
664            .collect();
665        assert_eq!(
666            calls,
667            vec!["alpha", "beta", "gamma"],
668            "each group's local `_call` capture should come from that row only"
669        );
670    }
671
672    #[gpui::test]
673    fn test_grouped_match_resolver_returning_none_skips_group(cx: &mut TestAppContext) {
674        let source = indoc! {r#"
675            fn outer() {
676                alpha();
677                skip_me();
678                gamma();
679            }
680        "#};
681        let resolver: Arc<dyn RunnableResolver> = Arc::new(SkipByTextResolver {
682            skip_text: "skip_me",
683        });
684        let runnables = collect_runnables(cx, source, GROUPED_QUERY, Some(resolver));
685
686        let run_texts: Vec<String> = runnables
687            .iter()
688            .map(|range| source[range.run_range.clone()].to_string())
689            .collect();
690        assert_eq!(run_texts, vec!["alpha", "gamma"]);
691    }
692
693    #[gpui::test]
694    fn test_grouped_match_offset_range_filters_groups(cx: &mut TestAppContext) {
695        let resolver: Arc<dyn RunnableResolver> = Arc::new(FirstRunResolver);
696        let beta_offset = GROUPED_SOURCE
697            .find("beta()")
698            .expect("source should contain `beta()`");
699        let runnables = collect_runnables_in(
700            cx,
701            GROUPED_SOURCE,
702            GROUPED_QUERY,
703            Some(resolver),
704            Some(beta_offset..beta_offset + "beta".len()),
705        );
706
707        let run_texts: Vec<String> = runnables
708            .iter()
709            .map(|range| GROUPED_SOURCE[range.run_range.clone()].to_string())
710            .collect();
711        assert_eq!(
712            run_texts,
713            vec!["beta"],
714            "offset_range should restrict emitted groups to those overlapping it"
715        );
716    }
717
718    #[gpui::test]
719    fn test_grouped_match_zero_width_offset_at_group_start(cx: &mut TestAppContext) {
720        let resolver: Arc<dyn RunnableResolver> = Arc::new(FirstRunResolver);
721        let alpha_offset = GROUPED_SOURCE
722            .find("alpha()")
723            .expect("source should contain `alpha()`");
724        let runnables = collect_runnables_in(
725            cx,
726            GROUPED_SOURCE,
727            GROUPED_QUERY,
728            Some(resolver),
729            Some(alpha_offset..alpha_offset),
730        );
731
732        let run_texts: Vec<String> = runnables
733            .iter()
734            .map(|range| GROUPED_SOURCE[range.run_range.clone()].to_string())
735            .collect();
736        assert_eq!(
737            run_texts,
738            vec!["alpha"],
739            "zero-width offset_range at the start of a group should include that group"
740        );
741    }
742
743    #[gpui::test]
744    fn test_local_extras_override_shared_extras_with_same_key(cx: &mut TestAppContext) {
745        let resolver: Arc<dyn RunnableResolver> = Arc::new(OverrideSharedResolver);
746        let runnables = collect_runnables(cx, GROUPED_SOURCE, GROUPED_QUERY, Some(resolver));
747
748        for range in &runnables {
749            assert_eq!(
750                range.extra_captures.get("_outer").map(String::as_str),
751                Some("LOCAL"),
752                "local extras should override shared extras with the same key"
753            );
754        }
755        assert_eq!(runnables.len(), 3);
756    }
757}
758