gpu-trace-perf 1.8.2

use anyhow::{Context, Result, bail};
use regex::Regex;
use std::{
    ffi::OsStr,
    path::{Path, PathBuf},
    process::Command,
    time::Duration,
};

use crate::{ReplayOutput, TraceTool, replay_command, snapshot::SnapshotResult};

pub struct GfxreconstructTrace {
    file: PathBuf,
    name: String,
    extra_args: Vec<String>,
}

impl GfxreconstructTrace {
    pub fn new(root: &Path, file: &Path, extra_args: Vec<String>) -> GfxreconstructTrace {
        GfxreconstructTrace {
            file: file.to_owned(),
            name: crate::relative_test_name(root, file),
            extra_args,
        }
    }
}

impl TraceTool for GfxreconstructTrace {
    fn replay(&self, wrapper: Option<&str>, envs: &[(String, String)]) -> Result<ReplayOutput> {
        let mut gfxr_command: Vec<&OsStr> = vec![
            OsStr::new("gfxrecon-replay"),
            OsStr::new("--swapchain"),
            OsStr::new("offscreen"),
            OsStr::new("-m"),
            OsStr::new("remap"),
        ];
        for arg in &self.extra_args {
            gfxr_command.push(OsStr::new(arg));
        }
        gfxr_command.push(self.file.as_os_str());

        let command = replay_command(&gfxr_command, wrapper, envs);
        let output: ReplayOutput = self.run_replay_command(command);

        Ok(output)
    }

    fn fps(&self, output: &crate::ReplayOutput) -> Result<f64> {
        parse_gfxrecon_fps_output(&output.stdout)
            .with_context(|| format!("with stderr: {}", &output.stderr))
    }

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

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

    fn snapshot(&self, output_dir: &str, loops: u32, timeout: Duration) -> Result<SnapshotResult> {
        let output_dir_path = self.output_dir(output_dir)?.unwrap();
        let total_frames = gfxrecon_total_frames(&self.file)?;
        let total_frames_str = total_frames.to_string();
        let screenshot_name = format!("screenshot_frame_{total_frames}.png");

        let start_time = std::time::Instant::now();
        let mut last_output = None;
        let mut files = Vec::new();

        for i in 1..=loops {
            let mut gfxr_command: Vec<&OsStr> = vec![
                OsStr::new("gfxrecon-replay"),
                OsStr::new("--swapchain"),
                OsStr::new("offscreen"),
                OsStr::new("-m"),
                OsStr::new("remap"),
            ];
            for arg in &self.extra_args {
                gfxr_command.push(OsStr::new(arg));
            }
            gfxr_command.extend(&[
                OsStr::new("--screenshots"),
                OsStr::new(&total_frames_str),
                OsStr::new("--screenshot-format"),
                OsStr::new("png"),
                OsStr::new("--screenshot-dir"),
                output_dir_path.as_os_str(),
                self.file.as_os_str(),
            ]);

            let command = replay_command(&gfxr_command, None, &[]);
            last_output = Some(self.run_replay_command_with_timeout(command, None, timeout));

            let screenshot_path = output_dir_path.join(&screenshot_name);

            if last_output.as_ref().unwrap().exit_code != 0 {
                // Delete any screenshot file the tool may have written before
                // failing, so it doesn't accumulate in the output directory.
                let _ = std::fs::remove_file(&screenshot_path);
                break;
            }

            // Move the screenshot to our numbered snapshot filename.
            let snapshot_filename = format!("snapshot{i:04}.png");
            let snapshot_path = output_dir_path.join(&snapshot_filename);

            std::fs::rename(&screenshot_path, &snapshot_path).with_context(|| {
                format!(
                    "renaming {} to {}",
                    screenshot_path.display(),
                    snapshot_path.display()
                )
            })?;

            files.push(PathBuf::from(&snapshot_filename));
        }

        let last_output = last_output.context("Finding the output for the last command")?;
        Ok(SnapshotResult {
            files,
            output: last_output,
            runtime: start_time.elapsed(),
        })
    }
}

