locum-cli 0.1.3

//! Words for what is in a capture: resources, stacks, constants by name.

use std::collections::HashMap;

use locum_capture::format;
use locum_capture::rdef::{self, ConstantBuffer};

use crate::open::Capture;

/// Prints a line to standard output, and leaves quietly when it is closed (piped into `head`).
#[macro_export]
macro_rules! out {
    ($($arg:tt)*) => {{
        use std::io::Write as _;
        if writeln!(std::io::stdout(), $($arg)*).is_err() {
            std::process::exit(0);
        }
    }};
}

/// The order bindings are listed in: input assembly, the shader stages, stream output, output merger.
pub const STAGE_ORDER: &str = "CASE stage WHEN 'ia' THEN 0 WHEN 'vs' THEN 1 WHEN 'hs' THEN 2 WHEN 'ds' THEN 3 \
     WHEN 'gs' THEN 4 WHEN 'so' THEN 5 WHEN 'ps' THEN 6 WHEN 'cs' THEN 7 ELSE 8 END";

/// A resource's shape as the `resources` table has it.
#[derive(Clone, Debug, Default)]
pub struct Shape {
    pub id: i64,
    pub kind: String,
    pub width: u32,
    pub height: u32,
    pub depth: u32,
    pub mips: u32,
    pub array: u32,
    pub format: u32,
    pub samples: u32,
    pub bytes: u64,
    pub stride: u32,
    pub misc: u32,
}

pub const SHAPE_COLUMNS: &str =
    "id, kind, width, height, depth, mips, array, format, samples, bytes, stride, misc";

impl Shape {
    /// From a row whose first columns are [`SHAPE_COLUMNS`].
    pub fn from_row(r: &rusqlite::Row) -> rusqlite::Result<Shape> {
        let n = |i: usize| -> rusqlite::Result<u32> { Ok(r.get::<_, Option<i64>>(i)?.unwrap_or(0) as u32) };
        Ok(Shape {
            id: r.get(0)?,
            kind: r.get(1)?,
            width: n(2)?,
            height: n(3)?,
            depth: n(4)?,
            mips: n(5)?,
            array: n(6)?,
            format: n(7)?,
            samples: n(8)?,
            bytes: r.get::<_, i64>(9)? as u64,
            stride: n(10)?,
            misc: n(11)?,
        })
    }

    pub fn cube(&self) -> bool {
        self.misc & 4 != 0
    }

    /// The subresources of its contents in the canonical layout (as locum laid them out).
    pub fn subs(&self) -> Vec<locum_capture::layout::Sub> {
        if self.kind == "buffer" {
            return locum_capture::layout::buffer(self.width);
        }
        locum_capture::layout::texture(
            self.format,
            self.width,
            self.height,
            self.depth,
            self.mips,
            self.array,
        )
        .unwrap_or_default()
    }
}

impl std::fmt::Display for Shape {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self.kind.as_str() {
            "buffer" => {
                write!(f, "#{} buffer {} bytes", self.id, self.width)?;
                if self.stride > 0 {
                    write!(f, " stride {}", self.stride)?;
                }
                Ok(())
            }
            kind => {
                write!(f, "#{} {kind} {}", self.id, self.width)?;
                if kind != "texture1d" {
                    write!(f, "x{}", self.height)?;
                }
                if kind == "texture3d" {
                    write!(f, "x{}", self.depth)?;
                }
                write!(f, " {}", format::name(self.format))?;
                if self.mips > 1 {
                    write!(f, " {} mips", self.mips)?;
                }
                if self.cube() {
                    write!(f, " cube")?;
                    if self.array > 6 {
                        write!(f, " x{}", self.array / 6)?;
                    }
                } else if self.array > 1 {
                    write!(f, " [{}]", self.array)?;
                }
                if self.samples > 1 {
                    write!(f, " {}x MSAA", self.samples)?;
                }
                Ok(())
            }
        }
    }
}

pub fn shape(c: &Capture, id: i64) -> Option<Shape> {
    c.row(
        &format!("SELECT {SHAPE_COLUMNS} FROM resources WHERE id = ?1"),
        [id],
        Shape::from_row,
    )
}

