gpu-trace-perf 1.9.0

use anyhow::{Context as _, Result, bail};
use std::{
    fmt,
    io::{self, BufReader, Read, Seek, SeekFrom},
    path::Path,
};

/// Rendering API used in a .rdc capture file, corresponding to RenderDoc's RDCDriver enum.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum RdcDriver {
    Unknown,
    D3D11,
    OpenGL,
    Mantle,
    D3D12,
    D3D10,
    D3D9,
    Image,
    Vulkan,
    OpenGLES,
    D3D8,
    Metal,
    Other(u32),
}

impl RdcDriver {
    pub fn is_directx(self) -> bool {
        matches!(
            self,
            RdcDriver::D3D8
                | RdcDriver::D3D9
                | RdcDriver::D3D10
                | RdcDriver::D3D11
                | RdcDriver::D3D12
        )
    }
}

impl fmt::Display for RdcDriver {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            RdcDriver::Unknown => write!(f, "Unknown"),
            RdcDriver::D3D8 => write!(f, "D3D8"),
            RdcDriver::D3D9 => write!(f, "D3D9"),
            RdcDriver::D3D10 => write!(f, "D3D10"),
            RdcDriver::D3D11 => write!(f, "D3D11"),
            RdcDriver::D3D12 => write!(f, "D3D12"),
            RdcDriver::OpenGL => write!(f, "OpenGL"),
            RdcDriver::OpenGLES => write!(f, "OpenGLES"),
            RdcDriver::Vulkan => write!(f, "Vulkan"),
            RdcDriver::Metal => write!(f, "Metal"),
            RdcDriver::Mantle => write!(f, "Mantle"),
            RdcDriver::Image => write!(f, "Image"),
            RdcDriver::Other(id) => write!(f, "Other({id})"),
        }
    }
}

const RDC_MAGIC: u32 = u32::from_le_bytes([b'R', b'D', b'O', b'C']);

fn read_u16_le(r: &mut impl Read) -> Result<u16> {
    let mut buf = [0u8; 2];
    r.read_exact(&mut buf)?;
    Ok(u16::from_le_bytes(buf))
}

fn read_u32_le(r: &mut impl Read) -> Result<u32> {
    let mut buf = [0u8; 4];
    r.read_exact(&mut buf)?;
    Ok(u32::from_le_bytes(buf))
}

fn read_u64_le(r: &mut impl Read) -> Result<u64> {
    let mut buf = [0u8; 8];
    r.read_exact(&mut buf)?;
    Ok(u64::from_le_bytes(buf))
}

/// Parse the rendering API from a RenderDoc .rdc capture file header.
pub fn parse_rdc_driver(path: &Path) -> Result<RdcDriver> {
    let mut f = std::fs::File::open(path).context("opening .rdc file")?;

    // FileHeader: magic (u64), version (u32), headerLength (u32), progVersion ([u8;16])
    // headerLength covers the entire preamble (FileHeader + BinaryThumbnail + CaptureMetaData).
    let magic_lo = read_u32_le(&mut f).context("reading magic")?;
    if magic_lo != RDC_MAGIC {
        bail!("not a valid .rdc file: bad magic");
    }
    let magic_hi = read_u32_le(&mut f).context("reading magic hi")?;
    if magic_hi != 0 {
        bail!("not a valid .rdc file: unexpected magic high bytes {magic_hi:#x}");
    }
    let _version = read_u32_le(&mut f).context("reading version")?;
    let _header_length = read_u32_le(&mut f).context("reading headerLength")?;
    // Skip progVersion[16].
    f.seek(SeekFrom::Current(16))
        .context("seeking past progVersion")?;

    // BinaryThumbnail: width (u16), height (u16), length (u32), data[length]
    let _thumb_width = read_u16_le(&mut f).context("reading thumbnail width")?;
    let _thumb_height = read_u16_le(&mut f).context("reading thumbnail height")?;
    let thumb_length = read_u32_le(&mut f).context("reading thumbnail length")?;
    f.seek(SeekFrom::Current(thumb_length as i64))
        .context("seeking past thumbnail data")?;

    // CaptureMetaData: machineIdent (u64), driverID (u32), driverNameLength (u8), driverName[...]
    let _machine_ident = read_u64_le(&mut f).context("reading machineIdent")?;
    let driver_id = read_u32_le(&mut f).context("reading driverID")?;

    Ok(match driver_id {
        0 => RdcDriver::Unknown,
        1 => RdcDriver::D3D11,
        2 => RdcDriver::OpenGL,
        3 => RdcDriver::Mantle,
        4 => RdcDriver::D3D12,
        5 => RdcDriver::D3D10,
        6 => RdcDriver::D3D9,
        7 => RdcDriver::Image,
        8 => RdcDriver::Vulkan,
        9 => RdcDriver::OpenGLES,
        10 => RdcDriver::D3D8,
        11 => RdcDriver::Metal,
        other => RdcDriver::Other(other),
    })
}

