benchmark.rs

// SPDX-License-Identifier: MPL-2.0
// Copyright (c) 2026 Jonathan D.A. Jewell <j.d.a.jewell@open.ac.uk>
//
// Benchmark comparison generator for julianiser.
//
// Produces scripts that measure execution time of the original Python/R
// code alongside the generated Julia equivalent, then computes and
// reports the speedup factor. This is the key deliverable of julianiser:
// demonstrating that Julia JIT compilation achieves C-like performance
// on data-science workloads.
//
// Generated artifacts:
//   benchmark.jl         — Julia benchmark script using BenchmarkTools
//   benchmark_runner.sh  — Shell script that runs both original and Julia
//   results.toml         — Template for recording benchmark results

use crate::abi::{BenchmarkResult, SourceLanguage, TranslationUnit};
use crate::manifest::Manifest;
use anyhow::Result;
use std::fmt::Write as FmtWrite;
use std::fs;
use std::path::Path;

/// Generate all benchmark artifacts for a set of translation units.
///
/// Creates:
/// 1. A Julia benchmark script that uses BenchmarkTools.jl
/// 2. A shell runner script that executes original and Julia code
/// 3. A TOML template for recording results
///
/// Returns the list of generated file paths.
pub fn generate_benchmarks(
    manifest: &Manifest,
    units: &[TranslationUnit],
    output_dir: &Path,
) -> Result<Vec<String>> {
    let bench_dir = output_dir.join("benchmarks");
    fs::create_dir_all(&bench_dir)?;

    let mut generated = Vec::new();

    // Generate Julia benchmark script.
    let julia_bench = generate_julia_benchmark(manifest, units);
    let julia_path = bench_dir.join("benchmark.jl");
    fs::write(&julia_path, &julia_bench)?;
    generated.push(julia_path.display().to_string());
    println!("  [bench] Generated {}", julia_path.display());

    // Generate shell runner.
    let runner = generate_benchmark_runner(manifest, units);
    let runner_path = bench_dir.join("benchmark_runner.sh");
    fs::write(&runner_path, &runner)?;
    // Make executable.
    #[cfg(unix)]
    {
        use std::os::unix::fs::PermissionsExt;
        let mut perms = fs::metadata(&runner_path)?.permissions();
        perms.set_mode(0o755);
        fs::set_permissions(&runner_path, perms)?;
    }
    generated.push(runner_path.display().to_string());
    println!("  [bench] Generated {}", runner_path.display());

    // Generate results template.
    let results = generate_results_template(units);
    let results_path = bench_dir.join("results.toml");
    fs::write(&results_path, &results)?;
    generated.push(results_path.display().to_string());
    println!("  [bench] Generated {}", results_path.display());

    Ok(generated)
}

