dxcontainer 0.5.2

//! Reader for the DXBC container, the file format Direct3D shader programs are stored in: Shader Model 4 and 5 programs
//! (`SHDR`, `SHEX`, signatures, `RDEF`) and Shader Model 6 programs (`DXIL`, `PSV0`, `RDAT`, ...) alike. [`chunks`] reads
//! the container's part directory and [`checksum`] computes the digest the compiler stores in it. [`feature_flags`]
//! reads the `SFI0` part, the shader's required features, with names for the bits ([`feature_name`]). [`serialize`]
//! writes parts as a container with a correct digest, [`repair_digest`] fixes the digest of one edited in place, and
//! [`replace_part`] swaps one part of a container for new bytes.
//!
//! The layout: a header (`DXBC`, a 16-byte digest, version 1, the container's size, the number of parts), the directory of
//! the parts' offsets right after it, then the parts one after another in the directory's order, each its four-character
//! kind, its length and its bytes, the last ending where the container does. Shader Model 6 parts are not padded to 4
//! bytes, so an offset need not be a multiple of 4. This holds for every fixture committed in `dxbc-reader` and `dxil-reader`
//! (the tests read each to its exact end and compare [`checksum`] with the stored digest).
//!
//! The reader reads to the exact end: a part anywhere else, a gap, a byte after the last part, a kind twice, a size other
//! than the container's or a version other than 1 is an [`Error`] naming the byte. The digest is not verified (it is not an
//! authenticity check); [`checksum`] is for a caller that wants to compare it.
//!
//! ```
//! // A container with one empty `ISGN` part: header, one directory entry, the part.
//! let mut bytes = b"DXBC".to_vec();
//! bytes.extend_from_slice(&[0; 16]); // digest
//! for word in [1u32, 44, 1, 36] {
//!     bytes.extend_from_slice(&word.to_le_bytes()); // version, size, part count, the part's offset
//! }
//! bytes.extend_from_slice(b"ISGN");
//! bytes.extend_from_slice(&0u32.to_le_bytes()); // the part's length
//!
//! let parts = dxcontainer::chunks(&bytes).expect("a well-formed container");
//! assert_eq!(&parts[0].kind, b"ISGN");
//! assert!(parts[0].data.is_empty());
//! assert!(dxcontainer::checksum(&bytes).is_some());
//! ```
#![deny(missing_docs)]
#![deny(clippy::disallowed_methods, clippy::disallowed_types)]
#![deny(
    clippy::arithmetic_side_effects,
    clippy::indexing_slicing,
    clippy::cast_possible_truncation
)]

use primase_foundations::bytes::View;

/// What is wrong with a container, and where (a byte offset in it).
#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)]
#[error("DXBC container at byte {offset}: {detail}")]
pub struct Error {
    /// The byte offset in the container where the problem is.
    pub offset: usize,
    /// What was expected there, or what was found instead.
    pub detail: String,
}

fn error(offset: usize, detail: impl Into<String>) -> Error {
    Error {
        offset,
        detail: detail.into(),
    }
}

/// A part of a container: its four-character kind and its bytes.
#[derive(Clone, Copy, Debug)]
pub struct Chunk<'a> {
    /// The part's four-character kind (`DXIL`, `RDEF`, ...).
    pub kind: [u8; 4],
    /// The part's bytes, without its kind and length.
    pub data: &'a [u8],
}

