[wgsl-in] Handle modf and frexp

src/back/glsl/keywords.rs

@@ -477,4 +477,9 @@ pub const RESERVED_KEYWORDS: &[&str] = &[
     // entry point name (should not be shadowed)
     //
     "main",
+    // Naga utilities:
+    super::FREXP_FUNCTION,
+    crate::front::type_gen::FREXP_RESULT_F32_STRUCT_NAME,
+    super::MODF_FUNCTION,
+    crate::front::type_gen::MODF_RESULT_F32_STRUCT_NAME,
 ];

src/back/glsl/mod.rs

@@ -72,6 +72,9 @@ pub const SUPPORTED_ES_VERSIONS: &[u16] = &[300, 310, 320];
 /// of detail for bounds checking in `ImageLoad`
 const CLAMPED_LOD_SUFFIX: &str = "_clamped_lod";
 
+pub(crate) const FREXP_FUNCTION: &str = "naga_frexp";
+pub(crate) const MODF_FUNCTION: &str = "naga_modf";
+
 /// Mapping between resources and bindings.
 pub type BindingMap = std::collections::BTreeMap<crate::ResourceBinding, u8>;
 
@@ -625,6 +628,34 @@ impl<'a, W: Write> Writer<'a, W> {
             }
         }
 
+        // Write special
+        if let Some(ty_struct) = self.module.special_types.frexp_result {
+            let struct_name = &self.names[&NameKey::Type(ty_struct)];
+            writeln!(self.out)?;
+            writeln!(
+                self.out,
+                "{struct_name} {FREXP_FUNCTION}(float arg) {{
+    int exp;
+    float fract = frexp(arg, exp);
+    return {struct_name}(fract, exp);
+}}"
+            )?;
+        }
+
+        if let Some(ty_struct) = self.module.special_types.modf_result {
+            let struct_name = &self.names[&NameKey::Type(ty_struct)];
+            writeln!(self.out)?;
+            writeln!(
+                self.out,
+                "{} {MODF_FUNCTION}(float arg) {{
+    float whole;
+    float fract = modf(arg, whole);
+    return {}(fract, whole);
+}}",
+                struct_name, struct_name
+            )?;
+        }
+
         // Write all named constants
         let mut constants = self
             .module
