gpu-trace-perf 1.8.2

#[macro_use]
extern crate lazy_static;
extern crate assert_approx_eq;
extern crate statrs;
extern crate walkdir;

mod angle;
mod apitrace;
mod gfxreconstruct;
mod graphs;
mod log_error;
mod mock_trace;
mod process_watcher;
mod renderdoc;
mod replay;
mod shader_analyze;
mod shader_parser;
mod shader_parser_ir3;
mod snapshot;
mod stats;
mod system_monitor;
mod trace_downloader;
mod u_trace;

use anyhow::{Context, Result, bail};
use apitrace::*;
use clap::{Args, Parser, Subcommand};
use gfxreconstruct::*;
use log::{debug, error, info, warn};
use log_error::LogError;
use rand::seq::IndexedRandom;
use rand::{Rng, SeedableRng};
use renderdoc::*;
use serde::{Deserialize, Serialize};
use simplelog::CombinedLogger;
use statrs::statistics::Statistics;
use stats::{BootstrappedRelativeAndMaxChange, ResultStats};
use std::cmp::max;
use std::ffi::OsStr;
use std::fs::{File, OpenOptions, create_dir_all};
use std::io::{self, prelude::*};
use std::path::{Path, PathBuf};
use std::process::{Command, Stdio, exit};
use std::time::Duration;
use u_trace::UTraceParser;
use walkdir::WalkDir;

use crate::angle::AngleTrace;
use crate::shader_parser::find_shaders;
use crate::snapshot::{Snapshot, SnapshotResult};
use crate::u_trace::{UTRACE_SHADER_DRAW_EVENTS, UTRACE_SHADER_STAGES};

const BUILD_VERSION: &str = env!("GPU_TRACE_PERF_VERSION");

#[derive(Debug, Parser)]
#[clap(
    version = BUILD_VERSION,
    author = "Emma Anholt <emma@anholt.net>",
    about = "Plays a collection of GPU traces under different wrappers to evaluate driver changes on performance"
)]
struct Cli {
    /// A level of verbosity, and can be used multiple times
    #[arg(short, long, action = clap::ArgAction::Count)]
    verbose: u8,

    #[arg(long = "no-log-time", default_value = "false")]
    /// Number of iterations to attempt on each trace before exiting.  Default
    /// is to loop endlessly.
    no_log_time: bool,

    #[command(subcommand)]
    subcmd: SubCommand,
}

#[derive(Default, Debug, Serialize, Deserialize)]
struct ReplayOutput {
    pub exit_code: i32,
    pub stdout: String,
    pub stderr: String,
    pub cmdline: String,
    #[serde(default)]
    pub peak: process_watcher::ProcessPeakMem,
}

impl From<std::process::Output> for ReplayOutput {
    fn from(value: std::process::Output) -> Self {
        ReplayOutput {
            exit_code: value.status.code().unwrap_or(1),
            stdout: String::from_utf8_lossy(&value.stdout).to_string(),
            stderr: String::from_utf8_lossy(&value.stderr).to_string(),
            cmdline: "".to_string(),
            peak: process_watcher::ProcessPeakMem::default(),
        }
    }
}

impl std::fmt::Display for ReplayOutput {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        if !self.cmdline.is_empty() {
            writeln!(f, "command: {}", &self.cmdline)?;
        }
        writeln!(f, "exit code: {}", self.exit_code)?;
        if self.stdout.lines().count() > 1 {
            writeln!(f, "stdout:")?;
            writeln!(f, "{}", &self.stdout)?;
        } else {
            writeln!(f, r#"stdout: "{}""#, &self.stdout.trim_end())?;
        }
        if self.stderr.lines().count() > 1 {
            writeln!(f, "stderr:")?;
            writeln!(f, "{}", &self.stderr)?;
        } else {
            writeln!(f, r#"stderr: "{}""#, &self.stderr.trim_end())?;
        }
        if let Some(b) = self.peak.vram_bytes {
            writeln!(f, "peak GPU VRAM: {} MiB", b / (1024 * 1024))?;
        }
        if let Some(b) = self.peak.sys_bytes {
            writeln!(f, "peak GPU system memory: {} MiB", b / (1024 * 1024))?;
        }
        if let Some(b) = self.peak.rss_bytes {
            writeln!(f, "peak RSS: {} MiB", b / (1024 * 1024))?;
        }
        Ok(())
    }
}
fn create_dir(path: PathBuf) -> Result<PathBuf> {
    match create_dir_all(&path) {
        Ok(()) => {}

        Err(err) => {
            if err.kind() != io::ErrorKind::AlreadyExists {
                bail!("failed to create {}: {:?}", path.display(), err);
            }
        }
    }

    Ok(path)
}

trait TraceTool {
    fn replay(&self, wrapper: Option<&str>, envs: &[(String, String)]) -> Result<ReplayOutput>;

    fn fps(&self, output: &ReplayOutput) -> Result<f64>;

    fn name(&self) -> &str;

