iridia-dxbc 0.5.0

//! Lifting a Shader Model 4 or 5 program of any stage (vertex, pixel, hull, domain, geometry or compute) to an
//! [`iridia::Program`].
//!
//! A DXBC register is four untyped 32-bit lanes, so the body keeps every register as a `vec4<u32>` local and each
//! instruction bitcasts its operands to the type the instruction implies; the typed values of the module are the
//! instruction results. Structured flow lifts to structured flow. The conventions:
//!
//! - **Varyings** are unpacked: one entry argument per input signature element, one result member per output element
//!   ([`iridia::entry_io`] says which), the D3D register packing undone by the prologue and epilogue and
//!   kept in the interface's provenance.
//! - **Resources** are module globals: constant buffer `cb<n>` as a uniform `array<vec4<u32>, rows>` (group 0, binding
//!   n), textures and buffer views `t<n>` (group 1), samplers `s<n>` (group 2), typed texture views `u<n>` (group 3, a storage
//!   image of four 32-bit lanes). A name and a layout come from the
//!   container's reflection where it has one; a stripped program's resources are named by register.
//! - **Behaviour** that differs from naga's or other APIs' is written out as operations
//!   ([`iridia::build`]'s D3D functions), and [`iridia::D3dModes`] says so: division by zero,
//!   float to integer conversion, `min`/`max` and `saturate` with NaN, shift counts, and reads out of range of a
//!   dynamically indexed buffer (zero).
//! - **Unknowns** are refused with the instruction word and the reason by name (`store_uav_typed: unordered-access
//!   views are not lifted`); nothing is approximated silently, and what is approximated is a note on the program.

mod body;
mod compute;
mod io;
mod operands;
mod resources;
mod setup;
mod tessellation;

use std::collections::{BTreeMap, BTreeSet};

use iridia::build::{Builder, Width};
use iridia::{Conventions, D3dModes, Program, Provenance, ShaderFormat, Stage};
use naga::{Handle, LocalVariable};

use crate::{Result, error};
use dxbc_reader::decode::{Decoded, decode};

/// The version of the lifter: bumped when what it produces for the same bytecode changes, so a cache keyed by the
/// bytecode's hash knows what is stale; gated by `tests/versions.rs` (the lifted fixtures hashed under it, `iridia::gate`).
pub const LIFTER_VERSION: u32 = 10;

/// Lift a DXBC container's program: Shader Model 4 or 5, any stage. Every instruction or declaration the lifter does
/// not carry (and every container that holds no such program) is an `Err` naming the word and the reason.
pub fn lift(bytes: &[u8]) -> Result<Program> {
    lift_with(bytes, &iridia::BoundSizes::new())
}

/// [`lift`] with the constant buffers `bound` names bound at the size it states: their uniform arrays have the rows of
/// that size, so a read past the declaration and inside the bound is the buffer's, and the interface's layout is the
/// bound size.
pub fn lift_with(bytes: &[u8], bound: &iridia::BoundSizes) -> Result<Program> {
    let mut program = lift_program(bytes, bound)?;
    program.provenance.source_hash = Some(iridia::source_hash(bytes));
    Ok(program)
}

fn lift_program(bytes: &[u8], bound: &iridia::BoundSizes) -> Result<Program> {
    let directory = crate::parts(bytes)?;
    let code = directory
        .iter()
        .find(|c| matches!(&c.kind, b"SHDR" | b"SHEX"))
        .ok_or_else(|| error(0, "missing shader code chunk"))?;
    if code.data.len() % 4 != 0 {
        return Err(error(0, "unaligned code length"));
    }
    let words: Vec<u32> = code
        .data
        .as_chunks::<4>()
        .0
        .iter()
        .map(|b| u32::from_le_bytes(*b))
        .collect();
    if words.len() > 1_048_576 {
        return Err(error(1, "program exceeds word budget"));
    }
    let decoded = decode(&words).map_err(crate::read)?;
    let stage = crate::reflection::stage_of(decoded.kind);
    let mut setup = setup::Setup::new(bytes, &decoded, stage)?;
    setup.bind_sizes(bound)?;
    let mut lift = Lift::new(&decoded, stage, setup)?;
    if let Some(sigs) = lift.setup.sigs.clone() {
        for sig in &sigs {
            lift.run_function(sig)?;
        }
        return lift.finish_functions(decoded.target());
    }
    if stage == Stage::Compute {
        lift.b.workgroup_size = lift
            .setup
            .workgroup
            .ok_or_else(|| error(0, "a compute program that declares no thread group size"))?;
    }
    lift.prologue()?;
    lift.body()?;
    lift.finish(decoded.target())
}

/// A conditional's `if` or a `loop`, `switch` being lifted.
struct Frame {
    kind: FrameKind,
    /// The word of the instruction that opened it, for the refusal of a frame never closed.
    word: usize,
    condition: Option<Handle<naga::Expression>>,
    accept: Option<naga::Block>,
    selector: Option<Handle<naga::Expression>>,
    cases: Vec<naga::SwitchCase>,
    current: Option<naga::SwitchValue>,
}
#[derive(PartialEq, Eq)]
enum FrameKind {
    If,
    Loop,
    Switch,
}

