diff options
Diffstat (limited to 'third_party/rust/naga/src/back/hlsl/writer.rs')
| -rw-r--r-- | third_party/rust/naga/src/back/hlsl/writer.rs | 3366 |
1 files changed, 3366 insertions, 0 deletions
diff --git a/third_party/rust/naga/src/back/hlsl/writer.rs b/third_party/rust/naga/src/back/hlsl/writer.rs new file mode 100644 index 0000000000..43f7212837 --- /dev/null +++ b/third_party/rust/naga/src/back/hlsl/writer.rs @@ -0,0 +1,3366 @@ +use super::{ + help::{WrappedArrayLength, WrappedConstructor, WrappedImageQuery, WrappedStructMatrixAccess}, + storage::StoreValue, + BackendResult, Error, Options, +}; +use crate::{ + back, + proc::{self, NameKey}, + valid, Handle, Module, ScalarKind, ShaderStage, TypeInner, +}; +use std::{fmt, mem}; + +const LOCATION_SEMANTIC: &str = "LOC"; +const SPECIAL_CBUF_TYPE: &str = "NagaConstants"; +const SPECIAL_CBUF_VAR: &str = "_NagaConstants"; +const SPECIAL_FIRST_VERTEX: &str = "first_vertex"; +const SPECIAL_FIRST_INSTANCE: &str = "first_instance"; +const SPECIAL_OTHER: &str = "other"; + +pub(crate) const MODF_FUNCTION: &str = "naga_modf"; +pub(crate) const FREXP_FUNCTION: &str = "naga_frexp"; + +struct EpStructMember { + name: String, + ty: Handle<crate::Type>, + // technically, this should always be `Some` + binding: Option<crate::Binding>, + index: u32, +} + +/// Structure contains information required for generating +/// wrapped structure of all entry points arguments +struct EntryPointBinding { + /// Name of the fake EP argument that contains the struct + /// with all the flattened input data. + arg_name: String, + /// Generated structure name + ty_name: String, + /// Members of generated structure + members: Vec<EpStructMember>, +} + +pub(super) struct EntryPointInterface { + /// If `Some`, the input of an entry point is gathered in a special + /// struct with members sorted by binding. + /// The `EntryPointBinding::members` array is sorted by index, + /// so that we can walk it in `write_ep_arguments_initialization`. + input: Option<EntryPointBinding>, + /// If `Some`, the output of an entry point is flattened. + /// The `EntryPointBinding::members` array is sorted by binding, + /// So that we can walk it in `Statement::Return` handler. + output: Option<EntryPointBinding>, +} + +#[derive(Clone, Eq, PartialEq, PartialOrd, Ord)] +enum InterfaceKey { + Location(u32), + BuiltIn(crate::BuiltIn), + Other, +} + +impl InterfaceKey { + const fn new(binding: Option<&crate::Binding>) -> Self { + match binding { + Some(&crate::Binding::Location { location, .. }) => Self::Location(location), + Some(&crate::Binding::BuiltIn(built_in)) => Self::BuiltIn(built_in), + None => Self::Other, + } + } +} + +#[derive(Copy, Clone, PartialEq)] +enum Io { + Input, + Output, +} + +impl<'a, W: fmt::Write> super::Writer<'a, W> { + pub fn new(out: W, options: &'a Options) -> Self { + Self { + out, + names: crate::FastHashMap::default(), + namer: proc::Namer::default(), + options, + entry_point_io: Vec::new(), + named_expressions: crate::NamedExpressions::default(), + wrapped: super::Wrapped::default(), + temp_access_chain: Vec::new(), + need_bake_expressions: Default::default(), + } + } + + fn reset(&mut self, module: &Module) { + self.names.clear(); + self.namer.reset( + module, + super::keywords::RESERVED, + super::keywords::TYPES, + super::keywords::RESERVED_CASE_INSENSITIVE, + &[], + &mut self.names, + ); + self.entry_point_io.clear(); + self.named_expressions.clear(); + self.wrapped.clear(); + self.need_bake_expressions.clear(); + } + + /// Helper method used to find which expressions of a given function require baking + /// + /// # Notes + /// Clears `need_bake_expressions` set before adding to it + fn update_expressions_to_bake( + &mut self, + module: &Module, + func: &crate::Function, + info: &valid::FunctionInfo, + ) { + use crate::Expression; + self.need_bake_expressions.clear(); + for (fun_handle, expr) in func.expressions.iter() { + let expr_info = &info[fun_handle]; + let min_ref_count = func.expressions[fun_handle].bake_ref_count(); + if min_ref_count <= expr_info.ref_count { + self.need_bake_expressions.insert(fun_handle); + } + + if let Expression::Math { + fun, + arg, + arg1, + arg2, + arg3, + } = *expr + { + match fun { + crate::MathFunction::Asinh + | crate::MathFunction::Acosh + | crate::MathFunction::Atanh + | crate::MathFunction::Unpack2x16float + | crate::MathFunction::Unpack2x16snorm + | crate::MathFunction::Unpack2x16unorm + | crate::MathFunction::Unpack4x8snorm + | crate::MathFunction::Unpack4x8unorm + | crate::MathFunction::Pack2x16float + | crate::MathFunction::Pack2x16snorm + | crate::MathFunction::Pack2x16unorm + | crate::MathFunction::Pack4x8snorm + | crate::MathFunction::Pack4x8unorm => { + self.need_bake_expressions.insert(arg); + } + crate::MathFunction::ExtractBits => { + self.need_bake_expressions.insert(arg); + self.need_bake_expressions.insert(arg1.unwrap()); + self.need_bake_expressions.insert(arg2.unwrap()); + } + crate::MathFunction::InsertBits => { + self.need_bake_expressions.insert(arg); + self.need_bake_expressions.insert(arg1.unwrap()); + self.need_bake_expressions.insert(arg2.unwrap()); + self.need_bake_expressions.insert(arg3.unwrap()); + } + crate::MathFunction::CountLeadingZeros => { + let inner = info[fun_handle].ty.inner_with(&module.types); + if let Some(crate::ScalarKind::Sint) = inner.scalar_kind() { + self.need_bake_expressions.insert(arg); + } + } + _ => {} + } + } + + if let Expression::Derivative { axis, ctrl, expr } = *expr { + use crate::{DerivativeAxis as Axis, DerivativeControl as Ctrl}; + if axis == Axis::Width && (ctrl == Ctrl::Coarse || ctrl == Ctrl::Fine) { + self.need_bake_expressions.insert(expr); + } + } + } + } + + pub fn write( + &mut self, + module: &Module, + module_info: &valid::ModuleInfo, + ) -> Result<super::ReflectionInfo, Error> { + self.reset(module); + + // Write special constants, if needed + if let Some(ref bt) = self.options.special_constants_binding { + writeln!(self.out, "struct {SPECIAL_CBUF_TYPE} {{")?; + writeln!(self.out, "{}int {};", back::INDENT, SPECIAL_FIRST_VERTEX)?; + writeln!(self.out, "{}int {};", back::INDENT, SPECIAL_FIRST_INSTANCE)?; + writeln!(self.out, "{}uint {};", back::INDENT, SPECIAL_OTHER)?; + writeln!(self.out, "}};")?; + write!( + self.out, + "ConstantBuffer<{}> {}: register(b{}", + SPECIAL_CBUF_TYPE, SPECIAL_CBUF_VAR, bt.register + )?; + if bt.space != 0 { + write!(self.out, ", space{}", bt.space)?; + } + writeln!(self.out, ");")?; + + // Extra newline for readability + writeln!(self.out)?; + } + + // Save all entry point output types + let ep_results = module + .entry_points + .iter() + .map(|ep| (ep.stage, ep.function.result.clone())) + .collect::<Vec<(ShaderStage, Option<crate::FunctionResult>)>>(); + + self.write_all_mat_cx2_typedefs_and_functions(module)?; + + // Write all structs + for (handle, ty) in module.types.iter() { + if let TypeInner::Struct { ref members, span } = ty.inner { + if module.types[members.last().unwrap().ty] + .inner + .is_dynamically_sized(&module.types) + { + // unsized arrays can only be in storage buffers, + // for which we use `ByteAddressBuffer` anyway. + continue; + } + + let ep_result = ep_results.iter().find(|e| { + if let Some(ref result) = e.1 { + result.ty == handle + } else { + false + } + }); + + self.write_struct( + module, + handle, + members, + span, + ep_result.map(|r| (r.0, Io::Output)), + )?; + writeln!(self.out)?; + } + } + + self.write_special_functions(module)?; + + self.write_wrapped_compose_functions(module, &module.const_expressions)?; + + // Write all named constants + let mut constants = module + .constants + .iter() + .filter(|&(_, c)| c.name.is_some()) + .peekable(); + while let Some((handle, _)) = constants.next() { + self.write_global_constant(module, handle)?; + // Add extra newline for readability on last iteration + if constants.peek().is_none() { + writeln!(self.out)?; + } + } + + // Write all globals + for (ty, _) in module.global_variables.iter() { + self.write_global(module, ty)?; + } + + if !module.global_variables.is_empty() { + // Add extra newline for readability + writeln!(self.out)?; + } + + // Write all entry points wrapped structs + for (index, ep) in module.entry_points.iter().enumerate() { + let ep_name = self.names[&NameKey::EntryPoint(index as u16)].clone(); + let ep_io = self.write_ep_interface(module, &ep.function, ep.stage, &ep_name)?; + self.entry_point_io.push(ep_io); + } + + // Write all regular functions + for (handle, function) in module.functions.iter() { + let info = &module_info[handle]; + + // Check if all of the globals are accessible + if !self.options.fake_missing_bindings { + if let Some((var_handle, _)) = + module + .global_variables + .iter() + .find(|&(var_handle, var)| match var.binding { + Some(ref binding) if !info[var_handle].is_empty() => { + self.options.resolve_resource_binding(binding).is_err() + } + _ => false, + }) + { + log::info!( + "Skipping function {:?} (name {:?}) because global {:?} is inaccessible", + handle, + function.name, + var_handle + ); + continue; + } + } + + let ctx = back::FunctionCtx { + ty: back::FunctionType::Function(handle), + info, + expressions: &function.expressions, + named_expressions: &function.named_expressions, + }; + let name = self.names[&NameKey::Function(handle)].clone(); + + self.write_wrapped_functions(module, &ctx)?; + + self.write_function(module, name.as_str(), function, &ctx, info)?; + + writeln!(self.out)?; + } + + let mut entry_point_names = Vec::with_capacity(module.entry_points.len()); + + // Write all entry points + for (index, ep) in module.entry_points.iter().enumerate() { + let info = module_info.get_entry_point(index); + + if !self.options.fake_missing_bindings { + let mut ep_error = None; + for (var_handle, var) in module.global_variables.iter() { + match var.binding { + Some(ref binding) if !info[var_handle].is_empty() => { + if let Err(err) = self.options.resolve_resource_binding(binding) { + ep_error = Some(err); + break; + } + } + _ => {} + } + } + if let Some(err) = ep_error { + entry_point_names.push(Err(err)); + continue; + } + } + + let ctx = back::FunctionCtx { + ty: back::FunctionType::EntryPoint(index as u16), + info, + expressions: &ep.function.expressions, + named_expressions: &ep.function.named_expressions, + }; + + self.write_wrapped_functions(module, &ctx)?; + + if ep.stage == ShaderStage::Compute { + // HLSL is calling workgroup size "num threads" + let num_threads = ep.workgroup_size; + writeln!( + self.out, + "[numthreads({}, {}, {})]", + num_threads[0], num_threads[1], num_threads[2] + )?; + } + + let name = self.names[&NameKey::EntryPoint(index as u16)].clone(); + self.write_function(module, &name, &ep.function, &ctx, info)?; + + if index < module.entry_points.len() - 1 { + writeln!(self.out)?; + } + + entry_point_names.push(Ok(name)); + } + + Ok(super::ReflectionInfo { entry_point_names }) + } + + fn write_modifier(&mut self, binding: &crate::Binding) -> BackendResult { + match *binding { + crate::Binding::BuiltIn(crate::BuiltIn::Position { invariant: true }) => { + write!(self.out, "precise ")?; + } + crate::Binding::Location { + interpolation, + sampling, + .. + } => { + if let Some(interpolation) = interpolation { + if let Some(string) = interpolation.to_hlsl_str() { + write!(self.out, "{string} ")? + } + } + + if let Some(sampling) = sampling { + if let Some(string) = sampling.to_hlsl_str() { + write!(self.out, "{string} ")? + } + } + } + crate::Binding::BuiltIn(_) => {} + } + + Ok(()) + } + + //TODO: we could force fragment outputs to always go through `entry_point_io.output` path + // if they are struct, so that the `stage` argument here could be omitted. + fn write_semantic( + &mut self, + binding: &crate::Binding, + stage: Option<(ShaderStage, Io)>, + ) -> BackendResult { + match *binding { + crate::Binding::BuiltIn(builtin) => { + let builtin_str = builtin.to_hlsl_str()?; + write!(self.out, " : {builtin_str}")?; + } + crate::Binding::Location { + second_blend_source: true, + .. + } => { + write!(self.out, " : SV_Target1")?; + } + crate::Binding::Location { + location, + second_blend_source: false, + .. + } => { + if stage == Some((crate::ShaderStage::Fragment, Io::Output)) { + write!(self.out, " : SV_Target{location}")?; + } else { + write!(self.out, " : {LOCATION_SEMANTIC}{location}")?; + } + } + } + + Ok(()) + } + + fn write_interface_struct( + &mut self, + module: &Module, + shader_stage: (ShaderStage, Io), + struct_name: String, + mut members: Vec<EpStructMember>, + ) -> Result<EntryPointBinding, Error> { + // Sort the members so that first come the user-defined varyings + // in ascending locations, and then built-ins. This allows VS and FS + // interfaces to match with regards to order. + members.sort_by_key(|m| InterfaceKey::new(m.binding.as_ref())); + + write!(self.out, "struct {struct_name}")?; + writeln!(self.out, " {{")?; + for m in members.iter() { + write!(self.out, "{}", back::INDENT)?; + if let Some(ref binding) = m.binding { + self.write_modifier(binding)?; + } + self.write_type(module, m.ty)?; + write!(self.out, " {}", &m.name)?; + if let Some(ref binding) = m.binding { + self.write_semantic(binding, Some(shader_stage))?; + } + writeln!(self.out, ";")?; + } + writeln!(self.out, "}};")?; + writeln!(self.out)?; + + match shader_stage.1 { + Io::Input => { + // bring back the original order + members.sort_by_key(|m| m.index); + } + Io::Output => { + // keep it sorted by binding + } + } + + Ok(EntryPointBinding { + arg_name: self.namer.call(struct_name.to_lowercase().as_str()), + ty_name: struct_name, + members, + }) + } + + /// Flatten all entry point arguments into a single struct. + /// This is needed since we need to re-order them: first placing user locations, + /// then built-ins. + fn write_ep_input_struct( + &mut self, + module: &Module, + func: &crate::Function, + stage: ShaderStage, + entry_point_name: &str, + ) -> Result<EntryPointBinding, Error> { + let struct_name = format!("{stage:?}Input_{entry_point_name}"); + + let mut fake_members = Vec::new(); + for arg in func.arguments.iter() { + match module.types[arg.ty].inner { + TypeInner::Struct { ref members, .. } => { + for member in members.iter() { + let name = self.namer.call_or(&member.name, "member"); + let index = fake_members.len() as u32; + fake_members.push(EpStructMember { + name, + ty: member.ty, + binding: member.binding.clone(), + index, + }); + } + } + _ => { + let member_name = self.namer.call_or(&arg.name, "member"); + let index = fake_members.len() as u32; + fake_members.push(EpStructMember { + name: member_name, + ty: arg.ty, + binding: arg.binding.clone(), + index, + }); + } + } + } + + self.write_interface_struct(module, (stage, Io::Input), struct_name, fake_members) + } + + /// Flatten all entry point results into a single struct. + /// This is needed since we need to re-order them: first placing user locations, + /// then built-ins. + fn write_ep_output_struct( + &mut self, + module: &Module, + result: &crate::FunctionResult, + stage: ShaderStage, + entry_point_name: &str, + ) -> Result<EntryPointBinding, Error> { + let struct_name = format!("{stage:?}Output_{entry_point_name}"); + + let mut fake_members = Vec::new(); + let empty = []; + let members = match module.types[result.ty].inner { + TypeInner::Struct { ref members, .. } => members, + ref other => { + log::error!("Unexpected {:?} output type without a binding", other); + &empty[..] + } + }; + + for member in members.iter() { + let member_name = self.namer.call_or(&member.name, "member"); + let index = fake_members.len() as u32; + fake_members.push(EpStructMember { + name: member_name, + ty: member.ty, + binding: member.binding.clone(), + index, + }); + } + + self.write_interface_struct(module, (stage, Io::Output), struct_name, fake_members) + } + + /// Writes special interface structures for an entry point. The special structures have + /// all the fields flattened into them and sorted by binding. They are only needed for + /// VS outputs and FS inputs, so that these interfaces match. + fn write_ep_interface( + &mut self, + module: &Module, + func: &crate::Function, + stage: ShaderStage, + ep_name: &str, + ) -> Result<EntryPointInterface, Error> { + Ok(EntryPointInterface { + input: if !func.arguments.is_empty() && stage == ShaderStage::Fragment { + Some(self.write_ep_input_struct(module, func, stage, ep_name)?) + } else { + None + }, + output: match func.result { + Some(ref fr) if fr.binding.is_none() && stage == ShaderStage::Vertex => { + Some(self.write_ep_output_struct(module, fr, stage, ep_name)?) + } + _ => None, + }, + }) + } + + /// Write an entry point preface that initializes the arguments as specified in IR. + fn write_ep_arguments_initialization( + &mut self, + module: &Module, + func: &crate::Function, + ep_index: u16, + ) -> BackendResult { + let ep_input = match self.entry_point_io[ep_index as usize].input.take() { + Some(ep_input) => ep_input, + None => return Ok(()), + }; + let mut fake_iter = ep_input.members.iter(); + for (arg_index, arg) in func.arguments.iter().enumerate() { + write!(self.out, "{}", back::INDENT)?; + self.write_type(module, arg.ty)?; + let arg_name = &self.names[&NameKey::EntryPointArgument(ep_index, arg_index as u32)]; + write!(self.out, " {arg_name}")?; + match module.types[arg.ty].inner { + TypeInner::Array { base, size, .. } => { + self.write_array_size(module, base, size)?; + let fake_member = fake_iter.next().unwrap(); + writeln!(self.out, " = {}.{};", ep_input.arg_name, fake_member.name)?; + } + TypeInner::Struct { ref members, .. } => { + write!(self.out, " = {{ ")?; + for index in 0..members.len() { + if index != 0 { + write!(self.out, ", ")?; + } + let fake_member = fake_iter.next().unwrap(); + write!(self.out, "{}.{}", ep_input.arg_name, fake_member.name)?; + } + writeln!(self.out, " }};")?; + } + _ => { + let fake_member = fake_iter.next().unwrap(); + writeln!(self.out, " = {}.{};", ep_input.arg_name, fake_member.name)?; + } + } + } + assert!(fake_iter.next().is_none()); + Ok(()) + } + + /// Helper method used to write global variables + /// # Notes + /// Always adds a newline + fn write_global( + &mut self, + module: &Module, + handle: Handle<crate::GlobalVariable>, + ) -> BackendResult { + let global = &module.global_variables[handle]; + let inner = &module.types[global.ty].inner; + + if let Some(ref binding) = global.binding { + if let Err(err) = self.options.resolve_resource_binding(binding) { + log::info!( + "Skipping global {:?} (name {:?}) for being inaccessible: {}", + handle, + global.name, + err, + ); + return Ok(()); + } + } + + // https://docs.microsoft.com/en-us/windows/win32/direct3dhlsl/dx-graphics-hlsl-variable-register + let register_ty = match global.space { + crate::AddressSpace::Function => unreachable!("Function address space"), + crate::AddressSpace::Private => { + write!(self.out, "static ")?; + self.write_type(module, global.ty)?; + "" + } + crate::AddressSpace::WorkGroup => { + write!(self.out, "groupshared ")?; + self.write_type(module, global.ty)?; + "" + } + crate::AddressSpace::Uniform => { + // constant buffer declarations are expected to be inlined, e.g. + // `cbuffer foo: register(b0) { field1: type1; }` + write!(self.out, "cbuffer")?; + "b" + } + crate::AddressSpace::Storage { access } => { + let (prefix, register) = if access.contains(crate::StorageAccess::STORE) { + ("RW", "u") + } else { + ("", "t") + }; + write!(self.out, "{prefix}ByteAddressBuffer")?; + register + } + crate::AddressSpace::Handle => { + let handle_ty = match *inner { + TypeInner::BindingArray { ref base, .. } => &module.types[*base].inner, + _ => inner, + }; + + let register = match *handle_ty { + TypeInner::Sampler { .. } => "s", + // all storage textures are UAV, unconditionally + TypeInner::Image { + class: crate::ImageClass::Storage { .. }, + .. + } => "u", + _ => "t", + }; + self.write_type(module, global.ty)?; + register + } + crate::AddressSpace::PushConstant => { + // The type of the push constants will be wrapped in `ConstantBuffer` + write!(self.out, "ConstantBuffer<")?; + "b" + } + }; + + // If the global is a push constant write the type now because it will be a + // generic argument to `ConstantBuffer` + if global.space == crate::AddressSpace::PushConstant { + self.write_global_type(module, global.ty)?; + + // need to write the array size if the type was emitted with `write_type` + if let TypeInner::Array { base, size, .. } = module.types[global.ty].inner { + self.write_array_size(module, base, size)?; + } + + // Close the angled brackets for the generic argument + write!(self.out, ">")?; + } + + let name = &self.names[&NameKey::GlobalVariable(handle)]; + write!(self.out, " {name}")?; + + // Push constants need to be assigned a binding explicitly by the consumer + // since naga has no way to know the binding from the shader alone + if global.space == crate::AddressSpace::PushConstant { + let target = self + .options + .push_constants_target + .as_ref() + .expect("No bind target was defined for the push constants block"); + write!(self.out, ": register(b{}", target.register)?; + if target.space != 0 { + write!(self.out, ", space{}", target.space)?; + } + write!(self.out, ")")?; + } + + if let Some(ref binding) = global.binding { + // this was already resolved earlier when we started evaluating an entry point. + let bt = self.options.resolve_resource_binding(binding).unwrap(); + + // need to write the binding array size if the type was emitted with `write_type` + if let TypeInner::BindingArray { base, size, .. } = module.types[global.ty].inner { + if let Some(overridden_size) = bt.binding_array_size { + write!(self.out, "[{overridden_size}]")?; + } else { + self.write_array_size(module, base, size)?; + } + } + + write!(self.out, " : register({}{}", register_ty, bt.register)?; + if bt.space != 0 { + write!(self.out, ", space{}", bt.space)?; + } + write!(self.out, ")")?; + } else { + // need to write the array size if the type was emitted with `write_type` + if let TypeInner::Array { base, size, .. } = module.types[global.ty].inner { + self.write_array_size(module, base, size)?; + } + if global.space == crate::AddressSpace::Private { + write!(self.out, " = ")?; + if let Some(init) = global.init { + self.write_const_expression(module, init)?; + } else { + self.write_default_init(module, global.ty)?; + } + } + } + + if global.space == crate::AddressSpace::Uniform { + write!(self.out, " {{ ")?; + + self.write_global_type(module, global.ty)?; + + write!( + self.out, + " {}", + &self.names[&NameKey::GlobalVariable(handle)] + )?; + + // need to write the array size if the type was emitted with `write_type` + if let TypeInner::Array { base, size, .. } = module.types[global.ty].inner { + self.write_array_size(module, base, size)?; + } + + writeln!(self.out, "; }}")?; + } else { + writeln!(self.out, ";")?; + } + + Ok(()) + } + + /// Helper method used to write global constants + /// + /// # Notes + /// Ends in a newline + fn write_global_constant( + &mut self, + module: &Module, + handle: Handle<crate::Constant>, + ) -> BackendResult { + write!(self.out, "static const ")?; + let constant = &module.constants[handle]; + self.write_type(module, constant.ty)?; + let name = &self.names[&NameKey::Constant(handle)]; + write!(self.out, " {}", name)?; + // Write size for array type + if let TypeInner::Array { base, size, .. } = module.types[constant.ty].inner { + self.write_array_size(module, base, size)?; + } + write!(self.out, " = ")?; + self.write_const_expression(module, constant.init)?; + writeln!(self.out, ";")?; + Ok(()) + } + + pub(super) fn write_array_size( + &mut self, + module: &Module, + base: Handle<crate::Type>, + size: crate::ArraySize, + ) -> BackendResult { + write!(self.out, "[")?; + + match size { + crate::ArraySize::Constant(size) => { + write!(self.out, "{size}")?; + } + crate::ArraySize::Dynamic => unreachable!(), + } + + write!(self.out, "]")?; + + if let TypeInner::Array { + base: next_base, + size: next_size, + .. + } = module.types[base].inner + { + self.write_array_size(module, next_base, next_size)?; + } + + Ok(()) + } + + /// Helper method used to write structs + /// + /// # Notes + /// Ends in a newline + fn write_struct( + &mut self, + module: &Module, + handle: Handle<crate::Type>, + members: &[crate::StructMember], + span: u32, + shader_stage: Option<(ShaderStage, Io)>, + ) -> BackendResult { + // Write struct name + let struct_name = &self.names[&NameKey::Type(handle)]; + writeln!(self.out, "struct {struct_name} {{")?; + + let mut last_offset = 0; + for (index, member) in members.iter().enumerate() { + if member.binding.is_none() && member.offset > last_offset { + // using int as padding should work as long as the backend + // doesn't support a type that's less than 4 bytes in size + // (Error::UnsupportedScalar catches this) + let padding = (member.offset - last_offset) / 4; + for i in 0..padding { + writeln!(self.out, "{}int _pad{}_{};", back::INDENT, index, i)?; + } + } + let ty_inner = &module.types[member.ty].inner; + last_offset = member.offset + ty_inner.size_hlsl(module.to_ctx()); + + // The indentation is only for readability + write!(self.out, "{}", back::INDENT)?; + + match module.types[member.ty].inner { + TypeInner::Array { base, size, .. } => { + // HLSL arrays are written as `type name[size]` + + self.write_global_type(module, member.ty)?; + + // Write `name` + write!( + self.out, + " {}", + &self.names[&NameKey::StructMember(handle, index as u32)] + )?; + // Write [size] + self.write_array_size(module, base, size)?; + } + // We treat matrices of the form `matCx2` as a sequence of C `vec2`s. + // See the module-level block comment in mod.rs for details. + TypeInner::Matrix { + rows, + columns, + scalar, + } if member.binding.is_none() && rows == crate::VectorSize::Bi => { + let vec_ty = crate::TypeInner::Vector { size: rows, scalar }; + let field_name_key = NameKey::StructMember(handle, index as u32); + + for i in 0..columns as u8 { + if i != 0 { + write!(self.out, "; ")?; + } + self.write_value_type(module, &vec_ty)?; + write!(self.out, " {}_{}", &self.names[&field_name_key], i)?; + } + } + _ => { + // Write modifier before type + if let Some(ref binding) = member.binding { + self.write_modifier(binding)?; + } + + // Even though Naga IR matrices are column-major, we must describe + // matrices passed from the CPU as being in row-major order. + // See the module-level block comment in mod.rs for details. + if let TypeInner::Matrix { .. } = module.types[member.ty].inner { + write!(self.out, "row_major ")?; + } + + // Write the member type and name + self.write_type(module, member.ty)?; + write!( + self.out, + " {}", + &self.names[&NameKey::StructMember(handle, index as u32)] + )?; + } + } + + if let Some(ref binding) = member.binding { + self.write_semantic(binding, shader_stage)?; + }; + writeln!(self.out, ";")?; + } + + // add padding at the end since sizes of types don't get rounded up to their alignment in HLSL + if members.last().unwrap().binding.is_none() && span > last_offset { + let padding = (span - last_offset) / 4; + for i in 0..padding { + writeln!(self.out, "{}int _end_pad_{};", back::INDENT, i)?; + } + } + + writeln!(self.out, "}};")?; + Ok(()) + } + + /// Helper method used to write global/structs non image/sampler types + /// + /// # Notes + /// Adds no trailing or leading whitespace + pub(super) fn write_global_type( + &mut self, + module: &Module, + ty: Handle<crate::Type>, + ) -> BackendResult { + let matrix_data = get_inner_matrix_data(module, ty); + + // We treat matrices of the form `matCx2` as a sequence of C `vec2`s. + // See the module-level block comment in mod.rs for details. + if let Some(MatrixType { + columns, + rows: crate::VectorSize::Bi, + width: 4, + }) = matrix_data + { + write!(self.out, "__mat{}x2", columns as u8)?; + } else { + // Even though Naga IR matrices are column-major, we must describe + // matrices passed from the CPU as being in row-major order. + // See the module-level block comment in mod.rs for details. + if matrix_data.is_some() { + write!(self.out, "row_major ")?; + } + + self.write_type(module, ty)?; + } + + Ok(()) + } + + /// Helper method used to write non image/sampler types + /// + /// # Notes + /// Adds no trailing or leading whitespace + pub(super) fn write_type(&mut self, module: &Module, ty: Handle<crate::Type>) -> BackendResult { + let inner = &module.types[ty].inner; + match *inner { + TypeInner::Struct { .. } => write!(self.out, "{}", self.names[&NameKey::Type(ty)])?, + // hlsl array has the size separated from the base type + TypeInner::Array { base, .. } | TypeInner::BindingArray { base, .. } => { + self.write_type(module, base)? + } + ref other => self.write_value_type(module, other)?, + } + + Ok(()) + } + + /// Helper method used to write value types + /// + /// # Notes + /// Adds no trailing or leading whitespace + pub(super) fn write_value_type(&mut self, module: &Module, inner: &TypeInner) -> BackendResult { + match *inner { + TypeInner::Scalar(scalar) | TypeInner::Atomic(scalar) => { + write!(self.out, "{}", scalar.to_hlsl_str()?)?; + } + TypeInner::Vector { size, scalar } => { + write!( + self.out, + "{}{}", + scalar.to_hlsl_str()?, + back::vector_size_str(size) + )?; + } + TypeInner::Matrix { + columns, + rows, + scalar, + } => { + // The IR supports only float matrix + // https://docs.microsoft.com/en-us/windows/win32/direct3dhlsl/dx-graphics-hlsl-matrix + + // Because of the implicit transpose all matrices have in HLSL, we need to transpose the size as well. + write!( + self.out, + "{}{}x{}", + scalar.to_hlsl_str()?, + back::vector_size_str(columns), + back::vector_size_str(rows), + )?; + } + TypeInner::Image { + dim, + arrayed, + class, + } => { + self.write_image_type(dim, arrayed, class)?; + } + TypeInner::Sampler { comparison } => { + let sampler = if comparison { + "SamplerComparisonState" + } else { + "SamplerState" + }; + write!(self.out, "{sampler}")?; + } + // HLSL arrays are written as `type name[size]` + // Current code is written arrays only as `[size]` + // Base `type` and `name` should be written outside + TypeInner::Array { base, size, .. } | TypeInner::BindingArray { base, size } => { + self.write_array_size(module, base, size)?; + } + _ => return Err(Error::Unimplemented(format!("write_value_type {inner:?}"))), + } + + Ok(()) + } + + /// Helper method used to write functions + /// # Notes + /// Ends in a newline + fn write_function( + &mut self, + module: &Module, + name: &str, + func: &crate::Function, + func_ctx: &back::FunctionCtx<'_>, + info: &valid::FunctionInfo, + ) -> BackendResult { + // Function Declaration Syntax - https://docs.microsoft.com/en-us/windows/win32/direct3dhlsl/dx-graphics-hlsl-function-syntax + + self.update_expressions_to_bake(module, func, info); + + // Write modifier + if let Some(crate::FunctionResult { + binding: + Some( + ref binding @ crate::Binding::BuiltIn(crate::BuiltIn::Position { + invariant: true, + }), + ), + .. + }) = func.result + { + self.write_modifier(binding)?; + } + + // Write return type + if let Some(ref result) = func.result { + match func_ctx.ty { + back::FunctionType::Function(_) => { + self.write_type(module, result.ty)?; + } + back::FunctionType::EntryPoint(index) => { + if let Some(ref ep_output) = self.entry_point_io[index as usize].output { + write!(self.out, "{}", ep_output.ty_name)?; + } else { + self.write_type(module, result.ty)?; + } + } + } + } else { + write!(self.out, "void")?; + } + + // Write function name + write!(self.out, " {name}(")?; + + let need_workgroup_variables_initialization = + self.need_workgroup_variables_initialization(func_ctx, module); + + // Write function arguments for non entry point functions + match func_ctx.ty { + back::FunctionType::Function(handle) => { + for (index, arg) in func.arguments.iter().enumerate() { + if index != 0 { + write!(self.out, ", ")?; + } + // Write argument type + let arg_ty = match module.types[arg.ty].inner { + // pointers in function arguments are expected and resolve to `inout` + TypeInner::Pointer { base, .. } => { + //TODO: can we narrow this down to just `in` when possible? + write!(self.out, "inout ")?; + base + } + _ => arg.ty, + }; + self.write_type(module, arg_ty)?; + + let argument_name = + &self.names[&NameKey::FunctionArgument(handle, index as u32)]; + + // Write argument name. Space is important. + write!(self.out, " {argument_name}")?; + if let TypeInner::Array { base, size, .. } = module.types[arg_ty].inner { + self.write_array_size(module, base, size)?; + } + } + } + back::FunctionType::EntryPoint(ep_index) => { + if let Some(ref ep_input) = self.entry_point_io[ep_index as usize].input { + write!(self.out, "{} {}", ep_input.ty_name, ep_input.arg_name,)?; + } else { + let stage = module.entry_points[ep_index as usize].stage; + for (index, arg) in func.arguments.iter().enumerate() { + if index != 0 { + write!(self.out, ", ")?; + } + self.write_type(module, arg.ty)?; + + let argument_name = + &self.names[&NameKey::EntryPointArgument(ep_index, index as u32)]; + + write!(self.out, " {argument_name}")?; + if let TypeInner::Array { base, size, .. } = module.types[arg.ty].inner { + self.write_array_size(module, base, size)?; + } + + if let Some(ref binding) = arg.binding { + self.write_semantic(binding, Some((stage, Io::Input)))?; + } + } + + if need_workgroup_variables_initialization { + if !func.arguments.is_empty() { + write!(self.out, ", ")?; + } + write!(self.out, "uint3 __local_invocation_id : SV_GroupThreadID")?; + } + } + } + } + // Ends of arguments + write!(self.out, ")")?; + + // Write semantic if it present + if let back::FunctionType::EntryPoint(index) = func_ctx.ty { + let stage = module.entry_points[index as usize].stage; + if let Some(crate::FunctionResult { + binding: Some(ref binding), + .. + }) = func.result + { + self.write_semantic(binding, Some((stage, Io::Output)))?; + } + } + + // Function body start + writeln!(self.out)?; + writeln!(self.out, "{{")?; + + if need_workgroup_variables_initialization { + self.write_workgroup_variables_initialization(func_ctx, module)?; + } + + if let back::FunctionType::EntryPoint(index) = func_ctx.ty { + self.write_ep_arguments_initialization(module, func, index)?; + } + + // Write function local variables + for (handle, local) in func.local_variables.iter() { + // Write indentation (only for readability) + write!(self.out, "{}", back::INDENT)?; + + // Write the local name + // The leading space is important + self.write_type(module, local.ty)?; + write!(self.out, " {}", self.names[&func_ctx.name_key(handle)])?; + // Write size for array type + if let TypeInner::Array { base, size, .. } = module.types[local.ty].inner { + self.write_array_size(module, base, size)?; + } + + write!(self.out, " = ")?; + // Write the local initializer if needed + if let Some(init) = local.init { + self.write_expr(module, init, func_ctx)?; + } else { + // Zero initialize local variables + self.write_default_init(module, local.ty)?; + } + + // Finish the local with `;` and add a newline (only for readability) + writeln!(self.out, ";")? + } + + if !func.local_variables.is_empty() { + writeln!(self.out)?; + } + + // Write the function body (statement list) + for sta in func.body.iter() { + // The indentation should always be 1 when writing the function body + self.write_stmt(module, sta, func_ctx, back::Level(1))?; + } + + writeln!(self.out, "}}")?; + + self.named_expressions.clear(); + + Ok(()) + } + + fn need_workgroup_variables_initialization( + &mut self, + func_ctx: &back::FunctionCtx, + module: &Module, + ) -> bool { + self.options.zero_initialize_workgroup_memory + && func_ctx.ty.is_compute_entry_point(module) + && module.global_variables.iter().any(|(handle, var)| { + !func_ctx.info[handle].is_empty() && var.space == crate::AddressSpace::WorkGroup + }) + } + + fn write_workgroup_variables_initialization( + &mut self, + func_ctx: &back::FunctionCtx, + module: &Module, + ) -> BackendResult { + let level = back::Level(1); + + writeln!( + self.out, + "{level}if (all(__local_invocation_id == uint3(0u, 0u, 0u))) {{" + )?; + + let vars = module.global_variables.iter().filter(|&(handle, var)| { + !func_ctx.info[handle].is_empty() && var.space == crate::AddressSpace::WorkGroup + }); + + for (handle, var) in vars { + let name = &self.names[&NameKey::GlobalVariable(handle)]; + write!(self.out, "{}{} = ", level.next(), name)?; + self.write_default_init(module, var.ty)?; + writeln!(self.out, ";")?; + } + + writeln!(self.out, "{level}}}")?; + self.write_barrier(crate::Barrier::WORK_GROUP, level) + } + + /// Helper method used to write statements + /// + /// # Notes + /// Always adds a newline + fn write_stmt( + &mut self, + module: &Module, + stmt: &crate::Statement, + func_ctx: &back::FunctionCtx<'_>, + level: back::Level, + ) -> BackendResult { + use crate::Statement; + + match *stmt { + Statement::Emit(ref range) => { + for handle in range.clone() { + let ptr_class = func_ctx.resolve_type(handle, &module.types).pointer_space(); + let expr_name = if ptr_class.is_some() { + // HLSL can't save a pointer-valued expression in a variable, + // but we shouldn't ever need to: they should never be named expressions, + // and none of the expression types flagged by bake_ref_count can be pointer-valued. + None + } else if let Some(name) = func_ctx.named_expressions.get(&handle) { + // Front end provides names for all variables at the start of writing. + // But we write them to step by step. We need to recache them + // Otherwise, we could accidentally write variable name instead of full expression. + // Also, we use sanitized names! It defense backend from generating variable with name from reserved keywords. + Some(self.namer.call(name)) + } else if self.need_bake_expressions.contains(&handle) { + Some(format!("_expr{}", handle.index())) + } else { + None + }; + + if let Some(name) = expr_name { + write!(self.out, "{level}")?; + self.write_named_expr(module, handle, name, handle, func_ctx)?; + } + } + } + // TODO: copy-paste from glsl-out + Statement::Block(ref block) => { + write!(self.out, "{level}")?; + writeln!(self.out, "{{")?; + for sta in block.iter() { + // Increase the indentation to help with readability + self.write_stmt(module, sta, func_ctx, level.next())? + } + writeln!(self.out, "{level}}}")? + } + // TODO: copy-paste from glsl-out + Statement::If { + condition, + ref accept, + ref reject, + } => { + write!(self.out, "{level}")?; + write!(self.out, "if (")?; + self.write_expr(module, condition, func_ctx)?; + writeln!(self.out, ") {{")?; + + let l2 = level.next(); + for sta in accept { + // Increase indentation to help with readability + self.write_stmt(module, sta, func_ctx, l2)?; + } + + // If there are no statements in the reject block we skip writing it + // This is only for readability + if !reject.is_empty() { + writeln!(self.out, "{level}}} else {{")?; + + for sta in reject { + // Increase indentation to help with readability + self.write_stmt(module, sta, func_ctx, l2)?; + } + } + + writeln!(self.out, "{level}}}")? + } + // TODO: copy-paste from glsl-out + Statement::Kill => writeln!(self.out, "{level}discard;")?, + Statement::Return { value: None } => { + writeln!(self.out, "{level}return;")?; + } + Statement::Return { value: Some(expr) } => { + let base_ty_res = &func_ctx.info[expr].ty; + let mut resolved = base_ty_res.inner_with(&module.types); + if let TypeInner::Pointer { base, space: _ } = *resolved { + resolved = &module.types[base].inner; + } + + if let TypeInner::Struct { .. } = *resolved { + // We can safely unwrap here, since we now we working with struct + let ty = base_ty_res.handle().unwrap(); + let struct_name = &self.names[&NameKey::Type(ty)]; + let variable_name = self.namer.call(&struct_name.to_lowercase()); + write!(self.out, "{level}const {struct_name} {variable_name} = ",)?; + self.write_expr(module, expr, func_ctx)?; + writeln!(self.out, ";")?; + + // for entry point returns, we may need to reshuffle the outputs into a different struct + let ep_output = match func_ctx.ty { + back::FunctionType::Function(_) => None, + back::FunctionType::EntryPoint(index) => { + self.entry_point_io[index as usize].output.as_ref() + } + }; + let final_name = match ep_output { + Some(ep_output) => { + let final_name = self.namer.call(&variable_name); + write!( + self.out, + "{}const {} {} = {{ ", + level, ep_output.ty_name, final_name, + )?; + for (index, m) in ep_output.members.iter().enumerate() { + if index != 0 { + write!(self.out, ", ")?; + } + let member_name = &self.names[&NameKey::StructMember(ty, m.index)]; + write!(self.out, "{variable_name}.{member_name}")?; + } + writeln!(self.out, " }};")?; + final_name + } + None => variable_name, + }; + writeln!(self.out, "{level}return {final_name};")?; + } else { + write!(self.out, "{level}return ")?; + self.write_expr(module, expr, func_ctx)?; + writeln!(self.out, ";")? + } + } + Statement::Store { pointer, value } => { + let ty_inner = func_ctx.resolve_type(pointer, &module.types); + if let Some(crate::AddressSpace::Storage { .. }) = ty_inner.pointer_space() { + let var_handle = self.fill_access_chain(module, pointer, func_ctx)?; + self.write_storage_store( + module, + var_handle, + StoreValue::Expression(value), + func_ctx, + level, + )?; + } else { + // We treat matrices of the form `matCx2` as a sequence of C `vec2`s. + // See the module-level block comment in mod.rs for details. + // + // We handle matrix Stores here directly (including sub accesses for Vectors and Scalars). + // Loads are handled by `Expression::AccessIndex` (since sub accesses work fine for Loads). + struct MatrixAccess { + base: Handle<crate::Expression>, + index: u32, + } + enum Index { + Expression(Handle<crate::Expression>), + Static(u32), + } + + let get_members = |expr: Handle<crate::Expression>| { + let resolved = func_ctx.resolve_type(expr, &module.types); + match *resolved { + TypeInner::Pointer { base, .. } => match module.types[base].inner { + TypeInner::Struct { ref members, .. } => Some(members), + _ => None, + }, + _ => None, + } + }; + + let mut matrix = None; + let mut vector = None; + let mut scalar = None; + + let mut current_expr = pointer; + for _ in 0..3 { + let resolved = func_ctx.resolve_type(current_expr, &module.types); + + match (resolved, &func_ctx.expressions[current_expr]) { + ( + &TypeInner::Pointer { base: ty, .. }, + &crate::Expression::AccessIndex { base, index }, + ) if matches!( + module.types[ty].inner, + TypeInner::Matrix { + rows: crate::VectorSize::Bi, + .. + } + ) && get_members(base) + .map(|members| members[index as usize].binding.is_none()) + == Some(true) => + { + matrix = Some(MatrixAccess { base, index }); + break; + } + ( + &TypeInner::ValuePointer { + size: Some(crate::VectorSize::Bi), + .. + }, + &crate::Expression::Access { base, index }, + ) => { + vector = Some(Index::Expression(index)); + current_expr = base; + } + ( + &TypeInner::ValuePointer { + size: Some(crate::VectorSize::Bi), + .. + }, + &crate::Expression::AccessIndex { base, index }, + ) => { + vector = Some(Index::Static(index)); + current_expr = base; + } + ( + &TypeInner::ValuePointer { size: None, .. }, + &crate::Expression::Access { base, index }, + ) => { + scalar = Some(Index::Expression(index)); + current_expr = base; + } + ( + &TypeInner::ValuePointer { size: None, .. }, + &crate::Expression::AccessIndex { base, index }, + ) => { + scalar = Some(Index::Static(index)); + current_expr = base; + } + _ => break, + } + } + + write!(self.out, "{level}")?; + + if let Some(MatrixAccess { index, base }) = matrix { + let base_ty_res = &func_ctx.info[base].ty; + let resolved = base_ty_res.inner_with(&module.types); + let ty = match *resolved { + TypeInner::Pointer { base, .. } => base, + _ => base_ty_res.handle().unwrap(), + }; + + if let Some(Index::Static(vec_index)) = vector { + self.write_expr(module, base, func_ctx)?; + write!( + self.out, + ".{}_{}", + &self.names[&NameKey::StructMember(ty, index)], + vec_index + )?; + + if let Some(scalar_index) = scalar { + write!(self.out, "[")?; + match scalar_index { + Index::Static(index) => { + write!(self.out, "{index}")?; + } + Index::Expression(index) => { + self.write_expr(module, index, func_ctx)?; + } + } + write!(self.out, "]")?; + } + + write!(self.out, " = ")?; + self.write_expr(module, value, func_ctx)?; + writeln!(self.out, ";")?; + } else { + let access = WrappedStructMatrixAccess { ty, index }; + match (&vector, &scalar) { + (&Some(_), &Some(_)) => { + self.write_wrapped_struct_matrix_set_scalar_function_name( + access, + )?; + } + (&Some(_), &None) => { + self.write_wrapped_struct_matrix_set_vec_function_name(access)?; + } + (&None, _) => { + self.write_wrapped_struct_matrix_set_function_name(access)?; + } + } + + write!(self.out, "(")?; + self.write_expr(module, base, func_ctx)?; + write!(self.out, ", ")?; + self.write_expr(module, value, func_ctx)?; + + if let Some(Index::Expression(vec_index)) = vector { + write!(self.out, ", ")?; + self.write_expr(module, vec_index, func_ctx)?; + + if let Some(scalar_index) = scalar { + write!(self.out, ", ")?; + match scalar_index { + Index::Static(index) => { + write!(self.out, "{index}")?; + } + Index::Expression(index) => { + self.write_expr(module, index, func_ctx)?; + } + } + } + } + writeln!(self.out, ");")?; + } + } else { + // We handle `Store`s to __matCx2 column vectors and scalar elements via + // the previously injected functions __set_col_of_matCx2 / __set_el_of_matCx2. + struct MatrixData { + columns: crate::VectorSize, + base: Handle<crate::Expression>, + } + + enum Index { + Expression(Handle<crate::Expression>), + Static(u32), + } + + let mut matrix = None; + let mut vector = None; + let mut scalar = None; + + let mut current_expr = pointer; + for _ in 0..3 { + let resolved = func_ctx.resolve_type(current_expr, &module.types); + match (resolved, &func_ctx.expressions[current_expr]) { + ( + &TypeInner::ValuePointer { + size: Some(crate::VectorSize::Bi), + .. + }, + &crate::Expression::Access { base, index }, + ) => { + vector = Some(index); + current_expr = base; + } + ( + &TypeInner::ValuePointer { size: None, .. }, + &crate::Expression::Access { base, index }, + ) => { + scalar = Some(Index::Expression(index)); + current_expr = base; + } + ( + &TypeInner::ValuePointer { size: None, .. }, + &crate::Expression::AccessIndex { base, index }, + ) => { + scalar = Some(Index::Static(index)); + current_expr = base; + } + _ => { + if let Some(MatrixType { + columns, + rows: crate::VectorSize::Bi, + width: 4, + }) = get_inner_matrix_of_struct_array_member( + module, + current_expr, + func_ctx, + true, + ) { + matrix = Some(MatrixData { + columns, + base: current_expr, + }); + } + + break; + } + } + } + + if let (Some(MatrixData { columns, base }), Some(vec_index)) = + (matrix, vector) + { + if scalar.is_some() { + write!(self.out, "__set_el_of_mat{}x2", columns as u8)?; + } else { + write!(self.out, "__set_col_of_mat{}x2", columns as u8)?; + } + write!(self.out, "(")?; + self.write_expr(module, base, func_ctx)?; + write!(self.out, ", ")?; + self.write_expr(module, vec_index, func_ctx)?; + + if let Some(scalar_index) = scalar { + write!(self.out, ", ")?; + match scalar_index { + Index::Static(index) => { + write!(self.out, "{index}")?; + } + Index::Expression(index) => { + self.write_expr(module, index, func_ctx)?; + } + } + } + + write!(self.out, ", ")?; + self.write_expr(module, value, func_ctx)?; + + writeln!(self.out, ");")?; + } else { + self.write_expr(module, pointer, func_ctx)?; + write!(self.out, " = ")?; + + // We cast the RHS of this store in cases where the LHS + // is a struct member with type: + // - matCx2 or + // - a (possibly nested) array of matCx2's + if let Some(MatrixType { + columns, + rows: crate::VectorSize::Bi, + width: 4, + }) = get_inner_matrix_of_struct_array_member( + module, pointer, func_ctx, false, + ) { + let mut resolved = func_ctx.resolve_type(pointer, &module.types); + if let TypeInner::Pointer { base, .. } = *resolved { + resolved = &module.types[base].inner; + } + + write!(self.out, "(__mat{}x2", columns as u8)?; + if let TypeInner::Array { base, size, .. } = *resolved { + self.write_array_size(module, base, size)?; + } + write!(self.out, ")")?; + } + + self.write_expr(module, value, func_ctx)?; + writeln!(self.out, ";")? + } + } + } + } + Statement::Loop { + ref body, + ref continuing, + break_if, + } => { + let l2 = level.next(); + if !continuing.is_empty() || break_if.is_some() { + let gate_name = self.namer.call("loop_init"); + writeln!(self.out, "{level}bool {gate_name} = true;")?; + writeln!(self.out, "{level}while(true) {{")?; + writeln!(self.out, "{l2}if (!{gate_name}) {{")?; + let l3 = l2.next(); + for sta in continuing.iter() { + self.write_stmt(module, sta, func_ctx, l3)?; + } + if let Some(condition) = break_if { + write!(self.out, "{l3}if (")?; + self.write_expr(module, condition, func_ctx)?; + writeln!(self.out, ") {{")?; + writeln!(self.out, "{}break;", l3.next())?; + writeln!(self.out, "{l3}}}")?; + } + writeln!(self.out, "{l2}}}")?; + writeln!(self.out, "{l2}{gate_name} = false;")?; + } else { + writeln!(self.out, "{level}while(true) {{")?; + } + + for sta in body.iter() { + self.write_stmt(module, sta, func_ctx, l2)?; + } + writeln!(self.out, "{level}}}")? + } + Statement::Break => writeln!(self.out, "{level}break;")?, + Statement::Continue => writeln!(self.out, "{level}continue;")?, + Statement::Barrier(barrier) => { + self.write_barrier(barrier, level)?; + } + Statement::ImageStore { + image, + coordinate, + array_index, + value, + } => { + write!(self.out, "{level}")?; + self.write_expr(module, image, func_ctx)?; + + write!(self.out, "[")?; + if let Some(index) = array_index { + // Array index accepted only for texture_storage_2d_array, so we can safety use int3(coordinate, array_index) here + write!(self.out, "int3(")?; + self.write_expr(module, coordinate, func_ctx)?; + write!(self.out, ", ")?; + self.write_expr(module, index, func_ctx)?; + write!(self.out, ")")?; + } else { + self.write_expr(module, coordinate, func_ctx)?; + } + write!(self.out, "]")?; + + write!(self.out, " = ")?; + self.write_expr(module, value, func_ctx)?; + writeln!(self.out, ";")?; + } + Statement::Call { + function, + ref arguments, + result, + } => { + write!(self.out, "{level}")?; + if let Some(expr) = result { + write!(self.out, "const ")?; + let name = format!("{}{}", back::BAKE_PREFIX, expr.index()); + let expr_ty = &func_ctx.info[expr].ty; + match *expr_ty { + proc::TypeResolution::Handle(handle) => self.write_type(module, handle)?, + proc::TypeResolution::Value(ref value) => { + self.write_value_type(module, value)? + } + }; + write!(self.out, " {name} = ")?; + self.named_expressions.insert(expr, name); + } + let func_name = &self.names[&NameKey::Function(function)]; + write!(self.out, "{func_name}(")?; + for (index, argument) in arguments.iter().enumerate() { + if index != 0 { + write!(self.out, ", ")?; + } + self.write_expr(module, *argument, func_ctx)?; + } + writeln!(self.out, ");")? + } + Statement::Atomic { + pointer, + ref fun, + value, + result, + } => { + write!(self.out, "{level}")?; + let res_name = format!("{}{}", back::BAKE_PREFIX, result.index()); + match func_ctx.info[result].ty { + proc::TypeResolution::Handle(handle) => self.write_type(module, handle)?, + proc::TypeResolution::Value(ref value) => { + self.write_value_type(module, value)? + } + }; + + // Validation ensures that `pointer` has a `Pointer` type. + let pointer_space = func_ctx + .resolve_type(pointer, &module.types) + .pointer_space() + .unwrap(); + + let fun_str = fun.to_hlsl_suffix(); + write!(self.out, " {res_name}; ")?; + match pointer_space { + crate::AddressSpace::WorkGroup => { + write!(self.out, "Interlocked{fun_str}(")?; + self.write_expr(module, pointer, func_ctx)?; + } + crate::AddressSpace::Storage { .. } => { + let var_handle = self.fill_access_chain(module, pointer, func_ctx)?; + // The call to `self.write_storage_address` wants + // mutable access to all of `self`, so temporarily take + // ownership of our reusable access chain buffer. + let chain = mem::take(&mut self.temp_access_chain); + let var_name = &self.names[&NameKey::GlobalVariable(var_handle)]; + write!(self.out, "{var_name}.Interlocked{fun_str}(")?; + self.write_storage_address(module, &chain, func_ctx)?; + self.temp_access_chain = chain; + } + ref other => { + return Err(Error::Custom(format!( + "invalid address space {other:?} for atomic statement" + ))) + } + } + write!(self.out, ", ")?; + // handle the special cases + match *fun { + crate::AtomicFunction::Subtract => { + // we just wrote `InterlockedAdd`, so negate the argument + write!(self.out, "-")?; + } + crate::AtomicFunction::Exchange { compare: Some(_) } => { + return Err(Error::Unimplemented("atomic CompareExchange".to_string())); + } + _ => {} + } + self.write_expr(module, value, func_ctx)?; + writeln!(self.out, ", {res_name});")?; + self.named_expressions.insert(result, res_name); + } + Statement::WorkGroupUniformLoad { pointer, result } => { + self.write_barrier(crate::Barrier::WORK_GROUP, level)?; + write!(self.out, "{level}")?; + let name = format!("_expr{}", result.index()); + self.write_named_expr(module, pointer, name, result, func_ctx)?; + + self.write_barrier(crate::Barrier::WORK_GROUP, level)?; + } + Statement::Switch { + selector, + ref cases, + } => { + // Start the switch + write!(self.out, "{level}")?; + write!(self.out, "switch(")?; + self.write_expr(module, selector, func_ctx)?; + writeln!(self.out, ") {{")?; + + // Write all cases + let indent_level_1 = level.next(); + let indent_level_2 = indent_level_1.next(); + + for (i, case) in cases.iter().enumerate() { + match case.value { + crate::SwitchValue::I32(value) => { + write!(self.out, "{indent_level_1}case {value}:")? + } + crate::SwitchValue::U32(value) => { + write!(self.out, "{indent_level_1}case {value}u:")? + } + crate::SwitchValue::Default => { + write!(self.out, "{indent_level_1}default:")? + } + } + + // The new block is not only stylistic, it plays a role here: + // We might end up having to write the same case body + // multiple times due to FXC not supporting fallthrough. + // Therefore, some `Expression`s written by `Statement::Emit` + // will end up having the same name (`_expr<handle_index>`). + // So we need to put each case in its own scope. + let write_block_braces = !(case.fall_through && case.body.is_empty()); + if write_block_braces { + writeln!(self.out, " {{")?; + } else { + writeln!(self.out)?; + } + + // Although FXC does support a series of case clauses before + // a block[^yes], it does not support fallthrough from a + // non-empty case block to the next[^no]. If this case has a + // non-empty body with a fallthrough, emulate that by + // duplicating the bodies of all the cases it would fall + // into as extensions of this case's own body. This makes + // the HLSL output potentially quadratic in the size of the + // Naga IR. + // + // [^yes]: ```hlsl + // case 1: + // case 2: do_stuff() + // ``` + // [^no]: ```hlsl + // case 1: do_this(); + // case 2: do_that(); + // ``` + if case.fall_through && !case.body.is_empty() { + let curr_len = i + 1; + let end_case_idx = curr_len + + cases + .iter() + .skip(curr_len) + .position(|case| !case.fall_through) + .unwrap(); + let indent_level_3 = indent_level_2.next(); + for case in &cases[i..=end_case_idx] { + writeln!(self.out, "{indent_level_2}{{")?; + let prev_len = self.named_expressions.len(); + for sta in case.body.iter() { + self.write_stmt(module, sta, func_ctx, indent_level_3)?; + } + // Clear all named expressions that were previously inserted by the statements in the block + self.named_expressions.truncate(prev_len); + writeln!(self.out, "{indent_level_2}}}")?; + } + + let last_case = &cases[end_case_idx]; + if last_case.body.last().map_or(true, |s| !s.is_terminator()) { + writeln!(self.out, "{indent_level_2}break;")?; + } + } else { + for sta in case.body.iter() { + self.write_stmt(module, sta, func_ctx, indent_level_2)?; + } + if !case.fall_through + && case.body.last().map_or(true, |s| !s.is_terminator()) + { + writeln!(self.out, "{indent_level_2}break;")?; + } + } + + if write_block_braces { + writeln!(self.out, "{indent_level_1}}}")?; + } + } + + writeln!(self.out, "{level}}}")? + } + Statement::RayQuery { .. } => unreachable!(), + } + + Ok(()) + } + + fn write_const_expression( + &mut self, + module: &Module, + expr: Handle<crate::Expression>, + ) -> BackendResult { + self.write_possibly_const_expression( + module, + expr, + &module.const_expressions, + |writer, expr| writer.write_const_expression(module, expr), + ) + } + + fn write_possibly_const_expression<E>( + &mut self, + module: &Module, + expr: Handle<crate::Expression>, + expressions: &crate::Arena<crate::Expression>, + write_expression: E, + ) -> BackendResult + where + E: Fn(&mut Self, Handle<crate::Expression>) -> BackendResult, + { + use crate::Expression; + + match expressions[expr] { + Expression::Literal(literal) => match literal { + // Floats are written using `Debug` instead of `Display` because it always appends the + // decimal part even it's zero + crate::Literal::F64(value) => write!(self.out, "{value:?}L")?, + crate::Literal::F32(value) => write!(self.out, "{value:?}")?, + crate::Literal::U32(value) => write!(self.out, "{}u", value)?, + crate::Literal::I32(value) => write!(self.out, "{}", value)?, + crate::Literal::I64(value) => write!(self.out, "{}L", value)?, + crate::Literal::Bool(value) => write!(self.out, "{}", value)?, + crate::Literal::AbstractInt(_) | crate::Literal::AbstractFloat(_) => { + return Err(Error::Custom( + "Abstract types should not appear in IR presented to backends".into(), + )); + } + }, + Expression::Constant(handle) => { + let constant = &module.constants[handle]; + if constant.name.is_some() { + write!(self.out, "{}", self.names[&NameKey::Constant(handle)])?; + } else { + self.write_const_expression(module, constant.init)?; + } + } + Expression::ZeroValue(ty) => self.write_default_init(module, ty)?, + Expression::Compose { ty, ref components } => { + match module.types[ty].inner { + TypeInner::Struct { .. } | TypeInner::Array { .. } => { + self.write_wrapped_constructor_function_name( + module, + WrappedConstructor { ty }, + )?; + } + _ => { + self.write_type(module, ty)?; + } + }; + write!(self.out, "(")?; + for (index, component) in components.iter().enumerate() { + if index != 0 { + write!(self.out, ", ")?; + } + write_expression(self, *component)?; + } + write!(self.out, ")")?; + } + Expression::Splat { size, value } => { + // hlsl is not supported one value constructor + // if we write, for example, int4(0), dxc returns error: + // error: too few elements in vector initialization (expected 4 elements, have 1) + let number_of_components = match size { + crate::VectorSize::Bi => "xx", + crate::VectorSize::Tri => "xxx", + crate::VectorSize::Quad => "xxxx", + }; + write!(self.out, "(")?; + write_expression(self, value)?; + write!(self.out, ").{number_of_components}")? + } + _ => unreachable!(), + } + + Ok(()) + } + + /// Helper method to write expressions + /// + /// # Notes + /// Doesn't add any newlines or leading/trailing spaces + pub(super) fn write_expr( + &mut self, + module: &Module, + expr: Handle<crate::Expression>, + func_ctx: &back::FunctionCtx<'_>, + ) -> BackendResult { + use crate::Expression; + + // Handle the special semantics of vertex_index/instance_index + let ff_input = if self.options.special_constants_binding.is_some() { + func_ctx.is_fixed_function_input(expr, module) + } else { + None + }; + let closing_bracket = match ff_input { + Some(crate::BuiltIn::VertexIndex) => { + write!(self.out, "({SPECIAL_CBUF_VAR}.{SPECIAL_FIRST_VERTEX} + ")?; + ")" + } + Some(crate::BuiltIn::InstanceIndex) => { + write!(self.out, "({SPECIAL_CBUF_VAR}.{SPECIAL_FIRST_INSTANCE} + ",)?; + ")" + } + Some(crate::BuiltIn::NumWorkGroups) => { + // Note: despite their names (`FIRST_VERTEX` and `FIRST_INSTANCE`), + // in compute shaders the special constants contain the number + // of workgroups, which we are using here. + write!( + self.out, + "uint3({SPECIAL_CBUF_VAR}.{SPECIAL_FIRST_VERTEX}, {SPECIAL_CBUF_VAR}.{SPECIAL_FIRST_INSTANCE}, {SPECIAL_CBUF_VAR}.{SPECIAL_OTHER})", + )?; + return Ok(()); + } + _ => "", + }; + + if let Some(name) = self.named_expressions.get(&expr) { + write!(self.out, "{name}{closing_bracket}")?; + return Ok(()); + } + + let expression = &func_ctx.expressions[expr]; + + match *expression { + Expression::Literal(_) + | Expression::Constant(_) + | Expression::ZeroValue(_) + | Expression::Compose { .. } + | Expression::Splat { .. } => { + self.write_possibly_const_expression( + module, + expr, + func_ctx.expressions, + |writer, expr| writer.write_expr(module, expr, func_ctx), + )?; + } + // All of the multiplication can be expressed as `mul`, + // except vector * vector, which needs to use the "*" operator. + Expression::Binary { + op: crate::BinaryOperator::Multiply, + left, + right, + } if func_ctx.resolve_type(left, &module.types).is_matrix() + || func_ctx.resolve_type(right, &module.types).is_matrix() => + { + // We intentionally flip the order of multiplication as our matrices are implicitly transposed. + write!(self.out, "mul(")?; + self.write_expr(module, right, func_ctx)?; + write!(self.out, ", ")?; + self.write_expr(module, left, func_ctx)?; + write!(self.out, ")")?; + } + + // TODO: handle undefined behavior of BinaryOperator::Modulo + // + // sint: + // if right == 0 return 0 + // if left == min(type_of(left)) && right == -1 return 0 + // if sign(left) != sign(right) return result as defined by WGSL + // + // uint: + // if right == 0 return 0 + // + // float: + // if right == 0 return ? see https://github.com/gpuweb/gpuweb/issues/2798 + + // While HLSL supports float operands with the % operator it is only + // defined in cases where both sides are either positive or negative. + Expression::Binary { + op: crate::BinaryOperator::Modulo, + left, + right, + } if func_ctx.resolve_type(left, &module.types).scalar_kind() + == Some(crate::ScalarKind::Float) => + { + write!(self.out, "fmod(")?; + self.write_expr(module, left, func_ctx)?; + write!(self.out, ", ")?; + self.write_expr(module, right, func_ctx)?; + write!(self.out, ")")?; + } + Expression::Binary { op, left, right } => { + write!(self.out, "(")?; + self.write_expr(module, left, func_ctx)?; + write!(self.out, " {} ", crate::back::binary_operation_str(op))?; + self.write_expr(module, right, func_ctx)?; + write!(self.out, ")")?; + } + Expression::Access { base, index } => { + if let Some(crate::AddressSpace::Storage { .. }) = + func_ctx.resolve_type(expr, &module.types).pointer_space() + { + // do nothing, the chain is written on `Load`/`Store` + } else { + // We use the function __get_col_of_matCx2 here in cases + // where `base`s type resolves to a matCx2 and is part of a + // struct member with type of (possibly nested) array of matCx2's. + // + // Note that this only works for `Load`s and we handle + // `Store`s differently in `Statement::Store`. + if let Some(MatrixType { + columns, + rows: crate::VectorSize::Bi, + width: 4, + }) = get_inner_matrix_of_struct_array_member(module, base, func_ctx, true) + { + write!(self.out, "__get_col_of_mat{}x2(", columns as u8)?; + self.write_expr(module, base, func_ctx)?; + write!(self.out, ", ")?; + self.write_expr(module, index, func_ctx)?; + write!(self.out, ")")?; + return Ok(()); + } + + let resolved = func_ctx.resolve_type(base, &module.types); + + let non_uniform_qualifier = match *resolved { + TypeInner::BindingArray { .. } => { + let uniformity = &func_ctx.info[index].uniformity; + + uniformity.non_uniform_result.is_some() + } + _ => false, + }; + + self.write_expr(module, base, func_ctx)?; + write!(self.out, "[")?; + if non_uniform_qualifier { + write!(self.out, "NonUniformResourceIndex(")?; + } + self.write_expr(module, index, func_ctx)?; + if non_uniform_qualifier { + write!(self.out, ")")?; + } + write!(self.out, "]")?; + } + } + Expression::AccessIndex { base, index } => { + if let Some(crate::AddressSpace::Storage { .. }) = + func_ctx.resolve_type(expr, &module.types).pointer_space() + { + // do nothing, the chain is written on `Load`/`Store` + } else { + fn write_access<W: fmt::Write>( + writer: &mut super::Writer<'_, W>, + resolved: &TypeInner, + base_ty_handle: Option<Handle<crate::Type>>, + index: u32, + ) -> BackendResult { + match *resolved { + // We specifically lift the ValuePointer to this case. While `[0]` is valid + // HLSL for any vector behind a value pointer, FXC completely miscompiles + // it and generates completely nonsensical DXBC. + // + // See https://github.com/gfx-rs/naga/issues/2095 for more details. + TypeInner::Vector { .. } | TypeInner::ValuePointer { .. } => { + // Write vector access as a swizzle + write!(writer.out, ".{}", back::COMPONENTS[index as usize])? + } + TypeInner::Matrix { .. } + | TypeInner::Array { .. } + | TypeInner::BindingArray { .. } => write!(writer.out, "[{index}]")?, + TypeInner::Struct { .. } => { + // This will never panic in case the type is a `Struct`, this is not true + // for other types so we can only check while inside this match arm + let ty = base_ty_handle.unwrap(); + + write!( + writer.out, + ".{}", + &writer.names[&NameKey::StructMember(ty, index)] + )? + } + ref other => { + return Err(Error::Custom(format!("Cannot index {other:?}"))) + } + } + Ok(()) + } + + // We write the matrix column access in a special way since + // the type of `base` is our special __matCx2 struct. + if let Some(MatrixType { + rows: crate::VectorSize::Bi, + width: 4, + .. + }) = get_inner_matrix_of_struct_array_member(module, base, func_ctx, true) + { + self.write_expr(module, base, func_ctx)?; + write!(self.out, "._{index}")?; + return Ok(()); + } + + let base_ty_res = &func_ctx.info[base].ty; + let mut resolved = base_ty_res.inner_with(&module.types); + let base_ty_handle = match *resolved { + TypeInner::Pointer { base, .. } => { + resolved = &module.types[base].inner; + Some(base) + } + _ => base_ty_res.handle(), + }; + + // We treat matrices of the form `matCx2` as a sequence of C `vec2`s. + // See the module-level block comment in mod.rs for details. + // + // We handle matrix reconstruction here for Loads. + // Stores are handled directly by `Statement::Store`. + if let TypeInner::Struct { ref members, .. } = *resolved { + let member = &members[index as usize]; + + match module.types[member.ty].inner { + TypeInner::Matrix { + rows: crate::VectorSize::Bi, + .. + } if member.binding.is_none() => { + let ty = base_ty_handle.unwrap(); + self.write_wrapped_struct_matrix_get_function_name( + WrappedStructMatrixAccess { ty, index }, + )?; + write!(self.out, "(")?; + self.write_expr(module, base, func_ctx)?; + write!(self.out, ")")?; + return Ok(()); + } + _ => {} + } + } + + self.write_expr(module, base, func_ctx)?; + write_access(self, resolved, base_ty_handle, index)?; + } + } + Expression::FunctionArgument(pos) => { + let key = func_ctx.argument_key(pos); + let name = &self.names[&key]; + write!(self.out, "{name}")?; + } + Expression::ImageSample { + image, + sampler, + gather, + coordinate, + array_index, + offset, + level, + depth_ref, + } => { + use crate::SampleLevel as Sl; + const COMPONENTS: [&str; 4] = ["", "Green", "Blue", "Alpha"]; + + let (base_str, component_str) = match gather { + Some(component) => ("Gather", COMPONENTS[component as usize]), + None => ("Sample", ""), + }; + let cmp_str = match depth_ref { + Some(_) => "Cmp", + None => "", + }; + let level_str = match level { + Sl::Zero if gather.is_none() => "LevelZero", + Sl::Auto | Sl::Zero => "", + Sl::Exact(_) => "Level", + Sl::Bias(_) => "Bias", + Sl::Gradient { .. } => "Grad", + }; + + self.write_expr(module, image, func_ctx)?; + write!(self.out, ".{base_str}{cmp_str}{component_str}{level_str}(")?; + self.write_expr(module, sampler, func_ctx)?; + write!(self.out, ", ")?; + self.write_texture_coordinates( + "float", + coordinate, + array_index, + None, + module, + func_ctx, + )?; + + if let Some(depth_ref) = depth_ref { + write!(self.out, ", ")?; + self.write_expr(module, depth_ref, func_ctx)?; + } + + match level { + Sl::Auto | Sl::Zero => {} + Sl::Exact(expr) => { + write!(self.out, ", ")?; + self.write_expr(module, expr, func_ctx)?; + } + Sl::Bias(expr) => { + write!(self.out, ", ")?; + self.write_expr(module, expr, func_ctx)?; + } + Sl::Gradient { x, y } => { + write!(self.out, ", ")?; + self.write_expr(module, x, func_ctx)?; + write!(self.out, ", ")?; + self.write_expr(module, y, func_ctx)?; + } + } + + if let Some(offset) = offset { + write!(self.out, ", ")?; + write!(self.out, "int2(")?; // work around https://github.com/microsoft/DirectXShaderCompiler/issues/5082#issuecomment-1540147807 + self.write_const_expression(module, offset)?; + write!(self.out, ")")?; + } + + write!(self.out, ")")?; + } + Expression::ImageQuery { image, query } => { + // use wrapped image query function + if let TypeInner::Image { + dim, + arrayed, + class, + } = *func_ctx.resolve_type(image, &module.types) + { + let wrapped_image_query = WrappedImageQuery { + dim, + arrayed, + class, + query: query.into(), + }; + + self.write_wrapped_image_query_function_name(wrapped_image_query)?; + write!(self.out, "(")?; + // Image always first param + self.write_expr(module, image, func_ctx)?; + if let crate::ImageQuery::Size { level: Some(level) } = query { + write!(self.out, ", ")?; + self.write_expr(module, level, func_ctx)?; + } + write!(self.out, ")")?; + } + } + Expression::ImageLoad { + image, + coordinate, + array_index, + sample, + level, + } => { + // https://docs.microsoft.com/en-us/windows/win32/direct3dhlsl/dx-graphics-hlsl-to-load + self.write_expr(module, image, func_ctx)?; + write!(self.out, ".Load(")?; + + self.write_texture_coordinates( + "int", + coordinate, + array_index, + level, + module, + func_ctx, + )?; + + if let Some(sample) = sample { + write!(self.out, ", ")?; + self.write_expr(module, sample, func_ctx)?; + } + + // close bracket for Load function + write!(self.out, ")")?; + + // return x component if return type is scalar + if let TypeInner::Scalar(_) = *func_ctx.resolve_type(expr, &module.types) { + write!(self.out, ".x")?; + } + } + Expression::GlobalVariable(handle) => match module.global_variables[handle].space { + crate::AddressSpace::Storage { .. } => {} + _ => { + let name = &self.names[&NameKey::GlobalVariable(handle)]; + write!(self.out, "{name}")?; + } + }, + Expression::LocalVariable(handle) => { + write!(self.out, "{}", self.names[&func_ctx.name_key(handle)])? + } + Expression::Load { pointer } => { + match func_ctx + .resolve_type(pointer, &module.types) + .pointer_space() + { + Some(crate::AddressSpace::Storage { .. }) => { + let var_handle = self.fill_access_chain(module, pointer, func_ctx)?; + let result_ty = func_ctx.info[expr].ty.clone(); + self.write_storage_load(module, var_handle, result_ty, func_ctx)?; + } + _ => { + let mut close_paren = false; + + // We cast the value loaded to a native HLSL floatCx2 + // in cases where it is of type: + // - __matCx2 or + // - a (possibly nested) array of __matCx2's + if let Some(MatrixType { + rows: crate::VectorSize::Bi, + width: 4, + .. + }) = get_inner_matrix_of_struct_array_member( + module, pointer, func_ctx, false, + ) + .or_else(|| get_inner_matrix_of_global_uniform(module, pointer, func_ctx)) + { + let mut resolved = func_ctx.resolve_type(pointer, &module.types); + if let TypeInner::Pointer { base, .. } = *resolved { + resolved = &module.types[base].inner; + } + + write!(self.out, "((")?; + if let TypeInner::Array { base, size, .. } = *resolved { + self.write_type(module, base)?; + self.write_array_size(module, base, size)?; + } else { + self.write_value_type(module, resolved)?; + } + write!(self.out, ")")?; + close_paren = true; + } + + self.write_expr(module, pointer, func_ctx)?; + + if close_paren { + write!(self.out, ")")?; + } + } + } + } + Expression::Unary { op, expr } => { + // https://docs.microsoft.com/en-us/windows/win32/direct3dhlsl/dx-graphics-hlsl-operators#unary-operators + let op_str = match op { + crate::UnaryOperator::Negate => "-", + crate::UnaryOperator::LogicalNot => "!", + crate::UnaryOperator::BitwiseNot => "~", + }; + write!(self.out, "{op_str}(")?; + self.write_expr(module, expr, func_ctx)?; + write!(self.out, ")")?; + } + Expression::As { + expr, + kind, + convert, + } => { + let inner = func_ctx.resolve_type(expr, &module.types); + match convert { + Some(dst_width) => { + let scalar = crate::Scalar { + kind, + width: dst_width, + }; + match *inner { + TypeInner::Vector { size, .. } => { + write!( + self.out, + "{}{}(", + scalar.to_hlsl_str()?, + back::vector_size_str(size) + )?; + } + TypeInner::Scalar(_) => { + write!(self.out, "{}(", scalar.to_hlsl_str()?,)?; + } + TypeInner::Matrix { columns, rows, .. } => { + write!( + self.out, + "{}{}x{}(", + scalar.to_hlsl_str()?, + back::vector_size_str(columns), + back::vector_size_str(rows) + )?; + } + _ => { + return Err(Error::Unimplemented(format!( + "write_expr expression::as {inner:?}" + ))); + } + }; + } + None => { + write!(self.out, "{}(", kind.to_hlsl_cast(),)?; + } + } + self.write_expr(module, expr, func_ctx)?; + write!(self.out, ")")?; + } + Expression::Math { + fun, + arg, + arg1, + arg2, + arg3, + } => { + use crate::MathFunction as Mf; + + enum Function { + Asincosh { is_sin: bool }, + Atanh, + ExtractBits, + InsertBits, + Pack2x16float, + Pack2x16snorm, + Pack2x16unorm, + Pack4x8snorm, + Pack4x8unorm, + Unpack2x16float, + Unpack2x16snorm, + Unpack2x16unorm, + Unpack4x8snorm, + Unpack4x8unorm, + Regular(&'static str), + MissingIntOverload(&'static str), + MissingIntReturnType(&'static str), + CountTrailingZeros, + CountLeadingZeros, + } + + let fun = match fun { + // comparison + Mf::Abs => Function::Regular("abs"), + Mf::Min => Function::Regular("min"), + Mf::Max => Function::Regular("max"), + Mf::Clamp => Function::Regular("clamp"), + Mf::Saturate => Function::Regular("saturate"), + // trigonometry + Mf::Cos => Function::Regular("cos"), + Mf::Cosh => Function::Regular("cosh"), + Mf::Sin => Function::Regular("sin"), + Mf::Sinh => Function::Regular("sinh"), + Mf::Tan => Function::Regular("tan"), + Mf::Tanh => Function::Regular("tanh"), + Mf::Acos => Function::Regular("acos"), + Mf::Asin => Function::Regular("asin"), + Mf::Atan => Function::Regular("atan"), + Mf::Atan2 => Function::Regular("atan2"), + Mf::Asinh => Function::Asincosh { is_sin: true }, + Mf::Acosh => Function::Asincosh { is_sin: false }, + Mf::Atanh => Function::Atanh, + Mf::Radians => Function::Regular("radians"), + Mf::Degrees => Function::Regular("degrees"), + // decomposition + Mf::Ceil => Function::Regular("ceil"), + Mf::Floor => Function::Regular("floor"), + Mf::Round => Function::Regular("round"), + Mf::Fract => Function::Regular("frac"), + Mf::Trunc => Function::Regular("trunc"), + Mf::Modf => Function::Regular(MODF_FUNCTION), + Mf::Frexp => Function::Regular(FREXP_FUNCTION), + Mf::Ldexp => Function::Regular("ldexp"), + // exponent + Mf::Exp => Function::Regular("exp"), + Mf::Exp2 => Function::Regular("exp2"), + Mf::Log => Function::Regular("log"), + Mf::Log2 => Function::Regular("log2"), + Mf::Pow => Function::Regular("pow"), + // geometry + Mf::Dot => Function::Regular("dot"), + //Mf::Outer => , + Mf::Cross => Function::Regular("cross"), + Mf::Distance => Function::Regular("distance"), + Mf::Length => Function::Regular("length"), + Mf::Normalize => Function::Regular("normalize"), + Mf::FaceForward => Function::Regular("faceforward"), + Mf::Reflect => Function::Regular("reflect"), + Mf::Refract => Function::Regular("refract"), + // computational + Mf::Sign => Function::Regular("sign"), + Mf::Fma => Function::Regular("mad"), + Mf::Mix => Function::Regular("lerp"), + Mf::Step => Function::Regular("step"), + Mf::SmoothStep => Function::Regular("smoothstep"), + Mf::Sqrt => Function::Regular("sqrt"), + Mf::InverseSqrt => Function::Regular("rsqrt"), + //Mf::Inverse =>, + Mf::Transpose => Function::Regular("transpose"), + Mf::Determinant => Function::Regular("determinant"), + // bits + Mf::CountTrailingZeros => Function::CountTrailingZeros, + Mf::CountLeadingZeros => Function::CountLeadingZeros, + Mf::CountOneBits => Function::MissingIntOverload("countbits"), + Mf::ReverseBits => Function::MissingIntOverload("reversebits"), + Mf::FindLsb => Function::MissingIntReturnType("firstbitlow"), + Mf::FindMsb => Function::MissingIntReturnType("firstbithigh"), + Mf::ExtractBits => Function::ExtractBits, + Mf::InsertBits => Function::InsertBits, + // Data Packing + Mf::Pack2x16float => Function::Pack2x16float, + Mf::Pack2x16snorm => Function::Pack2x16snorm, + Mf::Pack2x16unorm => Function::Pack2x16unorm, + Mf::Pack4x8snorm => Function::Pack4x8snorm, + Mf::Pack4x8unorm => Function::Pack4x8unorm, + // Data Unpacking + Mf::Unpack2x16float => Function::Unpack2x16float, + Mf::Unpack2x16snorm => Function::Unpack2x16snorm, + Mf::Unpack2x16unorm => Function::Unpack2x16unorm, + Mf::Unpack4x8snorm => Function::Unpack4x8snorm, + Mf::Unpack4x8unorm => Function::Unpack4x8unorm, + _ => return Err(Error::Unimplemented(format!("write_expr_math {fun:?}"))), + }; + + match fun { + Function::Asincosh { is_sin } => { + write!(self.out, "log(")?; + self.write_expr(module, arg, func_ctx)?; + write!(self.out, " + sqrt(")?; + self.write_expr(module, arg, func_ctx)?; + write!(self.out, " * ")?; + self.write_expr(module, arg, func_ctx)?; + match is_sin { + true => write!(self.out, " + 1.0))")?, + false => write!(self.out, " - 1.0))")?, + } + } + Function::Atanh => { + write!(self.out, "0.5 * log((1.0 + ")?; + self.write_expr(module, arg, func_ctx)?; + write!(self.out, ") / (1.0 - ")?; + self.write_expr(module, arg, func_ctx)?; + write!(self.out, "))")?; + } + Function::ExtractBits => { + // e: T, + // offset: u32, + // count: u32 + // T is u32 or i32 or vecN<u32> or vecN<i32> + if let (Some(offset), Some(count)) = (arg1, arg2) { + let scalar_width: u8 = 32; + // Works for signed and unsigned + // (count == 0 ? 0 : (e << (32 - count - offset)) >> (32 - count)) + write!(self.out, "(")?; + self.write_expr(module, count, func_ctx)?; + write!(self.out, " == 0 ? 0 : (")?; + self.write_expr(module, arg, func_ctx)?; + write!(self.out, " << ({scalar_width} - ")?; + self.write_expr(module, count, func_ctx)?; + write!(self.out, " - ")?; + self.write_expr(module, offset, func_ctx)?; + write!(self.out, ")) >> ({scalar_width} - ")?; + self.write_expr(module, count, func_ctx)?; + write!(self.out, "))")?; + } + } + Function::InsertBits => { + // e: T, + // newbits: T, + // offset: u32, + // count: u32 + // returns T + // T is i32, u32, vecN<i32>, or vecN<u32> + if let (Some(newbits), Some(offset), Some(count)) = (arg1, arg2, arg3) { + let scalar_width: u8 = 32; + let scalar_max: u32 = 0xFFFFFFFF; + // mask = ((0xFFFFFFFFu >> (32 - count)) << offset) + // (count == 0 ? e : ((e & ~mask) | ((newbits << offset) & mask))) + write!(self.out, "(")?; + self.write_expr(module, count, func_ctx)?; + write!(self.out, " == 0 ? ")?; + self.write_expr(module, arg, func_ctx)?; + write!(self.out, " : ")?; + write!(self.out, "(")?; + self.write_expr(module, arg, func_ctx)?; + write!(self.out, " & ~")?; + // mask + write!(self.out, "(({scalar_max}u >> ({scalar_width}u - ")?; + self.write_expr(module, count, func_ctx)?; + write!(self.out, ")) << ")?; + self.write_expr(module, offset, func_ctx)?; + write!(self.out, ")")?; + // end mask + write!(self.out, ") | ((")?; + self.write_expr(module, newbits, func_ctx)?; + write!(self.out, " << ")?; + self.write_expr(module, offset, func_ctx)?; + write!(self.out, ") & ")?; + // // mask + write!(self.out, "(({scalar_max}u >> ({scalar_width}u - ")?; + self.write_expr(module, count, func_ctx)?; + write!(self.out, ")) << ")?; + self.write_expr(module, offset, func_ctx)?; + write!(self.out, ")")?; + // // end mask + write!(self.out, "))")?; + } + } + Function::Pack2x16float => { + write!(self.out, "(f32tof16(")?; + self.write_expr(module, arg, func_ctx)?; + write!(self.out, "[0]) | f32tof16(")?; + self.write_expr(module, arg, func_ctx)?; + write!(self.out, "[1]) << 16)")?; + } + Function::Pack2x16snorm => { + let scale = 32767; + + write!(self.out, "uint((int(round(clamp(")?; + self.write_expr(module, arg, func_ctx)?; + write!( + self.out, + "[0], -1.0, 1.0) * {scale}.0)) & 0xFFFF) | ((int(round(clamp(" + )?; + self.write_expr(module, arg, func_ctx)?; + write!(self.out, "[1], -1.0, 1.0) * {scale}.0)) & 0xFFFF) << 16))",)?; + } + Function::Pack2x16unorm => { + let scale = 65535; + + write!(self.out, "(uint(round(clamp(")?; + self.write_expr(module, arg, func_ctx)?; + write!(self.out, "[0], 0.0, 1.0) * {scale}.0)) | uint(round(clamp(")?; + self.write_expr(module, arg, func_ctx)?; + write!(self.out, "[1], 0.0, 1.0) * {scale}.0)) << 16)")?; + } + Function::Pack4x8snorm => { + let scale = 127; + + write!(self.out, "uint((int(round(clamp(")?; + self.write_expr(module, arg, func_ctx)?; + write!( + self.out, + "[0], -1.0, 1.0) * {scale}.0)) & 0xFF) | ((int(round(clamp(" + )?; + self.write_expr(module, arg, func_ctx)?; + write!( + self.out, + "[1], -1.0, 1.0) * {scale}.0)) & 0xFF) << 8) | ((int(round(clamp(" + )?; + self.write_expr(module, arg, func_ctx)?; + write!( + self.out, + "[2], -1.0, 1.0) * {scale}.0)) & 0xFF) << 16) | ((int(round(clamp(" + )?; + self.write_expr(module, arg, func_ctx)?; + write!(self.out, "[3], -1.0, 1.0) * {scale}.0)) & 0xFF) << 24))",)?; + } + Function::Pack4x8unorm => { + let scale = 255; + + write!(self.out, "(uint(round(clamp(")?; + self.write_expr(module, arg, func_ctx)?; + write!(self.out, "[0], 0.0, 1.0) * {scale}.0)) | uint(round(clamp(")?; + self.write_expr(module, arg, func_ctx)?; + write!( + self.out, + "[1], 0.0, 1.0) * {scale}.0)) << 8 | uint(round(clamp(" + )?; + self.write_expr(module, arg, func_ctx)?; + write!( + self.out, + "[2], 0.0, 1.0) * {scale}.0)) << 16 | uint(round(clamp(" + )?; + self.write_expr(module, arg, func_ctx)?; + write!(self.out, "[3], 0.0, 1.0) * {scale}.0)) << 24)")?; + } + + Function::Unpack2x16float => { + write!(self.out, "float2(f16tof32(")?; + self.write_expr(module, arg, func_ctx)?; + write!(self.out, "), f16tof32((")?; + self.write_expr(module, arg, func_ctx)?; + write!(self.out, ") >> 16))")?; + } + Function::Unpack2x16snorm => { + let scale = 32767; + + write!(self.out, "(float2(int2(")?; + self.write_expr(module, arg, func_ctx)?; + write!(self.out, " << 16, ")?; + self.write_expr(module, arg, func_ctx)?; + write!(self.out, ") >> 16) / {scale}.0)")?; + } + Function::Unpack2x16unorm => { + let scale = 65535; + + write!(self.out, "(float2(")?; + self.write_expr(module, arg, func_ctx)?; + write!(self.out, " & 0xFFFF, ")?; + self.write_expr(module, arg, func_ctx)?; + write!(self.out, " >> 16) / {scale}.0)")?; + } + Function::Unpack4x8snorm => { + let scale = 127; + + write!(self.out, "(float4(int4(")?; + self.write_expr(module, arg, func_ctx)?; + write!(self.out, " << 24, ")?; + self.write_expr(module, arg, func_ctx)?; + write!(self.out, " << 16, ")?; + self.write_expr(module, arg, func_ctx)?; + write!(self.out, " << 8, ")?; + self.write_expr(module, arg, func_ctx)?; + write!(self.out, ") >> 24) / {scale}.0)")?; + } + Function::Unpack4x8unorm => { + let scale = 255; + + write!(self.out, "(float4(")?; + self.write_expr(module, arg, func_ctx)?; + write!(self.out, " & 0xFF, ")?; + self.write_expr(module, arg, func_ctx)?; + write!(self.out, " >> 8 & 0xFF, ")?; + self.write_expr(module, arg, func_ctx)?; + write!(self.out, " >> 16 & 0xFF, ")?; + self.write_expr(module, arg, func_ctx)?; + write!(self.out, " >> 24) / {scale}.0)")?; + } + Function::Regular(fun_name) => { + write!(self.out, "{fun_name}(")?; + self.write_expr(module, arg, func_ctx)?; + if let Some(arg) = arg1 { + write!(self.out, ", ")?; + self.write_expr(module, arg, func_ctx)?; + } + if let Some(arg) = arg2 { + write!(self.out, ", ")?; + self.write_expr(module, arg, func_ctx)?; + } + if let Some(arg) = arg3 { + write!(self.out, ", ")?; + self.write_expr(module, arg, func_ctx)?; + } + write!(self.out, ")")? + } + Function::MissingIntOverload(fun_name) => { + let scalar_kind = func_ctx.resolve_type(arg, &module.types).scalar_kind(); + if let Some(ScalarKind::Sint) = scalar_kind { + write!(self.out, "asint({fun_name}(asuint(")?; + self.write_expr(module, arg, func_ctx)?; + write!(self.out, ")))")?; + } else { + write!(self.out, "{fun_name}(")?; + self.write_expr(module, arg, func_ctx)?; + write!(self.out, ")")?; + } + } + Function::MissingIntReturnType(fun_name) => { + let scalar_kind = func_ctx.resolve_type(arg, &module.types).scalar_kind(); + if let Some(ScalarKind::Sint) = scalar_kind { + write!(self.out, "asint({fun_name}(")?; + self.write_expr(module, arg, func_ctx)?; + write!(self.out, "))")?; + } else { + write!(self.out, "{fun_name}(")?; + self.write_expr(module, arg, func_ctx)?; + write!(self.out, ")")?; + } + } + Function::CountTrailingZeros => { + match *func_ctx.resolve_type(arg, &module.types) { + TypeInner::Vector { size, scalar } => { + let s = match size { + crate::VectorSize::Bi => ".xx", + crate::VectorSize::Tri => ".xxx", + crate::VectorSize::Quad => ".xxxx", + }; + + if let ScalarKind::Uint = scalar.kind { + write!(self.out, "min((32u){s}, firstbitlow(")?; + self.write_expr(module, arg, func_ctx)?; + write!(self.out, "))")?; + } else { + write!(self.out, "asint(min((32u){s}, firstbitlow(")?; + self.write_expr(module, arg, func_ctx)?; + write!(self.out, ")))")?; + } + } + TypeInner::Scalar(scalar) => { + if let ScalarKind::Uint = scalar.kind { + write!(self.out, "min(32u, firstbitlow(")?; + self.write_expr(module, arg, func_ctx)?; + write!(self.out, "))")?; + } else { + write!(self.out, "asint(min(32u, firstbitlow(")?; + self.write_expr(module, arg, func_ctx)?; + write!(self.out, ")))")?; + } + } + _ => unreachable!(), + } + + return Ok(()); + } + Function::CountLeadingZeros => { + match *func_ctx.resolve_type(arg, &module.types) { + TypeInner::Vector { size, scalar } => { + let s = match size { + crate::VectorSize::Bi => ".xx", + crate::VectorSize::Tri => ".xxx", + crate::VectorSize::Quad => ".xxxx", + }; + + if let ScalarKind::Uint = scalar.kind { + write!(self.out, "((31u){s} - firstbithigh(")?; + self.write_expr(module, arg, func_ctx)?; + write!(self.out, "))")?; + } else { + write!(self.out, "(")?; + self.write_expr(module, arg, func_ctx)?; + write!( + self.out, + " < (0){s} ? (0){s} : (31){s} - asint(firstbithigh(" + )?; + self.write_expr(module, arg, func_ctx)?; + write!(self.out, ")))")?; + } + } + TypeInner::Scalar(scalar) => { + if let ScalarKind::Uint = scalar.kind { + write!(self.out, "(31u - firstbithigh(")?; + self.write_expr(module, arg, func_ctx)?; + write!(self.out, "))")?; + } else { + write!(self.out, "(")?; + self.write_expr(module, arg, func_ctx)?; + write!(self.out, " < 0 ? 0 : 31 - asint(firstbithigh(")?; + self.write_expr(module, arg, func_ctx)?; + write!(self.out, ")))")?; + } + } + _ => unreachable!(), + } + + return Ok(()); + } + } + } + Expression::Swizzle { + size, + vector, + pattern, + } => { + self.write_expr(module, vector, func_ctx)?; + write!(self.out, ".")?; + for &sc in pattern[..size as usize].iter() { + self.out.write_char(back::COMPONENTS[sc as usize])?; + } + } + Expression::ArrayLength(expr) => { + let var_handle = match func_ctx.expressions[expr] { + Expression::AccessIndex { base, index: _ } => { + match func_ctx.expressions[base] { + Expression::GlobalVariable(handle) => handle, + _ => unreachable!(), + } + } + Expression::GlobalVariable(handle) => handle, + _ => unreachable!(), + }; + + let var = &module.global_variables[var_handle]; + let (offset, stride) = match module.types[var.ty].inner { + TypeInner::Array { stride, .. } => (0, stride), + TypeInner::Struct { ref members, .. } => { + let last = members.last().unwrap(); + let stride = match module.types[last.ty].inner { + TypeInner::Array { stride, .. } => stride, + _ => unreachable!(), + }; + (last.offset, stride) + } + _ => unreachable!(), + }; + + let storage_access = match var.space { + crate::AddressSpace::Storage { access } => access, + _ => crate::StorageAccess::default(), + }; + let wrapped_array_length = WrappedArrayLength { + writable: storage_access.contains(crate::StorageAccess::STORE), + }; + + write!(self.out, "((")?; + self.write_wrapped_array_length_function_name(wrapped_array_length)?; + let var_name = &self.names[&NameKey::GlobalVariable(var_handle)]; + write!(self.out, "({var_name}) - {offset}) / {stride})")? + } + Expression::Derivative { axis, ctrl, expr } => { + use crate::{DerivativeAxis as Axis, DerivativeControl as Ctrl}; + if axis == Axis::Width && (ctrl == Ctrl::Coarse || ctrl == Ctrl::Fine) { + let tail = match ctrl { + Ctrl::Coarse => "coarse", + Ctrl::Fine => "fine", + Ctrl::None => unreachable!(), + }; + write!(self.out, "abs(ddx_{tail}(")?; + self.write_expr(module, expr, func_ctx)?; + write!(self.out, ")) + abs(ddy_{tail}(")?; + self.write_expr(module, expr, func_ctx)?; + write!(self.out, "))")? + } else { + let fun_str = match (axis, ctrl) { + (Axis::X, Ctrl::Coarse) => "ddx_coarse", + (Axis::X, Ctrl::Fine) => "ddx_fine", + (Axis::X, Ctrl::None) => "ddx", + (Axis::Y, Ctrl::Coarse) => "ddy_coarse", + (Axis::Y, Ctrl::Fine) => "ddy_fine", + (Axis::Y, Ctrl::None) => "ddy", + (Axis::Width, Ctrl::Coarse | Ctrl::Fine) => unreachable!(), + (Axis::Width, Ctrl::None) => "fwidth", + }; + write!(self.out, "{fun_str}(")?; + self.write_expr(module, expr, func_ctx)?; + write!(self.out, ")")? + } + } + Expression::Relational { fun, argument } => { + use crate::RelationalFunction as Rf; + + let fun_str = match fun { + Rf::All => "all", + Rf::Any => "any", + Rf::IsNan => "isnan", + Rf::IsInf => "isinf", + }; + write!(self.out, "{fun_str}(")?; + self.write_expr(module, argument, func_ctx)?; + write!(self.out, ")")? + } + Expression::Select { + condition, + accept, + reject, + } => { + write!(self.out, "(")?; + self.write_expr(module, condition, func_ctx)?; + write!(self.out, " ? ")?; + self.write_expr(module, accept, func_ctx)?; + write!(self.out, " : ")?; + self.write_expr(module, reject, func_ctx)?; + write!(self.out, ")")? + } + // Not supported yet + Expression::RayQueryGetIntersection { .. } => unreachable!(), + // Nothing to do here, since call expression already cached + Expression::CallResult(_) + | Expression::AtomicResult { .. } + | Expression::WorkGroupUniformLoadResult { .. } + | Expression::RayQueryProceedResult => {} + } + + if !closing_bracket.is_empty() { + write!(self.out, "{closing_bracket}")?; + } + Ok(()) + } + + fn write_named_expr( + &mut self, + module: &Module, + handle: Handle<crate::Expression>, + name: String, + // The expression which is being named. + // Generally, this is the same as handle, except in WorkGroupUniformLoad + named: Handle<crate::Expression>, + ctx: &back::FunctionCtx, + ) -> BackendResult { + match ctx.info[named].ty { + proc::TypeResolution::Handle(ty_handle) => match module.types[ty_handle].inner { + TypeInner::Struct { .. } => { + let ty_name = &self.names[&NameKey::Type(ty_handle)]; + write!(self.out, "{ty_name}")?; + } + _ => { + self.write_type(module, ty_handle)?; + } + }, + proc::TypeResolution::Value(ref inner) => { + self.write_value_type(module, inner)?; + } + } + + let resolved = ctx.resolve_type(named, &module.types); + + write!(self.out, " {name}")?; + // If rhs is a array type, we should write array size + if let TypeInner::Array { base, size, .. } = *resolved { + self.write_array_size(module, base, size)?; + } + write!(self.out, " = ")?; + self.write_expr(module, handle, ctx)?; + writeln!(self.out, ";")?; + self.named_expressions.insert(named, name); + + Ok(()) + } + + /// Helper function that write default zero initialization + fn write_default_init(&mut self, module: &Module, ty: Handle<crate::Type>) -> BackendResult { + write!(self.out, "(")?; + self.write_type(module, ty)?; + if let TypeInner::Array { base, size, .. } = module.types[ty].inner { + self.write_array_size(module, base, size)?; + } + write!(self.out, ")0")?; + Ok(()) + } + + fn write_barrier(&mut self, barrier: crate::Barrier, level: back::Level) -> BackendResult { + if barrier.contains(crate::Barrier::STORAGE) { + writeln!(self.out, "{level}DeviceMemoryBarrierWithGroupSync();")?; + } + if barrier.contains(crate::Barrier::WORK_GROUP) { + writeln!(self.out, "{level}GroupMemoryBarrierWithGroupSync();")?; + } + Ok(()) + } +} + +pub(super) struct MatrixType { + pub(super) columns: crate::VectorSize, + pub(super) rows: crate::VectorSize, + pub(super) width: crate::Bytes, +} + +pub(super) fn get_inner_matrix_data( + module: &Module, + handle: Handle<crate::Type>, +) -> Option<MatrixType> { + match module.types[handle].inner { + TypeInner::Matrix { + columns, + rows, + scalar, + } => Some(MatrixType { + columns, + rows, + width: scalar.width, + }), + TypeInner::Array { base, .. } => get_inner_matrix_data(module, base), + _ => None, + } +} + +/// Returns the matrix data if the access chain starting at `base`: +/// - starts with an expression with resolved type of [`TypeInner::Matrix`] if `direct = true` +/// - contains one or more expressions with resolved type of [`TypeInner::Array`] of [`TypeInner::Matrix`] +/// - ends at an expression with resolved type of [`TypeInner::Struct`] +pub(super) fn get_inner_matrix_of_struct_array_member( + module: &Module, + base: Handle<crate::Expression>, + func_ctx: &back::FunctionCtx<'_>, + direct: bool, +) -> Option<MatrixType> { + let mut mat_data = None; + let mut array_base = None; + + let mut current_base = base; + loop { + let mut resolved = func_ctx.resolve_type(current_base, &module.types); + if let TypeInner::Pointer { base, .. } = *resolved { + resolved = &module.types[base].inner; + }; + + match *resolved { + TypeInner::Matrix { + columns, + rows, + scalar, + } => { + mat_data = Some(MatrixType { + columns, + rows, + width: scalar.width, + }) + } + TypeInner::Array { base, .. } => { + array_base = Some(base); + } + TypeInner::Struct { .. } => { + if let Some(array_base) = array_base { + if direct { + return mat_data; + } else { + return get_inner_matrix_data(module, array_base); + } + } + + break; + } + _ => break, + } + + current_base = match func_ctx.expressions[current_base] { + crate::Expression::Access { base, .. } => base, + crate::Expression::AccessIndex { base, .. } => base, + _ => break, + }; + } + None +} + +/// Returns the matrix data if the access chain starting at `base`: +/// - starts with an expression with resolved type of [`TypeInner::Matrix`] +/// - contains zero or more expressions with resolved type of [`TypeInner::Array`] of [`TypeInner::Matrix`] +/// - ends with an [`Expression::GlobalVariable`](crate::Expression::GlobalVariable) in [`AddressSpace::Uniform`](crate::AddressSpace::Uniform) +fn get_inner_matrix_of_global_uniform( + module: &Module, + base: Handle<crate::Expression>, + func_ctx: &back::FunctionCtx<'_>, +) -> Option<MatrixType> { + let mut mat_data = None; + let mut array_base = None; + + let mut current_base = base; + loop { + let mut resolved = func_ctx.resolve_type(current_base, &module.types); + if let TypeInner::Pointer { base, .. } = *resolved { + resolved = &module.types[base].inner; + }; + + match *resolved { + TypeInner::Matrix { + columns, + rows, + scalar, + } => { + mat_data = Some(MatrixType { + columns, + rows, + width: scalar.width, + }) + } + TypeInner::Array { base, .. } => { + array_base = Some(base); + } + _ => break, + } + + current_base = match func_ctx.expressions[current_base] { + crate::Expression::Access { base, .. } => base, + crate::Expression::AccessIndex { base, .. } => base, + crate::Expression::GlobalVariable(handle) + if module.global_variables[handle].space == crate::AddressSpace::Uniform => + { + return mat_data.or_else(|| { + array_base.and_then(|array_base| get_inner_matrix_data(module, array_base)) + }) + } + _ => break, + }; + } + None +} |