tenant.openagents/omega

1use criterion::{
2    BatchSize, BenchmarkId, Criterion, Throughput, black_box, criterion_group, criterion_main,
3};
4use rand::{Rng as _, SeedableRng as _, rngs::StdRng};
5use streaming_diff::StreamingDiff;
6
7const SEED: u64 = 0x5EED_5EED;
8const CHUNK_SIZE: usize = 512;
9
10#[derive(Clone)]
11struct EditFixture {
12    name: &'static str,
13    old_text: String,
14    new_text: String,
15}
16
17fn streaming_diff_push_new(criterion: &mut Criterion) {
18    let fixtures = fixtures();
19    let mut group = criterion.benchmark_group("streaming_diff_push_new");
20    group.sample_size(10);
21
22    for fixture in fixtures {
23        group.throughput(Throughput::Bytes(fixture.new_text.len() as u64));
24        group.bench_with_input(
25            BenchmarkId::new(fixture.name, fixture.old_text.len()),
26            &fixture,
27            |bench, fixture| {
28                bench.iter_batched(
29                    || StreamingDiff::new(fixture.old_text.clone()),
30                    |mut diff| {
31                        let mut operation_count = 0;
32                        for chunk in chunk_text(&fixture.new_text, CHUNK_SIZE) {
33                            operation_count += black_box(diff.push_new(chunk)).len();
34                        }
35                        black_box(operation_count);
36                    },
37                    BatchSize::SmallInput,
38                );
39            },
40        );
41    }
42
43    group.finish();
44}
45
46fn streaming_diff_finish(criterion: &mut Criterion) {
47    let fixtures = fixtures();
48    let mut group = criterion.benchmark_group("streaming_diff_finish");
49    group.sample_size(10);
50
51    for fixture in fixtures {
52        group.throughput(Throughput::Bytes(fixture.new_text.len() as u64));
53        group.bench_with_input(
54            BenchmarkId::new(fixture.name, fixture.old_text.len()),
55            &fixture,
56            |bench, fixture| {
57                bench.iter_batched(
58                    || {
59                        let mut diff = StreamingDiff::new(fixture.old_text.clone());
60                        for chunk in chunk_text(&fixture.new_text, CHUNK_SIZE) {
61                            black_box(diff.push_new(chunk));
62                        }
63                        diff
64                    },
65                    |diff| {
66                        black_box(diff.finish());
67                    },
68                    BatchSize::SmallInput,
69                );
70            },
71        );
72    }
73
74    group.finish();
75}
76
77fn fixtures() -> Vec<EditFixture> {
78    // Keep fixtures modest because `StreamingDiff` is intentionally stressed here and
79    // can become very slow on tens of kilobytes of replacement text. These sizes still
80    // represent realistic `edit_file` old/new text blocks and are large enough to cross
81    // frame-budget-sized CPU work.
82    vec![
83        make_fixture(
84            "tiny_function_rewrite",
85            2,
86            EditPattern::LocalizedRewrite {
87                start_line: 12,
88                line_count: 6,
89            },
90            SEED,
91        ),
92        make_fixture(
93            "small_function_rewrite",
94            5,
95            EditPattern::LocalizedRewrite {
96                start_line: 22,
97                line_count: 12,
98            },
99            SEED + 1,
100        ),
101        make_fixture(
102            "medium_many_small_changes",
103            8,
104            EditPattern::ManySmallChanges { every_nth_line: 7 },
105            SEED + 2,
106        ),
107        make_fixture(
108            "medium_insertions",
109            8,
110            EditPattern::InsertHelperBlocks { every_nth_line: 9 },
111            SEED + 3,
112        ),
113    ]
114}
115
116enum EditPattern {
117    LocalizedRewrite {
118        start_line: usize,
119        line_count: usize,
120    },
121    ManySmallChanges {
122        every_nth_line: usize,
123    },
124    InsertHelperBlocks {
125        every_nth_line: usize,
126    },
127}
128
129fn make_fixture(
130    name: &'static str,
131    function_count: usize,
132    pattern: EditPattern,
133    seed: u64,
134) -> EditFixture {
135    let mut rng = StdRng::seed_from_u64(seed);
136    let mut lines = random_rust_module(&mut rng, function_count);
137    let old_text = lines.join("\n");
138
139    match pattern {
140        EditPattern::LocalizedRewrite {
141            start_line,
142            line_count,
143        } => rewrite_local_block(&mut lines, start_line, line_count, &mut rng),
144        EditPattern::ManySmallChanges { every_nth_line } => {
145            rewrite_many_small_lines(&mut lines, every_nth_line, &mut rng)
146        }
147        EditPattern::InsertHelperBlocks { every_nth_line } => {
148            insert_helper_blocks(&mut lines, every_nth_line, &mut rng)
149        }
150    }
151
152    EditFixture {
153        name,
154        old_text,
155        new_text: lines.join("\n"),
156    }
157}
158
159fn random_rust_module(rng: &mut StdRng, function_count: usize) -> Vec<String> {
160    let mut lines = vec![
161        "use anyhow::{Context as _, Result};".to_string(),
162        "use collections::HashMap;".to_string(),
163        "".to_string(),
164        "#[derive(Clone, Debug)]".to_string(),
165        "pub struct WorkspaceSnapshot {".to_string(),
166        "    buffers: HashMap<String, usize>,".to_string(),
167        "    version: usize,".to_string(),
168        "}".to_string(),
169        "".to_string(),
170        "impl WorkspaceSnapshot {".to_string(),
171    ];
172
173    for function_index in 0..function_count {
174        let function_name = identifier(rng, function_index);
175        let argument_name = identifier(rng, function_index + 1_000);
176        let local_name = identifier(rng, function_index + 2_000);
177        let branch_name = identifier(rng, function_index + 3_000);
178        let multiplier = rng.random_range(2..17);
179        let offset = rng.random_range(1..128);
180
181        lines.extend([
182            format!(
183                "    pub fn {function_name}(&mut self, {argument_name}: usize) -> Result<usize> {{"
184            ),
185            format!("        let mut {local_name} = {argument_name}.saturating_mul({multiplier});"),
186            format!("        if {local_name} % 2 == 0 {{"),
187            format!(
188                "            {local_name} = {local_name}.saturating_add(self.version + {offset});"
189            ),
190            "        } else {".to_string(),
191            format!("            {local_name} = {local_name}.saturating_sub({offset});"),
192            "        }".to_string(),
193            format!("        let {branch_name} = self.buffers.len().saturating_add({local_name});"),
194            format!("        self.version = self.version.saturating_add({branch_name});"),
195            format!("        Ok({branch_name})"),
196            "    }".to_string(),
197            "".to_string(),
198        ]);
199    }
200
201    lines.push("}".to_string());
202    lines.push("".to_string());
203    lines.push("pub fn normalize_path(path: &str) -> String {".to_string());
204    lines.push("    path.replace('\\\\', \"/\")".to_string());
205    lines.push("}".to_string());
206    lines
207}
208
209fn rewrite_local_block(
210    lines: &mut [String],
211    start_line: usize,
212    line_count: usize,
213    rng: &mut StdRng,
214) {
215    let end_line = (start_line + line_count).min(lines.len());
216    for (relative_index, line) in lines[start_line..end_line].iter_mut().enumerate() {
217        let suffix = identifier(rng, relative_index + 10_000);
218        if line.contains("saturating_add") {
219            *line = format!(
220                "        let {suffix} = self.version.checked_add({relative_index}).context(\"version overflow\")?;"
221            );
222        } else if line.contains("saturating_sub") {
223            *line = format!(
224                "            {suffix}.saturating_sub({});",
225                rng.random_range(8..256)
226            );
227        } else if line.trim().is_empty() {
228            *line = format!(
229                "        tracing::trace!(target: \"agent_bench\", value = {relative_index});"
230            );
231        } else {
232            *line = format!("{line} // updated {suffix}");
233        }
234    }
235}
236
237fn rewrite_many_small_lines(lines: &mut [String], every_nth_line: usize, rng: &mut StdRng) {
238    for (line_index, line) in lines.iter_mut().enumerate() {
239        if line_index % every_nth_line != 0 || line.trim().is_empty() {
240            continue;
241        }
242
243        if line.contains("let mut") {
244            *line = line.replace("let mut", "let mut updated");
245        } else if line.contains("Ok(") {
246            *line = line.replace("Ok(", "Ok(black_box_value(");
247        } else if line.ends_with('{') {
248            *line = format!("{line} // scenario {}", identifier(rng, line_index));
249        } else {
250            *line = format!("{line} // touched {}", identifier(rng, line_index));
251        }
252    }
253}
254
255fn insert_helper_blocks(lines: &mut Vec<String>, every_nth_line: usize, rng: &mut StdRng) {
256    let mut line_index = every_nth_line;
257    while line_index < lines.len() {
258        if lines[line_index].trim() == "}" {
259            let helper_name = identifier(rng, line_index + 20_000);
260            lines.splice(
261                line_index..line_index,
262                [
263                    format!("        let {helper_name} = self.buffers.len();"),
264                    format!("        tracing::trace!(target: \"agent_bench\", {helper_name});"),
265                ],
266            );
267            line_index += 2;
268        }
269        line_index += every_nth_line;
270    }
271}
272
273fn identifier(rng: &mut StdRng, salt: usize) -> String {
274    const WORDS: &[&str] = &[
275        "buffer",
276        "workspace",
277        "snapshot",
278        "version",
279        "project",
280        "entry",
281        "path",
282        "cursor",
283        "anchor",
284        "edit",
285        "thread",
286        "message",
287        "context",
288        "store",
289        "diff",
290        "range",
291        "token",
292        "parser",
293        "semantic",
294        "format",
295        "completion",
296        "diagnostic",
297        "terminal",
298        "channel",
299    ];
300
301    let first = WORDS[(rng.random_range(0..WORDS.len()) + salt) % WORDS.len()];
302    let second = WORDS[(rng.random_range(0..WORDS.len()) + salt / 3) % WORDS.len()];
303    format!("{first}_{second}_{salt}")
304}
305
306fn chunk_text(text: &str, max_chunk_size: usize) -> Vec<&str> {
307    let mut chunks = Vec::new();
308    let mut start = 0;
309    while start < text.len() {
310        let mut end = (start + max_chunk_size).min(text.len());
311        while end < text.len() && !text.is_char_boundary(end) {
312            end += 1;
313        }
314        chunks.push(&text[start..end]);
315        start = end;
316    }
317    chunks
318}
319
320criterion_group!(benches, streaming_diff_push_new, streaming_diff_finish);
321criterion_main!(benches);
322