/// Generate a Julia benchmark script using BenchmarkTools.jl.
///
/// The script imports each generated module, wraps the `run_pipeline()`
/// function in `@benchmark`, and collects timing statistics. Output is
/// printed as a human-readable table with median, mean, and allocation info.
fn generate_julia_benchmark(manifest: &Manifest, units: &[TranslationUnit]) -> String {
    let mut out = String::new();

    writeln!(out, "# SPDX-License-Identifier: MPL-2.0").expect("TODO: handle error");
    writeln!(out, "# Auto-generated by julianiser — benchmark script").expect("TODO: handle error");
    writeln!(out, "# Project: {}", manifest.project.name).expect("TODO: handle error");
    writeln!(out, "#").expect("TODO: handle error");
    writeln!(out, "# Run with: julia --project=. benchmarks/benchmark.jl")
        .expect("TODO: handle error");
    writeln!(out).expect("TODO: handle error");

    writeln!(out, "using BenchmarkTools").expect("TODO: handle error");
    writeln!(out, "using Printf").expect("TODO: handle error");
    writeln!(out).expect("TODO: handle error");

    // Include each generated module.
    for unit in units {
        writeln!(
            out,
            "include(\"../{}\")  # from {}",
            unit.output_path, unit.source_path
        )
        .expect("TODO: handle error");
        writeln!(out, "using .{}", unit.module_name).expect("TODO: handle error");
    }
    writeln!(out).expect("TODO: handle error");

    writeln!(out, "println(\"=\"^60)").expect("TODO: handle error");
    writeln!(
        out,
        "println(\"julianiser Benchmark: {}\")  ",
        manifest.project.name
    )
    .expect("TODO: handle error");
    writeln!(out, "println(\"=\"^60)").expect("TODO: handle error");
    writeln!(out).expect("TODO: handle error");

    writeln!(out, "# Warm up Julia's JIT compiler before benchmarking.")
        .expect("TODO: handle error");
    writeln!(out, "println(\"Warming up JIT...\")").expect("TODO: handle error");
    for unit in units {
        writeln!(out, "try").expect("TODO: handle error");
        writeln!(out, "    {}.run_pipeline()", unit.module_name).expect("TODO: handle error");
        writeln!(out, "catch e").expect("TODO: handle error");
        writeln!(
            out,
            "    @warn \"Warm-up failed for {}: $e\"",
            unit.module_name
        )
        .expect("TODO: handle error");
        writeln!(out, "end").expect("TODO: handle error");
    }
    writeln!(out).expect("TODO: handle error");

    writeln!(out, "println(\"\\nRunning benchmarks...\")").expect("TODO: handle error");
    writeln!(out, "println(\"-\"^60)").expect("TODO: handle error");
    writeln!(out).expect("TODO: handle error");

    // Benchmark each module.
    for unit in units {
        writeln!(
            out,
            "println(\"\\nBenchmark: {} (from {})\")  ",
            unit.module_name, unit.source_path
        )
        .expect("TODO: handle error");
        writeln!(
            out,
            "bench_{} = @benchmark {}.run_pipeline()",
            to_snake(&unit.module_name),
            unit.module_name
        )
        .expect("TODO: handle error");
        writeln!(out, "display(bench_{})", to_snake(&unit.module_name))
            .expect("TODO: handle error");
        writeln!(out, "println()").expect("TODO: handle error");
        writeln!(out).expect("TODO: handle error");
    }

    // Summary.
    writeln!(out, "println(\"=\"^60)").expect("TODO: handle error");
    writeln!(out, "println(\"Summary\")").expect("TODO: handle error");
    writeln!(out, "println(\"=\"^60)").expect("TODO: handle error");
    for unit in units {
        let snake = to_snake(&unit.module_name);
        writeln!(
            out,
            "@printf(\"  {}: median = %.3f ms, allocs = %d\\n\", median(bench_{}).time / 1e6, bench_{}.allocs)",
            unit.module_name, snake, snake
        ).expect("TODO: handle error");
    }
    writeln!(out).expect("TODO: handle error");
    writeln!(
        out,
        "println(\"\\nCompare these with your original {} timings.\")  ",
        units
            .first()
            .map(|u| u.language.to_string())
            .unwrap_or_default()
    )
    .expect("TODO: handle error");
    writeln!(
        out,
        "println(\"Julia's JIT typically delivers 10-100x speedup on array/DataFrame operations.\")"
    )
    .expect("TODO: handle error");

    out
}

