iridia-dxbc 0.3.0

//! The program's edges: locals for the registers, the prologue that unpacks the entry arguments into input registers,
//! and the epilogue that packs output registers into the result.

use iridia::build::{Builder, Kind, Val, Width};
use iridia::{IoType, ScalarType, Stage, SystemValue};
use naga::{Expression as E, Handle, Statement, TypeInner};

use super::resources::{Resources, lanes_of};
use super::setup::Setup;
use super::{Frame, Lift};
use crate::Result;
use dxbc_reader::decode::{Declared, Decoded, Register};

fn kind_of(s: ScalarType) -> Kind {
    Kind::of(s)
}

impl<'a> Lift<'a> {
    pub(super) fn new(decoded: &'a Decoded, stage: Stage, setup: Setup) -> Result<Self> {
        let mut lift = Lift {
            b: Builder::new(),
            decoded,
            stage,
            temps: Vec::new(),
            inputs: Default::default(),
            outputs: Default::default(),
            indexable: Default::default(),
            ids: Default::default(),
            sig: None,
            control_points: None,
            output_points: None,
            patch_in: None,
            domain_point: None,
            point_id: None,
            fork_id: None,
            join_id: None,
            gs: None,
            gs_instance: None,
            in_phase: false,
            pending_decl: None,
            primitive_id: None,
            input_coverage: None,
            depth: None,
            coverage: None,
            resources: Resources::default(),
            frames: Vec::<Frame>::new(),
            notes: setup.notes.clone(),
            word: 0,
            returned: false,
            arguments: Vec::new(),
            result_type: None,
            setup,
        };
        lift.declare_resources()?;
        if lift.setup.sigs.is_none() {
            lift.declare_locals();
        }
        Ok(lift)
    }

    /// The function's locals: the registers it uses, four raw lanes each. A program of several functions makes them
    /// again for each.
    pub(super) fn declare_locals(&mut self) {
        let decoded = self.decoded;
        let lift = self;
        lift.temps.clear();
        lift.inputs.clear();
        lift.outputs.clear();
        lift.indexable.clear();
        lift.ids.clear();
        lift.primitive_id = None;
        lift.control_points = None;
        lift.output_points = None;
        lift.patch_in = None;
        lift.domain_point = None;
        lift.point_id = None;
        lift.fork_id = None;
        lift.join_id = None;
        lift.gs = None;
        lift.gs_instance = None;
        lift.input_coverage = None;
        lift.depth = None;
        lift.coverage = None;
        let u4 = lift.u4();
        for n in 0..lift.setup.temps {
            let local = lift.b.local(&format!("r{n}"), u4);
            lift.temps.push(local);
        }
        for (register, size) in lift.setup.indexable.clone() {
            let Some(len) = std::num::NonZeroU32::new(size) else {
                continue;
            };
            let ty = lift.b.ty(TypeInner::Array {
                base: u4,
                size: naga::ArraySize::Constant(len),
                stride: 16,
            });
            let local = lift.b.local(&format!("x{register}"), ty);
            lift.indexable.insert(register, (local, size));
        }
        for register in lift.setup.declared_inputs.keys().copied().collect::<Vec<_>>() {
            let local = lift.b.local(&format!("v{register}"), u4);
            lift.inputs.insert(register, local);
        }
        // declared special inputs
        for d in &decoded.declarations {
            if let Declared::Input { operand, .. } = &d.kind {
                match operand.register {
                    Register::PrimitiveId if lift.primitive_id.is_none() => {
                        lift.primitive_id = Some(lift.b.local("vPrim", u4));
                    }
                    Register::InputCoverage if lift.input_coverage.is_none() => {
                        lift.input_coverage = Some(lift.b.local("vCoverage", u4));
                    }
                    Register::ThreadId
                    | Register::ThreadGroupId
                    | Register::ThreadIdInGroup
                    | Register::ThreadIdInGroupFlattened => {
                        let (system, name) = match operand.register {
                            Register::ThreadId => (SystemValue::DispatchThreadId, "vThreadID"),
                            Register::ThreadGroupId => (SystemValue::GroupId, "vThreadGroupID"),
                            Register::ThreadIdInGroup => (SystemValue::GroupThreadId, "vThreadIDInGroup"),
                            _ => (SystemValue::GroupIndex, "vThreadIDInGroupFlattened"),
                        };
                        lift.ids.entry(system).or_insert_with(|| lift.b.local(name, u4));
                    }
                    _ => {}
                }
            }
            if let Declared::Output { operand, .. } = &d.kind {
                match operand.register {
                    Register::Output => {
                        let register = operand.indices.first().map_or(0, |i| i.base);
                        if !lift.outputs.contains_key(&register) {
                            let local = lift.b.local(&format!("o{register}"), u4);
                            lift.outputs.insert(register, local);
                        }
                    }
                    Register::OutputDepth
                    | Register::OutputDepthGreaterEqual
                    | Register::OutputDepthLessEqual
                        if lift.depth.is_none() =>
                    {
                        lift.depth = Some(lift.b.local("oDepth", u4));
                    }
                    Register::OutputCoverage if lift.coverage.is_none() => {
                        lift.coverage = Some(lift.b.local("oMask", u4));
                    }
                    _ => {}
                }
            }
        }
        // the registers of a range addressed by a computed index
        for (output, first, count) in lift.setup.index_ranges.clone() {
            for r in first..first.saturating_add(count) {
                if output && !lift.outputs.contains_key(&r) {
                    let local = lift.b.local(&format!("o{r}"), u4);
                    lift.outputs.insert(r, local);
                }
            }
        }
        // registers the output signature (or a function's results) places elements in
        for v in &lift.result_varyings() {
            if let Some(d3d) = v.provenance.d3d
                && d3d.register != u32::MAX
                && !lift.outputs.contains_key(&d3d.register)
            {
                let local = lift.b.local(&format!("o{}", d3d.register), u4);
                lift.outputs.insert(d3d.register, local);
            }
        }
    }

