1use criterion::{BenchmarkId, Criterion, black_box, criterion_group, criterion_main};
  2use edit_prediction_cli::kept_rate::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);