/// A stack's frames as `module+0xoffset`, innermost first.
pub fn stack(c: &Capture, id: i64) -> Vec<String> {
    let Some(frames) = c.row("SELECT frames FROM stacks WHERE id = ?1", [id], |r| {
        r.get::<_, Vec<u8>>(0)
    }) else {
        return Vec::new();
    };
    let modules: Vec<(String, u64, u64)> = c.rows("SELECT path, base, size FROM modules", [], |r| {
        Ok((r.get(0)?, r.get::<_, i64>(1)? as u64, r.get::<_, i64>(2)? as u64))
    });
    frames
        .as_chunks::<8>()
        .0
        .iter()
        .map(|&b| {
            let a = u64::from_le_bytes(b);
            match modules
                .iter()
                .find(|(_, base, size)| (*base..base + size).contains(&a))
            {
                Some((path, base, _)) => {
                    let name = path.rsplit(['\\', '/']).next().unwrap_or(path);
                    format!("{name}+0x{:x}", a - base)
                }
                None => format!("0x{a:x}"),
            }
        })
        .collect()
}

/// What a program's reflection says, wherever it is kept.
pub struct Program {
    /// Its name: its bytecode's debug name, or the name the index's record gives it.
    pub name: Option<String>,
    pub constant_buffers: Option<Vec<ConstantBuffer>>,
    pub bindings: Option<Vec<rdef::Binding>>,
    /// Where the names come from: `bytecode`, `index` (a game's own record of a program built without reflection,
    /// `index --programs`) or `registers` (none: the registers its code declares, named by register).
    pub source: &'static str,
}

/// Program `code`'s reflection: its bytecode's own (DXBC `RDEF`, DXIL), else what an index of `names` keeps for its
/// key over the registers its code declares (the record names those it knows; a register it leaves out keeps its
/// register's name), else those registers alone.
pub fn reflected(names: &crate::join::Names, code: &[u8]) -> Program {
    let own = rdef::reflects(code);
    let key = locum_capture::hash::key(code);
    if !own && let Some(r) = names.reflection(&key) {
        let mut bindings = r.bindings;
        for d in rdef::declared(code).unwrap_or_default() {
            let named = bindings
                .iter()
                .any(|b| b.class() == d.class() && b.space == d.space && b.name_at(d.slot).is_some());
            if !named {
                bindings.push(d);
            }
        }
        // Variants compiled to the same code are each named in the index: the first, and how many more.
        let name = match names.program(&key) {
            Some(m) if m.others > 0 => format!("{} (+{} with the same bytes)", m.file.name, m.others),
            Some(m) => m.file.name,
            None => r.name,
        };
        return Program {
            name: Some(name),
            constant_buffers: Some(r.constant_buffers),
            bindings: Some(bindings),
            source: "index",
        };
    }
    Program {
        name: rdef::debug_name(code),
        constant_buffers: rdef::constant_buffers(code),
        bindings: rdef::bindings(code),
        source: if own { "bytecode" } else { "registers" },
    }
}

/// Shader reflection by shader hash, read once.
#[derive(Default)]
pub struct Reflection {
    programs: HashMap<String, Option<Program>>,
    /// The file an index names a program's bytecode by, when its reflection gives no name.
    files: HashMap<String, Option<String>>,
}

/// What a shader's reflection says of a register a draw bound.
pub enum Slot {
    /// The resource the program reads or writes there, by its name in the source.
    Named(String),
    /// The program does not use it (the compiler keeps only what the code reaches).
    Unused,
    /// No reflection to say.
    Unknown,
}

/// The blob key of shader `shader`'s bytecode, when the capture kept it.
fn blob_key(c: &Capture, shader: &str) -> Option<String> {
    c.row(
        "SELECT blob FROM shaders WHERE hash = ?1 AND blob IS NOT NULL",
        [shader],
        |r| r.get(0),
    )
}