    fn capture_draw_times(
        &self,
        output_dir: &str,
        run_name: &str,
        utrace: &mut UTraceParser,
    ) -> Result<()> {
        let path = self
            .output_dir(output_dir)
            .unwrap()
            .unwrap()
            .join(format!("{run_name}-draws.txt"));

        let mut file = std::fs::OpenOptions::new()
            .create(true)
            .append(true)
            .open(&path)
            .with_context(|| format!("appending to {}", path.display()))?;

        for frame in utrace.results()? {
            for events in UTRACE_SHADER_DRAW_EVENTS
                .iter()
                .flat_map(|event| frame.event_times(event))
            {
                for (time, params) in events.times {
                    for param in UTRACE_SHADER_STAGES {
                        if let Some(param) = params.get(*param) {
                            writeln!(file, "{param}: {time}")?;
                        }
                    }
                }
            }
        }

        Ok(())
    }

    fn can_snapshot(&self) -> bool {
        false
    }

    fn snapshot(
        &self,
        _output_dir: &str,
        _loops: u32,
        _timeout: Duration,
    ) -> Result<SnapshotResult> {
        anyhow::bail!("can't snapshot this trace type")
    }

    fn replay_get_fps(
        &self,
        wrapper: Option<&str>,
        output_dir: &str,
        u_trace_event: &str,
        utrace_use_csv: bool,
        capture_shaders_path: &str,
    ) -> Result<f64> {
        let u_trace_event = if u_trace_event.is_empty() && !capture_shaders_path.is_empty() {
            /* Make sure utrace is active if we're capturing draw times. */
            UTRACE_SHADER_DRAW_EVENTS.join(",")
        } else {
            u_trace_event.to_owned()
        };
        let mut utrace = UTraceParser::new(&u_trace_event, utrace_use_csv);
        let mut envs = utrace.env();

        envs.push(("NIR_VALIDATE".to_string(), "0".to_string()));

        let output = self.replay(wrapper, &envs)?;
        if output.exit_code != 0 {
            bail!("Command failed for {}:\n{output}", self.name());
        }
        if utrace.active() {
            if !capture_shaders_path.is_empty() {
                self.capture_draw_times(output_dir, capture_shaders_path, &mut utrace)?;
            }
            utrace.middle_frame_fps()
        } else {
            self.fps(&output)
        }
    }

    /// Runs the trace under the `a` and `b` wrappers, returning the fps.
    fn collect_fps(&self, run_opts: &Run, baseline_first: bool) -> Result<Vec<f64>> {
        info!("Running {}", self.name());

        let do_run = |name, wrapper| {
            self.replay_get_fps(
                Some(wrapper),
                &run_opts.output,
                &run_opts.utrace,
                run_opts.utrace_use_csv,
                if run_opts.capture_shaders { name } else { "" },
            )
        };

        let envs = run_opts.envs();
        let mut fps_results = Vec::new();

        // Determine the order: alternate between baseline-first and baseline-last
        let order: Vec<usize> = if baseline_first {
            (0..envs.len()).collect()
        } else {
            let mut order: Vec<usize> = (1..envs.len()).collect();
            order.push(0);
            order
        };

        for &idx in &order {
            let fps = do_run(&envs[idx].name, &envs[idx].wrapper)
                .context(format!("getting '{}' fps", envs[idx].name))?;
            fps_results.push((idx, fps));
        }

        // Sort back to original order (baseline first)
        fps_results.sort_by_key(|(idx, _)| *idx);
        let fps: Vec<f64> = fps_results.into_iter().map(|(_, fps)| fps).collect();

        // Log results to output files
        if let Some(output_dir) = self
            .output_dir(&run_opts.output)
            .expect("creating output directory")
        {
            for (i, &fps_val) in fps.iter().enumerate() {
                append_fps(&output_dir, &envs[i].name, fps_val).log_error();
            }
        }

        Ok(fps)
    }

    fn check_debug_filter(
        &self,
        a: &str,
        b: &str,
        envs: &FilterEnv,
        output_dir: &str,
        capture_shaders: bool,
    ) -> Result<bool> {
        let mut envs: Vec<_> = envs.to_vec();

        if capture_shaders {
            envs.push((
                "IR3_SHADER_DEBUG".to_owned(),
                "vs,fs,gs,tes,tcs,cs".to_owned(),
            ));
            envs.push(("MESA_SHADER_CACHE_DISABLE".to_owned(), "1".to_owned()));
            // Serialize zink so that the dumped shaders don't stomp each other.
            envs.push(("ZINK_DEBUG".to_owned(), "nobgc".to_owned()));
        };

        let runs = [
            ("a", self.replay(Some(a), &envs)?),
            ("b", self.replay(Some(b), &envs)?),
        ];

        if !output_dir.is_empty() {
            for (name, output) in &runs {
                let _ = self.write_output(
                    output_dir,
                    &format!("stdout-{name}.txt"),
                    output.stdout.as_bytes(),
                );
                let _ = self.write_output(
                    output_dir,
                    &format!("stderr-{name}.txt"),
                    output.stderr.as_bytes(),
                );

                if capture_shaders {
                    let shaders_dir = create_dir(Path::new(output_dir).join("shaders"))?;
                    for shader in find_shaders(&output.stderr)? {
                        std::fs::write(shaders_dir.join(&shader.sha1), shader.code.as_bytes())
                            .with_context(|| format!("writing {} shader", shader.sha1))?;
                    }
                }
            }
        }
        Ok(runs[0].1.stderr != runs[1].1.stderr)
    }