const SECTION_TYPE_FRAME_CAPTURE: u32 = 1;
const SECTION_FLAG_LZ4: u32 = 0x2;
const SECTION_FLAG_ZSTD: u32 = 0x4;

// D3D11 chunk IDs: SystemChunk::FirstDriverChunk=1000, SetResourceName=1001, CreateSwapBuffer=1002
const D3D11_CREATE_SWAP_BUFFER: u32 = 1002;

const CHUNK_INDEX_MASK: u32 = 0x0000_ffff;
const CHUNK_FLAG_CALLSTACK: u32 = 0x0001_0000;
const CHUNK_FLAG_THREAD_ID: u32 = 0x0002_0000;
const CHUNK_FLAG_DURATION: u32 = 0x0004_0000;
const CHUNK_FLAG_TIMESTAMP: u32 = 0x0008_0000;
const CHUNK_FLAG_64BIT_SIZE: u32 = 0x0010_0000;

// RenderDoc LZ4 block size
const RDC_LZ4_BLOCK_SIZE: usize = 1024 * 1024;

/// Streaming reader over RenderDoc's custom LZ4 block format.
///
/// Each block is an i32 compressed size followed by that many bytes of raw LZ4
/// block data.  Blocks are linked via `LZ4_decompress_safe_continue`: each can
/// reference the last 64 KB of the previous block's output.  We alternate
/// between two 1-MB page buffers (matching RenderDoc's compressor) so the
/// streaming context can resolve cross-block back-references.
struct RdcLz4BlockReader<R: Read> {
    reader: R,
    comp_remaining: usize,
    stream: *mut lz4_sys::LZ4StreamDecode,
    /// Two alternating output pages.  Both stay alive so the LZ4 streaming
    /// context's internal pointer into the previous page remains valid.
    pages: [Vec<u8>; 2],
    current_page: usize,
    block_len: usize,
    block_pos: usize,
}

impl<R: Read> RdcLz4BlockReader<R> {
    fn new(reader: R, comp_remaining: usize) -> io::Result<Self> {
        let stream = unsafe { lz4_sys::LZ4_createStreamDecode() };
        if stream.is_null() {
            return Err(io::Error::other("LZ4_createStreamDecode failed"));
        }
        unsafe { lz4_sys::LZ4_setStreamDecode(stream, std::ptr::null(), 0) };
        Ok(Self {
            reader,
            comp_remaining,
            stream,
            pages: [vec![0u8; RDC_LZ4_BLOCK_SIZE], vec![0u8; RDC_LZ4_BLOCK_SIZE]],
            current_page: 0,
            block_len: 0,
            block_pos: 0,
        })
    }

    fn next_block(&mut self) -> io::Result<bool> {
        if self.comp_remaining < 4 {
            return Ok(false);
        }
        let mut size_buf = [0u8; 4];
        self.reader.read_exact(&mut size_buf)?;
        self.comp_remaining -= 4;

        let block_comp_size = i32::from_le_bytes(size_buf);
        if block_comp_size <= 0 {
            return Ok(false);
        }
        let block_comp_size = block_comp_size as usize;
        if block_comp_size > self.comp_remaining {
            return Ok(false);
        }

        let mut compressed = vec![0u8; block_comp_size];
        self.reader.read_exact(&mut compressed)?;
        self.comp_remaining -= block_comp_size;

        self.current_page ^= 1;
        let page = &mut self.pages[self.current_page];

        let n = unsafe {
            lz4_sys::LZ4_decompress_safe_continue(
                self.stream,
                compressed.as_ptr(),
                page.as_mut_ptr(),
                block_comp_size as i32,
                RDC_LZ4_BLOCK_SIZE as i32,
            )
        };
        if n < 0 {
            return Err(io::Error::new(
                io::ErrorKind::InvalidData,
                format!("LZ4_decompress_safe_continue returned {n}"),
            ));
        }

        self.block_len = n as usize;
        self.block_pos = 0;
        Ok(true)
    }
}