/// The part directory of the container `bytes`, in the container's order. See the module documentation for what is refused.
pub fn chunks(bytes: &[u8]) -> Result<Vec<Chunk<'_>>, Error> {
    let view = View::new(bytes);
    let word = |at: usize| -> Result<usize, Error> {
        let word = view.array::<4>(at).map_err(|e| error(at, e.to_string()))?;
        usize::try_from(u32::from_le_bytes(word)).map_err(|_| error(at, "a word past the address space"))
    };
    if bytes.get(..4) != Some(b"DXBC") {
        return Err(error(0, "the magic is not DXBC"));
    }
    if word(20)? != 1 {
        return Err(error(20, "the container's version is not 1"));
    }
    if word(24)? != bytes.len() {
        return Err(error(
            24,
            format!(
                "the container states a size of {}, it has {} bytes",
                word(24)?,
                bytes.len()
            ),
        ));
    }
    let count = word(28)?;
    let directory_end = count
        .checked_mul(4)
        .and_then(|n| n.checked_add(32))
        .filter(|end| *end <= bytes.len())
        .ok_or_else(|| {
            error(
                28,
                format!("a directory of {count} parts does not fit the container"),
            )
        })?;
    let mut parts: Vec<Chunk<'_>> = Vec::new();
    let mut next = directory_end;
    for i in 0..count {
        let entry = i
            .checked_mul(4)
            .and_then(|n| n.checked_add(32))
            .ok_or_else(|| error(28, "a directory past the address space"))?;
        let start = word(entry)?;
        if start != next {
            let shown = view.array::<4>(start).map_or_else(
                |_| "past the end".to_owned(),
                |kind| String::from_utf8_lossy(&kind).into_owned(),
            );
            return Err(error(
                entry,
                if start < next {
                    format!(
                        "overlapping parts: part {i} ({shown}) starts at {start}, before the previous part ends at {next}"
                    )
                } else {
                    format!(
                        "{} bytes between parts: part {i} ({shown}) starts at {start}, where the previous part ends at {next}",
                        start.saturating_sub(next)
                    )
                },
            ));
        }
        let kind = view.array::<4>(start).map_err(|e| error(start, e.to_string()))?;
        let shown = String::from_utf8_lossy(&kind).into_owned();
        let length = word(
            start
                .checked_add(4)
                .ok_or_else(|| error(start, "a part past the address space"))?,
        )?;
        let data_at = start
            .checked_add(8)
            .ok_or_else(|| error(start, "a part past the address space"))?;
        let data = view
            .slice(data_at, length)
            .map_err(|e| error(start, format!("part {i} ({shown}): {e}")))?;
        if parts.iter().any(|p| p.kind == kind) {
            return Err(error(start, format!("a part kind twice: {shown}")));
        }
        next = data_at
            .checked_add(length)
            .ok_or_else(|| error(start, "a part past the address space"))?;
        parts.push(Chunk { kind, data });
    }
    if next != bytes.len() {
        return Err(error(
            next,
            format!("{} bytes after the last part", bytes.len().saturating_sub(next)),
        ));
    }
    Ok(parts)
}

/// The digest a container stores at bytes 4 to 20: MD5's block function over its bytes from 20 on, the last block padded as the
/// compiler does (the bit count first, then the rest, `0x80`, zeros, and at 60 the bit count over 4, or 1; a rest of 56 bytes
/// or more padded in a block of its own first). `None` for fewer than 20 bytes.
pub fn checksum(container: &[u8]) -> Option<[u8; 16]> {
    let data = container.get(20..)?;
    let bits = u32::try_from(data.len()).ok()?.wrapping_mul(8);
    let (blocks, rest) = data.as_chunks::<64>();
    let mut state = primase_foundations::hash::MD5_INIT;
    for block in blocks {
        primase_foundations::hash::md5_block(&mut state, block);
    }
    let mut last = [0u8; 64];
    if rest.len() < 56 {
        last.get_mut(4..)?.get_mut(..rest.len())?.copy_from_slice(rest);
        *last.get_mut(rest.len().checked_add(4)?)? = 0x80;
    } else {
        let mut first = [0u8; 64];
        first.get_mut(..rest.len())?.copy_from_slice(rest);
        *first.get_mut(rest.len())? = 0x80;
        primase_foundations::hash::md5_block(&mut state, &first);
    }
    last.get_mut(..4)?.copy_from_slice(&bits.to_le_bytes());
    last.get_mut(60..)?
        .copy_from_slice(&((bits >> 2) | 1).to_le_bytes());
    primase_foundations::hash::md5_block(&mut state, &last);
    Some(primase_foundations::hash::md5_bytes(state))
}

/// A container that cannot be written: what is wrong with the parts it was asked to hold.
#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)]
#[error("cannot write a DXBC container: {detail}")]
pub struct WriteError {
    /// What cannot be written.
    pub detail: String,
}