    fn output_subdir(&self) -> String {
        // ':' breaks image references for snapshot comparisons.
        self.name()
            .trim_start_matches('/')
            .replace(['/', ':', '\''], "-")
    }

    fn output_dir(&self, output_dir: &str) -> Result<Option<PathBuf>> {
        if output_dir.is_empty() {
            return Ok(None);
        }

        create_dir(Path::new(output_dir).join(self.output_subdir())).map(Some)
    }

    fn write_output(&self, output_dir: &str, filename: &str, data: &[u8]) -> Result<()> {
        if let Some(output_dir) = self
            .output_dir(output_dir)
            .expect("creating output directory")
        {
            let output_path = output_dir.join(filename);

            let mut file = std::fs::File::create(&output_path)
                .with_context(|| format!("creating {}", output_path.display()))?;

            file.write_all(data).context("failed to write")?;
        }

        Ok(())
    }

    fn run_replay_command(&self, command: Command) -> ReplayOutput {
        debug!("Running replay for {}", self.name());
        let mut command = command;
        let output = command
            .output()
            .with_context(|| format!("Running {:?}", &command))
            .unwrap();
        debug!("Finished replay for {}", self.name());

        let mut output = ReplayOutput::from(output);

        output.cmdline = format!("{command:?}");

        output
    }

    /// Runs a command with a timeout, optionally writing `stdin_data` to the
    /// process's stdin before waiting.  Returns the collected output.  If the
    /// process is killed due to the timeout, a message is appended to stderr and
    /// the exit code is set to non-zero.
    ///
    /// Stdout and stderr are drained in background threads so that pipe buffers
    /// cannot block the child process.
    fn run_replay_command_with_timeout(
        &self,
        command: Command,
        stdin_data: Option<&[u8]>,
        timeout: Duration,
    ) -> ReplayOutput {
        let mut command = command;

        let cmdline = format!("{command:?}");
        command.stdout(Stdio::piped()).stderr(Stdio::piped());
        if stdin_data.is_some() {
            command.stdin(Stdio::piped());
        }

        debug!("Running replay for {}", self.name());

        let env_marker = format!("{:016x}", rand::random::<u64>());
        command.env("GPU_TRACE_PERF_ID", &env_marker);

        let mut child = match command.spawn() {
            Ok(c) => c,
            Err(e) => {
                return ReplayOutput {
                    exit_code: 1,
                    stderr: format!("Failed to spawn: {e}"),
                    cmdline,
                    ..Default::default()
                };
            }
        };

        if let Some(data) = stdin_data {
            if let Some(mut stdin) = child.stdin.take() {
                let _ = stdin.write_all(data);
            }
        }

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

        let stdout_thread = std::thread::spawn(move || {
            let mut buf = Vec::new();
            std::io::BufReader::new(stdout).read_to_end(&mut buf).ok();
            buf
        });
        let stderr_thread = std::thread::spawn(move || {
            let mut buf = Vec::new();
            std::io::BufReader::new(stderr).read_to_end(&mut buf).ok();
            buf
        });

        // Poll /proc/<pid>/process_watcher/ for peak GPU memory while the child runs.
        let process_watcher_watcher =
            crate::process_watcher::ProcessWatcher::watch(child.id(), env_marker);

        use wait_timeout::ChildExt;
        let timed_out = match child.wait_timeout(timeout).ok().flatten() {
            Some(_) => false,
            None => {
                let _ = child.kill();
                true
            }
        };

        let exit_code = child.wait().ok().and_then(|s| s.code()).unwrap_or(1);
        let peak = process_watcher_watcher.stop();

        debug!("Finished replay for {}", self.name());
        if let Some(b) = peak.vram_bytes {
            debug!(
                "Peak GPU VRAM for {}: {} MiB",
                self.name(),
                b / (1024 * 1024)
            );
        }
        if let Some(b) = peak.sys_bytes {
            debug!(
                "Peak GPU system memory for {}: {} MiB",
                self.name(),
                b / (1024 * 1024)
            );
        }
        if let Some(b) = peak.rss_bytes {
            debug!("Peak RSS for {}: {} MiB", self.name(), b / (1024 * 1024));
        }
        let stdout = String::from_utf8_lossy(&stdout_thread.join().unwrap_or_default()).to_string();
        let mut stderr =
            String::from_utf8_lossy(&stderr_thread.join().unwrap_or_default()).to_string();

        let exit_code = if timed_out {
            warn!("Snapshot replay killed after timeout: {cmdline}");
            stderr.push_str("\nKilled due to snapshot timeout.");
            if exit_code == 0 { 1 } else { exit_code }
        } else {
            exit_code
        };

        ReplayOutput {
            exit_code,
            stdout,
            stderr,
            cmdline,
            peak,
        }
    }
}

pub fn replay_command<S>(args: &[S], wrapper: Option<&str>, envs: &[(String, String)]) -> Command
where
    S: AsRef<OsStr>,
{
    let mut command = if let Some(wrapper) = wrapper {
        let mut command = Command::new(wrapper);
        command.arg(&args[0]);
        command
    } else {
        Command::new(&args[0])
    };
    for arg in &args[1..] {
        command.arg(arg.as_ref());
    }

    for env in envs {
        command.env(&env.0, &env.1);
    }

    command
}

type FilterEnv<'a> = Vec<(String, String)>;

#[derive(Debug, Args)]
#[command(name = "replay")]
/// Runs CI snapshot replay from a .toml config
struct Replay {
    #[arg(long)]
    /// Path to the .toml config file
    config: PathBuf,