/// Generate a shell script that runs both original and Julia code for comparison.
///
/// The script:
/// 1. Times the original Python/R script using `time`
/// 2. Times the Julia equivalent (including precompilation)
/// 3. Times the Julia equivalent again (JIT-warmed)
/// 4. Prints a comparison table
fn generate_benchmark_runner(manifest: &Manifest, units: &[TranslationUnit]) -> String {
    let mut out = String::new();

    writeln!(out, "#!/usr/bin/env bash").expect("TODO: handle error");
    writeln!(out, "# SPDX-License-Identifier: MPL-2.0").expect("TODO: handle error");
    writeln!(out, "# Auto-generated by julianiser — benchmark runner").expect("TODO: handle error");
    writeln!(out, "# Project: {}", manifest.project.name).expect("TODO: handle error");
    writeln!(out, "set -euo pipefail").expect("TODO: handle error");
    writeln!(out).expect("TODO: handle error");

    writeln!(
        out,
        "SCRIPT_DIR=\"$(cd \"$(dirname \"${{BASH_SOURCE[0]}}\")\" && pwd)\""
    )
    .expect("TODO: handle error");
    writeln!(out, "PROJECT_DIR=\"$(dirname \"$SCRIPT_DIR\")\"").expect("TODO: handle error");
    writeln!(out).expect("TODO: handle error");

    writeln!(
        out,
        "echo \"============================================================\""
    )
    .expect("TODO: handle error");
    writeln!(
        out,
        "echo \"julianiser Benchmark Runner: {}\"",
        manifest.project.name
    )
    .expect("TODO: handle error");
    writeln!(
        out,
        "echo \"============================================================\""
    )
    .expect("TODO: handle error");
    writeln!(out).expect("TODO: handle error");

    // Run original source files.
    for unit in units {
        writeln!(out, "echo \"\"").expect("TODO: handle error");
        writeln!(
            out,
            "echo \"--- Original: {} ({}) ---\"",
            unit.source_path, unit.language
        )
        .expect("TODO: handle error");

        match unit.language {
            SourceLanguage::Python => {
                writeln!(out, "if command -v python3 &> /dev/null; then")
                    .expect("TODO: handle error");
                writeln!(out, "    echo \"Timing: python3 {}\"", unit.source_path)
                    .expect("TODO: handle error");
                writeln!(
                    out,
                    "    time python3 \"$PROJECT_DIR/../{}\" 2>&1 || echo \"  (Python script failed or not found)\"",
                    unit.source_path
                )
                .expect("TODO: handle error");
                writeln!(out, "else").expect("TODO: handle error");
                writeln!(
                    out,
                    "    echo \"  python3 not found — skipping original benchmark\""
                )
                .expect("TODO: handle error");
                writeln!(out, "fi").expect("TODO: handle error");
            }
            SourceLanguage::R => {
                writeln!(out, "if command -v Rscript &> /dev/null; then")
                    .expect("TODO: handle error");
                writeln!(out, "    echo \"Timing: Rscript {}\"", unit.source_path)
                    .expect("TODO: handle error");
                writeln!(
                    out,
                    "    time Rscript \"$PROJECT_DIR/../{}\" 2>&1 || echo \"  (R script failed or not found)\"",
                    unit.source_path
                )
                .expect("TODO: handle error");
                writeln!(out, "else").expect("TODO: handle error");
                writeln!(
                    out,
                    "    echo \"  Rscript not found — skipping original benchmark\""
                )
                .expect("TODO: handle error");
                writeln!(out, "fi").expect("TODO: handle error");
            }
        }
    }

    writeln!(out).expect("TODO: handle error");
    writeln!(out, "echo \"\"").expect("TODO: handle error");
    writeln!(out, "echo \"--- Julia (generated) ---\"").expect("TODO: handle error");
    writeln!(out, "if command -v julia &> /dev/null; then").expect("TODO: handle error");
    writeln!(
        out,
        "    echo \"Running Julia benchmark (includes JIT compilation)...\""
    )
    .expect("TODO: handle error");
    writeln!(
        out,
        "    time julia --project=\"$PROJECT_DIR\" \"$SCRIPT_DIR/benchmark.jl\" 2>&1"
    )
    .expect("TODO: handle error");
    writeln!(out, "else").expect("TODO: handle error");
    writeln!(
        out,
        "    echo \"  julia not found — install Julia >= {} to run benchmarks\"",
        manifest.julia.version
    )
    .expect("TODO: handle error");
    writeln!(out, "fi").expect("TODO: handle error");
    writeln!(out).expect("TODO: handle error");

    writeln!(out, "echo \"\"").expect("TODO: handle error");
    writeln!(
        out,
        "echo \"============================================================\""
    )
    .expect("TODO: handle error");
    writeln!(out, "echo \"Benchmark complete.\"").expect("TODO: handle error");
    writeln!(
        out,
        "echo \"For accurate Julia timings, run benchmark.jl directly\""
    )
    .expect("TODO: handle error");
    writeln!(
        out,
        "echo \"(the first run includes precompilation overhead).\""
    )
    .expect("TODO: handle error");
    writeln!(
        out,
        "echo \"============================================================\""
    )
    .expect("TODO: handle error");

    out
}