fn write_error(detail: impl Into<String>) -> WriteError {
    WriteError {
        detail: detail.into(),
    }
}

/// Writes `parts` as a container: the header, the directory of the parts' offsets, then the parts in the order given,
/// each its kind, its length and its bytes with nothing between them, and the digest [`checksum`] computes over the
/// result. A part's bytes are written exactly as given, so a part whose producer padded it (or did not) stays so, and
/// `serialize(&chunks(bytes)?)` is `bytes` again whenever `bytes` stores the digest [`checksum`] computes, as the
/// games' shipped containers do. Refuses a kind twice (the container [`chunks`] reads has each kind once) and a size or
/// count that does not fit the header's 32 bits.
///
/// ```
/// let part = dxcontainer::Chunk { kind: *b"ISGN", data: &[] };
/// let bytes = dxcontainer::serialize(&[part])?;
/// assert_eq!(bytes.len(), 44);
/// assert_eq!(dxcontainer::checksum(&bytes).unwrap()[..], bytes[4..20]);
/// assert_eq!(&dxcontainer::chunks(&bytes).expect("a container reads back")[0].kind, b"ISGN");
/// # Ok::<(), dxcontainer::WriteError>(())
/// ```
pub fn serialize(parts: &[Chunk<'_>]) -> Result<Vec<u8>, WriteError> {
    let too_big = || write_error("the container is larger than the header's 32-bit size can state");
    let mut offset = parts
        .len()
        .checked_mul(4)
        .and_then(|n| n.checked_add(32))
        .ok_or_else(too_big)?;
    let mut directory = Vec::with_capacity(parts.len());
    for (i, part) in parts.iter().enumerate() {
        if let Some(first) = parts.iter().take(i).position(|p| p.kind == part.kind) {
            return Err(write_error(format!(
                "parts {first} and {i} are both {}",
                String::from_utf8_lossy(&part.kind)
            )));
        }
        directory.push(u32::try_from(offset).map_err(|_| too_big())?);
        offset = offset
            .checked_add(8)
            .and_then(|n| n.checked_add(part.data.len()))
            .ok_or_else(too_big)?;
    }
    let size = u32::try_from(offset).map_err(|_| too_big())?;
    let count = u32::try_from(parts.len()).map_err(|_| too_big())?;
    let mut out = Vec::with_capacity(offset);
    out.extend_from_slice(b"DXBC");
    out.extend_from_slice(&[0; 16]);
    for word in [1, size, count] {
        out.extend_from_slice(&u32::to_le_bytes(word));
    }
    for at in directory {
        out.extend_from_slice(&at.to_le_bytes());
    }
    for part in parts {
        out.extend_from_slice(&part.kind);
        out.extend_from_slice(
            &u32::try_from(part.data.len())
                .map_err(|_| too_big())?
                .to_le_bytes(),
        );
        out.extend_from_slice(part.data);
    }
    repair_digest(&mut out)?;
    Ok(out)
}

/// Rewrites the digest at bytes 4 to 20 of `container` to the one [`checksum`] computes over the rest of it: for a
/// container whose bytes were edited in place. Refuses bytes that are not a container's (fewer than 20, or not starting
/// `DXBC`), or too long for the digest's 32-bit bit count.
///
/// ```
/// let mut bytes = dxcontainer::serialize(&[dxcontainer::Chunk { kind: *b"ISGN", data: &[0; 8] }])?;
/// bytes[36 + 8] ^= 1; // edit a byte of the part
/// assert_ne!(dxcontainer::checksum(&bytes).unwrap()[..], bytes[4..20]);
/// dxcontainer::repair_digest(&mut bytes)?;
/// assert_eq!(dxcontainer::checksum(&bytes).unwrap()[..], bytes[4..20]);
/// # Ok::<(), dxcontainer::WriteError>(())
/// ```
pub fn repair_digest(container: &mut [u8]) -> Result<(), WriteError> {
    if container.get(..4) != Some(b"DXBC") || container.len() < 20 {
        return Err(write_error(
            "the bytes are not a DXBC container (no DXBC magic and 16-byte digest)",
        ));
    }
    let digest = checksum(container)
        .ok_or_else(|| write_error("the container is too long for the digest's bit count"))?;
    container
        .get_mut(4..20)
        .ok_or_else(|| write_error("the digest field is out of range"))?
        .copy_from_slice(&digest);
    Ok(())
}

/// `container` with the part of kind `kind` replaced by `data` (a part not in it is appended after the others), every
/// other part kept in its place and as it was, and the digest repaired: the edit a shader patch makes. Refuses a
/// container [`chunks`] refuses.
pub fn replace_part(container: &[u8], kind: [u8; 4], data: &[u8]) -> Result<Vec<u8>, ReplaceError> {
    let mut parts = chunks(container).map_err(ReplaceError::Read)?;
    match parts.iter_mut().find(|p| p.kind == kind) {
        Some(part) => part.data = data,
        None => parts.push(Chunk { kind, data }),
    }
    serialize(&parts).map_err(ReplaceError::Write)
}

/// Why [`replace_part`] failed: the container it was given, or the container it would write.
#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)]
pub enum ReplaceError {
    /// The container is not one [`chunks`] reads.
    #[error(transparent)]
    Read(Error),
    /// The edited container cannot be written.
    #[error(transparent)]
    Write(WriteError),
}

/// Which bytecode a container's program part holds. It decides what the bits of the container's `SFI0` part name: the
/// two compilers share the first bits and then diverge.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Bytecode {
    /// Shader Model 4 or 5 (`fxc`): an `SHDR` or `SHEX` part.
    Dxbc,
    /// Shader Model 6 (`dxc`): a `DXIL` part.
    Dxil,
}

