tenant.openagents/omega

1use aho_corasick::{AhoCorasick, AhoCorasickBuilder};
2use anyhow::{Ok, Result};
3use client::proto;
4use fancy_regex::{Captures, Regex, RegexBuilder};
5use gpui::Entity;
6use itertools::Itertools as _;
7use language::{Buffer, BufferSnapshot, CharKind};
8use smol::future::yield_now;
9use std::{
10    borrow::Cow,
11    io::{BufRead, BufReader, Read},
12    ops::Range,
13    sync::{Arc, LazyLock},
14};
15use text::Anchor;
16use util::{
17    paths::{PathMatcher, PathStyle},
18    rel_path::RelPath,
19};
20
21#[derive(Debug)]
22pub enum SearchResult {
23    Buffer {
24        buffer: Entity<Buffer>,
25        ranges: Vec<Range<Anchor>>,
26    },
27    LimitReached,
28    WaitingForScan,
29    Searching,
30}
31
32#[derive(Clone, Copy, PartialEq)]
33pub enum SearchInputKind {
34    Query,
35    Include,
36    Exclude,
37}
38
39#[derive(Clone, Debug)]
40pub struct SearchInputs {
41    query: Arc<str>,
42    files_to_include: PathMatcher,
43    files_to_exclude: PathMatcher,
44    match_full_paths: bool,
45    buffers: Option<Vec<Entity<Buffer>>>,
46}
47
48pub type LineHint = u32;
49
50impl SearchInputs {
51    pub fn as_str(&self) -> &str {
52        self.query.as_ref()
53    }
54    pub fn files_to_include(&self) -> &PathMatcher {
55        &self.files_to_include
56    }
57    pub fn files_to_exclude(&self) -> &PathMatcher {
58        &self.files_to_exclude
59    }
60    pub fn buffers(&self) -> &Option<Vec<Entity<Buffer>>> {
61        &self.buffers
62    }
63}
64#[derive(Clone, Debug)]
65pub enum SearchQuery {
66    Text {
67        search: AhoCorasick,
68        replacement: Option<String>,
69        whole_word: bool,
70        case_sensitive: bool,
71        include_ignored: bool,
72        inner: SearchInputs,
73    },
74    Regex {
75        regex: Regex,
76        replacement: Option<String>,
77        multiline: bool,
78        whole_word: bool,
79        case_sensitive: bool,
80        include_ignored: bool,
81        one_match_per_line: bool,
82        inner: SearchInputs,
83        escaped: bool,
84    },
85}
86
87static WORD_MATCH_TEST: LazyLock<Regex> = LazyLock::new(|| {
88    RegexBuilder::new(r"\B")
89        .build()
90        .expect("Failed to create WORD_MATCH_TEST")
91});
92
93impl SearchQuery {
94    /// Create a text query
95    ///
96    /// If `match_full_paths` is true, include/exclude patterns will always be matched against fully qualified project paths beginning with a project root.
97    /// If `match_full_paths` is false, patterns will be matched against worktree-relative paths.
98    pub fn text(
99        query: impl ToString,
100        whole_word: bool,
101        case_sensitive: bool,
102        include_ignored: bool,
103        files_to_include: PathMatcher,
104        files_to_exclude: PathMatcher,
105        match_full_paths: bool,
106        buffers: Option<Vec<Entity<Buffer>>>,
107    ) -> Result<Self> {
108        let mut query = query.to_string();
109        text::LineEnding::normalize(&mut query);
110        if !case_sensitive && !query.is_ascii() {
111            // AhoCorasickBuilder doesn't support case-insensitive search with unicode characters
112            // Fallback to regex search as recommended by
113            // https://docs.rs/aho-corasick/1.1/aho_corasick/struct.AhoCorasickBuilder.html#method.ascii_case_insensitive
114            return Self::escaped_regex(
115                query,
116                whole_word,
117                case_sensitive,
118                include_ignored,
119                files_to_include,
120                files_to_exclude,
121                match_full_paths,
122                buffers,
123            );
124        }
125        let search = AhoCorasickBuilder::new()
126            .ascii_case_insensitive(!case_sensitive)
127            .build([&query])?;
128        let inner = SearchInputs {
129            query: query.into(),
130            files_to_exclude,
131            files_to_include,
132            match_full_paths,
133            buffers,
134        };
135        Ok(Self::Text {
136            search,
137            replacement: None,
138            whole_word,
139            case_sensitive,
140            include_ignored,
141            inner,
142        })
143    }
144
145    /// Create a regex query
146    ///
147    /// If `match_full_paths` is true, include/exclude patterns will be matched against fully qualified project paths
148    /// beginning with a project root name. If false, they will be matched against project-relative paths (which don't start
149    /// with their respective project root).
150    pub fn regex(
151        query: impl ToString,
152        whole_word: bool,
153        case_sensitive: bool,
154        include_ignored: bool,
155        one_match_per_line: bool,
156        files_to_include: PathMatcher,
157        files_to_exclude: PathMatcher,
158        match_full_paths: bool,
159        buffers: Option<Vec<Entity<Buffer>>>,
160    ) -> Result<Self> {
161        let query = query.to_string();
162        let inner = SearchInputs {
163            query: Arc::from(query.as_str()),
164            files_to_include,
165            files_to_exclude,
166            match_full_paths,
167            buffers,
168        };
169        Self::build_regex(
170            query,
171            whole_word,
172            case_sensitive,
173            include_ignored,
174            one_match_per_line,
175            inner,
176            false,
177        )
178    }
179
180    /// Create a regex query from a literal string, escaping any regex
181    /// metacharacters so that the resulting query matches the literal text.
182    ///
183    /// Unlike `regex`, the query stored on the resulting `SearchQuery` is the
184    /// original unescaped text, so `as_str` returns what the user typed.
185    pub fn escaped_regex(
186        query: impl ToString,
187        whole_word: bool,
188        case_sensitive: bool,
189        include_ignored: bool,
190        files_to_include: PathMatcher,
191        files_to_exclude: PathMatcher,
192        match_full_paths: bool,
193        buffers: Option<Vec<Entity<Buffer>>>,
194    ) -> Result<Self> {
195        let mut query = query.to_string();
196        text::LineEnding::normalize(&mut query);
197        let inner = SearchInputs {
198            query: Arc::from(query.as_str()),
199            files_to_include,
200            files_to_exclude,
201            match_full_paths,
202            buffers,
203        };
204        Self::build_regex(
205            regex::escape(&query),
206            whole_word,
207            case_sensitive,
208            include_ignored,
209            false,
210            inner,
211            true,
212        )
213    }
214
215    fn build_regex(
216        mut pattern: String,
217        whole_word: bool,
218        mut case_sensitive: bool,
219        include_ignored: bool,
220        one_match_per_line: bool,
221        inner: SearchInputs,
222        escaped: bool,
223    ) -> Result<Self> {
224        if let Some((case_sensitive_from_pattern, new_pattern)) =
225            Self::case_sensitive_from_pattern(&pattern)
226        {
227            case_sensitive = case_sensitive_from_pattern;
228            pattern = new_pattern
229        }
230
231        if whole_word {
232            let mut word_pattern = String::new();
233            if let Some(first) = pattern.get(0..1)
234                && WORD_MATCH_TEST.is_match(first).is_ok_and(|x| !x)
235            {
236                word_pattern.push_str("\\b");
237            }
238            word_pattern.push_str(&pattern);
239            if let Some(last) = pattern.get(pattern.len() - 1..)
240                && WORD_MATCH_TEST.is_match(last).is_ok_and(|x| !x)
241            {
242                word_pattern.push_str("\\b");
243            }
244            pattern = word_pattern
245        }
246
247        let multiline = pattern.contains('\n') || pattern.contains("\\n");
248        if multiline {
249            pattern.insert_str(0, "(?m)");
250        }
251
252        let regex = RegexBuilder::new(&pattern)
253            .case_insensitive(!case_sensitive)
254            .crlf(true)
255            .build()?;
256        Ok(Self::Regex {
257            regex,
258            replacement: None,
259            multiline,
260            whole_word,
261            case_sensitive,
262            include_ignored,
263            inner,
264            one_match_per_line,
265            escaped,
266        })
267    }
268
269    /// Extracts case sensitivity settings from pattern items in the provided
270    /// query and returns the same query, with the pattern items removed.
271    ///
272    /// The following pattern modifiers are supported:
273    ///
274    /// - `\c` (case_sensitive: false)
275    /// - `\C` (case_sensitive: true)
276    ///
277    /// If no pattern item were found, `None` will be returned.
278    fn case_sensitive_from_pattern(query: &str) -> Option<(bool, String)> {
279        if !(query.contains("\\c") || query.contains("\\C")) {
280            return None;
281        }
282
283        let mut was_escaped = false;
284        let mut new_query = String::new();
285        let mut is_case_sensitive = None;
286
287        for c in query.chars() {
288            if was_escaped {
289                if c == 'c' {
290                    is_case_sensitive = Some(false);
291                } else if c == 'C' {
292                    is_case_sensitive = Some(true);
293                } else {
294                    new_query.push('\\');
295                    new_query.push(c);
296                }
297                was_escaped = false
298            } else if c == '\\' {
299                was_escaped = true
300            } else {
301                new_query.push(c);
302            }
303        }
304
305        is_case_sensitive.map(|c| (c, new_query))
306    }
307
308    pub fn from_proto(message: proto::SearchQuery, path_style: PathStyle) -> Result<Self> {
309        let files_to_include = if message.files_to_include.is_empty() {
310            message
311                .files_to_include_legacy
312                .split(',')
313                .map(str::trim)
314                .filter(|&glob_str| !glob_str.is_empty())
315                .map(|s| s.to_string())
316                .collect()
317        } else {
318            message.files_to_include
319        };
320
321        let files_to_exclude = if message.files_to_exclude.is_empty() {
322            message
323                .files_to_exclude_legacy
324                .split(',')
325                .map(str::trim)
326                .filter(|&glob_str| !glob_str.is_empty())
327                .map(|s| s.to_string())
328                .collect()
329        } else {
330            message.files_to_exclude
331        };
332
333        if message.regex {
334            Self::regex(
335                message.query,
336                message.whole_word,
337                message.case_sensitive,
338                message.include_ignored,
339                false,
340                PathMatcher::new(files_to_include, path_style)?,
341                PathMatcher::new(files_to_exclude, path_style)?,
342                message.match_full_paths,
343                None, // search opened only don't need search remote
344            )
345        } else {
346            Self::text(
347                message.query,
348                message.whole_word,
349                message.case_sensitive,
350                message.include_ignored,
351                PathMatcher::new(files_to_include, path_style)?,
352                PathMatcher::new(files_to_exclude, path_style)?,
353                message.match_full_paths,
354                None, // search opened only don't need search remote
355            )
356        }
357    }
358
359    pub fn with_replacement(mut self, new_replacement: String) -> Self {
360        match self {
361            Self::Text {
362                ref mut replacement,
363                ..
364            }
365            | Self::Regex {
366                ref mut replacement,
367                ..
368            } => {
369                *replacement = Some(new_replacement);
370                self
371            }
372        }
373    }
374
375    pub fn to_proto(&self) -> proto::SearchQuery {
376        let mut files_to_include = self.files_to_include().sources();
377        let mut files_to_exclude = self.files_to_exclude().sources();
378        proto::SearchQuery {
379            query: self.as_str().to_string(),
380            regex: self.is_regex(),
381            whole_word: self.whole_word(),
382            case_sensitive: self.case_sensitive(),
383            include_ignored: self.include_ignored(),
384            files_to_include: files_to_include.clone().map(ToOwned::to_owned).collect(),
385            files_to_exclude: files_to_exclude.clone().map(ToOwned::to_owned).collect(),
386            match_full_paths: self.match_full_paths(),
387            // Populate legacy fields for backwards compatibility
388            files_to_include_legacy: files_to_include.join(","),
389            files_to_exclude_legacy: files_to_exclude.join(","),
390        }
391    }
392
393    pub(crate) async fn detect(
394        &self,
395        mut reader: BufReader<Box<dyn Read + Send + Sync>>,
396    ) -> Result<Option<LineHint>> {
397        let query_str = self.as_str();
398        if query_str.is_empty() {
399            return Ok(None);
400        }
401
402        // Yield from this function every 20KB scanned.
403        const YIELD_THRESHOLD: usize = 20 * 1024;
404
405        match self {
406            Self::Text { search, .. } => {
407                let mut text = String::new();
408                if query_str.contains('\n') {
409                    reader.read_to_string(&mut text)?;
410                    text::LineEnding::normalize(&mut text);
411                    if search.is_match(&text) {
412                        Ok(Some(LineHint::default()))
413                    } else {
414                        Ok(None)
415                    }
416                } else {
417                    let mut bytes_read = 0;
418                    let mut line_number: LineHint = LineHint::default();
419                    while reader.read_line(&mut text)? > 0 {
420                        if search.is_match(&text) {
421                            return Ok(Some(line_number));
422                        }
423                        bytes_read += text.len();
424                        if bytes_read >= YIELD_THRESHOLD {
425                            bytes_read = 0;
426                            smol::future::yield_now().await;
427                        }
428                        text.clear();
429                        line_number += 1;
430                    }
431                    Ok(None)
432                }
433            }
434            Self::Regex {
435                regex, multiline, ..
436            } => {
437                let mut text = String::new();
438                if *multiline {
439                    reader.read_to_string(&mut text)?;
440                    text::LineEnding::normalize(&mut text);
441                    if regex.is_match(&text)? {
442                        Ok(Some(LineHint::default()))
443                    } else {
444                        Ok(None)
445                    }
446                } else {
447                    let mut bytes_read = 0;
448                    let mut line_number: LineHint = LineHint::default();
449                    while reader.read_line(&mut text)? > 0 {
450                        if regex.is_match(&text)? {
451                            return Ok(Some(line_number));
452                        }
453                        bytes_read += text.len();
454                        if bytes_read >= YIELD_THRESHOLD {
455                            bytes_read = 0;
456                            smol::future::yield_now().await;
457                        }
458                        text.clear();
459                        line_number += 1;
460                    }
461                    Ok(None)
462                }
463            }
464        }
465    }
466    /// Returns the replacement text for this `SearchQuery`.
467    pub fn replacement(&self) -> Option<&str> {
468        match self {
469            SearchQuery::Text { replacement, .. } | SearchQuery::Regex { replacement, .. } => {
470                replacement.as_deref()
471            }
472        }
473    }
474    /// Replaces search hits if replacement is set. `text` is assumed to be a string that matches this `SearchQuery` exactly, without any leftovers on either side.
475    pub fn replacement_for<'a>(&self, text: &'a str) -> Option<Cow<'a, str>> {
476        match self {
477            SearchQuery::Text { replacement, .. }
478            | SearchQuery::Regex {
479                replacement,
480                escaped: true,
481                ..
482            } => replacement.clone().map(Cow::from),
483
484            SearchQuery::Regex {
485                regex,
486                replacement: Some(replacement),
487                escaped: false,
488                ..
489            } => {
490                static TEXT_REPLACEMENT_SPECIAL_CHARACTERS_REGEX: LazyLock<Regex> =
491                    LazyLock::new(|| Regex::new(r"\\\\|\\n|\\t").unwrap());
492                let replacement = TEXT_REPLACEMENT_SPECIAL_CHARACTERS_REGEX.replace_all(
493                    replacement,
494                    |c: &Captures| match c.get(0).unwrap().as_str() {
495                        r"\\" => "\\",
496                        r"\n" => "\n",
497                        r"\t" => "\t",
498                        x => unreachable!("Unexpected escape sequence: {}", x),
499                    },
500                );
501                Some(regex.replace(text, replacement))
502            }
503
504            SearchQuery::Regex {
505                replacement: None, ..
506            } => None,
507        }
508    }
509
510    pub async fn search(
511        &self,
512        buffer: &BufferSnapshot,
513        subrange: Option<Range<usize>>,
514    ) -> Vec<Range<usize>> {
515        const YIELD_INTERVAL: usize = 20000;
516
517        if self.as_str().is_empty() {
518            return Default::default();
519        }
520
521        let range_offset = subrange.as_ref().map(|r| r.start).unwrap_or(0);
522        let rope = if let Some(range) = subrange {
523            buffer.as_rope().slice(range)
524        } else {
525            buffer.as_rope().clone()
526        };
527
528        let mut matches = Vec::new();
529        match self {
530            Self::Text {
531                search, whole_word, ..
532            } => {
533                for (ix, mat) in search
534                    .stream_find_iter(rope.bytes_in_range(0..rope.len()))
535                    .enumerate()
536                {
537                    if (ix + 1) % YIELD_INTERVAL == 0 {
538                        yield_now().await;
539                    }
540
541                    let mat = mat.unwrap();
542                    if *whole_word {
543                        let classifier = buffer.char_classifier_at(range_offset + mat.start());
544
545                        let prev_kind = rope
546                            .reversed_chars_at(mat.start())
547                            .next()
548                            .map(|c| classifier.kind(c));
549                        let start_kind =
550                            classifier.kind(rope.chars_at(mat.start()).next().unwrap());
551                        let end_kind =
552                            classifier.kind(rope.reversed_chars_at(mat.end()).next().unwrap());
553                        let next_kind = rope.chars_at(mat.end()).next().map(|c| classifier.kind(c));
554                        if (Some(start_kind) == prev_kind && start_kind == CharKind::Word)
555                            || (Some(end_kind) == next_kind && end_kind == CharKind::Word)
556                        {
557                            continue;
558                        }
559                    }
560                    matches.push(mat.start()..mat.end())
561                }
562            }
563
564            Self::Regex {
565                regex, multiline, ..
566            } => {
567                if *multiline {
568                    let text = rope.to_string();
569                    for (ix, mat) in regex.find_iter(&text).enumerate() {
570                        if (ix + 1) % YIELD_INTERVAL == 0 {
571                            yield_now().await;
572                        }
573
574                        if let std::result::Result::Ok(mat) = mat {
575                            matches.push(mat.start()..mat.end());
576                        }
577                    }
578                } else {
579                    let mut line = String::new();
580                    let mut line_offset = 0;
581                    for (chunk_ix, chunk) in rope.chunks().chain(["\n"]).enumerate() {
582                        if (chunk_ix + 1) % YIELD_INTERVAL == 0 {
583                            yield_now().await;
584                        }
585
586                        for (newline_ix, text) in chunk.split('\n').enumerate() {
587                            if newline_ix > 0 {
588                                for mat in regex.find_iter(&line).flatten() {
589                                    let start = line_offset + mat.start();
590                                    let end = line_offset + mat.end();
591                                    matches.push(start..end);
592                                    if self.one_match_per_line() == Some(true) {
593                                        break;
594                                    }
595                                }
596
597                                line_offset += line.len() + 1;
598                                line.clear();
599                            }
600                            line.push_str(text);
601                        }
602                    }
603                }
604            }
605        }
606
607        matches
608    }
609
610    pub fn is_empty(&self) -> bool {
611        self.as_str().is_empty()
612    }
613
614    pub fn as_str(&self) -> &str {
615        self.as_inner().as_str()
616    }
617
618    pub fn whole_word(&self) -> bool {
619        match self {
620            Self::Text { whole_word, .. } => *whole_word,
621            Self::Regex { whole_word, .. } => *whole_word,
622        }
623    }
624
625    pub fn case_sensitive(&self) -> bool {
626        match self {
627            Self::Text { case_sensitive, .. } => *case_sensitive,
628            Self::Regex { case_sensitive, .. } => *case_sensitive,
629        }
630    }
631
632    pub fn include_ignored(&self) -> bool {
633        match self {
634            Self::Text {
635                include_ignored, ..
636            } => *include_ignored,
637            Self::Regex {
638                include_ignored, ..
639            } => *include_ignored,
640        }
641    }
642
643    pub fn is_regex(&self) -> bool {
644        matches!(self, Self::Regex { .. })
645    }
646
647    pub fn files_to_include(&self) -> &PathMatcher {
648        self.as_inner().files_to_include()
649    }
650
651    pub fn files_to_exclude(&self) -> &PathMatcher {
652        self.as_inner().files_to_exclude()
653    }
654
655    pub fn buffers(&self) -> Option<&Vec<Entity<Buffer>>> {
656        self.as_inner().buffers.as_ref()
657    }
658
659    pub fn is_opened_only(&self) -> bool {
660        self.as_inner().buffers.is_some()
661    }
662
663    pub fn filters_path(&self) -> bool {
664        !(self.files_to_exclude().sources().next().is_none()
665            && self.files_to_include().sources().next().is_none())
666    }
667
668    pub fn match_full_paths(&self) -> bool {
669        self.as_inner().match_full_paths
670    }
671
672    /// Check match full paths to determine whether you're required to pass a fully qualified
673    /// project path (starts with a project root).
674    pub fn match_path(&self, file_path: &RelPath) -> bool {
675        let mut path = file_path.to_rel_path_buf();
676        loop {
677            if self.files_to_exclude().is_match(&path) {
678                return false;
679            } else if self.files_to_include().sources().next().is_none()
680                || self.files_to_include().is_match(&path)
681            {
682                return true;
683            } else if !path.pop() {
684                return false;
685            }
686        }
687    }
688    pub fn as_inner(&self) -> &SearchInputs {
689        match self {
690            Self::Regex { inner, .. } | Self::Text { inner, .. } => inner,
691        }
692    }
693
694    /// Whether this search should replace only one match per line, instead of
695    /// all matches.
696    /// Returns `None` for text searches, as only regex searches support this
697    /// option.
698    pub fn one_match_per_line(&self) -> Option<bool> {
699        match self {
700            Self::Regex {
701                one_match_per_line, ..
702            } => Some(*one_match_per_line),
703            Self::Text { .. } => None,
704        }
705    }
706
707    pub fn search_str(&self, text: &str) -> Vec<Range<usize>> {
708        if self.as_str().is_empty() {
709            return Vec::new();
710        }
711
712        let is_word_char = |c: char| c.is_alphanumeric() || c == '_';
713
714        let mut matches = Vec::new();
715        match self {
716            Self::Text {
717                search, whole_word, ..
718            } => {
719                for mat in search.find_iter(text.as_bytes()) {
720                    if *whole_word {
721                        let prev_char = text[..mat.start()].chars().last();
722                        let next_char = text[mat.end()..].chars().next();
723                        if prev_char.is_some_and(&is_word_char)
724                            || next_char.is_some_and(&is_word_char)
725                        {
726                            continue;
727                        }
728                    }
729                    matches.push(mat.start()..mat.end());
730                }
731            }
732            Self::Regex {
733                regex,
734                multiline,
735                one_match_per_line,
736                ..
737            } => {
738                if *multiline {
739                    for mat in regex.find_iter(text).flatten() {
740                        matches.push(mat.start()..mat.end());
741                    }
742                } else {
743                    let mut line_offset = 0;
744                    for line in text.split('\n') {
745                        for mat in regex.find_iter(line).flatten() {
746                            matches.push((line_offset + mat.start())..(line_offset + mat.end()));
747                            if *one_match_per_line {
748                                break;
749                            }
750                        }
751                        line_offset += line.len() + 1;
752                    }
753                }
754            }
755        }
756        matches
757    }
758}
759