src/main.rs - main - am/tt

src/main.rs
657 строк22 KB
Andrey Manzharov
добавлена возможность запустить на одном конкретном тесте
13 июл 2026, 15:22
13 июл 2026, 15:22fb5e570
use clap::Parser;
use colored::*;
use serde::Deserialize;
use similar::{ChangeTag, TextDiff};
use std::collections::HashMap;
use std::io::{Read, Write};
use std::path::{Path, PathBuf};
use std::process::{Command, Stdio};
use std::time::{Duration, Instant};

#[derive(Parser)]
#[command(name = "tester", about = "Тестировщик олимпиадных решений", version)]
struct Args {
    /// Команда запуска или файл решения (например "main.rs", "sol.py", "./binary")
    solution: String,

    /// Файл с тестами (вход и выход разделены '---', тесты — двумя пустыми строками)
    tests_file: String,

    /// Номер конкретного теста для запуска (опционально)
    test_number: Option<usize>,

    /// Таймаут на один тест, секунд
    #[arg(short = 'T', long, default_value = "5")]
    timeout: u64,

    /// Показывать stderr для пройденных (PASS) тестов
    #[arg(short = 'v', long, default_value_t = false)]
    verbose: bool,
}

#[derive(Deserialize, Clone, Debug)]
struct Config {
    languages: HashMap<String, LanguageConfig>,
}

#[derive(Deserialize, Clone, Debug)]
struct LanguageConfig {
    /// Команда компиляции. Плейсхолдеры: {src}, {out}
    compile: Option<Vec<String>>,
    /// Команда запуска. Плейсхолдеры: {src}, {out}
    run: Vec<String>,
}

impl Config {
    /// Загружает конфиг из текущей директории или XDG_CONFIG_HOME
    fn load() -> Option<Config> {
        let mut paths = Vec::new();
        paths.push(PathBuf::from("tt.toml"));

        if let Ok(xdg_config) = std::env::var("XDG_CONFIG_HOME") {
            paths.push(PathBuf::from(xdg_config).join("tt/config.toml"));
        } else if let Ok(home) = std::env::var("HOME") {
            paths.push(PathBuf::from(home).join(".config/tt/config.toml"));
        }

        for path in paths {
            if path.exists() {
                match std::fs::read_to_string(&path) {
                    Ok(content) => match toml::from_str::<Config>(&content) {
                        Ok(cfg) => return Some(cfg),
                        Err(e) => {
                            eprintln!(
                                "{} Ошибка разбора конфига {}: {}",
                                "!".yellow().bold(),
                                path.display(),
                                e
                            );
                        }
                    },
                    Err(e) => {
                        eprintln!(
                            "{} Не удалось прочитать конфиг {}: {}",
                            "!".yellow().bold(),
                            path.display(),
                            e
                        );
                    }
                }
            }
        }
        None
    }
}

/// Разбивает файл на тесты, а каждый тест на вход и выход.
fn parse_tests(content: &str) -> Vec<(String, String)> {
    let normalized: String = content
        .lines()
        .map(|l| l.trim_end())
        .collect::<Vec<_>>()
        .join("\n");

    normalized
        .split("\n\n\n")
        .filter_map(|block| {
            let block = block.trim();
            if block.is_empty() {
                return None;
            }

            let mut input_parts = Vec::new();
            let mut expected_parts = Vec::new();
            let mut in_expected = false;

            for line in block.lines() {
                if line.trim() == "---" {
                    in_expected = true;
                    continue;
                }
                if in_expected {
                    expected_parts.push(line);
                } else {
                    input_parts.push(line);
                }
            }

            Some((input_parts.join("\n"), expected_parts.join("\n")))
        })
        .collect()
}

/// Нормализует вывод: убирает trailing-пробелы и пустые строки в конце.
fn normalize(s: &str) -> String {
    let mut lines: Vec<&str> = s.lines().map(|l| l.trim_end()).collect();
    while lines.last().map_or(false, |l| l.is_empty()) {
        lines.pop();
    }
    lines.join("\n")
}

