deflake-rs 0.1.0

use std::{
    collections::HashMap,
    env::current_dir,
    fmt::{Display, Write},
    fs::{create_dir, remove_dir_all, remove_file},
    path::{self},
    process::Output,
    rc::Rc,
    str::FromStr,
    thread::sleep,
    time::{Duration, Instant},
};

use cargo_toml::Manifest;
use duct::{cmd, Handle};
use serde::{Deserialize, Serialize};

use crate::{ast::ItemFn, cli::Cli, git::Edit};

pub static COVERAGE_FOLDER: &str = "deflaker_coverage";
pub static LOGS_FOLDER: &str = "deflake_logs";
pub static IGNORE_FILE: &str = "0ignore.profraw";

#[derive(Deserialize, Debug)]
struct CompilerArtifactProfile {
    opt_level: String,
    debug_assertions: bool,
    overflow_checks: bool,
    test: bool,
}

#[derive(Deserialize, Debug)]
struct CompilerArtifactTarget {
    kind: Vec<String>,
    crate_types: Vec<String>,
    name: String,
    src_path: String,
    doc: bool,
    doctest: bool,
    test: bool,
}

#[derive(Deserialize, Debug)]
struct CompilerArtifact {
    reason: String,
    package_id: String,
    manifest_path: Option<String>,
    executable: Option<String>,
    filenames: Vec<String>,
    features: Vec<String>,
    profile: CompilerArtifactProfile,
    target: CompilerArtifactTarget,
    fresh: bool,
}

#[derive(Deserialize, Debug)]
#[serde(tag = "reason")]
enum ArtifactReason {
    // WARN: may be missing options, hence the Unknown value
    #[serde(rename(deserialize = "compiler-artifact"))]
    CompilerArtifact,
    #[serde(rename(deserialize = "build-script-executed"))]
    BuildScript,
    #[serde(rename(deserialize = "build-finished"))]
    Finished,

    #[serde(other)]
    Unknown,
}

// TODO: work out how the types differ
#[derive(Debug)]
pub enum BinaryType {
    Test,
    Bin,
    Lib,
    Bench,
    Proc,
    Unknown,
}

#[derive(Debug)]
pub struct TestBinary {
    pub path: String,
    pub name: String,
    pub test_type: BinaryType,
    pub workspace: Option<String>,
}

impl TestBinary {
    fn test_cases(self) -> Vec<TestCase> {
        let bin = Rc::new(self);

        // Cases will be outout "xxxx: test"
        let stdout = match cmd!(bin.path.clone(), "--list", "--format", "terse")
            // Ignoring coverage which gets collected when running
            .env(
                "LLVM_PROFILE_FILE",
                format!("{}/{}", COVERAGE_FOLDER, IGNORE_FILE),
            )
            .stderr_capture()
            .stdout_capture()
            .read()
        {
            Ok(s) => s,
            Err(_) => panic!("ERR: could not get tests for binary {}", bin.path),
        };

        let cases = stdout.split("\n");

        let mut test_cases = vec![];
        for case in cases {
            if case == "" {
                continue;
            }
            let mut parts = case.split(": ");
            let name = parts.next();
            let test_type = parts.next();

            if name.is_some() && test_type.is_some() {
                let name = name.unwrap();
                let test_type = test_type.unwrap();
                test_cases.push(TestCase {
                    name: name.to_string(),
                    case_type: CaseType::from_str(test_type).unwrap(),
                    binary: Rc::clone(&bin),
                })
            } else {
                println!(
                    "WARN: Couldn't parse test case output. Line = \"{}\" ",
                    case
                );
            }
        }

        test_cases
    }
}

#[derive(Debug)]
pub enum CaseType {
    Test,
    Bench,
}

impl FromStr for CaseType {
    type Err = ();

    fn from_str(case_type: &str) -> Result<Self, Self::Err> {
        match case_type {
            "test" => Ok(Self::Test),
            "bench" => Ok(Self::Bench),
            _ => {
                eprintln!("Unexpected test case type: {}", case_type);
                Err(())
            }
        }
    }
}

pub type TestId = String;

#[derive(Debug)]
pub struct TestCase {
    pub name: String,
    pub case_type: CaseType,
    pub binary: Rc<TestBinary>,
}

impl TestCase {
    // Start test case execution in the background
    pub fn exec_child(&self) -> Handle {
        let test_exec = cmd!(self.binary.path.clone(), self.name.clone(), "--exact")
            .dir(
                current_dir()
                    .unwrap()
                    .join(self.binary.workspace.clone().unwrap_or(".".to_string())),
            )
            .env("LLVM_PROFILE_FILE", self.profraw_filename())
            .unchecked()
            .stderr_to_stdout()
            .stdout_path(self.stdout_path())
            .start();

        let handle = test_exec.unwrap();

        handle
    }