impl Bytecode {
    /// The bytecode `parts` hold: `None` for a container with no program part or with both kinds.
    pub fn of(parts: &[Chunk<'_>]) -> Option<Self> {
        let has = |kinds: &[&[u8; 4]]| parts.iter().any(|p| kinds.contains(&&p.kind));
        match (has(&[b"SHDR", b"SHEX"]), has(&[b"DXIL"])) {
            (true, false) => Some(Self::Dxbc),
            (false, true) => Some(Self::Dxil),
            _ => None,
        }
    }
}

/// The name of bit `bit` of the `SFI0` shader feature flags of `bytecode`, `None` for a bit with no name here.
///
/// A name is the description the producing compiler prints for the flag. For `fxc` every named bit (0 to 13) was read
/// off the compiler's own output: programs that need one feature each were compiled with `fxc` (Shader Model 4.0 to
/// 5.0), and each container's flags were compared with the compiler's listing of what the program requires (bit 1 only
/// with a Shader Model 4 pixel program that uses a raw buffer; there it also sets the program's
/// `enableRawAndStructuredBuffers` global flag). For `dxc` the names are those its listing (`dxc -dumpbin`, 1.9.2609)
/// gives a container with that bit alone set: bits 0 to 33, and 40 and 41, the two flags it lists as optional
/// ("Function uses derivatives", "Function requires visible group"). Programs compiled to need one feature each (Shader
/// Model 6.0 to 6.7) set bits 0, 2 to 6, 8 to 20, 22 to 26 and 28 to 30, with those names.
///
/// In the games' stock containers (the 69,728 `DXIL` containers of skate. and Elden Ring, and 21 of Unity's 4,854
/// `DXBC` ones; Sekiro's and Riders Republic's have no `SFI0` part) the bits that occur are, for `DXIL`, 2, 3, 4, 11,
/// 14, 15, 20 and 25, and for `DXBC` 1 and 4. Against the program itself: bit 14 is set in every program that uses
/// wave or quad operations and in no other, bit 3 in every program that declares more than 8 UAV slots and no other,
/// bit 25 in every program that creates a handle from a descriptor heap and no other, bit 20 in every program that
/// uses a ray query (and in 16 that do not).
pub fn feature_name(bytecode: Bytecode, bit: u32) -> Option<&'static str> {
    // the bits both compilers name, alike
    let shared = match bit {
        0 => Some("Double-precision floating point"),
        1 => Some("Raw and Structured buffers"),
        2 => Some("UAVs at every shader stage"),
        3 => Some("64 UAV slots"),
        4 => Some("Minimum-precision data types"),
        5 => Some("Double-precision extensions for 11.1"),
        6 => Some("Shader extensions for 11.1"),
        7 => Some("Comparison filtering for feature level 9"),
        8 => Some("Tiled resources"),
        9 => Some("PS Output Stencil Ref"),
        10 => Some("PS Inner Coverage"),
        11 => Some("Typed UAV Load Additional Formats"),
        12 => Some("Raster Ordered UAVs"),
        13 => Some("SV_RenderTargetArrayIndex or SV_ViewportArrayIndex from any shader feeding rasterizer"),
        _ => None,
    };
    if shared.is_some() || bytecode == Bytecode::Dxbc {
        return shared;
    }
    Some(match bit {
        14 => "Wave level operations",
        15 => "64-Bit integer",
        16 => "View Instancing",
        17 => "Barycentrics",
        18 => "Use native low precision",
        19 => "Shading Rate",
        20 => "Raytracing tier 1.1 features",
        21 => "Sampler feedback",
        22 => "64-bit Atomics on Typed Resources",
        23 => "64-bit Atomics on Group Shared",
        24 => "Derivatives in mesh and amplification shaders",
        25 => "Resource descriptor heap indexing",
        26 => "Sampler descriptor heap indexing",
        27 => "Wave Matrix",
        28 => "64-bit Atomics on Heap Resources",
        29 => "Advanced Texture Ops",
        30 => "Writeable MSAA Textures",
        31 => "SampleCmp with gradient or bias",
        32 => "Extended command info",
        33 => "Linear Algebra",
        40 => "Function uses derivatives",
        41 => "Function requires visible group",
        _ => return None,
    })
}