#[cfg(unix)]
fn wait_with_timeout_and_rusage(
    pid: libc::pid_t,
    timeout: Duration,
) -> Result<(libc::c_int, libc::rusage), String> {
    let start = Instant::now();
    loop {
        let mut status: libc::c_int = 0;
        let mut rusage: libc::rusage = unsafe { std::mem::zeroed() };

        let ret = unsafe { libc::wait4(pid, &mut status, libc::WNOHANG, &mut rusage) };

        if ret == pid {
            return Ok((status, rusage));
        } else if ret == -1 {
            return Err(format!("wait4 failed: {}", std::io::Error::last_os_error()));
        }

        if start.elapsed() > timeout {
            unsafe { libc::kill(pid, libc::SIGKILL) };
            let mut s = 0;
            unsafe { libc::wait4(pid, &mut s, 0, std::ptr::null_mut()) };
            return Err("Time Limit Exceeded".to_string());
        }

        std::thread::sleep(Duration::from_millis(5));
    }
}

struct Metrics {
    wall_ms: u64,
    usr_ms: u64,
    sys_ms: u64,
    mem_kb: u64,
}

/// Запускает решение, кормит его входом, возвращает (stdout, stderr, Метрики) либо ошибку.
fn run_solution(
    cmd: &[String],
    input: &str,
    timeout: Duration,
) -> (Result<String, String>, String, Metrics) {
    if cmd.is_empty() {
        return (
            Err("пустая команда решения".to_string()),
            String::new(),
            Metrics {
                wall_ms: 0,
                usr_ms: 0,
                sys_ms: 0,
                mem_kb: 0,
            },
        );
    }

    let program = &cmd[0];
    let args = &cmd[1..];

    let mut child = match Command::new(program)
        .args(args)
        .stdin(Stdio::piped())
        .stdout(Stdio::piped())
        .stderr(Stdio::piped())
        .spawn()
    {
        Ok(c) => c,
        Err(e) => {
            return (
                Err(format!("не удалось запустить решение: {e}")),
                String::new(),
                Metrics {
                    wall_ms: 0,
                    usr_ms: 0,
                    sys_ms: 0,
                    mem_kb: 0,
                },
            );
        }
    };

    let stdin = child.stdin.take().unwrap();
    let stdout = child.stdout.take().unwrap();
    let stderr = child.stderr.take().unwrap();

    let input_owned = input.to_string();
    let write_thread = std::thread::spawn(move || {
        let mut stdin = stdin;
        let _ = stdin.write_all(input_owned.as_bytes());
    });

    let read_stdout = std::thread::spawn(move || {
        let mut buf = String::new();
        let mut stdout = stdout;
        let _ = stdout.read_to_string(&mut buf);
        buf
    });
    let read_stderr = std::thread::spawn(move || {
        let mut buf = String::new();
        let mut stderr = stderr;
        let _ = stderr.read_to_string(&mut buf);
        buf
    });

    let start = Instant::now();

    #[cfg(unix)]
    let pid = child.id() as libc::pid_t;

    #[cfg(unix)]
    let wait_result = wait_with_timeout_and_rusage(pid, timeout);

    #[cfg(not(unix))]
    let wait_result: Result<(i32, (u64, u64, u64)), String> = {
        match child.wait_timeout(timeout) {
            Ok(Some(status)) => Ok((status.code().unwrap_or(-1), (0, 0, 0))),
            Ok(None) => {
                let _ = child.kill();
                let _ = child.wait();
                Err("Time Limit Exceeded".to_string())
            }
            Err(e) => Err(e.to_string()),
        }
    };

    let wall_ms = start.elapsed().as_millis() as u64;

    #[cfg(unix)]
    {
        let _ = child.wait();
    }

    let out = read_stdout.join().unwrap_or_default();
    let err = read_stderr.join().unwrap_or_default();
    let _ = write_thread.join();

    match wait_result {
        Ok((status, rusage)) => {
            #[cfg(unix)]
            let (usr_ms, sys_ms, mem_kb) = {
                let usr = (rusage.ru_utime.tv_sec as u64) * 1000
                    + (rusage.ru_utime.tv_usec as u64) / 1000;
                let sys = (rusage.ru_stime.tv_sec as u64) * 1000
                    + (rusage.ru_stime.tv_usec as u64) / 1000;
                #[cfg(target_os = "macos")]
                let mem = rusage.ru_maxrss / 1024;
                #[cfg(not(target_os = "macos"))]
                let mem = rusage.ru_maxrss as u64;
                (usr, sys, mem)
            };
            #[cfg(not(unix))]
            let (usr_ms, sys_ms, mem_kb) = (0, 0, 0);

            let metrics = Metrics {
                wall_ms,
                usr_ms,
                sys_ms,
                mem_kb,
            };

            if status == 0 {
                (Ok(out), err, metrics)
            } else {
                let code = if status == 0 {
                    "0".into()
                } else {
                    format!("{}", status)
                };
                (Err(format!("Runtime Error (код {code})")), err, metrics)
            }
        }
        Err(e) => (
            Err(e),
            err,
            Metrics {
                wall_ms,
                usr_ms: 0,
                sys_ms: 0,
                mem_kb: 0,
            },
        ),
    }
}