/// The state of one lift.
struct Lift<'a> {
    b: Builder,
    decoded: &'a Decoded,
    stage: Stage,
    setup: setup::Setup,
    /// `r<n>`, `v<n>`, `o<n>`: four raw lanes each.
    temps: Vec<Handle<LocalVariable>>,
    inputs: BTreeMap<u32, Handle<LocalVariable>>,
    outputs: BTreeMap<u32, Handle<LocalVariable>>,
    /// `x<n>[i]`: the local and its length.
    indexable: BTreeMap<u32, (Handle<LocalVariable>, u32)>,
    /// `vThreadID`, `vThreadGroupID`, `vThreadIDInGroup` and `vThreadIDInGroupFlattened`, by the system value each is.
    ids: BTreeMap<iridia::SystemValue, Handle<LocalVariable>>,
    /// A hull, domain or position-less vertex program's function being lifted (see [`iridia::stage_functions`]).
    sig: Option<iridia::FunctionSig>,
    /// The control points of a hull or domain program: `v[point][register]` and a hull program's `vocp`, as arrays of
    /// registers, `registers` of them to a point.
    control_points: Option<(Handle<LocalVariable>, u32)>,
    output_points: Option<(Handle<LocalVariable>, u32)>,
    /// A domain program's patch constants `vpc[register]`, `vDomain`, and the phase numbers of a hull program.
    patch_in: Option<Handle<LocalVariable>>,
    domain_point: Option<Handle<LocalVariable>>,
    point_id: Option<Handle<LocalVariable>>,
    fork_id: Option<Handle<LocalVariable>>,
    join_id: Option<Handle<LocalVariable>>,
    /// A geometry program's emitted vertices so far (see [`tessellation::GsState`]) and its instance number.
    gs: Option<tessellation::GsState>,
    gs_instance: Option<Handle<LocalVariable>>,
    /// Lifting a phase of a hull program: a `ret` ends the phase, not the function.
    in_phase: bool,
    /// The declared arguments and result of the function being lifted, until it is closed.
    pending_decl: Option<iridia::build::FunctionDecl>,
    /// `vPrim`, `vCoverage`, `oDepth` and `oMask`, lane 0.
    primitive_id: Option<Handle<LocalVariable>>,
    input_coverage: Option<Handle<LocalVariable>>,
    depth: Option<Handle<LocalVariable>>,
    coverage: Option<Handle<LocalVariable>>,
    resources: resources::Resources,
    frames: Vec<Frame>,
    notes: BTreeSet<String>,
    /// The word of the instruction being lifted, for refusals.
    word: usize,
    /// A `ret` ended the program at the top level: what follows is unreachable.
    returned: bool,
    arguments: Vec<naga::FunctionArgument>,
    result_type: Option<Handle<naga::Type>>,
}

impl<'a> Lift<'a> {
    fn refuse<T>(&self, detail: impl Into<String>) -> Result<T> {
        Err(error(self.word, detail))
    }

    fn u4(&mut self) -> Handle<naga::Type> {
        self.b.ty_of(iridia::build::Kind::U, Width::Four)
    }

    /// Refuse a program that ends with an `if`, `loop` or `switch` still open.
    fn check_closed(&self) -> Result<()> {
        match self.frames.last() {
            None => Ok(()),
            Some(frame) => {
                let (opener, closer) = match frame.kind {
                    FrameKind::If => ("if", "endif"),
                    FrameKind::Loop => ("loop", "endloop"),
                    FrameKind::Switch => ("switch", "endswitch"),
                };
                Err(error(
                    frame.word,
                    format!("{opener} is not closed: the program ends without its {closer}"),
                ))
            }
        }
    }

    fn finish(mut self, target: String) -> Result<Program> {
        self.check_closed()?;
        if !self.returned {
            self.epilogue()?;
        }
        self.b.check().map_err(|detail| error(self.word, detail))?;
        let stage = match self.stage {
            Stage::Vertex => naga::ShaderStage::Vertex,
            Stage::Compute => naga::ShaderStage::Compute,
            _ => naga::ShaderStage::Fragment,
        };
        let word = self.word;
        let module = self
            .b
            .finish(
                "main",
                stage,
                self.arguments,
                self.result_type,
                self.setup.early_depth_test,
            )
            .map_err(|detail| error(word, detail))?;
        let mut limit_notes = iridia::compute_limit_notes(&module);
        let interface = self.setup.interface.clone();
        limit_notes.extend(iridia::interface_notes(&interface));
        Ok(Program {
            module,
            stage: self.stage,
            entry: "main".into(),
            interface: self.setup.interface,
            bindings: self.resources.bound,
            conventions: Conventions::D3D,
            modes: D3dModes::ALL,
            provenance: Provenance {
                format: ShaderFormat::Dxbc,
                target,
                lifter: "iridia-dxbc".into(),
                lifter_version: LIFTER_VERSION,
                // the hash of the bytes, which `lift` sets
                source_hash: None,
                instruction_count: Some(iridia::InstructionCount {
                    count: self.decoded.instructions.len() as u64,
                    unit: iridia::InstructionUnit::DxbcInstructions,
                }),
            },
            notes: self.notes.into_iter().chain(limit_notes).collect(),
        })
    }
}