tenant.openagents/omega

1use crate::{CharClassifier, CharKind, CharScopeContext, LanguageScope};
2use anyhow::{Context, anyhow};
3use imara_diff::{Algorithm, Diff, InternedInput, Interner, Token, sources::lines};
4use std::{fmt::Write, iter, ops::Range, sync::Arc};
5
6const MAX_WORD_DIFF_LEN: usize = 512;
7const MAX_WORD_DIFF_LINE_COUNT: usize = 8;
8
9/// Computes a diff between two strings, returning a unified diff string.
10pub fn unified_diff(old_text: &str, new_text: &str) -> String {
11    unified_diff_with_offsets(old_text, new_text, 0, 0)
12}
13
14/// Computes a diff between two strings, returning a unified diff string with
15/// hunk headers adjusted to reflect the given starting line numbers (zero-indexed).
16pub fn unified_diff_with_offsets(
17    old_text: &str,
18    new_text: &str,
19    old_start_line: u32,
20    new_start_line: u32,
21) -> String {
22    unified_diff_with_context(old_text, new_text, old_start_line, new_start_line, 3)
23}
24
25/// Computes a diff between two strings, returning a unified diff string with
26/// hunk headers adjusted to reflect the given starting line numbers (zero-indexed),
27/// and a configurable number of context lines around changes.
28pub fn unified_diff_with_context(
29    old_text: &str,
30    new_text: &str,
31    old_start_line: u32,
32    new_start_line: u32,
33    context_lines: u32,
34) -> String {
35    // The builder appends its own line terminators, so tokenize without them.
36    let mut input = InternedInput::default();
37    input.update_before(old_text.lines());
38    input.update_after(new_text.lines());
39    let diff = Diff::compute(Algorithm::Histogram, &input);
40    let mut builder =
41        OffsetUnifiedDiffBuilder::new(&input, old_start_line, new_start_line, context_lines);
42    for hunk in diff.hunks() {
43        builder.process_change(hunk.before, hunk.after);
44    }
45    builder.finish()
46}
47
48/// A unified diff builder that applies line number offsets to hunk headers.
49struct OffsetUnifiedDiffBuilder<'a> {
50    before: &'a [Token],
51    after: &'a [Token],
52    interner: &'a Interner<&'a str>,
53
54    pos: u32,
55    before_hunk_start: u32,
56    after_hunk_start: u32,
57    before_hunk_len: u32,
58    after_hunk_len: u32,
59
60    old_line_offset: u32,
61    new_line_offset: u32,
62    context_lines: u32,
63
64    buffer: String,
65    dst: String,
66}
67
68impl<'a> OffsetUnifiedDiffBuilder<'a> {
69    fn new(
70        input: &'a InternedInput<&'a str>,
71        old_line_offset: u32,
72        new_line_offset: u32,
73        context_lines: u32,
74    ) -> Self {
75        Self {
76            before_hunk_start: 0,
77            after_hunk_start: 0,
78            before_hunk_len: 0,
79            after_hunk_len: 0,
80            old_line_offset,
81            new_line_offset,
82            context_lines,
83            buffer: String::with_capacity(8),
84            dst: String::new(),
85            interner: &input.interner,
86            before: &input.before,
87            after: &input.after,
88            pos: 0,
89        }
90    }
91
92    fn print_tokens(&mut self, tokens: &[Token], prefix: char) {
93        for &token in tokens {
94            writeln!(&mut self.buffer, "{prefix}{}", self.interner[token]).unwrap();
95        }
96    }
97
98    fn flush(&mut self) {
99        if self.before_hunk_len == 0 && self.after_hunk_len == 0 {
100            return;
101        }
102
103        let end = (self.pos + self.context_lines).min(self.before.len() as u32);
104        self.update_pos(end, end);
105
106        writeln!(
107            &mut self.dst,
108            "@@ -{},{} +{},{} @@",
109            self.before_hunk_start + 1 + self.old_line_offset,
110            self.before_hunk_len,
111            self.after_hunk_start + 1 + self.new_line_offset,
112            self.after_hunk_len,
113        )
114        .unwrap();
115        write!(&mut self.dst, "{}", &self.buffer).unwrap();
116        self.buffer.clear();
117        self.before_hunk_len = 0;
118        self.after_hunk_len = 0;
119    }
120
121    fn update_pos(&mut self, print_to: u32, move_to: u32) {
122        self.print_tokens(&self.before[self.pos as usize..print_to as usize], ' ');
123        let len = print_to - self.pos;
124        self.pos = move_to;
125        self.before_hunk_len += len;
126        self.after_hunk_len += len;
127    }
128}
129
130impl OffsetUnifiedDiffBuilder<'_> {
131    fn process_change(&mut self, before: Range<u32>, after: Range<u32>) {
132        if before.start - self.pos > self.context_lines * 2 {
133            self.flush();
134        }
135        if self.before_hunk_len == 0 && self.after_hunk_len == 0 {
136            self.pos = before.start.saturating_sub(self.context_lines);
137            self.before_hunk_start = self.pos;
138            self.after_hunk_start = after.start.saturating_sub(self.context_lines);
139        }
140        self.update_pos(before.start, before.end);
141        self.before_hunk_len += before.end - before.start;
142        self.after_hunk_len += after.end - after.start;
143        self.print_tokens(
144            &self.before[before.start as usize..before.end as usize],
145            '-',
146        );
147        self.print_tokens(&self.after[after.start as usize..after.end as usize], '+');
148    }
149
150    fn finish(mut self) -> String {
151        self.flush();
152        self.dst
153    }
154}
155
156/// Computes a diff between two strings, returning a vector of old and new row
157/// ranges.
158pub fn line_diff(old_text: &str, new_text: &str) -> Vec<(Range<u32>, Range<u32>)> {
159    let mut edits = Vec::new();
160    let input = InternedInput::new(lines(old_text), lines(new_text));
161    diff_internal(&input, &mut |_, _, old_rows, new_rows| {
162        edits.push((old_rows, new_rows));
163    });
164    edits
165}
166
167/// Computes a diff between two strings, returning a vector of edits.
168///
169/// The edits are represented as tuples of byte ranges and replacement strings.
170///
171/// Internally, this function first performs a line-based diff, and then performs a second
172/// word-based diff within hunks that replace small numbers of lines.
173pub fn text_diff(old_text: &str, new_text: &str) -> Vec<(Range<usize>, Arc<str>)> {
174    text_diff_with_options(old_text, new_text, DiffOptions::default())
175}
176
177/// Computes word-level diff ranges between two strings.
178///
179/// Returns a tuple of (old_ranges, new_ranges) where each vector contains
180/// the byte ranges of changed words in the respective text.
181pub fn word_diff_ranges(
182    old_text: &str,
183    new_text: &str,
184    options: DiffOptions,
185) -> (Vec<Range<usize>>, Vec<Range<usize>>) {
186    let mut input: InternedInput<&str> = InternedInput::default();
187    input.update_before(tokenize(old_text, options.language_scope.clone()));
188    input.update_after(tokenize(new_text, options.language_scope));
189
190    let mut old_ranges: Vec<Range<usize>> = Vec::new();
191    let mut new_ranges: Vec<Range<usize>> = Vec::new();
192
193    diff_internal(&input, &mut |old_byte_range, new_byte_range, _, _| {
194        if !old_byte_range.is_empty() {
195            if let Some(last) = old_ranges.last_mut()
196                && last.end >= old_byte_range.start
197            {
198                last.end = old_byte_range.end;
199            } else {
200                old_ranges.push(old_byte_range);
201            }
202        }
203
204        if !new_byte_range.is_empty() {
205            if let Some(last) = new_ranges.last_mut()
206                && last.end >= new_byte_range.start
207            {
208                last.end = new_byte_range.end;
209            } else {
210                new_ranges.push(new_byte_range);
211            }
212        }
213    });
214
215    (old_ranges, new_ranges)
216}
217
218/// Computes character-level diff between two strings.
219///
220/// Usually, you should use `text_diff`, which performs a word-wise diff.
221pub fn char_diff<'a>(old_text: &'a str, new_text: &'a str) -> Vec<(Range<usize>, &'a str)> {
222    let mut input: InternedInput<&str> = InternedInput::default();
223    input.update_before(tokenize_chars(old_text));
224    input.update_after(tokenize_chars(new_text));
225    let mut edits: Vec<(Range<usize>, &str)> = Vec::new();
226    diff_internal(&input, &mut |old_byte_range, new_byte_range, _, _| {
227        let replacement = if new_byte_range.is_empty() {
228            ""
229        } else {
230            &new_text[new_byte_range]
231        };
232        edits.push((old_byte_range, replacement));
233    });
234    edits
235}
236
237pub struct DiffOptions {
238    pub language_scope: Option<LanguageScope>,
239    pub max_word_diff_len: usize,
240    pub max_word_diff_line_count: usize,
241}
242
243impl Default for DiffOptions {
244    fn default() -> Self {
245        Self {
246            language_scope: Default::default(),
247            max_word_diff_len: MAX_WORD_DIFF_LEN,
248            max_word_diff_line_count: MAX_WORD_DIFF_LINE_COUNT,
249        }
250    }
251}
252
253/// Computes a diff between two strings, using a specific language scope's
254/// word characters for word-level diffing.
255pub fn text_diff_with_options(
256    old_text: &str,
257    new_text: &str,
258    options: DiffOptions,
259) -> Vec<(Range<usize>, Arc<str>)> {
260    let empty: Arc<str> = Arc::default();
261    let mut edits = Vec::new();
262    let mut hunk_input = InternedInput::default();
263    let input = InternedInput::new(lines(old_text), lines(new_text));
264    diff_internal(&input, &mut |old_byte_range,
265                                new_byte_range,
266                                old_rows,
267                                new_rows| {
268        if should_perform_word_diff_within_hunk(
269            &old_rows,
270            &old_byte_range,
271            &new_rows,
272            &new_byte_range,
273            &options,
274        ) {
275            let old_offset = old_byte_range.start;
276            let new_offset = new_byte_range.start;
277            hunk_input.clear();
278            hunk_input.update_before(tokenize(
279                &old_text[old_byte_range],
280                options.language_scope.clone(),
281            ));
282            hunk_input.update_after(tokenize(
283                &new_text[new_byte_range],
284                options.language_scope.clone(),
285            ));
286            diff_internal(&hunk_input, &mut |old_byte_range, new_byte_range, _, _| {
287                let old_byte_range =
288                    old_offset + old_byte_range.start..old_offset + old_byte_range.end;
289                let new_byte_range =
290                    new_offset + new_byte_range.start..new_offset + new_byte_range.end;
291                let replacement_text = if new_byte_range.is_empty() {
292                    empty.clone()
293                } else {
294                    new_text[new_byte_range].into()
295                };
296                edits.push((old_byte_range, replacement_text));
297            });
298        } else {
299            let replacement_text = if new_byte_range.is_empty() {
300                empty.clone()
301            } else {
302                new_text[new_byte_range].into()
303            };
304            edits.push((old_byte_range, replacement_text));
305        }
306    });
307    edits
308}
309
310pub fn apply_diff_patch(base_text: &str, patch: &str) -> Result<String, anyhow::Error> {
311    let patch = diffy::Patch::from_str(patch).context("Failed to parse patch")?;
312    let result = diffy::apply(base_text, &patch);
313    result.map_err(|err| anyhow!(err))
314}
315
316pub fn apply_reversed_diff_patch(base_text: &str, patch: &str) -> Result<String, anyhow::Error> {
317    let patch = diffy::Patch::from_str(patch).context("Failed to parse patch")?;
318    let reversed = patch.reverse();
319    diffy::apply(base_text, &reversed).map_err(|err| anyhow!(err))
320}
321
322fn should_perform_word_diff_within_hunk(
323    old_row_range: &Range<u32>,
324    old_byte_range: &Range<usize>,
325    new_row_range: &Range<u32>,
326    new_byte_range: &Range<usize>,
327    options: &DiffOptions,
328) -> bool {
329    !old_byte_range.is_empty()
330        && !new_byte_range.is_empty()
331        && old_byte_range.len() <= options.max_word_diff_len
332        && new_byte_range.len() <= options.max_word_diff_len
333        && old_row_range.len() <= options.max_word_diff_line_count
334        && new_row_range.len() <= options.max_word_diff_line_count
335}
336
337fn diff_internal(
338    input: &InternedInput<&str>,
339    on_change: &mut dyn FnMut(Range<usize>, Range<usize>, Range<u32>, Range<u32>),
340) {
341    let mut old_offset = 0;
342    let mut new_offset = 0;
343    let mut old_token_ix = 0;
344    let mut new_token_ix = 0;
345    let diff = Diff::compute(Algorithm::Histogram, input);
346    for hunk in diff.hunks() {
347        let old_tokens = hunk.before;
348        let new_tokens = hunk.after;
349        old_offset += token_len(
350            input,
351            &input.before[old_token_ix as usize..old_tokens.start as usize],
352        );
353        new_offset += token_len(
354            input,
355            &input.after[new_token_ix as usize..new_tokens.start as usize],
356        );
357        let old_len = token_len(
358            input,
359            &input.before[old_tokens.start as usize..old_tokens.end as usize],
360        );
361        let new_len = token_len(
362            input,
363            &input.after[new_tokens.start as usize..new_tokens.end as usize],
364        );
365        let old_byte_range = old_offset..old_offset + old_len;
366        let new_byte_range = new_offset..new_offset + new_len;
367        old_token_ix = old_tokens.end;
368        new_token_ix = new_tokens.end;
369        old_offset = old_byte_range.end;
370        new_offset = new_byte_range.end;
371        on_change(old_byte_range, new_byte_range, old_tokens, new_tokens);
372    }
373}
374
375fn tokenize_chars(text: &str) -> impl Iterator<Item = &str> {
376    let mut chars = text.char_indices().peekable();
377    iter::from_fn(move || {
378        let (start, c) = chars.next()?;
379        Some(&text[start..start + c.len_utf8()])
380    })
381}
382
383fn tokenize(text: &str, language_scope: Option<LanguageScope>) -> impl Iterator<Item = &str> {
384    let classifier =
385        CharClassifier::new(language_scope).scope_context(Some(CharScopeContext::Completion));
386    let mut chars = text.char_indices();
387    let mut prev = None;
388    let mut start_ix = 0;
389    iter::from_fn(move || {
390        for (ix, c) in chars.by_ref() {
391            let mut token = None;
392            let kind = classifier.kind(c);
393            if let Some((prev_char, prev_kind)) = prev
394                && (kind != prev_kind || (kind == CharKind::Punctuation && c != prev_char))
395            {
396                token = Some(&text[start_ix..ix]);
397                start_ix = ix;
398            }
399            prev = Some((c, kind));
400            if token.is_some() {
401                return token;
402            }
403        }
404        if start_ix < text.len() {
405            let token = &text[start_ix..];
406            start_ix = text.len();
407            return Some(token);
408        }
409        None
410    })
411}
412
413fn token_len(input: &InternedInput<&str>, tokens: &[Token]) -> usize {
414    tokens
415        .iter()
416        .map(|token| input.interner[*token].len())
417        .sum()
418}
419
420#[cfg(test)]
421mod tests {
422    use super::*;
423
424    #[test]
425    fn test_tokenize() {
426        let text = "";
427        assert_eq!(tokenize(text, None).collect::<Vec<_>>(), Vec::<&str>::new());
428
429        let text = " ";
430        assert_eq!(tokenize(text, None).collect::<Vec<_>>(), vec![" "]);
431
432        let text = "one";
433        assert_eq!(tokenize(text, None).collect::<Vec<_>>(), vec!["one"]);
434
435        let text = "one\n";
436        assert_eq!(tokenize(text, None).collect::<Vec<_>>(), vec!["one", "\n"]);
437
438        let text = "one.two(three)";
439        assert_eq!(
440            tokenize(text, None).collect::<Vec<_>>(),
441            vec!["one", ".", "two", "(", "three", ")"]
442        );
443
444        let text = "one two three()";
445        assert_eq!(
446            tokenize(text, None).collect::<Vec<_>>(),
447            vec!["one", " ", "two", " ", "three", "(", ")"]
448        );
449
450        let text = "   one\n two three";
451        assert_eq!(
452            tokenize(text, None).collect::<Vec<_>>(),
453            vec!["   ", "one", "\n ", "two", " ", "three"]
454        );
455    }
456
457    #[test]
458    fn test_text_diff() {
459        let old_text = "one two three";
460        let new_text = "one TWO three";
461        assert_eq!(text_diff(old_text, new_text), [(4..7, "TWO".into()),]);
462
463        let old_text = "one\ntwo\nthree\n";
464        let new_text = "one\ntwo\nAND\nTHEN\nthree\n";
465        assert_eq!(
466            text_diff(old_text, new_text),
467            [(8..8, "AND\nTHEN\n".into()),]
468        );
469
470        let old_text = "one two\nthree four five\nsix seven eight nine\nten\n";
471        let new_text = "one two\nthree FOUR five\nsix SEVEN eight nine\nten\nELEVEN\n";
472        assert_eq!(
473            text_diff(old_text, new_text),
474            [
475                (14..18, "FOUR".into()),
476                (28..33, "SEVEN".into()),
477                (49..49, "ELEVEN\n".into())
478            ]
479        );
480    }
481
482    #[test]
483    fn test_apply_diff_patch() {
484        let old_text = "one two\nthree four five\nsix seven eight nine\nten\n";
485        let new_text = "one two\nthree FOUR five\nsix SEVEN eight nine\nten\nELEVEN\n";
486        let patch = unified_diff(old_text, new_text);
487        assert_eq!(apply_diff_patch(old_text, &patch).unwrap(), new_text);
488    }
489
490    #[test]
491    fn test_apply_reversed_diff_patch() {
492        let old_text = "one two\nthree four five\nsix seven eight nine\nten\n";
493        let new_text = "one two\nthree FOUR five\nsix SEVEN eight nine\nten\nELEVEN\n";
494        let patch = unified_diff(old_text, new_text);
495        assert_eq!(
496            apply_reversed_diff_patch(new_text, &patch).unwrap(),
497            old_text
498        );
499    }
500
501    #[test]
502    fn test_char_diff() {
503        assert_eq!(char_diff("", ""), vec![]);
504        assert_eq!(char_diff("", "abc"), vec![(0..0, "abc")]);
505        assert_eq!(char_diff("abc", ""), vec![(0..3, "")]);
506        assert_eq!(char_diff("ac", "abc"), vec![(1..1, "b")]); // "b" inserted
507        assert_eq!(char_diff("abc", "ac"), vec![(1..2, "")]); // "b" deleted
508        assert_eq!(char_diff("abc", "adc"), vec![(1..2, "d")]); // "b" replaced with "d"
509        assert_eq!(char_diff("日", "日本語"), vec![(3..3, "本語")]); // "本語" inserted
510        assert_eq!(char_diff("日本語", "日"), vec![(3..9, "")]); // "本語" deleted
511        assert_eq!(char_diff("🎉", "🎉🎊🎈"), vec![(4..4, "🎊🎈")]); // "🎊🎈" inserted
512        assert_eq!(
513            char_diff("test日本", "test日本語です"),
514            vec![(10..10, "語です")]
515        );
516    }
517
518    #[test]
519    fn test_unified_diff_with_offsets() {
520        let old_text = "foo\nbar\nbaz\n";
521        let new_text = "foo\nBAR\nbaz\n";
522
523        let expected_diff_body = " foo\n-bar\n+BAR\n baz\n";
524
525        let diff_no_offset = unified_diff(old_text, new_text);
526        assert_eq!(
527            diff_no_offset,
528            format!("@@ -1,3 +1,3 @@\n{}", expected_diff_body)
529        );
530
531        let diff_with_offset = unified_diff_with_offsets(old_text, new_text, 9, 11);
532        assert_eq!(
533            diff_with_offset,
534            format!("@@ -10,3 +12,3 @@\n{}", expected_diff_body)
535        );
536
537        let diff_with_offset = unified_diff_with_offsets(old_text, new_text, 99, 104);
538        assert_eq!(
539            diff_with_offset,
540            format!("@@ -100,3 +105,3 @@\n{}", expected_diff_body)
541        );
542    }
543
544    #[test]
545    fn test_unified_diff_with_context() {
546        // Test that full context includes all lines from the start
547        let old_text = "line1\nline2\nline3\nline4\nline5\nCHANGE_ME\nline7\nline8\n";
548        let new_text = "line1\nline2\nline3\nline4\nline5\nCHANGED\nline7\nline8\n";
549
550        // With default 3 lines of context, the diff starts at line 3
551        let diff_default = unified_diff_with_offsets(old_text, new_text, 0, 0);
552        assert_eq!(
553            diff_default,
554            "@@ -3,6 +3,6 @@\n line3\n line4\n line5\n-CHANGE_ME\n+CHANGED\n line7\n line8\n"
555        );
556
557        // With full context (8 lines), the diff starts at line 1
558        let diff_full_context = unified_diff_with_context(old_text, new_text, 0, 0, 8);
559        assert_eq!(
560            diff_full_context,
561            "@@ -1,8 +1,8 @@\n line1\n line2\n line3\n line4\n line5\n-CHANGE_ME\n+CHANGED\n line7\n line8\n"
562        );
563
564        // With 0 context, only the changed line is shown
565        let diff_no_context = unified_diff_with_context(old_text, new_text, 0, 0, 0);
566        assert_eq!(diff_no_context, "@@ -6,1 +6,1 @@\n-CHANGE_ME\n+CHANGED\n");
567    }
568}
569