acr-cli 0.2.8

A CLI tool for AtCoder competitive programming in Rust
use std::path::{Path, PathBuf};
use std::time::Duration;

use anyhow::Context;
use colored::Colorize;

use super::TestResult;
use crate::atcoder::TestCase;

const TLE_TIMEOUT: Duration = Duration::from_secs(5);

/// Build the problem binary with `cargo build --release`.
/// Returns the path to the compiled binary.
pub fn build(problem_dir: &Path) -> anyhow::Result<PathBuf> {
    let output = std::process::Command::new("cargo")
        .args(["build", "--release"])
        .current_dir(problem_dir)
        .output()
        .context("Failed to run cargo build")?;

    if !output.status.success() {
        let stderr = String::from_utf8_lossy(&output.stderr);
        anyhow::bail!("Build failed:\n{}", stderr);
    }

    // Find the binary name from Cargo.toml
    let cargo_toml = std::fs::read_to_string(problem_dir.join("Cargo.toml"))
        .context("Failed to read problem Cargo.toml")?;
    let doc: toml::Value = toml::from_str(&cargo_toml)?;
    let name = doc
        .get("package")
        .and_then(|p| p.get("name"))
        .and_then(|n| n.as_str())
        .context("Could not find package name in Cargo.toml")?;

    // Binary path: find the workspace target dir
    // Walk up from problem_dir to find workspace root with target/
    let mut search = problem_dir.to_path_buf();
    loop {
        let target_bin = search.join("target/release").join(name);
        if target_bin.exists() {
            return Ok(target_bin);
        }
        if !search.pop() {
            break;
        }
    }

    anyhow::bail!("Could not find compiled binary for {}", name)
}

/// Run a single test case against the binary.
pub async fn run_test(binary: &Path, test_case: &TestCase) -> TestResult {
    use tokio::process::Command;

    let result = tokio::time::timeout(TLE_TIMEOUT, async {
        let child = Command::new(binary)
            .stdin(std::process::Stdio::piped())
            .stdout(std::process::Stdio::piped())
            .stderr(std::process::Stdio::piped())
            .spawn();

        let mut child = match child {
            Ok(c) => c,
            Err(e) => return TestResult::Re { stderr: e.to_string() },
        };

        // Write input to stdin
        if let Some(mut stdin) = child.stdin.take() {
            use tokio::io::AsyncWriteExt;
            let _ = stdin.write_all(test_case.input.as_bytes()).await;
            drop(stdin);
        }

        let output = match child.wait_with_output().await {
            Ok(o) => o,
            Err(e) => return TestResult::Re { stderr: e.to_string() },
        };

        if !output.status.success() {
            let stderr = String::from_utf8_lossy(&output.stderr).to_string();
            return TestResult::Re { stderr };
        }

        let actual = String::from_utf8_lossy(&output.stdout).to_string();
        if actual.trim_end() == test_case.expected.trim_end() {
            TestResult::Ac
        } else {
            TestResult::Wa {
                actual,
                expected: test_case.expected.clone(),
            }
        }
    })
    .await;

    match result {
        Ok(test_result) => test_result,
        Err(_) => TestResult::Tle,
    }
}

/// Run all test cases and return results.
pub async fn run_all(
    problem_dir: &Path,
    test_cases: &[TestCase],
) -> anyhow::Result<Vec<(usize, TestResult)>> {
    println!("{}", "Building...".dimmed());
    let binary = build(problem_dir)?;
    println!("{}", "Running tests...".dimmed());

    let mut results = Vec::new();
    for tc in test_cases {
        let result = run_test(&binary, tc).await;
        results.push((tc.index, result));
    }
    Ok(results)
}

/// Display test results with colored output.
pub fn display_results(results: &[(usize, TestResult)]) {
    for (index, result) in results {
        match result {
            TestResult::Ac => {
                println!("  Test {} ... {}", index, "AC".green().bold());
            }
            TestResult::Wa { actual, expected } => {
                println!("  Test {} ... {}", index, "WA".red().bold());
                println!("    Expected: {}", expected.trim_end());
                println!("    Actual:   {}", actual.trim_end());
            }
            TestResult::Re { stderr } => {
                println!("  Test {} ... {}", index, "RE".yellow().bold());
                if !stderr.is_empty() {
                    println!("    {}", stderr.trim_end());
                }
            }
            TestResult::Tle => {
                println!("  Test {} ... {}", index, "TLE".yellow().bold());
            }
        }
    }

    let total = results.len();
    let passed = results
        .iter()
        .filter(|(_, r)| matches!(r, TestResult::Ac))
        .count();
    let status = if passed == total {
        format!("All tests passed ({}/{})", passed, total).green()
    } else {
        format!("{}/{} passed", passed, total).red()
    };
    println!("\n  {}", status.bold());
}