impl<R: Read> Drop for RdcLz4BlockReader<R> {
    fn drop(&mut self) {
        unsafe { lz4_sys::LZ4_freeStreamDecode(self.stream) };
    }
}

impl<R: Read> Read for RdcLz4BlockReader<R> {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        if self.block_pos >= self.block_len && !self.next_block()? {
            return Ok(0);
        }
        let available = &self.pages[self.current_page][self.block_pos..self.block_len];
        let n = buf.len().min(available.len());
        buf[..n].copy_from_slice(&available[..n]);
        self.block_pos += n;
        Ok(n)
    }
}

/// Streaming iterator over frame-capture chunks from a .rdc file.
///
/// Internally reads from a decompressing stream (LZ4, Zstd, or raw) via a
/// `BufReader`, parsing each chunk header and yielding `(chunk_id, data)`.
/// Only one decompressed LZ4 block (~1 MB) plus a 64-KB back-reference
/// dictionary are held in memory at a time.
pub struct RdcFrameCaptureChunks {
    reader: BufReader<Box<dyn Read>>,
    /// Byte offset in the decompressed stream — used to compute 64-byte chunk alignment.
    stream_pos: usize,
    done: bool,
}

/// Seeks to the FrameCapture section in an .rdc reader and returns a chunk
/// iterator.  The reader must be positioned at the start of the file.
pub fn parse_rdc_frame_capture_chunks<R: Read + Seek + 'static>(
    mut reader: R,
) -> Result<RdcFrameCaptureChunks> {
    let magic_lo = read_u32_le(&mut reader).context("reading magic")?;
    if magic_lo != RDC_MAGIC {
        bail!("not a valid .rdc file: bad magic");
    }
    let _ = read_u32_le(&mut reader)?; // magic_hi
    let _ = read_u32_le(&mut reader)?; // version
    let header_length = read_u32_le(&mut reader).context("reading headerLength")? as u64;

    reader
        .seek(SeekFrom::Start(header_length))
        .context("seeking to sections")?;

    loop {
        let mut first = [0u8; 1];
        match reader.read_exact(&mut first) {
            Ok(_) => {}
            Err(e) if e.kind() == io::ErrorKind::UnexpectedEof => {
                bail!("no FrameCapture section found")
            }
            Err(e) => return Err(e.into()),
        }
        if first[0] == b'A' {
            bail!("ASCII .rdc sections not supported");
        }

        // BinarySectionHeader remainder: zero[3] + sectionType(4) + compLen(8) +
        //   uncompLen(8) + version(8) + flags(4) + nameLen(4)
        reader.seek(SeekFrom::Current(3))?;
        let section_type = read_u32_le(&mut reader)?;
        let comp_len = read_u64_le(&mut reader)? as usize;
        let _ = read_u64_le(&mut reader)?; // uncompressed length
        let _ = read_u64_le(&mut reader)?; // section version
        let section_flags = read_u32_le(&mut reader)?;
        let name_len = read_u32_le(&mut reader)?;
        reader.seek(SeekFrom::Current(name_len as i64))?;

        if section_type != SECTION_TYPE_FRAME_CAPTURE {
            reader.seek(SeekFrom::Current(comp_len as i64))?;
            continue;
        }

        let section_reader: Box<dyn Read> = if section_flags & SECTION_FLAG_LZ4 != 0 {
            Box::new(RdcLz4BlockReader::new(reader, comp_len).context("creating LZ4 reader")?)
        } else if section_flags & SECTION_FLAG_ZSTD != 0 {
            let take = reader.take(comp_len as u64);
            Box::new(zstd::Decoder::new(BufReader::new(take)).context("creating zstd decoder")?)
        } else {
            Box::new(reader.take(comp_len as u64))
        };

        return Ok(RdcFrameCaptureChunks {
            reader: BufReader::new(section_reader),
            stream_pos: 0,
            done: false,
        });
    }
}

impl RdcFrameCaptureChunks {
    fn read_u32(&mut self) -> io::Result<u32> {
        let mut buf = [0u8; 4];
        self.reader.read_exact(&mut buf)?;
        self.stream_pos += 4;
        Ok(u32::from_le_bytes(buf))
    }

