wgsl: handle RayDesc/RayIntersection at the type decl level instead o…

…f an AST constructor
gfx-rs · Feb 23, 2023 · a3787be · a3787be
1 parent 7f165ba
commit a3787be
Show file tree

Hide file tree

Showing 5 changed files with 38 additions and 29 deletions.
diff --git a/src/front/wgsl/lower/construction.rs b/src/front/wgsl/lower/construction.rs
@@ -660,10 +660,6 @@ impl<'source, 'temp> Lowerer<'source, 'temp> {
                 });
                 ConcreteConstructorHandle::Type(ty)
             }
-            ast::ConstructorType::RayDesc => {
-                let ty = ctx.module.generate_ray_desc_type();
-                ConcreteConstructorHandle::Type(ty)
-            }
             ast::ConstructorType::Type(ty) => ConcreteConstructorHandle::Type(ty),
         };
 

diff --git a/src/front/wgsl/lower/mod.rs b/src/front/wgsl/lower/mod.rs
@@ -637,14 +637,21 @@ impl<'source, 'temp> Lowerer<'source, 'temp> {
             module: &mut module,
         };
 
-        for decl in self.index.visit_ordered() {
-            let span = tu.decls.get_span(decl);
-            let decl = &tu.decls[decl];
+        for decl_handle in self.index.visit_ordered() {
+            let span = tu.decls.get_span(decl_handle);
+            let decl = &tu.decls[decl_handle];
+
+            //TODO: find a nicer way?
+            if let Some(dep) = decl.dependencies.iter().find(|dep| dep.ident == "RayDesc") {
+                let ty_handle = ctx.module.generate_ray_desc_type();
+                ctx.globals
+                    .insert(dep.ident, LoweredGlobalDecl::Type(ty_handle));
+            }
 
             match decl.kind {
                 ast::GlobalDeclKind::Fn(ref f) => {
-                    let decl = self.function(f, span, ctx.reborrow())?;
-                    ctx.globals.insert(f.name.name, decl);
+                    let lowered_decl = self.function(f, span, ctx.reborrow())?;
+                    ctx.globals.insert(f.name.name, lowered_decl);
                 }
                 ast::GlobalDeclKind::Var(ref v) => {
                     let ty = self.resolve_ast_type(v.ty, ctx.reborrow())?;
@@ -2302,6 +2309,12 @@ impl<'source, 'temp> Lowerer<'source, 'temp> {
                     },
                 }
             }
+            ast::Type::RayDesc => {
+                return Ok(ctx.module.generate_ray_desc_type());
+            }
+            ast::Type::RayIntersection => {
+                return Ok(ctx.module.generate_ray_intersection_type());
+            }
             ast::Type::User(ref ident) => {
                 return match ctx.globals.get(ident.name) {
                     Some(&LoweredGlobalDecl::Type(handle)) => Ok(handle),

diff --git a/src/front/wgsl/parse/ast.rs b/src/front/wgsl/parse/ast.rs
@@ -231,6 +231,8 @@ pub enum Type<'a> {
     },
     AccelerationStructure,
     RayQuery,
+    RayDesc,
+    RayIntersection,
     BindingArray {
         base: Handle<Type<'a>>,
         size: ArraySize<'a>,
@@ -370,9 +372,6 @@ pub enum ConstructorType<'a> {
         size: ArraySize<'a>,
     },
 
-    /// Ray description.
-    RayDesc,
-
     /// Constructing a value of a known Naga IR type.
     ///
     /// This variant is produced only during lowering, when we have Naga types

diff --git a/src/front/wgsl/parse/mod.rs b/src/front/wgsl/parse/mod.rs
@@ -441,7 +441,6 @@ impl Parser {
                 }))
             }
             "array" => ast::ConstructorType::PartialArray,
-            "RayDesc" => ast::ConstructorType::RayDesc,
             "atomic"
             | "binding_array"
             | "sampler"
@@ -1382,6 +1381,8 @@ impl Parser {
             }
             "acceleration_structure" => ast::Type::AccelerationStructure,
             "ray_query" => ast::Type::RayQuery,