/// Generate a TOML template for recording benchmark results.
///
/// Users fill this in after running the benchmarks. The results can then
/// be parsed by other tools (e.g. CI pipelines, dashboards) to track
/// performance over time.
fn generate_results_template(units: &[TranslationUnit]) -> String {
    let mut out = String::new();

    writeln!(out, "# SPDX-License-Identifier: MPL-2.0").expect("TODO: handle error");
    writeln!(out, "# julianiser benchmark results").expect("TODO: handle error");
    writeln!(
        out,
        "# Fill in timing data after running benchmark_runner.sh"
    )
    .expect("TODO: handle error");
    writeln!(out).expect("TODO: handle error");

    for unit in units {
        writeln!(out, "[[benchmarks]]").expect("TODO: handle error");
        writeln!(out, "name = \"{}\"", unit.module_name).expect("TODO: handle error");
        writeln!(out, "source = \"{}\"", unit.source_path).expect("TODO: handle error");
        writeln!(out, "language = \"{}\"", unit.language).expect("TODO: handle error");
        writeln!(out, "original_time_seconds = 0.0  # TODO: fill in").expect("TODO: handle error");
        writeln!(out, "julia_time_seconds = 0.0     # TODO: fill in").expect("TODO: handle error");
        writeln!(
            out,
            "speedup = 0.0                # TODO: compute original / julia"
        )
        .expect("TODO: handle error");
        writeln!(out, "iterations = 1000").expect("TODO: handle error");
        writeln!(out, "notes = \"\"").expect("TODO: handle error");
        writeln!(out).expect("TODO: handle error");
    }

    out
}

/// Create a BenchmarkResult placeholder for a translation unit.
///
/// Used by the library API to return structured results that can be
/// populated later with actual timing data.
pub fn create_benchmark_placeholder(unit: &TranslationUnit) -> BenchmarkResult {
    BenchmarkResult::new(&unit.module_name, unit.language)
}

/// Convert PascalCase to snake_case for Julia variable names.
fn to_snake(s: &str) -> String {
    let mut result = String::new();
    for (i, ch) in s.chars().enumerate() {
        if ch.is_uppercase() {
            if i > 0 {
                result.push('_');
            }
            result.push(ch.to_lowercase().next().unwrap_or(ch));
        } else {
            result.push(ch);
        }
    }
    result
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::abi::{DetectedCall, SourceLanguage};
    use crate::manifest::{JuliaConfig, ProjectConfig};

    /// Helper to create a test manifest.
    fn test_manifest() -> Manifest {
        Manifest {
            project: ProjectConfig {
                name: "bench-test".to_string(),
                version: "0.1.0".to_string(),
                description: "Test pipeline".to_string(),
            },
            sources: vec![],
            mappings: vec![],
            julia: JuliaConfig {
                version: "1.10".to_string(),
                packages: vec!["BenchmarkTools".to_string()],
                flags: vec![],
            },
            workload: None,
            data: None,
        }
    }

    /// Helper to create a test translation unit with one detected call.
    fn test_unit() -> TranslationUnit {
        let mut unit = TranslationUnit::new("analysis.py", SourceLanguage::Python);
        unit.detected_calls.push(DetectedCall {
            library: "pandas".to_string(),
            function: "read_csv".to_string(),
            language: SourceLanguage::Python,
            line_number: 3,
            source_line: "df = pd.read_csv(\"data.csv\")".to_string(),
            arguments: vec!["\"data.csv\"".to_string()],
        });
        unit
    }

    #[test]
    fn test_julia_benchmark_contains_module() {
        let manifest = test_manifest();
        let unit = test_unit();
        let bench = generate_julia_benchmark(&manifest, &[unit]);
        assert!(bench.contains("using BenchmarkTools"));
        assert!(bench.contains("Analysis"));
        assert!(bench.contains("@benchmark"));
    }

    #[test]
    fn test_benchmark_runner_contains_python() {
        let manifest = test_manifest();
        let unit = test_unit();
        let runner = generate_benchmark_runner(&manifest, &[unit]);
        assert!(runner.contains("python3"));
        assert!(runner.contains("julia"));
        assert!(runner.contains("set -euo pipefail"));
    }

    #[test]
    fn test_results_template_format() {
        let unit = test_unit();
        let results = generate_results_template(&[unit]);
        assert!(results.contains("[[benchmarks]]"));
        assert!(results.contains("name = \"Analysis\""));
        assert!(results.contains("language = \"python\""));
        assert!(results.contains("original_time_seconds"));
    }

    #[test]
    fn test_to_snake() {
        assert_eq!(to_snake("DataPipeline"), "data_pipeline");
        assert_eq!(to_snake("Analysis"), "analysis");
        assert_eq!(to_snake("MyCSVLoader"), "my_c_s_v_loader");
    }

    #[test]
    fn test_benchmark_placeholder() {
        let unit = test_unit();
        let result = create_benchmark_placeholder(&unit);
        assert_eq!(result.name, "Analysis");
        assert_eq!(result.language, SourceLanguage::Python);
        assert!(result.original_time_seconds.is_none());
        assert!(result.julia_time_seconds.is_none());
    }
}