    fn read_u64(&mut self) -> io::Result<u64> {
        let mut buf = [0u8; 8];
        self.reader.read_exact(&mut buf)?;
        self.stream_pos += 8;
        Ok(u64::from_le_bytes(buf))
    }

    fn skip(&mut self, n: usize) -> io::Result<()> {
        let mut remaining = n;
        let mut discard = [0u8; 64];
        while remaining > 0 {
            let to_read = remaining.min(discard.len());
            self.reader.read_exact(&mut discard[..to_read])?;
            remaining -= to_read;
        }
        self.stream_pos += n;
        Ok(())
    }
}

impl Iterator for RdcFrameCaptureChunks {
    /// `(chunk_id, chunk_data)` where chunk_data is this chunk's serialized bytes.
    type Item = (u32, Vec<u8>);

    fn next(&mut self) -> Option<(u32, Vec<u8>)> {
        if self.done {
            return None;
        }

        let word = match self.read_u32() {
            Ok(0) | Err(_) => {
                self.done = true;
                return None;
            }
            Ok(w) => w,
        };

        let chunk_id = word & CHUNK_INDEX_MASK;

        if word & CHUNK_FLAG_CALLSTACK != 0 {
            let n = self.read_u32().ok()? as usize;
            self.skip(n * 8).ok()?;
        }
        let meta_skip = (if word & CHUNK_FLAG_THREAD_ID != 0 {
            8
        } else {
            0
        }) + (if word & CHUNK_FLAG_DURATION != 0 {
            8
        } else {
            0
        }) + (if word & CHUNK_FLAG_TIMESTAMP != 0 {
            8
        } else {
            0
        });
        if meta_skip > 0 {
            self.skip(meta_skip).ok()?;
        }

        let chunk_size = if word & CHUNK_FLAG_64BIT_SIZE != 0 {
            self.read_u64().ok()? as usize
        } else {
            self.read_u32().ok()? as usize
        };

        let mut data = vec![0u8; chunk_size];
        if self.reader.read_exact(&mut data).is_err() {
            self.done = true;
            return None;
        }
        self.stream_pos += chunk_size;

        // EndChunk aligns the stream to ChunkAlignment (64) on both read and write paths.
        let padding = ((self.stream_pos + 63) & !63) - self.stream_pos;
        if padding > 0 {
            let _ = self.skip(padding);
        }

        Some((chunk_id, data))
    }
}

/// For a D3D11 `CreateSwapBuffer` chunk, returns `(width, height)` of the
/// back buffer.  Returns `None` for any other chunk type or truncated data.
pub fn d3d11_swapchain_size_from_chunk(chunk_id: u32, data: &[u8]) -> Option<(u32, u32)> {
    if chunk_id != D3D11_CREATE_SWAP_BUFFER {
        return None;
    }
    // Serialised layout: Buffer(u32=4) + SwapbufferID(u64=8) + Width(u32) + Height(u32)
    if data.len() < 20 {
        return None;
    }
    let width = u32::from_le_bytes(data[12..16].try_into().unwrap());
    let height = u32::from_le_bytes(data[16..20].try_into().unwrap());
    Some((width, height))
}

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

    #[test]
    fn test_parse_rdc_driver_vulkan() {
        assert_eq!(
            parse_rdc_driver(Path::new("src/test_data/vkcube.rdc")).unwrap(),
            RdcDriver::Vulkan
        );
    }

    #[test]
    fn test_parse_rdc_driver_d3d11() {
        assert_eq!(
            parse_rdc_driver(Path::new(
                "src/test_data/d3d11-humus-modernlightmapping.rdc"
            ))
            .unwrap(),
            RdcDriver::D3D11
        );
    }

    #[test]
    fn test_parse_rdc_d3d11_swapchain_size() {
        use std::{fs::File, io::BufReader};
        let f =
            BufReader::new(File::open("src/test_data/d3d11-humus-modernlightmapping.rdc").unwrap());
        assert_eq!(
            parse_rdc_frame_capture_chunks(f)
                .unwrap()
                .filter_map(|(id, data)| d3d11_swapchain_size_from_chunk(id, &data))
                .collect::<Vec<_>>(),
            vec![(640, 480)]
        );
    }
}