add raycast special types (backport to 0.17)

Author: robtfm

src/compose/mod.rs

@@ -134,7 +134,7 @@ use tracing::{debug, trace};
 
 use crate::{
     compose::preprocess::{PreprocessOutput, PreprocessorMetaData},
-    derive::DerivedModule,
+    derive::{DerivedModule, RequiredSpecialTypes},
     redirect::Redirector,
 };
 
@@ -246,6 +246,8 @@ pub struct ComposableModule {
     header_ir: naga::Module,
     // character offset of the start of the owned module string
     start_offset: usize,
+    // any required special types for this module
+    required_special_types: RequiredSpecialTypes,
 }
 
 // data used to build a ComposableModule
@@ -1095,12 +1097,22 @@ impl Composer {
         module_builder.set_shader_source(&source_ir, 0);
         header_builder.set_shader_source(&source_ir, 0);
 
+        let mut special_types: HashSet<&naga::Handle<naga::Type>> = HashSet::new();
+        special_types.extend(source_ir.special_types.predeclared_types.values());
+        special_types.extend(
+            [
+                source_ir.special_types.ray_desc.as_ref(),
+                source_ir.special_types.ray_intersection.as_ref(),
+            ]
+            .iter()
+            .flatten(),
+        );
+
         let mut owned_types = HashSet::new();
         for (h, ty) in source_ir.types.iter() {
             if let Some(name) = &ty.name {
-                // we need to exclude autogenerated struct names, i.e. those that begin with "__"
-                // "__" is a reserved prefix for naga so user variables cannot use it.
-                if !name.contains(DECORATION_PRE) && !name.starts_with("__") {
+                // we exclude any special types, these are added back later
+                if !name.contains(DECORATION_PRE) && !special_types.contains(&h) {
                     let name = format!("{name}{module_decoration}");
                     owned_types.insert(name.clone());
                     // copy and rename types
@@ -1165,10 +1177,12 @@ impl Composer {
             }
         }
 
+        let required_special_types = module_builder.get_required_special_types();
         let module_ir = module_builder.into_module_with_entrypoints();
         let mut header_ir: naga::Module = header_builder.into();
 
-        if self.validate && create_headers {
+        // note: we cannot validate when special types are used, as writeback isn't supported
+        if self.validate && create_headers && required_special_types.is_empty() {
             // check that identifiers haven't been renamed
             #[allow(clippy::single_element_loop)]
             for language in [
@@ -1202,6 +1216,7 @@ impl Composer {
             module_ir,
             header_ir,
             start_offset,
+            required_special_types,
         };
 
         Ok(composable_module)
@@ -1285,6 +1300,8 @@ impl Composer {
             }
         }
 
+        derived.add_required_special_types(composable.required_special_types);
+
         derived.clear_shader_source();
     }
 
@@ -1792,11 +1809,15 @@ impl Composer {
             });
         }
 
+        let required_special_types = derived.get_required_special_types();
+
         let mut naga_module = naga::Module {
             entry_points,
             ..derived.into()
         };
 
+        required_special_types.generate(&mut naga_module);
+
         // apply overrides
         if !composable.override_functions.is_empty() {
             let mut redirect = Redirector::new(naga_module);

src/compose/test.rs

@@ -1263,6 +1263,45 @@ mod test {
         output_eq!(wgsl, "tests/expected/atomics.txt");
     }
 
+    #[test]
+    fn test_raycasts() {
+        let mut composer =
+            Composer::default().with_capabilities(naga::valid::Capabilities::RAY_QUERY);
+
+        composer
+            .add_composable_module(ComposableModuleDescriptor {
+                source: include_str!("tests/raycast/mod.wgsl"),
+                file_path: "tests/raycast/mod.wgsl",
+                ..Default::default()
+            })
+            .unwrap();
+
+        let _module = composer
+            .make_naga_module(NagaModuleDescriptor {
+                source: include_str!("tests/raycast/top.wgsl"),
+                file_path: "tests/raycast/top.wgsl",
+                ..Default::default()
+            })
+            .unwrap();
+
+        // writeback doesn't work for raycast structures
+        // at least we can test that it compiles
+
+        // let info = composer.create_validator().validate(&module).unwrap();
+        // let wgsl = naga::back::wgsl::write_string(
+        //     &module,
+        //     &info,
+        //     naga::back::wgsl::WriterFlags::EXPLICIT_TYPES,
+        // )
+        // .unwrap();
+
+        // let mut f = std::fs::File::create("raycast.txt").unwrap();
+        // f.write_all(wgsl.as_bytes()).unwrap();
+        // drop(f);
+
+        // output_eq!(wgsl, "tests/expected/raycast.txt");
+    }
+
     #[test]
     #[should_panic] // Diagnostic filters not yet supported
     fn test_diagnostic_filters() {

src/compose/tests/raycast/mod.wgsl

@@ -0,0 +1,13 @@
+#define_import_path test_module
+
+@group(0) @binding(0) var tlas: acceleration_structure;
+
+const RAY_NO_CULL = 0xFFu;
+
+fn ray_func() -> RayIntersection {
+    let ray = RayDesc(0u, RAY_NO_CULL, 0.0001, 100000.0, vec3<f32>(0.0, 0.0, 0.0), vec3<f32>(1.0, 0.0, 0.0));
+    var rq: ray_query;
+    rayQueryInitialize(&rq, tlas, ray);
+    rayQueryProceed(&rq);
+    return rayQueryGetCommittedIntersection(&rq);
+}

src/compose/tests/raycast/top.wgsl

@@ -0,0 +1,6 @@
+#import test_module
+
+fn main() -> f32 {
+    let ray = test_module::ray_func();
+    return ray.t;
+}

src/derive.rs

@@ -6,6 +6,35 @@ use naga::{
 };
 use std::{cell::RefCell, rc::Rc};
 
+#[derive(Clone, Copy, Debug, Default)]
+pub struct RequiredSpecialTypes {
+    pub ray_query: bool,
+    pub ray_intersection: bool,
+}
+
+impl std::ops::BitOrAssign for RequiredSpecialTypes {
+    fn bitor_assign(&mut self, rhs: Self) {
+        self.ray_query |= rhs.ray_query;
+        self.ray_intersection |= rhs.ray_intersection;
+    }
+}
+
+impl RequiredSpecialTypes {
+    pub fn generate(&self, target: &mut naga::Module) {
+        if self.ray_query {
+            target.generate_ray_desc_type();
+        }
+
+        if self.ray_intersection {
+            target.generate_ray_intersection_type();
+        }
+    }
+
+    pub fn is_empty(&self) -> bool {
+        !self.ray_intersection && !self.ray_query
+    }
+}
+
 #[derive(Debug, Default)]
 pub struct DerivedModule<'a> {
     shader: Option<&'a Module>,
@@ -30,6 +59,7 @@ pub struct DerivedModule<'a> {
     globals: Arena<GlobalVariable>,
     functions: Arena<Function>,
     pipeline_overrides: Arena<Override>,
+    required_special_types: RequiredSpecialTypes,
 }
 
 impl<'a> DerivedModule<'a> {
@@ -397,10 +427,14 @@ impl<'a> DerivedModule<'a> {
                             naga::RayQueryFunction::Initialize {
                                 acceleration_structure,
                                 descriptor,
-                            } => naga::RayQueryFunction::Initialize {
-                                acceleration_structure: map_expr!(acceleration_structure),
-                                descriptor: map_expr!(descriptor),
-                            },
+                            } => {
+                                // record the use of ray queries, to later add to the final module
+                                self.required_special_types.ray_query = true;
+                                naga::RayQueryFunction::Initialize {
+                                    acceleration_structure: map_expr!(acceleration_structure),
+                                    descriptor: map_expr!(descriptor),
+                                }
+                            }
                             naga::RayQueryFunction::Proceed { result } => {
                                 naga::RayQueryFunction::Proceed {
                                     result: map_expr!(result),
@@ -689,6 +723,8 @@ impl<'a> DerivedModule<'a> {
             }
             Expression::RayQueryProceedResult => expr.clone(),
             Expression::RayQueryGetIntersection { query, committed } => {
+                // record use of the intersection type
+                self.required_special_types.ray_intersection = true;
                 Expression::RayQueryGetIntersection {
                     query: map_expr!(query),
                     committed: *committed,
@@ -831,6 +867,16 @@ impl<'a> DerivedModule<'a> {
         self.import_function(func, span)
     }
 
+    /// get any required special types for this module
+    pub fn get_required_special_types(&self) -> RequiredSpecialTypes {
+        self.required_special_types
+    }
+
+    /// add required special types for this module
+    pub fn add_required_special_types(&mut self, types: RequiredSpecialTypes) {
+        self.required_special_types |= types;
+    }
+
     pub fn into_module_with_entrypoints(mut self) -> naga::Module {
         let entry_points = self
             .shader