tenant.openagents/omega

1use criterion::{BenchmarkId, Criterion, black_box, criterion_group, criterion_main};
2use edit_prediction::metrics::compute_kept_rate;
3
4fn repeated_function_lines(line_count: usize) -> String {
5    let mut text = String::with_capacity(line_count * 32);
6    for index in 0..line_count {
7        text.push_str("fn helper_");
8        text.push_str(&(index % 16).to_string());
9        text.push_str("() { value += old_name + 1; }\n");
10    }
11    text
12}
13
14fn localized_rename_inputs(line_count: usize) -> (String, String, String) {
15    let base = repeated_function_lines(line_count);
16    let mut predicted = base.clone();
17    let mut final_text = base.clone();
18
19    let needle = "value += old_name + 1;";
20    let prediction = "value += very_long_predicted_name + 1;";
21    let accepted = "value += new_name + 1;";
22
23    let offset = base
24        .rfind(needle)
25        .expect("expected needle in synthetic input");
26    let end = offset + needle.len();
27
28    predicted.replace_range(offset..end, prediction);
29    final_text.replace_range(offset..end, accepted);
30
31    (base, predicted, final_text)
32}
33
34fn identical_new_content_inputs(line_count: usize) -> (String, String, String) {
35    let predicted = repeated_function_lines(line_count);
36    (String::new(), predicted.clone(), predicted)
37}
38
39fn repetitive_token_inputs(token_repetitions: usize) -> (String, String, String) {
40    let repeated_old = "foo + foo + foo + foo + foo\n".repeat(token_repetitions);
41    let repeated_predicted = "foo + foo + prediction_token + foo + foo\n".repeat(token_repetitions);
42    let repeated_final = "foo + foo + kept_token + foo + foo\n".repeat(token_repetitions);
43    (repeated_old, repeated_predicted, repeated_final)
44}
45
46fn kept_rate_benchmark(c: &mut Criterion) {
47    let mut no_change_group = c.benchmark_group("kept_rate/no_change");
48    for line_count in [128usize, 512, 2048] {
49        let text = repeated_function_lines(line_count);
50        no_change_group.bench_with_input(
51            BenchmarkId::new("lines", line_count),
52            &text,
53            |bench, text| {
54                bench.iter(|| {
55                    black_box(compute_kept_rate(
56                        black_box(text),
57                        black_box(text),
58                        black_box(text),
59                    ));
60                });
61            },
62        );
63    }
64    no_change_group.finish();
65
66    let mut localized_group = c.benchmark_group("kept_rate/localized_rename");
67    for line_count in [128usize, 512, 2048] {
68        let inputs = localized_rename_inputs(line_count);
69        localized_group.bench_with_input(
70            BenchmarkId::new("lines", line_count),
71            &inputs,
72            |bench, inputs| {
73                let (base, predicted, final_text) = inputs;
74                bench.iter(|| {
75                    black_box(compute_kept_rate(
76                        black_box(base),
77                        black_box(predicted),
78                        black_box(final_text),
79                    ));
80                });
81            },
82        );
83    }
84    localized_group.finish();
85
86    let mut addition_group = c.benchmark_group("kept_rate/identical_addition");
87    for line_count in [128usize, 512, 2048] {
88        let inputs = identical_new_content_inputs(line_count);
89        addition_group.bench_with_input(
90            BenchmarkId::new("lines", line_count),
91            &inputs,
92            |bench, inputs| {
93                let (base, predicted, final_text) = inputs;
94                bench.iter(|| {
95                    black_box(compute_kept_rate(
96                        black_box(base),
97                        black_box(predicted),
98                        black_box(final_text),
99                    ));
100                });
101            },
102        );
103    }
104    addition_group.finish();
105
106    let mut repetitive_group = c.benchmark_group("kept_rate/repetitive_tokens");
107    for token_repetitions in [64usize, 256, 1024] {
108        let inputs = repetitive_token_inputs(token_repetitions);
109        repetitive_group.bench_with_input(
110            BenchmarkId::new("repetitions", token_repetitions),
111            &inputs,
112            |bench, inputs| {
113                let (base, predicted, final_text) = inputs;
114                bench.iter(|| {
115                    black_box(compute_kept_rate(
116                        black_box(base),
117                        black_box(predicted),
118                        black_box(final_text),
119                    ));
120                });
121            },
122        );
123    }
124    repetitive_group.finish();
125}
126
127criterion_group!(benches, kept_rate_benchmark);
128criterion_main!(benches);
129