+            "RayDesc" => ast::Type::RayDesc,
+            "RayIntersection" => ast::Type::RayIntersection,
             _ => return Ok(None),
         }))
     }

diff --git a/tests/out/spv/ray-query.spvasm b/tests/out/spv/ray-query.spvasm
@@ -10,12 +10,12 @@ OpMemoryModel Logical GLSL450
 OpEntryPoint GLCompute %29 "main" %21 %23
 OpExecutionMode %29 LocalSize 1 1 1
 OpMemberDecorate %13 0 Offset 0
-OpMemberDecorate %16 0 Offset 0
-OpMemberDecorate %16 1 Offset 4
-OpMemberDecorate %16 2 Offset 8
-OpMemberDecorate %16 3 Offset 12
-OpMemberDecorate %16 4 Offset 16
-OpMemberDecorate %16 5 Offset 32
+OpMemberDecorate %15 0 Offset 0
+OpMemberDecorate %15 1 Offset 4
+OpMemberDecorate %15 2 Offset 8
+OpMemberDecorate %15 3 Offset 12
+OpMemberDecorate %15 4 Offset 16
+OpMemberDecorate %15 5 Offset 32
 OpMemberDecorate %20 0 Offset 0
 OpMemberDecorate %20 1 Offset 4
 OpMemberDecorate %20 2 Offset 8
@@ -49,19 +49,19 @@ OpMemberDecorate %24 0 Offset 0
 %11 = OpConstant  %4  0
 %12 = OpTypeAccelerationStructureNV
 %13 = OpTypeStruct %4
-%14 = OpTypeRayQueryKHR
-%15 = OpTypeVector %7 3
-%16 = OpTypeStruct %4 %4 %7 %7 %15 %15
+%14 = OpTypeVector %7 3
+%15 = OpTypeStruct %4 %4 %7 %7 %14 %14
+%16 = OpTypeRayQueryKHR
 %17 = OpTypeVector %7 2
 %18 = OpTypeBool
-%19 = OpTypeMatrix %15 4
+%19 = OpTypeMatrix %14 4
 %20 = OpTypeStruct %4 %7 %4 %4 %4 %4 %4 %17 %18 %19 %19
 %22 = OpTypePointer UniformConstant %12
 %21 = OpVariable  %22  UniformConstant
 %24 = OpTypeStruct %13
 %25 = OpTypePointer StorageBuffer %24
 %23 = OpVariable  %25  StorageBuffer
-%27 = OpTypePointer Function %14
+%27 = OpTypePointer Function %16
 %30 = OpTypeFunction %2
 %32 = OpTypePointer StorageBuffer %13
 %45 = OpConstant  %4  1
@@ -73,15 +73,15 @@ OpMemberDecorate %24 0 Offset 0
 %33 = OpAccessChain  %32  %23 %11
 OpBranch %34
 %34 = OpLabel
-%35 = OpCompositeConstruct  %15  %9 %9 %9
-%36 = OpCompositeConstruct  %15  %9 %10 %9
-%37 = OpCompositeConstruct  %16  %3 %5 %6 %8 %35 %36
+%35 = OpCompositeConstruct  %14  %9 %9 %9
+%36 = OpCompositeConstruct  %14  %9 %10 %9
+%37 = OpCompositeConstruct  %15  %3 %5 %6 %8 %35 %36
 %38 = OpCompositeExtract  %4  %37 0
 %39 = OpCompositeExtract  %4  %37 1
 %40 = OpCompositeExtract  %7  %37 2
 %41 = OpCompositeExtract  %7  %37 3
-%42 = OpCompositeExtract  %15  %37 4
-%43 = OpCompositeExtract  %15  %37 5
+%42 = OpCompositeExtract  %14  %37 4
+%43 = OpCompositeExtract  %14  %37 5
 OpRayQueryInitializeKHR %26 %31 %38 %39 %42 %40 %43 %41
 %44 = OpRayQueryProceedKHR  %18  %26
 %46 = OpRayQueryGetIntersectionTypeKHR  %4  %26 %45