/// A container's shader feature flags: its `SFI0` part (eight bytes, a little-endian `u64`) and the bytecode that says
/// what the bits mean.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct FeatureFlags {
    /// The bytecode of the container's program.
    pub bytecode: Bytecode,
    /// The flags as stored.
    pub bits: u64,
}

impl FeatureFlags {
    /// Whether bit `bit` is set.
    pub fn has(self, bit: u32) -> bool {
        self.bits.checked_shr(bit).is_some_and(|rest| rest & 1 == 1)
    }

    /// The names of the set bits that have one ([`feature_name`]), lowest bit first.
    pub fn names(self) -> Vec<&'static str> {
        (0..u64::BITS)
            .filter(|&b| self.has(b))
            .filter_map(|b| feature_name(self.bytecode, b))
            .collect()
    }

    /// The set bits that have no name, as a mask.
    pub fn unnamed(self) -> u64 {
        (0..u64::BITS)
            .filter(|&b| self.has(b) && feature_name(self.bytecode, b).is_none())
            .fold(0, |mask, b| mask | 1u64.checked_shl(b).unwrap_or(0))
    }
}

/// The shader feature flags of `container`: `None` for a container with no `SFI0` part (Sekiro's and Riders Republic's
/// `DXBC` containers have none; the part is there in all 69,728 `DXIL` containers of skate. and Elden Ring). Refuses a
/// container [`chunks`] refuses, an `SFI0` part of other than eight bytes, and an `SFI0` part in a container whose
/// bytecode is unknown (no program part, or both kinds): the bits mean nothing without it.
pub fn feature_flags(container: &[u8]) -> Result<Option<FeatureFlags>, Error> {
    let parts = chunks(container)?;
    let Some(part) = parts.iter().find(|p| &p.kind == b"SFI0") else {
        return Ok(None);
    };
    let at = part
        .data
        .as_ptr()
        .addr()
        .saturating_sub(container.as_ptr().addr());
    let bits = <[u8; 8]>::try_from(part.data)
        .map_err(|_| error(at, format!("the SFI0 part has {} bytes, not 8", part.data.len())))?;
    let bytecode = Bytecode::of(&parts).ok_or_else(|| {
        error(
            at,
            "an SFI0 part in a container with no single kind of program part",
        )
    })?;
    Ok(Some(FeatureFlags {
        bytecode,
        bits: u64::from_le_bytes(bits),
    }))
}