    #[arg(long)]
    /// Device name to test (e.g. freedreno-a530)
    device: String,

    #[arg(long = "output")]
    /// Output directory for results and HTML report
    output: PathBuf,

    #[arg(long)]
    /// Path to a file containing a JWT bearer token for S3 auth
    jwt: Option<PathBuf>,

    #[arg(long)]
    /// Local traces-db path (overrides S3 download; disables uploads)
    traces_db: Option<PathBuf>,

    #[arg(long, default_value = "5368709120")]
    /// Maximum disk space for trace cache in bytes (default 5 GiB)
    disk_limit: u64,
}

#[derive(Debug, Subcommand)]
enum SubCommand {
    Run(Run),
    ShaderAnalyze(ShaderAnalyze),
    Snapshot(Snapshot),
    Replay(replay::Replay),
}

#[derive(Debug, Clone)]
struct Environment {
    name: String,
    wrapper: String,
}

#[derive(Debug, Args)]
/// Runs a collection of traces, checking for performance differences between wrapper scripts
struct Run {
    #[arg(short, long, number_of_values = 1)]
    /// paths to directories containing traces (or traces themselves)
    traces: Vec<String>,

    #[arg(long, number_of_values = 1)]
    /// if present, only traces that differ in stderr when these env vars are
    /// set will be tested
    debug_filter: Vec<String>,

    #[arg(long, default_value = "")]
    /// The name of a Mesa utrace event to parse the times of rather than frame
    /// execution wall time.
    utrace: String,

    #[arg(long, default_value_t = false)]
    /// Use utrace's CSV output instead of the default JSON. May be more reliable for apps that
    /// don't properly close their GL context.
    utrace_use_csv: bool,

    #[arg(long = "output", default_value = "")]
    /// if present, path to a directory to store debug output and logs
    output: String,

    /// If set, capture per-draw, per-shader timings (where a whole draw's time
    /// is attributed to each shader stage) from utrace events and save them in
    /// the output dir.
    #[arg(long)]
    capture_shaders: bool,

    #[arg(long = "samples", default_value = "0")]
    /// Number of iterations to attempt on each trace before exiting.  Default
    /// is to loop endlessly.
    samples: usize,

    #[arg(long = "env", num_args = 2, value_names = &["NAME", "WRAPPER"])]
    /// Environment configuration: NAME is the label, WRAPPER is the command wrapper script.
    /// Can be specified multiple times. Example: --env baseline ./baseline.sh --env new ./new.sh
    env: Vec<String>,

    #[arg(long, number_of_values = 1)]
    /// Extra argument to pass to the ANGLE replay binary; may be repeated.
    /// Example: --angle-args --use-gl=native
    angle_args: Vec<String>,

    #[arg(long, number_of_values = 1)]
    /// Extra argument to pass to eglretrace; may be repeated.
    apitrace_args: Vec<String>,

    #[arg(long, number_of_values = 1)]
    /// Extra argument to pass to gfxrecon-replay; may be repeated.
    gfxreconstruct_args: Vec<String>,

    #[arg(skip)]
    environments: Vec<Environment>,
}

impl Run {
    fn parse_environments(&mut self) -> Result<()> {
        if self.env.len() % 2 != 0 {
            bail!("--env requires pairs of NAME and WRAPPER arguments");
        }
        if self.env.len() < 2 {
            bail!("At least one --env argument must be provided");
        }

        self.environments = self
            .env
            .chunks(2)
            .map(|chunk| Environment {
                name: chunk[0].clone(),
                wrapper: chunk[1].clone(),
            })
            .collect();

        Ok(())
    }

    fn envs(&self) -> &[Environment] {
        &self.environments
    }
}

#[derive(Debug, Args)]
/// Runs a collection of traces, checking for performance differences between two wrapper scripts
struct Thread64 {
    #[arg(short, long, number_of_values = 1)]
    /// paths to directories containing traces (or traces themselves)
    traces: Vec<String>,

