tenant.openagents/omega

1use language::LineIndent;
2use std::{cmp, iter};
3
4#[derive(Copy, Clone, Debug, PartialEq, Eq)]
5pub enum IndentDelta {
6    Spaces(isize),
7    Tabs(isize),
8}
9
10impl IndentDelta {
11    pub fn character(&self) -> char {
12        match self {
13            IndentDelta::Spaces(_) => ' ',
14            IndentDelta::Tabs(_) => '\t',
15        }
16    }
17
18    pub fn len(&self) -> isize {
19        match self {
20            IndentDelta::Spaces(n) => *n,
21            IndentDelta::Tabs(n) => *n,
22        }
23    }
24}
25
26pub fn compute_indent_delta(buffer_indent: LineIndent, query_indent: LineIndent) -> IndentDelta {
27    if buffer_indent.tabs > 0 {
28        IndentDelta::Tabs(buffer_indent.tabs as isize - query_indent.tabs as isize)
29    } else {
30        IndentDelta::Spaces(buffer_indent.spaces as isize - query_indent.spaces as isize)
31    }
32}
33
34/// Computes the indent delta for the lines after the first, given per-line
35/// `(buffer, query)` indents for those lines.
36///
37/// When the remaining lines agree on a consistent delta, that delta is
38/// returned even if it differs from `first_line_delta`. This handles queries
39/// where only the first line's indentation was stripped. When the remaining
40/// lines are inconsistent (or all blank), falls back to `first_line_delta`,
41/// preserving the uniform re-indentation behavior.
42pub fn compute_rest_indent_delta(
43    first_line_delta: IndentDelta,
44    indent_pairs: impl IntoIterator<Item = (LineIndent, LineIndent)>,
45) -> IndentDelta {
46    let mut rest_delta = None;
47    for (buffer_indent, query_indent) in indent_pairs {
48        if buffer_indent.line_blank || query_indent.line_blank {
49            continue;
50        }
51        let delta = compute_indent_delta(buffer_indent, query_indent);
52        match rest_delta {
53            None => rest_delta = Some(delta),
54            Some(existing) if existing == delta => {}
55            Some(_) => return first_line_delta,
56        }
57    }
58    rest_delta.unwrap_or(first_line_delta)
59}
60
61/// Synchronous re-indentation adapter. Buffers incomplete lines and applies
62/// an `IndentDelta` to each line's leading whitespace before emitting it.
63///
64/// Models sometimes omit the leading indentation only on the first line of
65/// `old_text`/`new_text` (e.g. when copying from mid-line context), so the
66/// first line and the remaining lines can require different deltas.
67pub struct Reindenter {
68    first_line_delta: IndentDelta,
69    rest_delta: IndentDelta,
70    buffer: String,
71    in_leading_whitespace: bool,
72    on_first_line: bool,
73}
74
75impl Reindenter {
76    #[cfg(test)]
77    fn uniform(delta: IndentDelta) -> Self {
78        Self::with_deltas(delta, delta)
79    }
80
81    pub fn with_deltas(first_line_delta: IndentDelta, rest_delta: IndentDelta) -> Self {
82        Self {
83            first_line_delta,
84            rest_delta,
85            buffer: String::new(),
86            in_leading_whitespace: true,
87            on_first_line: true,
88        }
89    }
90
91    /// Feed a chunk of text and return the re-indented portion that is
92    /// ready to emit. Incomplete trailing lines are buffered internally.
93    pub fn push(&mut self, chunk: &str) -> String {
94        self.buffer.push_str(chunk);
95        self.drain(false)
96    }
97
98    /// Flush any remaining buffered content (call when the stream is done).
99    pub fn finish(&mut self) -> String {
100        self.drain(true)
101    }
102
103    fn drain(&mut self, is_final: bool) -> String {
104        let mut indented = String::new();
105        let mut start_ix = 0;
106        let mut newlines = self.buffer.match_indices('\n');
107        loop {
108            let (line_end, is_pending_line) = match newlines.next() {
109                Some((ix, _)) => (ix, false),
110                None => (self.buffer.len(), true),
111            };
112            let line = &self.buffer[start_ix..line_end];
113            let delta = if self.on_first_line {
114                self.first_line_delta
115            } else {
116                self.rest_delta
117            };
118
119            if self.in_leading_whitespace {
120                if let Some(non_whitespace_ix) = line.find(|c| delta.character() != c) {
121                    // We found a non-whitespace character, adjust indentation
122                    // based on the delta.
123                    let new_indent_len =
124                        cmp::max(0, non_whitespace_ix as isize + delta.len()) as usize;
125                    indented.extend(iter::repeat(delta.character()).take(new_indent_len));
126                    indented.push_str(&line[non_whitespace_ix..]);
127                    self.in_leading_whitespace = false;
128                } else if is_pending_line && !is_final {
129                    // We're still in leading whitespace and this line is incomplete.
130                    // Stop processing until we receive more input.
131                    break;
132                } else {
133                    // This line is entirely whitespace. Push it without indentation.
134                    indented.push_str(line);
135                }
136            } else {
137                indented.push_str(line);
138            }
139
140            if is_pending_line {
141                start_ix = line_end;
142                break;
143            } else {
144                self.in_leading_whitespace = true;
145                self.on_first_line = false;
146                indented.push('\n');
147                start_ix = line_end + 1;
148            }
149        }
150        self.buffer.replace_range(..start_ix, "");
151        if is_final {
152            indented.push_str(&self.buffer);
153            self.buffer.clear();
154        }
155        indented
156    }
157}
158
159#[cfg(test)]
160mod tests {
161    use super::*;
162
163    #[test]
164    fn test_indent_single_chunk() {
165        let mut r = Reindenter::uniform(IndentDelta::Spaces(2));
166        let out = r.push("    abc\n  def\n      ghi");
167        // All three lines are emitted: "ghi" starts with spaces but
168        // contains non-whitespace, so it's processed immediately.
169        assert_eq!(out, "      abc\n    def\n        ghi");
170        let out = r.finish();
171        assert_eq!(out, "");
172    }
173
174    #[test]
175    fn test_outdent_tabs() {
176        let mut r = Reindenter::uniform(IndentDelta::Tabs(-2));
177        let out = r.push("\t\t\t\tabc\n\t\tdef\n\t\t\t\t\t\tghi");
178        assert_eq!(out, "\t\tabc\ndef\n\t\t\t\tghi");
179        let out = r.finish();
180        assert_eq!(out, "");
181    }
182
183    #[test]
184    fn test_incremental_chunks() {
185        let mut r = Reindenter::uniform(IndentDelta::Spaces(2));
186        // Feed "    ab" — the `a` is non-whitespace, so the line is
187        // processed immediately even without a trailing newline.
188        let out = r.push("    ab");
189        assert_eq!(out, "      ab");
190        // Feed "c\n" — appended to the already-processed line (no longer
191        // in leading whitespace).
192        let out = r.push("c\n");
193        assert_eq!(out, "c\n");
194        let out = r.finish();
195        assert_eq!(out, "");
196    }
197
198    #[test]
199    fn test_zero_delta() {
200        let mut r = Reindenter::uniform(IndentDelta::Spaces(0));
201        let out = r.push("  hello\n  world\n");
202        assert_eq!(out, "  hello\n  world\n");
203        let out = r.finish();
204        assert_eq!(out, "");
205    }
206
207    #[test]
208    fn test_clamp_negative_indent() {
209        let mut r = Reindenter::uniform(IndentDelta::Spaces(-10));
210        let out = r.push("  abc\n");
211        // max(0, 2 - 10) = 0, so no leading spaces.
212        assert_eq!(out, "abc\n");
213        let out = r.finish();
214        assert_eq!(out, "");
215    }
216
217    #[test]
218    fn test_whitespace_only_lines() {
219        let mut r = Reindenter::uniform(IndentDelta::Spaces(2));
220        let out = r.push("   \n  code\n");
221        // First line is all whitespace — emitted verbatim. Second line is indented.
222        assert_eq!(out, "   \n    code\n");
223        let out = r.finish();
224        assert_eq!(out, "");
225    }
226
227    #[test]
228    fn test_distinct_first_line_delta() {
229        // First line's indentation was stripped in the query (delta +8),
230        // while the remaining lines are already correct (delta 0). Chunks
231        // split mid-line and mid-indentation to exercise the streaming path,
232        // and the blank line is passed through verbatim.
233        let mut r = Reindenter::with_deltas(IndentDelta::Spaces(8), IndentDelta::Spaces(0));
234        let mut out = r.push("self.target_a = ");
235        out.push_str(&r.push("\"after\"\n    "));
236        out.push_str(&r.push("    self.target_b = \"after\"\n"));
237        out.push_str(&r.push("\n        self.target_c = \"after\""));
238        out.push_str(&r.finish());
239        assert_eq!(
240            out,
241            concat!(
242                "        self.target_a = \"after\"\n",
243                "        self.target_b = \"after\"\n",
244                "\n",
245                "        self.target_c = \"after\"",
246            )
247        );
248    }
249
250    fn line_indent(text: &str) -> LineIndent {
251        LineIndent::from_iter(text.chars())
252    }
253
254    #[test]
255    fn test_compute_rest_indent_delta() {
256        let first_line_delta = IndentDelta::Spaces(8);
257
258        // Remaining lines that agree on a delta override the first-line
259        // delta, and blank lines are skipped when forming the consensus.
260        assert_eq!(
261            compute_rest_indent_delta(
262                first_line_delta,
263                vec![
264                    (line_indent("        b"), line_indent("        b")),
265                    (line_indent(""), line_indent("")),
266                    (line_indent("        c"), line_indent("        c")),
267                ],
268            ),
269            IndentDelta::Spaces(0)
270        );
271        assert_eq!(
272            compute_rest_indent_delta(
273                first_line_delta,
274                vec![
275                    (line_indent("        b"), line_indent("    b")),
276                    (line_indent("   "), line_indent("")),
277                    (line_indent("        c"), line_indent("    c")),
278                ],
279            ),
280            IndentDelta::Spaces(4)
281        );
282        assert_eq!(
283            compute_rest_indent_delta(
284                first_line_delta,
285                vec![(line_indent("\t\tb"), line_indent("\tb"))],
286            ),
287            IndentDelta::Tabs(1)
288        );
289
290        // Inconsistent remaining lines fall back to the first-line delta...
291        assert_eq!(
292            compute_rest_indent_delta(
293                first_line_delta,
294                vec![
295                    (line_indent("        b"), line_indent("        b")),
296                    (line_indent("        c"), line_indent("    c")),
297                ],
298            ),
299            first_line_delta
300        );
301
302        // ...and so do all-blank and empty pairings.
303        assert_eq!(
304            compute_rest_indent_delta(
305                first_line_delta,
306                vec![(line_indent("   "), line_indent(""))],
307            ),
308            first_line_delta
309        );
310        assert_eq!(
311            compute_rest_indent_delta(first_line_delta, vec![]),
312            first_line_delta
313        );
314    }
315
316    #[test]
317    fn test_compute_indent_delta_spaces() {
318        let buffer = LineIndent {
319            tabs: 0,
320            spaces: 8,
321            line_blank: false,
322        };
323        let query = LineIndent {
324            tabs: 0,
325            spaces: 4,
326            line_blank: false,
327        };
328        let delta = compute_indent_delta(buffer, query);
329        assert_eq!(delta.len(), 4);
330        assert_eq!(delta.character(), ' ');
331    }
332
333    #[test]
334    fn test_compute_indent_delta_tabs() {
335        let buffer = LineIndent {
336            tabs: 2,
337            spaces: 0,
338            line_blank: false,
339        };
340        let query = LineIndent {
341            tabs: 3,
342            spaces: 0,
343            line_blank: false,
344        };
345        let delta = compute_indent_delta(buffer, query);
346        assert_eq!(delta.len(), -1);
347        assert_eq!(delta.character(), '\t');
348    }
349}
350