    pub fn stdout_path(&self) -> String {
        format!("{}/{}.txt", LOGS_FOLDER, self.id())
    }

    pub fn profraw_filename(&self) -> String {
        format!(
            "{}/{}_{}.profraw",
            COVERAGE_FOLDER, &self.binary.name, &self.name
        )
    }
    pub fn profdata_filename(&self) -> String {
        format!(
            "{}/{}_{}.profdata",
            COVERAGE_FOLDER, &self.binary.name, &self.name
        )
    }

    // Basic handling of test execution to get pass/fail
    // Also cleans up resources used by passing tests as not needed
    fn process_result(&self, output: &Output) -> TestResult {
        match output.status.code() {
            Some(code) => match code {
                0 => {
                    self.cleanup();
                    TestResult::Success
                }
                _ => TestResult::Error,
            },
            None => {
                eprintln!("Error: runner exited by signal");
                TestResult::Error
            }
        }
    }

    pub fn wait_handle(&self, handle: Handle) -> TestResult {
        let output = handle.into_output().unwrap();
        self.process_result(&output)
    }

    pub fn id(&self) -> TestId {
        format!("{}.{}", self.binary.name.clone(), self.name.clone())
    }

    // Remove files associated with a test from execution
    pub fn cleanup(&self) {
        remove_file(self.profraw_filename());
        remove_file(self.stdout_path()).unwrap();
    }
}

#[derive(Debug, Eq, PartialEq, Clone, Serialize, Deserialize)]
pub enum TestResult {
    Success,
    Error,
    Flaky,
    Timeout,
}

impl Display for TestResult {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_str(match self {
            TestResult::Success => "Pass",
            TestResult::Error => "Fail",
            TestResult::Flaky => "Fail (flaky)",
            TestResult::Timeout => "Fail (timeout)",
        });
        Ok(())
    }
}

#[derive(Debug, Serialize, Deserialize)]
pub struct TestMeta {
    pub execution_time: u128,
    pub deflaker_version: String,
    pub tests: usize,
    pub deflake_time: Option<u128>,
}

#[derive(Debug, Deserialize, Serialize)]
pub struct TestResultOutput {
    pub tests: HashMap<TestId, TestResult>,
    pub meta: TestMeta,
}