fn print_diff(expected: &str, actual: &str) {
    let diff = TextDiff::from_lines(expected, actual);
    const MAX_LINES: usize = 60;
    let mut count = 0usize;

    for change in diff.iter_all_changes() {
        if count >= MAX_LINES {
            println!("  {}", "… (diff обрезан)".dimmed());
            return;
        }
        let (sign, color) = match change.tag() {
            ChangeTag::Delete => ("−", Color::Red),
            ChangeTag::Insert => ("+", Color::Green),
            ChangeTag::Equal => (" ", Color::BrightBlack),
        };
        let value = change.value();
        let trimmed = value.trim_end_matches('\n');
        let display = if trimmed.is_empty() {
            "(пусто)".to_string()
        } else {
            trimmed.to_string()
        };
        println!("  {} {}", sign.color(color).bold(), display.color(color));
        count += 1;
    }
}

fn show_test_input(input: &str) {
    let lines: Vec<&str> = input.lines().collect();
    const MAX: usize = 6;
    println!("  {}", "── Входные данные ──".dimmed());
    for line in lines.iter().take(MAX) {
        println!("  {} {}", "›".dimmed(), line);
    }
    if lines.len() > MAX {
        println!("  {} … (ещё {} строк)", "›".dimmed(), lines.len() - MAX);
    }
}

fn print_stderr(stderr: &str) {
    if !stderr.trim().is_empty() {
        println!("  {}", "── stderr ──".dimmed());
        for line in stderr.lines() {
            println!("  {} {}", "↳".dimmed(), line);
        }
    }
}

fn format_metrics(m: &Metrics) -> String {
    let mem_str = if m.mem_kb >= 1024 {
        format!("{:.1} MB", m.mem_kb as f64 / 1024.0)
    } else {
        format!("{} KB", m.mem_kb)
    };
    format!(
        "wall: {}ms | usr: {}ms | sys: {}ms | mem: {}",
        m.wall_ms, m.usr_ms, m.sys_ms, mem_str
    )
}

/// Заменяет плейсхолдеры {src} и {out} в массиве аргументов
fn substitute_args(args: &[String], src: &str, out: &str) -> Vec<String> {
    args.iter()
        .map(|a| a.replace("{src}", src).replace("{out}", out))
        .collect()
}