    #[arg(long = "output")]
    /// if present, path to a directory to store debug output and logs
    output: String,
}

#[derive(Debug, Args)]
/// Runs a collection of traces, checking for performance differences between two wrapper scripts
struct ShaderAnalyze {
    /// path to a directory to store debug output and logs
    output: String,
}

type TraceList = Vec<Box<dyn TraceTool + Send>>;
type TraceResultList = Vec<(Box<dyn TraceTool + Send>, TraceResults)>;

/// Per-trace-type extra arguments to pass to the replay command.
#[derive(Default)]
pub(crate) struct TraceExtraArgs {
    pub angle: Vec<String>,
    pub apitrace: Vec<String>,
    pub gfxreconstruct: Vec<String>,
}

pub fn relative_test_name(root: &Path, file: &Path) -> String {
    let relative = if file.is_relative() || root == Path::new("/") {
        file
    } else {
        file.strip_prefix(root)
            .with_context(|| {
                format!(
                    "finding {}'s path relative to {}",
                    file.display(),
                    root.display()
                )
            })
            .unwrap()
    };
    relative.display().to_string()
}

fn path_as_angle_trace(path: &Path, angle_args: &[String]) -> Option<Box<dyn TraceTool + Send>> {
    // Allow specifying a subtest on angle_trace_tests.
    let parent = path.parent()?;
    if parent.file_name() == Some(OsStr::new("angle_trace_tests")) {
        return Some(Box::new(AngleTrace::new(
            parent,
            path.file_name()
                .context("finding angle_trace_tests subtest")
                .unwrap()
                .to_str()
                .unwrap(),
            angle_args.to_vec(),
        )));
    }
    None
}
/// Constructs the appropriate trace-tool backend from a local file path.
pub(crate) fn collect_traces_for_path(
    root: &Path,
    path: &Path,
    utrace: bool,
    nonloopable: bool,
    extra: &TraceExtraArgs,
) -> TraceList {
    let mut traces: TraceList = Vec::new();

    let suffix = path.extension();

    if path.ends_with("angle_trace_tests") {
        match AngleTrace::enumerate(path, &extra.angle) {
            Ok(angles) => {
                for trace in angles {
                    traces.push(Box::new(trace));
                }
            }
            Err(e) => error!("{}", e),
        }
    } else if suffix == Some(OsStr::new("trace")) || suffix == Some(OsStr::new("trace-dxgi")) {
        if utrace {
            traces.push(Box::new(ApitraceUtraceTrace::new(
                root,
                path,
                extra.apitrace.clone(),
            )));
        } else {
            traces.push(Box::new(ApitraceTrace::new(
                root,
                path,
                nonloopable,
                extra.apitrace.clone(),
            )));
        }
    } else if suffix == Some(OsStr::new("rdc")) {
        if utrace {
            traces.push(Box::new(RenderdocUtraceTrace::new(root, path)));
        } else {
            traces.push(Box::new(RenderdocPyTrace::new(root, path)));
        }
    } else if suffix == Some(OsStr::new("gfxr")) {
        traces.push(Box::new(GfxreconstructTrace::new(
            root,
            path,
            extra.gfxreconstruct.clone(),
        )));
    } else if suffix == Some(OsStr::new("mock-trace")) {
        traces.push(Box::new(mock_trace::MockTrace::new(root, path)));
    };

    if let Some(angle) = path_as_angle_trace(path, &extra.angle) {
        traces.push(angle);
    }

    traces
}

// Walk the directories specified in the command line args and make a vec of the actual traces to run
pub(crate) fn collect_traces(
    root: Option<&Path>,
    paths: &[String],
    utrace: bool,
    extra: &TraceExtraArgs,
) -> TraceList {
    let mut traces: TraceList = Vec::new();

    let root = root.unwrap_or(Path::new("/"));

    for file in paths {
        let walk = WalkDir::new(file).follow_links(true);

        for entry in walk.into_iter().filter_map(|e| e.ok()) {
            traces.append(&mut collect_traces_for_path(
                root,
                entry.path(),
                utrace,
                false,
                extra,
            ));
        }

        // Parse for angle_trace_tests subtest specifications
        if !Path::new(file).exists() {
            if let Some(angle) = path_as_angle_trace(Path::new(file), &extra.angle) {
                traces.push(angle);
            }
        }
    }

    traces
}

#[derive(Clone)]
struct TraceResults {
    /// Results for each environment. First element is the baseline.
    results: Vec<Vec<f64>>,
    logged: bool,
}

impl TraceResults {
    pub fn new(num_envs: usize) -> TraceResults {
        TraceResults {
            results: vec![Vec::new(); num_envs],
            logged: false,
        }
    }

