json-extractor 0.1.0

use std::hint::black_box;

use criterion::{Criterion, Throughput, criterion_group, criterion_main};
use json_extractor::StagedScanner;

fn generate_simple_json(count: usize) -> Vec<u8> {
    let mut result = Vec::new();
    for i in 0..count {
        result.extend_from_slice(format!(r#"{{"id":{},"name":"user{}"}}"#, i, i).as_bytes());
        result.push(b' ');
    }
    result
}

fn generate_nested_json(depth: usize) -> Vec<u8> {
    let mut json = String::new();
    json.push('{');
    for i in 0..depth {
        json.push_str(&format!(r#""level{}":{{"#, i));
    }
    json.push_str(r#""value":123"#);
    for _ in 0..=depth {
        json.push('}');
    }
    json.into_bytes()
}

fn generate_array_json(size: usize) -> Vec<u8> {
    let mut json = String::from("[");
    for i in 0..size {
        if i > 0 {
            json.push(',');
        }
        json.push_str(&i.to_string());
    }
    json.push(']');
    json.into_bytes()
}

fn bench_simple_fragments(c: &mut Criterion) {
    let mut group = c.benchmark_group("simple_fragments");

    for count in [10, 100, 1000] {
        let data = generate_simple_json(count);
        group.throughput(Throughput::Bytes(data.len() as u64));

        group.bench_function(format!("{}_objects", count), |b| {
            let mut scanner = StagedScanner::new();
            b.iter(|| {
                let fragments = scanner.scan_fragments(&data);
                black_box(fragments);
            });
        });
    }

    group.finish();
}

fn bench_nested_fragments(c: &mut Criterion) {
    let mut group = c.benchmark_group("nested_fragments");

    for depth in [5, 10, 20, 50] {
        let data = generate_nested_json(depth);
        group.throughput(Throughput::Bytes(data.len() as u64));

        group.bench_function(format!("depth_{}", depth), |b| {
            let mut scanner = StagedScanner::new();
            b.iter(|| {
                let fragments = scanner.scan_fragments(black_box(&data));
                black_box(fragments);
            });
        });
    }

    group.finish();
}

fn bench_array_fragments(c: &mut Criterion) {
    let mut group = c.benchmark_group("array_fragments");

    for size in [10, 100, 1000, 10000] {
        let data = generate_array_json(size);
        group.throughput(Throughput::Bytes(data.len() as u64));

        group.bench_function(format!("{}_elements", size), |b| {
            let mut scanner = StagedScanner::new();
            b.iter(|| {
                let fragments = scanner.scan_fragments(black_box(&data));
                black_box(fragments);
            });
        });
    }

    group.finish();
}

fn bench_document_sizes(c: &mut Criterion) {
    let mut group = c.benchmark_group("document_sizes");

    // Test performance across different document sizes
    for object_count in [10, 100, 1000, 10000] {
        let data = generate_simple_json(object_count);
        group.throughput(Throughput::Bytes(data.len() as u64));

        group.bench_function(format!("{}_objects", object_count), |b| {
            let mut scanner = StagedScanner::new();
            b.iter(|| {
                let fragments = scanner.scan_fragments(black_box(&data));
                black_box(fragments);
            });
        });
    }

    group.finish();
}

fn bench_string_heavy(c: &mut Criterion) {
    let mut group = c.benchmark_group("string_heavy");

    // JSON with lots of strings and escapes
    let json_simple = br#"{"text": "simple string value without escapes"}"#;
    let json_escapes =
        br#"{"text": "string with \"quotes\" and \\backslashes\\ and \n\t escapes"}"#;
    let json_long = format!(r#"{{"text": "{}"}}"#, "a".repeat(1000));

    group.throughput(Throughput::Bytes(json_simple.len() as u64));
    group.bench_function("simple_string", |b| {
        let mut scanner = StagedScanner::new();
        b.iter(|| {
            let fragments = scanner.scan_fragments(black_box(json_simple));
            black_box(fragments);
        });
    });

    group.throughput(Throughput::Bytes(json_escapes.len() as u64));
    group.bench_function("escaped_string", |b| {
        let mut scanner = StagedScanner::new();
        b.iter(|| {
            let fragments = scanner.scan_fragments(black_box(json_escapes));
            black_box(fragments);
        });
    });

    group.throughput(Throughput::Bytes(json_long.len() as u64));
    group.bench_function("long_string", |b| {
        let mut scanner = StagedScanner::new();
        b.iter(|| {
            let fragments = scanner.scan_fragments(black_box(json_long.as_bytes()));
            black_box(fragments);
        });
    });

    group.finish();
}

fn bench_error_cases(c: &mut Criterion) {
    let mut group = c.benchmark_group("error_cases");

    let invalid = br#"{"bad": !!!}"#;
    let mismatched = br#"{"test": [1, 2}"#;
    let incomplete = br#"{"incomplete": 123"#;

    group.throughput(Throughput::Bytes(invalid.len() as u64));
    group.bench_function("invalid_character", |b| {
        let mut scanner = StagedScanner::new();
        b.iter(|| {
            let fragments = scanner.scan_fragments(black_box(invalid));
            black_box(fragments);
        });
    });

    group.throughput(Throughput::Bytes(mismatched.len() as u64));
    group.bench_function("mismatched_bracket", |b| {
        let mut scanner = StagedScanner::new();
        b.iter(|| {
            let fragments = scanner.scan_fragments(black_box(mismatched));
            black_box(fragments);
        });
    });

    group.throughput(Throughput::Bytes(incomplete.len() as u64));
    group.bench_function("incomplete_fragment", |b| {
        let mut scanner = StagedScanner::new();
        b.iter(|| {
            let fragments = scanner.scan_fragments(black_box(incomplete));
            black_box(fragments);
        });
    });

    group.finish();
}

fn bench_real_world(c: &mut Criterion) {
    let mut group = c.benchmark_group("real_world");

    // Simulate a log file with mixed content and JSON fragments
    let log_data = br#"
[2024-01-01 10:00:00] INFO Starting application
[2024-01-01 10:00:01] DEBUG {"event":"startup","pid":1234,"memory_mb":128}
Some random log text without JSON
[2024-01-01 10:00:02] INFO {"event":"request","method":"GET","path":"/api/users","status":200}
More log text
[2024-01-01 10:00:03] ERROR {"event":"error","type":"connection_failed","details":{"host":"db.example.com","port":5432}}
[2024-01-01 10:00:04] INFO Application running
    "#;

    group.throughput(Throughput::Bytes(log_data.len() as u64));
    group.bench_function("mixed_log_file", |b| {
        let mut scanner = StagedScanner::new();
        b.iter(|| {
            let fragments = scanner.scan_fragments(black_box(log_data));
            black_box(fragments);
        });
    });

    group.finish();
}

criterion_group!(
    benches,
    bench_simple_fragments,
    bench_nested_fragments,
    bench_array_fragments,
    bench_string_heavy,
    bench_error_cases,
    bench_real_world,
    bench_document_sizes,
);
criterion_main!(benches);