    /// The entry arguments and result type, and the stores of each argument into its input register lanes.
    pub(super) fn prologue(&mut self) -> Result<()> {
        let io = self.setup.io.clone();
        let declared = match self.b.declare_entry(&io) {
            Ok(d) => d,
            Err(detail) => return self.refuse(detail),
        };
        self.arguments = declared.arguments;
        self.result_type = declared.result;
        for (i, item) in io.arguments.iter().enumerate() {
            let IoType::Value(v) = item.ty else {
                return self.refuse("an array-typed input");
            };
            let Ok(argument) = u32::try_from(i) else {
                return self.refuse("more than 4294967295 inputs");
            };
            let expr = self.b.app(E::FunctionArgument(argument));
            let mut value = Val {
                expr,
                kind: kind_of(v.scalar),
                width: match Width::try_from(v.width) {
                    Ok(width) => width,
                    Err(e) => return self.refuse(format!("an input of {e}")),
                },
            };
            let Some(varying) = item
                .varyings
                .first()
                .and_then(|&v| self.setup.interface.inputs.get(v))
                .cloned()
            else {
                return self.refuse(format!("the input {} names no varying", item.name));
            };
            // raw bits: a bool is D3D's all-ones or zero
            if value.kind == Kind::B {
                let ones = self.b.lit_u(u32::MAX);
                let zero = self.b.lit_u(0);
                value = self.b.select(value, ones, zero);
            } else {
                value = self.b.bitcast(value, Kind::U);
            }
            if let Some(d3d) = varying.provenance.d3d {
                let declared = self
                    .setup
                    .declared_inputs
                    .get(&d3d.register)
                    .copied()
                    .unwrap_or(0);
                if declared & d3d.mask != 0 {
                    let Some(&local) = self.inputs.get(&d3d.register) else {
                        return self.refuse(format!("input register v{} has no local", d3d.register));
                    };
                    let pointer = self.b.local_ptr(local);
                    for (j, lane) in lanes_of(d3d.mask).into_iter().enumerate() {
                        let component = self.b.lane(value, super::operands::lane_u8(j));
                        let target = self.b.app(E::AccessIndex {
                            base: pointer,
                            index: u32::from(lane),
                        });
                        self.b.store(target, component);
                    }
                }
            }
            let special = match varying.system {
                Some(SystemValue::PrimitiveIndex) => self.primitive_id,
                Some(SystemValue::Coverage) => self.input_coverage,
                Some(
                    SystemValue::DispatchThreadId
                    | SystemValue::GroupId
                    | SystemValue::GroupThreadId
                    | SystemValue::GroupIndex,
                ) => {
                    // an id's lanes, then zero in the lanes it does not have
                    let system = varying.system.expect("matched");
                    let Some(&local) = self.ids.get(&system) else {
                        continue;
                    };
                    let pointer = self.b.local_ptr(local);
                    for lane in value.width.lanes() {
                        let component = self.b.lane(value, lane);
                        let target = self.b.app(E::AccessIndex {
                            base: pointer,
                            index: u32::from(lane),
                        });
                        self.b.store(target, component);
                    }
                    None
                }
                _ => None,
            };
            if let Some(local) = special {
                let pointer = self.b.local_ptr(local);
                let target = self.b.app(E::AccessIndex {
                    base: pointer,
                    index: 0,
                });
                self.b.store(target, value);
            }
        }
        Ok(())
    }