fn main() {
    let args = Args::parse();

    println!();
    println!("{}", "  🧪  Tester".bold());
    println!("{}", "  ─────────────────────────────────".dimmed());

    let config = Config::load();
    let run_cmd: Vec<String>;
    let mut compiled_out_path: Option<PathBuf> = None;

    let path = Path::new(&args.solution);

    // Если передан файл (а не готовая команда типа "./a.out"), пытаемся найти его в конфиге
    if path.is_file() {
        if let Some(ext) = path.extension().and_then(|e| e.to_str()) {
            if let Some(cfg) = config.as_ref().and_then(|c| c.languages.get(ext)) {
                println!(
                    "  {} {}",
                    "🔍".dimmed(),
                    format!("Найден конфиг для .{}", ext).cyan()
                );

                if let Some(compile_cmd) = &cfg.compile {
                    let mut out_path =
                        std::env::temp_dir().join(format!("t_sol_{}", std::process::id()));
                    if cfg!(windows) {
                        out_path.set_extension("exe");
                    }
                    let out_str = out_path.to_string_lossy().to_string();

                    let final_cmd = substitute_args(compile_cmd, &args.solution, &out_str);
                    println!(
                        "  {} {}",
                        "🛠".dimmed(),
                        format!("Компиляция: {}", final_cmd.join(" ")).cyan()
                    );

                    let output = Command::new(&final_cmd[0]).args(&final_cmd[1..]).output();

                    match output {
                        Ok(out) if out.status.success() => {
                            println!("  {} {}", "✓".green().bold(), "Компиляция успешна.".green());
                            compiled_out_path = Some(out_path);
                        }
                        Ok(out) => {
                            println!(
                                "  {} {}",
                                "✗".red().bold(),
                                "Ошибка компиляции:".red().bold()
                            );
                            eprintln!("{}", String::from_utf8_lossy(&out.stderr).trim_end());
                            std::process::exit(1);
                        }
                        Err(e) => {
                            println!(
                                "  {} {}",
                                "✗".red().bold(),
                                format!("Не удалось запустить компилятор: {e}").red()
                            );
                            std::process::exit(1);
                        }
                    }
                }

                let out_str = compiled_out_path
                    .as_ref()
                    .map(|p| p.to_string_lossy().to_string())
                    .unwrap_or_default();

                run_cmd = substitute_args(&cfg.run, &args.solution, &out_str);
            } else {
                eprintln!(
                    "{} Ошибка: Файл '{}' имеет неизвестное расширение '.{}'.",
                    "✗".red().bold(),
                    args.solution,
                    ext
                );
                eprintln!(
                    "{} Добавьте конфигурацию для '.{}' в tt.toml или передайте готовый бинарь.",
                    "!".yellow(),
                    ext
                );
                std::process::exit(1);
            }
        } else {
            eprintln!(
                "{} Ошибка: Файл '{}' не имеет расширения.",
                "✗".red().bold(),
                args.solution
            );
            std::process::exit(1);
        }
    } else {
        // Если это не файл, предполагаем, что это команда (например, "./binary" или "python sol.py")
        run_cmd = args
            .solution
            .split_whitespace()
            .map(|s| s.to_string())
            .collect();

        // Проверяем, существует ли бинарник, если первый аргумент выглядит как путь
        if let Some(first_arg) = run_cmd.get(0) {
            if first_arg.contains('/') || first_arg.starts_with("./") {
                if !Path::new(first_arg).exists() {
                    eprintln!(
                        "{} Ошибка: Исполняемый файл '{}' не найден.",
                        "✗".red().bold(),
                        first_arg
                    );
                    std::process::exit(1);
                }
            }
        }
    }

    if run_cmd.is_empty() {
        eprintln!("{} Не указана команда для запуска.", "✗".red().bold());
        std::process::exit(1);
    }

    let content = std::fs::read_to_string(&args.tests_file).unwrap_or_else(|e| {
        eprintln!(
            "\n{} Не удалось прочитать файл тестов '{}': {e}",
            "✗".red().bold(),
            args.tests_file
        );
        std::process::exit(1);
    });

    let tests = parse_tests(&content);

    let total = tests.len();
    if total == 0 {
        eprintln!("{} Тесты не найдены в файле.", "✗".red().bold());
        std::process::exit(1);
    }

    // Проверка корректности переданного номера теста
    if let Some(num) = args.test_number {
        if num == 0 || num > total {
            eprintln!(
                "{} Ошибка: Теста с номером {} не существует (всего {} тестов).",
                "✗".red().bold(),
                num,
                total
            );
            std::process::exit(1);
        }
    }

    // Шапка
    println!("  {} {}", "Запуск  :".dimmed(), run_cmd.join(" ").cyan());
    println!("  {} {}", "Файл    :".dimmed(), args.tests_file.cyan());

    if let Some(num) = args.test_number {
        println!(
            "  {} {}",
            "Тестов  :".dimmed(),
            format!("1 (только #{})", num).cyan()
        );
    } else {
        println!("  {} {}", "Тестов  :".dimmed(), total.to_string().cyan());
    }

    println!(
        "  {} {}{}",
        "Таймаут:".dimmed(),
        args.timeout.to_string().cyan(),
        "s".dimmed()
    );
    println!();

    let timeout = Duration::from_secs(args.timeout);
    let mut passed = 0usize;

    for (i, (test, expected)) in tests.iter().enumerate() {
        let idx = i + 1;

        // Пропускаем тесты, если запрошен конкретный
        if let Some(num) = args.test_number {
            if idx != num {
                continue;
            }
        }

        let expected = normalize(expected);

        let (result, stderr, metrics) = run_solution(&run_cmd, test, timeout);
        let metrics_str = format_metrics(&metrics);

        match result {
            Ok(output) => {
                let actual = normalize(&output);
                if actual == expected {
                    passed += 1;
                    println!(
                        "  {} [{:>3}/{:>3}] {} {}",
                        "✓".green().bold(),
                        idx,
                        total,
                        "PASS".green().bold(),
                        format!("({})", metrics_str).dimmed()
                    );
                    // В режиме verbose выводим stderr даже для успешных тестов
                    if args.verbose {
                        print_stderr(&stderr);
                    }
                } else {
                    println!(
                        "  {} [{:>3}/{:>3}] {} {}",
                        "✗".red().bold(),
                        idx,
                        total,
                        "FAIL".red().bold(),
                        format!("({})", metrics_str).dimmed()
                    );
                    show_test_input(test);
                    println!("  {}", "── Diff: ожидаемый vs полученный ──".dimmed());
                    print_diff(&expected, &actual);
                    // При неверном ответе всегда выводим stderr
                    print_stderr(&stderr);
                    println!();
                }
            }
            Err(e) => {
                println!(
                    "  {} [{:>3}/{:>3}] {} {}",
                    "✗".red().bold(),
                    idx,
                    total,
                    "ERROR".red().bold(),
                    format!("({})", metrics_str).dimmed()
                );
                show_test_input(test);
                println!("  {} {}", "→".red(), e.red());
                // При ошибке выполнения всегда выводим stderr
                print_stderr(&stderr);
                println!();
            }
        }
    }

    // Очистка скомпилированного файла
    if let Some(out) = compiled_out_path {
        let _ = std::fs::remove_file(out);
    }

    // Итог
    // Если мы запускали только один тест, то всего в сводке тоже 1
    let total_run = if args.test_number.is_some() { 1 } else { total };

    println!();
    let line = "════════════════════════════════════════";
    if passed == total_run {
        println!("{}", line.green());
        println!(
            "  {} Все тесты пройдены: {}/{}",
            "🎉",
            passed.to_string().green().bold(),
            total_run.to_string().green().bold()
        );
        println!("{}", line.green());
        std::process::exit(0);
    } else {
        println!("{}", line.yellow());
        println!(
            "  {} Пройдено {} из {} тестов",
            "📊",
            passed.to_string().green().bold(),
            total_run.to_string().cyan().bold()
        );
        println!("{}", line.yellow());
        std::process::exit(1);
    }
}