fn parse_gfxrecon_total_frames(output: &str) -> Result<u32> {
    lazy_static! {
        static ref FRAMES_RE: Regex = Regex::new(r"Total frames:\s*([0-9]+)").unwrap();
    }
    for line in output.lines() {
        if let Some(cap) = FRAMES_RE.captures(line) {
            return cap[1].parse::<u32>().context("parsing total frame count");
        }
    }
    bail!("Could not find 'Total frames' in gfxrecon-info output")
}

fn gfxrecon_total_frames(file: &Path) -> Result<u32> {
    let output = Command::new("gfxrecon-info")
        .arg(file)
        .output()
        .context("calling gfxrecon-info")?;
    parse_gfxrecon_total_frames(&String::from_utf8_lossy(&output.stdout))
        .with_context(|| format!("getting frame count for {}: {:?}", file.display(), output))
}

// Returns the FPS for the frame from gpu-trace-perf-renderdoc-wrapper.py output
fn parse_gfxrecon_fps_output(output: &str) -> Result<f64> {
    lazy_static! {
        static ref OLD_RE: Regex = Regex::new("Replay FPS: ([0-9.]*) fps,").unwrap();
        static ref NEW_RE: Regex = Regex::new("Measured FPS: ([0-9.]*) fps,").unwrap();
    }

    for line in output.lines() {
        if let Some(cap) = OLD_RE.captures(line) {
            return cap[1]
                .parse::<f64>()
                .with_context(|| format!("Parsing gfxrecon FPS: {line}"));
        }
        if let Some(cap) = NEW_RE.captures(line) {
            return cap[1]
                .parse::<f64>()
                .with_context(|| format!("Parsing gfxrecon FPS: {line}"));
        }
    }

    bail!("Failed to find replay FPS line in {:?}", output);
}

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

    #[test]
    fn test_gfxrecon_total_frames() {
        // Actual gfxrecon-info output for a trimmed capture.
        let info_output = "
[gfxrecon] INFO - Loading state for captured frame 38
[gfxrecon] INFO - Finished loading state for captured frame 38
Exe info:
\tApplication exe name: vkcube

File info:
\tTotal frames: 10 (trimmed frame range 38-47)
";
        assert_eq!(parse_gfxrecon_total_frames(info_output).unwrap(), 10);
    }

    #[test]
    fn test_gfxrecon_parsing_old() {
        let gfxrecon_input = "
[gfxrecon] INFO - Loading state for captured frame 2
[gfxrecon] INFO - Finished loading state for captured frame 2
Load time:  1.853056 seconds
Total time: 2.019838 seconds
Replay FPS: 59.958586 fps, 0.166782 seconds, 10 frames, framerange 2-11
";
        assert_eq!(
            parse_gfxrecon_fps_output(gfxrecon_input).unwrap(),
            59.958586
        );
    }

    #[test]
    fn test_gfxrecon_parsing_new() {
        let gfxrecon_input = "
================== Start timer (Frame: 1) ==================
[gfxrecon] INFO - Loading state for captured frame 38
[gfxrecon] INFO - Replay adjusted the vkGetPhysicalDeviceSurfaceFormatsKHR array count: capture count = 2, replay count = 0
[gfxrecon] WARNING - OverrideCreatePipelineCache(): PipelineCache data was provided, but pipelineCacheUUIDs did not match. This requires a pipeline-recompilation and may cause unexpected delays.
[gfxrecon] INFO - Finished loading state for captured frame 38
================== Start timer (Frame: 2) ==================
================== End timer (Frame: 11) ==================
Load time:  0.189234 seconds (frame 38)
Total time: 0.282403 seconds
Measured FPS: 107.736265 fps, 0.092819 seconds, 10 frames, 1 loop, framerange [1-11)
";
        assert_eq!(
            parse_gfxrecon_fps_output(gfxrecon_input).unwrap(),
            107.736265
        );
    }
}