    /// Performs a resampling with replacement from the "from" instance.
    pub fn resample<R: Rng>(&mut self, from: &TraceResults, rng: &mut R) {
        assert_eq!(self.results.len(), from.results.len());
        for i in 0..self.results.len() {
            assert_eq!(self.results[i].len(), from.results[i].len());
        }

        for i in 0..self.results.len() {
            for val in &mut self.results[i] {
                *val = *from.results[i].choose(rng).unwrap();
            }
        }
    }
}

fn append_fps(path: &Path, wrapper: &str, fps: f64) -> Result<()> {
    let path = path.join(format!("fps-{wrapper}.txt"));
    let mut file = OpenOptions::new()
        .append(true)
        .create(true)
        .open(&path)
        .with_context(|| format!(" opening {}", path.display()))?;

    writeln!(file, "{fps}").with_context(|| format!("writing fps to {}", path.display()))?;

    Ok(())
}

fn run_each_trace(run_opts: &Run, results: &mut TraceResultList, baseline_first: bool) {
    for (trace, results) in results.iter_mut() {
        match trace.collect_fps(run_opts, baseline_first) {
            Ok(fps_vec) => {
                for (i, fps) in fps_vec.into_iter().enumerate() {
                    results.results[i].push(fps);
                }
            }
            Err(e) => {
                if !results.logged {
                    error!("  Failed: {:?}", e);
                    results.logged = true;
                } else {
                    error!("  Failed");
                }
            }
        }
    }
}

fn print_stats<R: Rng>(results: &TraceResultList, env_names: &[&str], rng: &mut R) {
    if env_names.is_empty() {
        return;
    }

    // Filter results to only include those with data for all environments
    let results: Vec<_> = results
        .iter()
        .filter(|(_, result)| result.results.iter().all(|r| !r.is_empty()))
        .collect();

    if results.is_empty() {
        return;
    }

    let namelen = results
        .iter()
        .map(|(trace, _)| trace.name().len())
        .max()
        .unwrap_or(0);

    // Handle single environment case (no comparison)
    if env_names.len() == 1 {
        // Sort by FPS
        let mut single_stats: Vec<_> = results
            .iter()
            .map(|(trace, results)| {
                let data = &results.results[0];
                (
                    trace.name(),
                    data.iter().mean(),
                    data.iter().std_dev(),
                    data.len(),
                )
            })
            .collect();

        single_stats
            .sort_by(|(_, a_fps, _, _), (_, b_fps, _, _)| a_fps.partial_cmp(b_fps).unwrap());

        for (trace, mean_fps, stddev, count) in &single_stats {
            println!("{trace:namelen$}: {mean_fps:6.1} fps (+/- {stddev:5.1}%) (n={count})");
        }

        return;
    }

    // Apply the Bonferroni correction for multiple hypothesis testing.
    let comparison_count = env_names.len() - 1;
    let trace_count = results.len();
    let alpha = 0.05 / ((comparison_count * trace_count) as f64);

    // For each comparison (baseline vs each other environment)
    for env_idx in 1..env_names.len() {
        let mut stats: Vec<_> = results
            .iter()
            .map(|(trace, results)| {
                (
                    trace.name(),
                    ResultStats::new(&results.results[0], &results.results[env_idx], alpha),
                )
            })
            .collect();

        stats.sort_by(|(_, a), (_, b)| a.change.total_cmp(&b.change));

        for (trace, stat) in &stats {
            let change = if stat.has_fps() {
                format!("{:>7.2}%", stat.change * 100.0)
            } else {
                "no time detected".to_string()
            };

            let error = if stat.n[0] > 1 && stat.has_fps() {
                format!(" (+/- {:5.1}%)", stat.error * 100.0)
            } else {
                "".to_string()
            };

            let (before, after) = if stat.has_fps() {
                (
                    format!("{:5.1}", stat.means[0]),
                    format!("{:5.1}", stat.means[1]),
                )
            } else {
                ("".to_string(), "".to_string())
            };

            let count = if stat.n[0] == stat.n[1] {
                format!("{}", stat.n[0])
            } else {
                format!("{}/{}", stat.n[0], stat.n[1])
            };

            println!("{trace:namelen$}: {before:6} -> {after:6} fps {change}{error} (n={count})");
        }

        let samples = stats.iter().map(|x| x.1.n[0]).max().unwrap_or(0);
        if samples > 5 {
            let summary_stats = BootstrappedRelativeAndMaxChange::new(
                &results
                    .iter()
                    .map(|(_, r)| {
                        let mut tr = TraceResults::new(2);
                        tr.results[0] = r.results[0].clone();
                        tr.results[1] = r.results[env_idx].clone();
                        tr
                    })
                    .collect::<Vec<TraceResults>>(),
                100,
                rng,
            );
            println!(
                "average fps {:+.1}% (+/- {:.1}%)",
                summary_stats.relative_mean_change * 100.0 - 100.0,
                summary_stats.relative_mean_error * 100.0
            );
            println!(
                "max fps {:+.1}% (+/- {:.1}%)",
                summary_stats.relative_max_change * 100.0 - 100.0,
                summary_stats.relative_max_error * 100.0
            );
        }
    }
}

fn run(run_opts: &Run, traces: TraceList) {
    let envs = run_opts.envs();
    let num_envs = envs.len();

    let mut results: Vec<_> = traces
        .into_iter()
        .map(|trace| (trace, TraceResults::new(num_envs)))
        .collect();

    // Use a consistently seeded rng for bootstrapping, so that we get reproducible output for testing.
    let mut rng = rand_chacha::ChaCha12Rng::seed_from_u64(0xabcdef12);

    let env_names: Vec<&str> = envs.iter().map(|e| e.name.as_str()).collect();
    let mut baseline_first = true;
    for samples in 1.. {
        run_each_trace(run_opts, &mut results, baseline_first);
        print_stats(&results, &env_names, &mut rng);

        baseline_first = !baseline_first;

        if run_opts.samples > 0 && samples >= run_opts.samples {
            break;
        }

        if samples == 1 {
            results.retain(|(_trace, result)| !result.results.iter().any(|x| x.is_empty()));
        }
    }
}

// Test the wrapper environment scripts people hand us before trying to start up an actual trace
fn test_run_wrapper(wrapper: &str) -> bool {
    let output = Command::new(wrapper).arg("true").output();
    match output {
        Ok(output) => {
            if !output.status.success() {
                error!(
                    "Test invocation of '{} true' failed: {}",
                    wrapper,
                    ReplayOutput::from(output)
                );
                false
            } else {
                true
            }
        }
        Err(err) => {
            error!(
                "Failed to spawn test invocation of '{} true': {:?}",
                wrapper, err
            );
            info!(
                "TIP: Exec format error here probably means you need #!/bin/sh in your shell script."
            );
            false
        }
    }
}

fn debug_filter_traces(traces: TraceList, run_opts: &Run) -> TraceList {
    let mut filters = Vec::new();
    for f in &run_opts.debug_filter {
        if let Some((k, v)) = f.split_once('=') {
            filters.push((k.to_string(), v.to_string()));
        } else {
            error!(
                "Debug filter '{}' should be a KEY=value environment assignment",
                f
            );
            exit(1);
        }
    }
    if filters.is_empty() && !run_opts.capture_shaders {
        return traces;
    }

    println!(
        "Checking for effects of debug filters{}:",
        if run_opts.capture_shaders {
            " and capturing shaders"
        } else {
            ""
        }
    );

    let mut traces = traces;
    traces.retain(|trace| {
        print!("  {}: ", trace.name());
        match trace.check_debug_filter(
            &run_opts.envs()[0].wrapper,
            &run_opts.envs()[1].wrapper,
            &filters,
            &run_opts.output,
            run_opts.capture_shaders,
        ) {
            Ok(true) => {
                println!("affected");
                true
            }
            Ok(false) => {
                println!("skipped");
                false
            }
            Err(e) => {
                println!("failed ({e})");
                false
            }
        }
    });

    traces
}

fn init_logging(cli: &Cli) -> Result<()> {
    let level_filter = match cli.verbose {
        0 => log::LevelFilter::Info,
        1 => log::LevelFilter::Debug,
        _ => log::LevelFilter::Trace,
    };
    let mut loggers: Vec<Box<dyn simplelog::SharedLogger>> = Vec::new();
    let log_time_level = if cli.no_log_time {
        log::LevelFilter::Off
    } else {
        log::LevelFilter::Error
    };

    let output_dir: Option<&Path> = match &cli.subcmd {
        SubCommand::Run(opts) if !opts.output.is_empty() => Some(Path::new(&opts.output)),
        SubCommand::Snapshot(opts) => Some(Path::new(&opts.output)),
        SubCommand::Replay(opts) => Some(opts.output.as_path()),
        _ => None,
    };

    let config = simplelog::ConfigBuilder::new()
        .set_time_level(log_time_level)
        .build();
    if let Some(output_dir) = output_dir {
        create_dir_all(output_dir).context("creating output dir")?;

        loggers.push(simplelog::WriteLogger::new(
            max(level_filter, log::LevelFilter::Debug),
            config.clone(),
            File::create(output_dir.join("log.txt")).context("creating output log file")?,
        ));
    }

    loggers.push(simplelog::SimpleLogger::new(level_filter, config));

    CombinedLogger::init(loggers).context("setting up logging")?;

    Ok(())
}

fn main() -> Result<()> {
    let mut cli = Cli::parse();

    init_logging(&cli)?;

    match &mut cli.subcmd {
        SubCommand::Run(run_opts) => {
            if let Err(e) = run_opts.parse_environments() {
                error!("{}", e);
                error!("Example: --env baseline ./baseline.sh --env new ./new.sh");
                std::process::exit(1);
            }

            for env in run_opts.envs() {
                if !test_run_wrapper(&env.wrapper) {
                    std::process::exit(1);
                }
            }

            let extra = TraceExtraArgs {
                angle: run_opts.angle_args.clone(),
                apitrace: run_opts.apitrace_args.clone(),
                gfxreconstruct: run_opts.gfxreconstruct_args.clone(),
            };
            let traces = collect_traces(
                None,
                &run_opts.traces,
                !run_opts.utrace.is_empty() || run_opts.capture_shaders,
                &extra,
            );
            if traces.is_empty() {
                error!("No traces found in the given directories:");
                for t in &run_opts.traces {
                    error!("  {}", t);
                }
                std::process::exit(1);
            }

            let traces = debug_filter_traces(traces, run_opts);
            if traces.is_empty() {
                warn!("All traces filtered out by the debug filter.");
                std::process::exit(0);
            }

            run(run_opts, traces);

            Ok(())
        }

        SubCommand::ShaderAnalyze(analyze) => shader_analyze::shader_analyze(analyze),

        SubCommand::Snapshot(opts) => snapshot::snapshot(opts),

        SubCommand::Replay(replay) => replay::replay(replay),
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    struct UnitTestTrace {
        fps: f64,
        name: String,
    }

    impl TraceTool for UnitTestTrace {
        fn replay(
            &self,
            wrapper: Option<&str>,
            _envs: &[(String, String)],
        ) -> Result<ReplayOutput> {
            Ok(ReplayOutput {
                stdout: wrapper.unwrap_or("bad").to_string(),
                ..Default::default()
            })
        }

        fn fps(&self, output: &ReplayOutput) -> Result<f64> {
            output
                .stdout
                .parse::<f64>()
                .context("parsing unit test wrapper as f64")
                .map(|x| x * self.fps)
        }

        fn name(&self) -> &str {
            &self.name
        }
    }

    impl UnitTestTrace {
        pub fn new(fps: f64) -> UnitTestTrace {
            UnitTestTrace {
                fps,
                name: format!("unit_test_{fps}"),
            }
        }
    }

    #[test]
    fn test_ordering() {
        let mut results: TraceResultList =
            vec![(Box::new(UnitTestTrace::new(5.0)), TraceResults::new(2))];

        let run = Run {
            traces: Vec::new(),
            debug_filter: Vec::new(),
            output: String::new(),
            utrace: "".to_string(),
            utrace_use_csv: false,
            capture_shaders: false,
            samples: 0,
            env: vec![],
            angle_args: vec![],
            apitrace_args: vec![],
            gfxreconstruct_args: vec![],
            environments: vec![
                Environment {
                    name: "env0".to_string(),
                    wrapper: "3.0".to_string(),
                },
                Environment {
                    name: "env1".to_string(),
                    wrapper: "4.0".to_string(),
                },
            ],
        };

        run_each_trace(&run, &mut results, true);
        run_each_trace(&run, &mut results, false);

        for (_, results) in results {
            for val in &results.results[0] {
                assert_eq!(*val, 5.0 * 3.0);
            }
            for val in &results.results[1] {
                assert_eq!(*val, 5.0 * 4.0);
            }
        }
    }

    #[test]
    fn verify_cli() {
        use clap::CommandFactory;
        Cli::command().debug_assert();
    }

    #[test]
    fn test_collect_traces() -> Result<()> {
        let manifest_path = env!("CARGO_MANIFEST_PATH");
        let top = Path::new(manifest_path).parent().unwrap();

        let traces = collect_traces(
            Some(top),
            &[
                top.join("src/test_data/walkdir_traces/")
                    .display()
                    .to_string(),
                top.join("src/test_data/vkcube.rdc").display().to_string(),
                "angle_trace_tests/b".to_string(),
            ],
            false,
            &TraceExtraArgs::default(),
        );
        let mut names: Vec<_> = traces.iter().map(|x| x.name()).collect();
        names.sort();

        assert_eq!(
            names,
            vec![
                "angle_trace_tests/b",
                "src/test_data/vkcube.rdc",
                "src/test_data/walkdir_traces/vkcube-10-frames.gfxr",
                "src/test_data/walkdir_traces/vkcube.rdc",
            ]
        );

        Ok(())
    }

    #[test]
    fn test_timeout_pipe() {
        struct TimeoutTest {}

        impl TraceTool for TimeoutTest {
            fn replay(
                &self,
                _wrapper: Option<&str>,
                _envs: &[(String, String)],
            ) -> Result<ReplayOutput> {
                let mut command = Command::new("sh");
                command.arg("-c");
                command.arg("yes extremely long output line | head -n 50000");

                Ok(self.run_replay_command_with_timeout(command, None, Duration::MAX))
            }

            fn fps(&self, _output: &ReplayOutput) -> Result<f64> {
                todo!()
            }

            fn name(&self) -> &str {
                todo!()
            }
        }
        let trace = TimeoutTest {};
        let output = trace.replay(None, &[]).unwrap();
        assert_eq!(output.exit_code, 0, "output: {output}");
        assert_eq!(output.stdout.lines().count(), 50000);
    }
}