impl Reflection {
    /// Shader `shader`'s reflection (`reflected`).
    pub fn program<'a>(&'a mut self, c: &Capture, shader: &str) -> Option<&'a Program> {
        self.programs
            .entry(shader.to_string())
            .or_insert_with(|| Some(reflected(&c.names, &c.blob(&blob_key(c, shader)?).ok()?)))
            .as_ref()
    }

    /// Whether shader `shader` declares a resource or constant buffer called `name` (any case).
    pub fn declares(&mut self, c: &Capture, shader: &str, name: &str) -> bool {
        self.program(c, shader)
            .and_then(|p| p.bindings.as_ref())
            .is_some_and(|bs| bs.iter().any(|b| b.name.eq_ignore_ascii_case(name)))
    }

    /// Shader `shader`'s name: its reflection's (`reflected`: a debug name, or an index's record), or else the file an
    /// index (`--index`) says the same bytecode was read from.
    pub fn name(&mut self, c: &Capture, shader: &str) -> Option<String> {
        if let Some(name) = self.program(c, shader).and_then(|p| p.name.clone()) {
            return Some(name);
        }
        self.files
            .entry(shader.to_string())
            .or_insert_with(|| c.names.program(&blob_key(c, shader)?).map(|m| m.file.name))
            .clone()
    }

    /// What shader `shader` calls register `slot` (space `space`) of class `class` (`srv`, `uav`, `sampler`, `cb`).
    pub fn slot(&mut self, c: &Capture, shader: &str, class: &str, slot: u32, space: u32) -> Slot {
        let Some(bindings) = self.program(c, shader).and_then(|p| p.bindings.as_ref()) else {
            return Slot::Unknown;
        };
        bindings
            .iter()
            .filter(|b| b.class() == class && b.space == space)
            .find_map(|b| b.name_at(slot))
            .map_or(Slot::Unused, Slot::Named)
    }

    /// The constant buffers a shader declares.
    fn cbs<'a>(&'a mut self, c: &Capture, shader: &str) -> Option<&'a Vec<ConstantBuffer>> {
        self.program(c, shader)?.constant_buffers.as_ref()
    }

    /// The constant buffer a shader declares at register `slot` (space `space`).
    pub fn cb<'a>(
        &'a mut self,
        c: &Capture,
        shader: &str,
        slot: u32,
        space: u32,
    ) -> Option<&'a ConstantBuffer> {
        self.cbs(c, shader)?
            .iter()
            .find(|b| b.slot == Some(slot) && b.space == space)
    }

    /// The constant buffer called `name` (any case) as the capture's programs declare it, and the program it is
    /// taken from: of those declaring it, the one declaring the most bytes.
    pub fn declared_as(&mut self, c: &Capture, name: &str) -> Option<(ConstantBuffer, String)> {
        let shaders: Vec<String> = c.rows(
            "SELECT hash FROM shaders WHERE blob IS NOT NULL ORDER BY hash",
            [],
            |r| r.get(0),
        );
        let mut best: Option<(ConstantBuffer, String)> = None;
        for shader in shaders {
            let Some(cbs) = self.cbs(c, &shader) else {
                continue;
            };
            for cb in cbs.iter().filter(|cb| cb.name.eq_ignore_ascii_case(name)) {
                if best.as_ref().is_none_or(|(b, _)| cb.size > b.size) {
                    best = Some((cb.clone(), shader.clone()));
                }
            }
        }
        best
    }
}

/// A 4-byte constant as its type says (float when unknown), at most 12 characters.
pub fn value(bytes: [u8; 4], base: Option<u16>) -> String {
    match base {
        Some(1) => (u32::from_le_bytes(bytes) != 0).to_string(),
        Some(2) => i32::from_le_bytes(bytes).to_string(),
        Some(19 | 20) => u32::from_le_bytes(bytes).to_string(),
        _ => float(f32::from_le_bytes(bytes)),
    }
}

/// A float in at most 12 characters: plainly when it fits, else in exponent form; denormals and NaNs (more likely
/// integers or packed bits) as hex.
pub fn float(f: f32) -> String {
    let bits = f.to_bits();
    if (f != 0.0 && f.abs() < 1e-30) || f.is_nan() {
        return format!("0x{bits:08x}");
    }
    let a = f.abs();
    let plain = format!("{f}");
    if plain.len() <= 12 && a < 1e7 {
        return plain;
    }
    if (1e-4..1e7).contains(&a) {
        // Round to fit: the integer digits, the point and as many decimals as are left.
        let digits = (a.log10().floor() as i32 + 1).max(1) as usize;
        let sign = usize::from(f < 0.0);
        let decimals = 12usize.saturating_sub(digits + 1 + sign).min(9);
        return format!("{f:.decimals$}");
    }
    format!("{f:.4e}")
}

/// The names of a register's four components, shortened to `name` when they are `name.x` to `name.w`.
pub fn register_name(names: &[Option<String>; 4]) -> String {
    if let Some(Some(first)) = names.first()
        && let Some(base) = first.strip_suffix(".x")
        && names
            .iter()
            .zip(["x", "y", "z", "w"])
            .all(|(n, c)| n.as_deref() == Some(&format!("{base}.{c}")))
    {
        return base.to_string();
    }
    if let Some(line) = matrix_line(names) {
        return line;
    }
    let mut distinct: Vec<&str> = Vec::new();
    for n in names.iter().flatten() {
        if !distinct.contains(&n.as_str()) {
            distinct.push(n);
        }
    }
    distinct.join(" ")
}