#[derive(Debug)]
pub struct TestResults<'a>(HashMap<TestId, (&'a TestCase, TestResult)>);
impl<'a> TestResults<'a> {
    pub fn new() -> TestResults<'a> {
        TestResults(HashMap::new())
    }

    pub fn push(&mut self, res: (&'a TestCase, TestResult)) {
        self.0.insert(res.0.id(), res);
    }

    pub fn has_failure(&self) -> bool {
        (&self.0).into_iter().any(|(_, (_, r))| {
            *r == TestResult::Error || *r == TestResult::Timeout || *r == TestResult::Flaky
        })
    }

    pub fn find_result(&self, res: TestResult) -> Vec<(&TestId, &(&'a TestCase, TestResult))> {
        (&self.0)
            .into_iter()
            .filter(|(id, (case, result))| *result == res)
            .collect()
    }

    pub fn failures(&self) -> Vec<(&TestId, &(&'a TestCase, TestResult))> {
        let mut v = vec![];
        v.append(self.find_result(TestResult::Error).as_mut());
        v.append(self.find_result(TestResult::Timeout).as_mut());

        v
    }
    pub fn failures_no_timeout(&self) -> Vec<(&TestId, &(&TestCase, TestResult))> {
        let mut v = vec![];
        v.append(self.find_result(TestResult::Error).as_mut());
        v
    }

    pub fn reclassify(&mut self, id: TestId, status: TestResult) {
        self.0.get_mut(&id).unwrap().1 = status;
    }

    pub fn flaky(&self) -> Vec<(&TestId, &(&'a TestCase, TestResult))> {
        self.find_result(TestResult::Flaky)
    }

    // Serialise the rest results to an string
    pub fn serialise_old(&self) -> String {
        (&self.0)
            .into_iter()
            .map(|(k, v)| format!("{} = {}\n", v.0.id(), v.1.to_string()))
            .fold(String::new(), |a, b| a + &b)
    }

    pub fn serialise(&self, meta: TestMeta) -> String {
        let mut results = HashMap::new();
        (&self.0).into_iter().for_each(|(id, (_, result))| {
            results.insert(id.clone(), result.clone());
        });
        let test_results = TestResultOutput {
            tests: results,
            meta,
        };

        serde_json::to_string(&test_results).unwrap()
    }
}

pub struct TestRunner {}

impl TestRunner {
    // Runs tests concurrently up to a max limit
    pub fn exec_batch(tests: &Vec<TestCase>, max: usize, timeout: u64) -> TestResults {
        let init = match max < tests.len() {
            true => max,
            false => tests.len(),
        };

        let mut results = TestResults::new();

        let mut children = vec![];

        let mut next_test = 0;
        for _ in 0..init {
            let test = tests.get(next_test).unwrap();

            let child = test.exec_child();
            let start_time = Instant::now();
            children.push((test, child, start_time));

            next_test += 1;
        }

        // While still tests running
        // TODO: broken?
        let mut completed = 0;
        while children.len() > 0 {
            let mut new_children = vec![];
            // Remove finished tests
            children.retain(|(test, child, start_time)| {
                // See if test has finished
                let result = match child.try_wait() {
                    Ok(o) => match o {
                        Some(output) => Some(test.process_result(output)),
                        None => None,
                    },
                    Err(_) => {
                        panic!("ERR: something went wrong running a test");
                    }
                };

                if let Some(result) = result {
                    completed += 1;
                    results.push((test, result));

                    if next_test < tests.len() {
                        let test = tests.get(next_test).unwrap();
                        let child = test.exec_child();
                        let start_time = Instant::now();
                        new_children.push((test, child, start_time));
                        next_test += 1;
                    }

                    return false;
                }

                if start_time.elapsed().as_secs() > timeout {
                    println!("Timeout hit, {} secs", start_time.elapsed().as_secs());
                    // TODO: do we need to kill?
                    child.kill();

                    results.push((test, TestResult::Timeout));
                    return false;
                }

                true
            });

            // Add newly run tests
            children.append(&mut new_children);

            print!(
                "Tests finished {}/{} ({} processes running)        \r",
                completed,
                tests.len(),
                children.len()
            );

            sleep(Duration::from_millis(50));
        }
        println!("");

        results
    }

    // Get the test cases which failed
    pub fn failures(
        results: Vec<(&TestCase, TestResult)>,
    ) -> HashMap<TestId, (&TestCase, TestResult)> {
        let mut failed = HashMap::new();
        for result in results {
            if result.1 == TestResult::Error {
                failed.insert(result.0.id(), result);
            }
        }
        return failed;
    }

    // Clear coverage and logs before execution
    // TODO: remove all of these in the workspace folders too?
    pub fn setup() {
        // Delete old coverage
        let path = path::Path::new(COVERAGE_FOLDER);
        if path.exists() {
            match remove_dir_all(path) {
                Err(_) => {
                    eprintln!("WARN: couldn't remove old coverage data")
                }
                _ => {
                    println!("[dbg] Cleared coverage folder")
                }
            };
        }

        // Add logs folder
        let path = path::Path::new(LOGS_FOLDER);
        match remove_dir_all(path) {
            Err(_) => {}
            _ => {
                println!("[dbg] Cleared logs folder")
            }
        }
        match create_dir(path) {
            Err(_) => {
                eprintln!("WARN: could not create logs folder")
            }
            _ => {}
        }
    }

    // All workspace folders are also cleaned as these can also contain coverage data
    // Must be removed so repeated builds don't cause a dirty fingerprint for the project
    pub fn cleanup() {
        let paths = get_workspaces();
        remove_file(format!("{}/{}", COVERAGE_FOLDER, IGNORE_FILE));
        if let Some(paths) = paths {
            for path in paths {
                remove_file(format!("{}/{}/{}", path, COVERAGE_FOLDER, IGNORE_FILE));
            }
        }
    }
}

// If none, then likely no workspaces
fn get_workspaces() -> Option<Vec<String>> {
    // TODO: allow specifying
    let cargo = Manifest::from_path("./Cargo.toml");

    let empty = None;
    match cargo {
        Ok(cargo) => match cargo.workspace {
            // TODO: use members or default_members
            Some(workspace) => Some(workspace.members),
            None => empty,
        },
        Err(_) => empty,
    }
}

#[derive(Debug, Clone)]
pub enum TrackReason {
    FuncAdded,
    FuncModified,
    LineAdded,
    LineDeleted,
}

#[allow(dead_code)]
#[derive(Debug, Clone)]
pub struct TrackMeta {
    pub func: ItemFn,
    pub reason: TrackReason,
    pub edit: Option<Edit>,
}

// A line to be checked in the coverage report
#[derive(Debug, Clone)]
pub enum Track {
    #[allow(dead_code)]
    Line(usize, Option<TrackMeta>), // line
    Span(usize, usize, usize, usize), // line start, line end, col start, col end
}

impl Track {
    pub fn covers(&self, l: usize) -> bool {
        match self {
            Track::Line(line, _) => l == *line,
            Track::Span(ls, le, _, _) => *ls <= l && *le >= l,
        }
    }
}

pub struct TestCollector {}

// TODO: look at edge cases such as workspaces and how to deal with it
// also maybe make this implementation of TestBinary or something
impl TestCollector {
    // NOTE: as a side effect, this will build the tests as well if not up to date so may be slow
    pub fn get_tests(args: &Cli) -> Vec<TestCase> {
        //let now = Instant::now();
        let test_binaries = TestCollector::get_test_binaries(!args.build_only, args.no_coverage);

        println!(
            "[dbg] {} binaries ({})",
            test_binaries.len(),
            (&test_binaries)
                .iter()
                .map(|b| b.name.as_str())
                .collect::<Vec<&str>>()
                .join(", ")
        );

        //let perf_time = now.elapsed().as_millis();
        //println!("PERF: Time getting binaries =  {}", perf_time);

        //let now = Instant::now();
        // Collect tests in each test binary
        let mut cases = vec![];
        for binary in test_binaries {
            // If filtering, skip if no provided binary equals current
            if let Some(only_run) = &args.filter_binaries {
                if !only_run.contains(&binary.name) {
                    println!("[dbg] Ignoring binary: {}", binary.name);
                    continue;
                }
            }
            let binary_cases = binary.test_cases();
            cases.extend(binary_cases);
        }
        //let perf_time = now.elapsed().as_millis();
        //println!("PERF: Time getting test cases =  {}", perf_time);

        return cases;
    }

    pub fn get_test_binaries(hide_logs: bool, no_coverage: bool) -> Vec<TestBinary> {
        let workspaces = get_workspaces();
        // tODO: if no workspace then get tests from current level
        // else, get tests from each workspace folder

        let workspaces = match workspaces {
            Some(v) => v,
            None => vec![".".to_string()],
        };
        //dbg!(&workspaces);

        let mut test_binaries = vec![];

        // TODO: run in parallel to speed up
        for member in workspaces {
            let mut cmd = cmd!(
                "cargo",
                "test",
                "--no-run",
                "--message-format",
                "json-render-diagnostics"
            )
            .dir(current_dir().unwrap().join(&member));
            if !no_coverage {
                cmd = cmd
                    .env("RUSTFLAGS", "-C instrument-coverage")
                    // NOTE: workaround as coverage is being generated at build time by dependencies
                    .env(
                        "LLVM_PROFILE_FILE",
                        format!("{}/{}", COVERAGE_FOLDER, IGNORE_FILE),
                    );
            }
            // Builds tests with coveraged integrated

            if hide_logs {
                cmd = cmd.stdout_capture().stderr_capture();
            }

            let stdout = match cmd.read() {
                Ok(s) => s,
                Err(e) => {
                    dbg!(e);
                    panic!(
                        "Cargo test command failed (you are probably not in the right directory)"
                    );
                }
            };

            let artifacts = stdout.split("\n");

            for artifact in artifacts {
                let member_name = match member == "." {
                    true => None,
                    false => Some(member.clone()),
                };
                if let Some(a) = TestCollector::parse_artifact(artifact, member_name) {
                    test_binaries.push(a);
                }
            }

            // TODO: cleanup between workspace members?
        }
        return test_binaries;
    }

    fn parse_artifact(artifact_log: &str, workspace: Option<String>) -> Option<TestBinary> {
        // Check reason field before parsing

        let artifact_type: ArtifactReason = serde_json::from_str(artifact_log).unwrap();

        // Ignoring logs that wont be related to test binaries
        match artifact_type {
            ArtifactReason::CompilerArtifact => {
                if let Some(compiler_artifact) =
                    TestCollector::process_compiler_artifact(artifact_log)
                {
                    Some(TestBinary {
                        name: compiler_artifact.target.name,
                        path: compiler_artifact.executable.unwrap(),
                        test_type: TestCollector::kind_to_type(compiler_artifact.target.kind),
                        workspace,
                    })
                } else {
                    None
                }
            }
            _ => None,
        }
    }

    fn kind_to_type(kind: Vec<String>) -> BinaryType {
        // TODO: handle other kinds cargo/src/doc/src/reference/external-tools.md
        match kind.first() {
            Some(k) => match k.as_str() {
                "bin" => BinaryType::Bin,
                "test" => BinaryType::Test,
                "bench" => BinaryType::Bench,
                _ => BinaryType::Unknown,
            },
            None => BinaryType::Unknown,
        }
    }

    fn process_compiler_artifact(artifact_log: &str) -> Option<CompilerArtifact> {
        match serde_json::from_str::<CompilerArtifact>(artifact_log) {
            Ok(a) => {
                // WARN: This may be inaccurate, but works for my small testing setup
                if a.target.test && a.profile.test {
                    Some(a)
                } else {
                    None
                }
            }
            Err(e) => {
                eprintln!("Failed to parse artifact");
                //dbg!(&artifact_log);
                //dbg!(e);
                None
            }
        }
    }
}