    /// Pack the output registers into the result and return it.
    pub(super) fn epilogue(&mut self) -> Result<()> {
        if self.sig.is_some() {
            return self.function_epilogue();
        }
        let Some(result) = self.result_type else {
            self.b.push(Statement::Return { value: None });
            return Ok(());
        };
        let io = self.setup.io.clone();
        let mut parts = Vec::new();
        for item in &io.results {
            let part = match item.ty {
                IoType::Array(..) | IoType::Memory => return self.refuse("an array-typed output"),
                IoType::ClipDistances(_) => {
                    let mut components = Vec::new();
                    for &i in &item.varyings {
                        let Some(v) = self.setup.interface.outputs.get(i).cloned() else {
                            return self.refuse("a clip distance past the interface's outputs");
                        };
                        let Some(d3d) = v.provenance.d3d else {
                            return self.refuse(format!("clip distance {} has no register", item.name));
                        };
                        let Some(&local) = self.outputs.get(&d3d.register) else {
                            return self.refuse(format!(
                                "clip distance {} has no register o{}",
                                item.name, d3d.register
                            ));
                        };
                        let p = self.b.local_ptr(local);
                        let row = self.b.load(p, Kind::U, Width::Four);
                        for lane in lanes_of(d3d.mask) {
                            let c = self.b.lane(row, lane);
                            components.push(self.b.bitcast(c, Kind::F).expr);
                        }
                    }
                    let ty = self.array_type(&item.ty);
                    self.b.app(E::Compose { ty, components })
                }
                IoType::Value(carried) => {
                    let Some(v) = item
                        .varyings
                        .first()
                        .and_then(|&v| self.setup.interface.outputs.get(v))
                        .cloned()
                    else {
                        return self.refuse(format!("the output {} names no varying", item.name));
                    };
                    let (local, lanes): (Handle<naga::LocalVariable>, Vec<u8>) = match v.system {
                        Some(
                            SystemValue::Depth | SystemValue::DepthGreaterEqual | SystemValue::DepthLessEqual,
                        ) => (self.depth.expect("a depth output has its local"), vec![0]),
                        Some(SystemValue::Coverage) => {
                            (self.coverage.expect("a coverage output has its local"), vec![0])
                        }
                        Some(SystemValue::ClipPosition | SystemValue::Target(_)) => {
                            let d3d = v.provenance.d3d.expect("a signature element has a register");
                            match self.outputs.get(&d3d.register) {
                                Some(&l) => (l, vec![0, 1, 2, 3]),
                                None => {
                                    return self.refuse(format!(
                                        "output register o{} is not declared",
                                        d3d.register
                                    ));
                                }
                            }
                        }
                        _ => {
                            let d3d = v.provenance.d3d.expect("a signature element has a register");
                            match self.outputs.get(&d3d.register) {
                                Some(&l) => (l, lanes_of(d3d.mask)),
                                None => {
                                    return self.refuse(format!(
                                        "output register o{} is not declared",
                                        d3d.register
                                    ));
                                }
                            }
                        }
                    };
                    let p = self.b.local_ptr(local);
                    let row = self.b.load(p, Kind::U, Width::Four);
                    let picked = self.b.swizzle(row, &lanes);
                    let value = self.b.bitcast(picked, kind_of(carried.scalar));
                    debug_assert_eq!(value.width.get(), carried.width);
                    value.expr
                }
            };
            parts.push(part);
        }
        let composed = self.b.app(E::Compose {
            ty: result,
            components: parts,
        });
        self.b.push(Statement::Return {
            value: Some(composed),
        });
        Ok(())
    }

    fn array_type(&mut self, ty: &IoType) -> Handle<naga::Type> {
        let IoType::ClipDistances(n) = ty else {
            unreachable!("an array type")
        };
        let f = self.b.ty_of(Kind::F, Width::One);
        self.b.ty(TypeInner::Array {
            base: f,
            size: naga::ArraySize::Constant(std::num::NonZeroU32::new(*n).expect("clip distances")),
            stride: 4,
        })
    }
}