/// A register holding one row or column of a matrix, named once: `m.VC_MatrixViewProj column 2 ([0..3][2])`. A
/// column-major matrix keeps a column in each register, so the register `mul(v, M)` dots with v for output .z is
/// column 2.
fn matrix_line(names: &[Option<String>; 4]) -> Option<String> {
    // `base[r][c]`: the matrix's name (an array element's included) and the element's row and column.
    let element = |n: &str| -> Option<(String, u32, u32)> {
        let (rest, c) = n.strip_suffix(']')?.rsplit_once('[')?;
        let (base, r) = rest.strip_suffix(']')?.rsplit_once('[')?;
        Some((base.to_string(), r.parse().ok()?, c.parse().ok()?))
    };
    let lanes: Vec<(String, u32, u32)> = names
        .iter()
        .flatten()
        .map(|n| element(n))
        .collect::<Option<_>>()?;
    let (base, r0, c0) = lanes.first()?.clone();
    if lanes.len() < 2 || lanes.iter().any(|(b, _, _)| *b != base) {
        return None;
    }
    let last = lanes.last()?;
    let steps = |f: fn(&(String, u32, u32)) -> u32, start: u32| {
        lanes.iter().enumerate().all(|(i, l)| f(l) == start + i as u32)
    };
    if lanes.iter().all(|l| l.2 == c0) && steps(|l| l.1, r0) {
        return Some(format!("{base} column {c0} ([{r0}..{}][{c0}])", last.1));
    }
    if lanes.iter().all(|l| l.1 == r0) && steps(|l| l.2, c0) {
        return Some(format!("{base} row {r0} ([{r0}][{c0}..{}])", last.2));
    }
    None
}

/// What a Direct3D 12 buffer's contents may lack for event `event`: contents taken from an upload heap as the GPU
/// reached the frame's first captured submission hold what the CPU had written by then, and an event of a later
/// submission can read what it wrote since (a ring of per-frame data). None for other contents and the first
/// submission's events.
pub fn later_upload(c: &Capture, resource: i64, event: i64) -> Option<String> {
    let source: String = c
        .row(
            "SELECT content_source FROM resources WHERE id = ?1",
            [resource],
            |r| r.get::<_, Option<String>>(0),
        )
        .flatten()?;
    if source != "upload-at-first-submission" {
        return None;
    }
    let (submission, first, nth, all): (i64, i64, i64, i64) = c.row(
        "SELECT x.submission, (SELECT min(submission) FROM d12_executions), \
         (SELECT count(DISTINCT submission) FROM d12_executions WHERE submission <= x.submission), \
         (SELECT count(DISTINCT submission) FROM d12_executions) \
         FROM d12_event_locations l JOIN d12_executions x ON x.id = l.execution WHERE l.event = ?1",
        [event],
        |r| Ok((r.get(0)?, r.get(1)?, r.get(2)?, r.get(3)?)),
    )?;
    (submission != first).then(|| {
        format!(
            "taken from the upload heap as the GPU reached the frame's first submission, and event {event} ran in \
             submission {nth} of {all}: what the CPU wrote to the buffer in between is not in them"
        )
    })
}

/// Short key for tables.
pub fn short(key: &str) -> &str {
    key.get(..12).unwrap_or(key)
}

/// A count of bytes in MiB with one decimal.
pub fn mib(bytes: u64) -> String {
    format!("{:.1} MiB", bytes as f64 / (1 << 20) as f64)
}

/// A draw's or dispatch's arguments by name (`count=714 start=885904 base=0`): the API puts the start index in a
/// different place in each call. Indirect calls name their argument buffer. Arguments of another shape are shown as
/// recorded.
pub fn call_args(call: &str, args: &str) -> String {
    let names: &[&str] = match call {
        "Draw" => &["count", "start"],
        "DrawIndexed" => &["count", "start", "base"],
        "DrawInstanced" => &["count", "instances", "start", "first_instance"],
        "DrawIndexedInstanced" => &["count", "instances", "start", "base", "first_instance"],
        "Dispatch" => &["x", "y", "z"],
        "DrawInstancedIndirect" | "DrawIndexedInstancedIndirect" | "DispatchIndirect" => &["args", "offset"],
        _ => &[],
    };
    let values: Vec<i64> = serde_json::from_str(args).unwrap_or_default();
    if names.is_empty() || values.len() != names.len() {
        return args.to_string();
    }
    names
        .iter()
        .zip(&values)
        .map(|(name, v)| match *name {
            "args" => format!("args=#{v}"),
            _ => format!("{name}={v}"),
        })
        .collect::<Vec<_>>()
        .join(" ")
}