@@ -2985,8 +3016,8 @@ impl<'a, W: Write> Writer<'a, W> {
                     Mf::Round => "roundEven",
                     Mf::Fract => "fract",
                     Mf::Trunc => "trunc",
-                    Mf::Modf => "modf",
-                    Mf::Frexp => "frexp",
+                    Mf::Modf => MODF_FUNCTION,
+                    Mf::Frexp => FREXP_FUNCTION,
                     Mf::Ldexp => "ldexp",
                     // exponent
                     Mf::Exp => "exp",

src/back/hlsl/help.rs

@@ -781,6 +781,42 @@ impl<'a, W: Write> super::Writer<'a, W> {
         Ok(())
     }
 
+    pub(super) fn write_special_functions(&mut self, module: &crate::Module) -> BackendResult {
+        if let Some(ty_struct) = module.special_types.frexp_result {
+            let struct_name = &self.names[&NameKey::Type(ty_struct)];
+            writeln!(
+                self.out,
+                "{struct_name} {}(in float arg) {{
+    float exp;
+    float fract = frexp(arg, exp);
+    {struct_name} result;
+    result.exp = exp;
+    result.fract = fract;
+    return result;
+}}",
+                super::writer::FREXP_FUNCTION
+            )?;
+            writeln!(self.out)?;
+        }
+        if let Some(ty_struct) = module.special_types.modf_result {
+            let struct_name = &self.names[&NameKey::Type(ty_struct)];
+            writeln!(
+                self.out,
+                "{struct_name} {}(in float arg) {{
+    float whole;
+    float fract = modf(arg, whole);
+    {struct_name} result;
+    result.whole = whole;
+    result.fract = fract;
+    return result;
+}}",
+                super::writer::MODF_FUNCTION,
+            )?;
+            writeln!(self.out)?;
+        }
+        Ok(())
+    }
+
     /// Helper function that writes compose wrapped functions
     pub(super) fn write_wrapped_compose_functions(
         &mut self,

src/back/hlsl/keywords.rs

@@ -814,6 +814,11 @@ pub const RESERVED: &[&str] = &[
     "TextureBuffer",
     "ConstantBuffer",
     "RayQuery",
+    // Naga utilities
+    super::writer::FREXP_FUNCTION,
+    crate::front::type_gen::FREXP_RESULT_F32_STRUCT_NAME,
+    super::writer::MODF_FUNCTION,
+    crate::front::type_gen::MODF_RESULT_F32_STRUCT_NAME,
 ];
 
 // DXC scalar types, from https://github.com/microsoft/DirectXShaderCompiler/blob/18c9e114f9c314f93e68fbc72ce207d4ed2e65ae/tools/clang/lib/AST/ASTContextHLSL.cpp#L48-L254

src/back/hlsl/writer.rs

@@ -17,6 +17,9 @@ const SPECIAL_BASE_VERTEX: &str = "base_vertex";
 const SPECIAL_BASE_INSTANCE: &str = "base_instance";
 const SPECIAL_OTHER: &str = "other";
 
+pub(crate) const FREXP_FUNCTION: &str = "naga_frexp";
+pub(crate) const MODF_FUNCTION: &str = "naga_modf";
+
 struct EpStructMember {
     name: String,
     ty: Handle<crate::Type>,
@@ -215,7 +218,10 @@ impl<'a, W: fmt::Write> super::Writer<'a, W> {
 
         // Write all structs
         for (handle, ty) in module.types.iter() {
-            if let TypeInner::Struct { ref members, span } = ty.inner {
+            if let TypeInner::Struct {
+                ref members, span, ..
+            } = ty.inner
+            {
                 if module.types[members.last().unwrap().ty]
                     .inner
                     .is_dynamically_sized(&module.types)
@@ -244,6 +250,8 @@ impl<'a, W: fmt::Write> super::Writer<'a, W> {
             }
         }
 
+        self.write_special_functions(module)?;
+
         self.write_wrapped_compose_functions(module, &module.const_expressions)?;
 
         // Write all named constants
@@ -2665,8 +2673,8 @@ impl<'a, W: fmt::Write> super::Writer<'a, W> {
                     Mf::Round => Function::Regular("round"),
                     Mf::Fract => Function::Regular("frac"),
                     Mf::Trunc => Function::Regular("trunc"),
-                    Mf::Modf => Function::Regular("modf"),
-                    Mf::Frexp => Function::Regular("frexp"),
+                    Mf::Modf => Function::Regular(MODF_FUNCTION),
+                    Mf::Frexp => Function::Regular(FREXP_FUNCTION),
                     Mf::Ldexp => Function::Regular("ldexp"),
                     // exponent
                     Mf::Exp => Function::Regular("exp"),

src/back/msl/keywords.rs

@@ -214,4 +214,8 @@ pub const RESERVED: &[&str] = &[
     // Naga utilities
     "DefaultConstructible",
     "clamped_lod_e",
+    super::writer::FREXP_FUNCTION,
+    crate::front::type_gen::FREXP_RESULT_F32_STRUCT_NAME,
+    super::writer::MODF_FUNCTION,
+    crate::front::type_gen::MODF_RESULT_F32_STRUCT_NAME,
 ];

src/back/msl/writer.rs

@@ -32,6 +32,9 @@ const RAY_QUERY_FIELD_INTERSECTION: &str = "intersection";
 const RAY_QUERY_FIELD_READY: &str = "ready";
 const RAY_QUERY_FUN_MAP_INTERSECTION: &str = "_map_intersection_type";
 
+pub(crate) const FREXP_FUNCTION: &str = "naga_frexp";
+pub(crate) const MODF_FUNCTION: &str = "naga_modf";
+
 /// Write the Metal name for a Naga numeric type: scalar, vector, or matrix.
 ///
 /// The `sizes` slice determines whether this function writes a
@@ -1678,8 +1681,8 @@ impl<W: Write> Writer<W> {
                     Mf::Round => "rint",
                     Mf::Fract => "fract",
                     Mf::Trunc => "trunc",
-                    Mf::Modf => "modf",
-                    Mf::Frexp => "frexp",
+                    Mf::Modf => MODF_FUNCTION,
+                    Mf::Frexp => FREXP_FUNCTION,
                     Mf::Ldexp => "ldexp",
                     // exponent
                     Mf::Exp => "exp",
@@ -1813,6 +1816,12 @@ impl<W: Write> Writer<W> {
                     write!(self.out, "((")?;
                     self.put_expression(arg, context, false)?;
                     write!(self.out, ") * 57.295779513082322865)")?;
+                } else if fun == Mf::Modf || fun == Mf::Frexp {
+                    write!(self.out, "{fun_name}")?;
+                    self.put_call_parameters(
+                        iter::once(arg).chain(arg1).chain(arg2).chain(arg3),
+                        context,
+                    )?;
                 } else {
                     write!(self.out, "{NAMESPACE}::{fun_name}")?;
                     self.put_call_parameters(
@@ -3236,6 +3245,31 @@ impl<W: Write> Writer<W> {
                 }
             }
         }
+
+        if let Some(struct_ty) = module.special_types.frexp_result {
+            let struct_name = &self.names[&NameKey::Type(struct_ty)];
+            writeln!(
+                self.out,
+                "struct {struct_name} {FREXP_FUNCTION}(float arg) {{
+    int exp;
+    float fract = {NAMESPACE}::frexp(arg, exp);
+    return {struct_name}{{ fract, exp }};
+}};"
+            )?;
+        }
+
+        if let Some(struct_ty) = module.special_types.modf_result {
+            let struct_name = &self.names[&NameKey::Type(struct_ty)];
+            writeln!(
+                self.out,
+                "struct {struct_name} {MODF_FUNCTION}(float arg) {{
+    float whole;
+    float fract = {NAMESPACE}::modf(arg, whole);
+    return {struct_name}{{ fract, whole }};
+}};"
+            )?;
+        }
+
         Ok(())
     }
 

src/back/spv/block.rs

@@ -787,8 +787,8 @@ impl<'w> BlockContext<'w> {
                     Mf::Floor => MathOp::Ext(spirv::GLOp::Floor),
                     Mf::Fract => MathOp::Ext(spirv::GLOp::Fract),
                     Mf::Trunc => MathOp::Ext(spirv::GLOp::Trunc),
-                    Mf::Modf => MathOp::Ext(spirv::GLOp::Modf),
-                    Mf::Frexp => MathOp::Ext(spirv::GLOp::Frexp),
+                    Mf::Modf => MathOp::Ext(spirv::GLOp::ModfStruct),
+                    Mf::Frexp => MathOp::Ext(spirv::GLOp::FrexpStruct),
                     Mf::Ldexp => MathOp::Ext(spirv::GLOp::Ldexp),
                     // geometry
                     Mf::Dot => match *self.fun_info[arg].ty.inner_with(&self.ir_module.types) {

src/back/spv/writer.rs

@@ -241,6 +241,16 @@ impl Writer {
         self.get_type_id(local_type.into())
     }
 
+    pub(super) fn get_int_type_id(&mut self) -> Word {
+        let local_type = LocalType::Value {
+            vector_size: None,
+            kind: crate::ScalarKind::Sint,
+            width: 4,
+            pointer_space: None,
+        };
+        self.get_type_id(local_type.into())
+    }
+
     pub(super) fn get_float_type_id(&mut self) -> Word {
         let local_type = LocalType::Value {
             vector_size: None,
@@ -1693,26 +1703,45 @@ impl Writer {
         index: usize,
         member: &crate::StructMember,
         arena: &UniqueArena<crate::Type>,
+    ) -> Result<(), Error> {
+        self.decorate_struct_member_raw(
+            struct_id,
+            index,
+            member.offset,
+            member.name.as_deref(),
+            &arena[member.ty].inner,
+            arena,
+        )
+    }
+
+    fn decorate_struct_member_raw(
+        &mut self,
+        struct_id: Word,
+        index: usize,
+        offset: Word,
+        member_name: Option<&str>,
+        member_ty_inner: &crate::TypeInner,
+        arena: &UniqueArena<crate::Type>,
     ) -> Result<(), Error> {
         use spirv::Decoration;
 
         self.annotations.push(Instruction::member_decorate(
             struct_id,
             index as u32,
             Decoration::Offset,
-            &[member.offset],
+            &[offset],
         ));
 
         if self.flags.contains(WriterFlags::DEBUG) {
-            if let Some(ref name) = member.name {
+            if let Some(name) = member_name {
                 self.debugs
                     .push(Instruction::member_name(struct_id, index as u32, name));
             }
         }
 
         // Matrices and arrays of matrices both require decorations,
         // so "see through" an array to determine if they're needed.
-        let member_array_subty_inner = match arena[member.ty].inner {
+        let member_array_subty_inner = match *member_ty_inner {
             crate::TypeInner::Array { base, .. } => &arena[base].inner,
             ref other => other,
         };

src/front/mod.rs

@@ -3,7 +3,7 @@ Frontend parsers that consume binary and text shaders and load them into [`Modul
 */
 
 mod interpolator;
-mod type_gen;
+pub(crate) mod type_gen;
 
 #[cfg(feature = "glsl-in")]
 pub mod glsl;