Ray query expressions and special types

src/back/dot/mod.rs

@@ -566,6 +566,12 @@ fn write_function_expressions(
                 edges.insert("", expr);
                 ("ArrayLength".into(), 7)
             }
+            E::RayQueryProceedResult => ("rayQueryProceedResult".into(), 4),
+            E::RayQueryGetIntersection { query, committed } => {
+                edges.insert("", query);
+                let ty = if committed { "Committed" } else { "Candidate" };
+                (format!("rayQueryGet{}Intersection", ty).into(), 4)
+            }
         };
 
         // give uniform expressions an outline

src/back/glsl/mod.rs

@@ -3172,13 +3172,17 @@ impl<'a, W: Write> Writer<'a, W> {
                 }
             }
             // These expressions never show up in `Emit`.
-            Expression::CallResult(_) | Expression::AtomicResult { .. } => unreachable!(),
+            Expression::CallResult(_)
+            | Expression::AtomicResult { .. }
+            | Expression::RayQueryProceedResult => unreachable!(),
             // `ArrayLength` is written as `expr.length()` and we convert it to a uint
             Expression::ArrayLength(expr) => {
                 write!(self.out, "uint(")?;
                 self.write_expr(expr, ctx)?;
                 write!(self.out, ".length())")?
             }
+            // not supported yet
+            Expression::RayQueryGetIntersection { .. } => unreachable!(),
         }
 
         Ok(())

src/back/hlsl/writer.rs

@@ -2825,8 +2825,12 @@ impl<'a, W: fmt::Write> super::Writer<'a, W> {
                 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::CallResult(_)
+            | Expression::AtomicResult { .. }
+            | Expression::RayQueryProceedResult => {}
         }
 
         if !closing_bracket.is_empty() {

src/back/msl/writer.rs

@@ -1799,7 +1799,9 @@ impl<W: Write> Writer<W> {
                 _ => return Err(Error::Validation),
             },
             // has to be a named expression
-            crate::Expression::CallResult(_) | crate::Expression::AtomicResult { .. } => {
+            crate::Expression::CallResult(_)
+            | crate::Expression::AtomicResult { .. }
+            | crate::Expression::RayQueryProceedResult => {
                 unreachable!()
             }
             crate::Expression::ArrayLength(expr) => {
@@ -1824,6 +1826,8 @@ impl<W: Write> Writer<W> {
                     write!(self.out, ")")?;
                 }
             }
+            // hot supported yet
+            crate::Expression::RayQueryGetIntersection { .. } => unreachable!(),
         }
         Ok(())
     }

src/back/spv/block.rs

@@ -1364,6 +1364,10 @@ impl<'w> BlockContext<'w> {
                 id
             }
             crate::Expression::ArrayLength(expr) => self.write_runtime_array_length(expr, block)?,
+            //TODO
+            crate::Expression::RayQueryProceedResult => unreachable!(),
+            //TODO
+            crate::Expression::RayQueryGetIntersection { .. } => unreachable!(),
         };
 
         self.cached[expr_handle] = id;

src/back/wgsl/writer.rs

@@ -1685,8 +1685,12 @@ impl<W: Write> Writer<W> {
 
                 write!(self.out, ")")?
             }
+            // Not supported yet
+            Expression::RayQueryGetIntersection { .. } => unreachable!(),
             // Nothing to do here, since call expression already cached
-            Expression::CallResult(_) | Expression::AtomicResult { .. } => {}
+            Expression::CallResult(_)
+            | Expression::AtomicResult { .. }
+            | Expression::RayQueryProceedResult => {}
         }
 
         Ok(())

src/front/glsl/constants.rs

@@ -37,6 +37,8 @@ pub enum ConstantSolvingError {
     Load,
     #[error("Constants don't support image expressions")]
     ImageExpression,
+    #[error("Constants don't support ray query expressions")]
+    RayQueryExpression,
     #[error("Cannot access the type")]
     InvalidAccessBase,
     #[error("Cannot access at the index")]
@@ -295,6 +297,9 @@ impl<'a> ConstantSolver<'a> {
             Expression::ImageSample { .. }
             | Expression::ImageLoad { .. }
             | Expression::ImageQuery { .. } => Err(ConstantSolvingError::ImageExpression),
+            Expression::RayQueryProceedResult | Expression::RayQueryGetIntersection { .. } => {
+                Err(ConstantSolvingError::RayQueryExpression)
+            }
         }
     }
 

src/front/glsl/types.rs

@@ -246,14 +246,7 @@ impl Frontend {
         expr: Handle<Expression>,
         meta: Span,
     ) -> Result<()> {
-        let resolve_ctx = ResolveContext {
-            constants: &self.module.constants,
-            types: &self.module.types,
-            global_vars: &self.module.global_variables,
-            local_vars: &ctx.locals,
-            functions: &self.module.functions,
-            arguments: &ctx.arguments,
-        };
+        let resolve_ctx = ResolveContext::with_locals(&self.module, &ctx.locals, &ctx.arguments);
 
         ctx.typifier
             .grow(expr, &ctx.expressions, &resolve_ctx)
@@ -312,14 +305,7 @@ impl Frontend {
         expr: Handle<Expression>,
         meta: Span,
     ) -> Result<()> {
-        let resolve_ctx = ResolveContext {
-            constants: &self.module.constants,
-            types: &self.module.types,
-            global_vars: &self.module.global_variables,
-            local_vars: &ctx.locals,
-            functions: &self.module.functions,
-            arguments: &ctx.arguments,
-        };
+        let resolve_ctx = ResolveContext::with_locals(&self.module, &ctx.locals, &ctx.arguments);
 
         ctx.typifier
             .invalidate(expr, &ctx.expressions, &resolve_ctx)

src/front/mod.rs

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

src/front/type_gen.rs

@@ -0,0 +1,153 @@
+/*!
+Type generators.
+*/
+
+use crate::{arena::Handle, span::Span};
+
+impl crate::Module {
+    pub(super) fn generate_ray_desc_type(&mut self) -> Handle<crate::Type> {
+        if let Some(handle) = self.special_types.ray_desc {
+            return handle;
+        }
+
+        let width = 4;
+        let ty_flag = self.types.insert(
+            crate::Type {
+                name: None,
+                inner: crate::TypeInner::Scalar {
+                    width,
+                    kind: crate::ScalarKind::Uint,
+                },
+            },
+            Span::UNDEFINED,
+        );
+        let ty_scalar = self.types.insert(
+            crate::Type {
+                name: None,
+                inner: crate::TypeInner::Scalar {
+                    width,
+                    kind: crate::ScalarKind::Float,
+                },
+            },
+            Span::UNDEFINED,
+        );
+        let ty_vector = self.types.insert(
+            crate::Type {
+                name: None,
+                inner: crate::TypeInner::Vector {
+                    size: crate::VectorSize::Tri,
+                    kind: crate::ScalarKind::Float,
+                    width,
+                },
+            },
+            Span::UNDEFINED,
+        );
+
+        let handle = self.types.insert(
+            crate::Type {
+                name: Some("RayDesc".to_string()),
+                inner: crate::TypeInner::Struct {
+                    members: vec![
+                        crate::StructMember {
+                            name: Some("flags".to_string()),
+                            ty: ty_flag,
+                            binding: None,
+                            offset: 0,
+                        },
+                        crate::StructMember {
+                            name: Some("cull_mask".to_string()),
+                            ty: ty_flag,
+                            binding: None,
+                            offset: 4,
+                        },
+                        crate::StructMember {
+                            name: Some("tmin".to_string()),
+                            ty: ty_scalar,
+                            binding: None,
+                            offset: 8,
+                        },
+                        crate::StructMember {
+                            name: Some("tmax".to_string()),
+                            ty: ty_scalar,
+                            binding: None,
+                            offset: 12,
+                        },
+                        crate::StructMember {
+                            name: Some("origin".to_string()),
+                            ty: ty_vector,
+                            binding: None,
+                            offset: 16,
+                        },
+                        crate::StructMember {
+                            name: Some("dir".to_string()),
+                            ty: ty_vector,
+                            binding: None,
+                            offset: 32,
+                        },
+                    ],
+                    span: 48,
+                },
+            },
+            Span::UNDEFINED,
+        );
+
+        self.special_types.ray_desc = Some(handle);
+        handle
+    }
+
+    pub(super) fn generate_ray_intersection_type(&mut self) -> Handle<crate::Type> {
+        if let Some(handle) = self.special_types.ray_intersection {
+            return handle;
+        }
+
+        let width = 4;
+        let ty_flag = self.types.insert(
+            crate::Type {
+                name: None,
+                inner: crate::TypeInner::Scalar {
+                    width,
+                    kind: crate::ScalarKind::Uint,
+                },
+            },
+            Span::UNDEFINED,
+        );
+        let ty_scalar = self.types.insert(
+            crate::Type {
+                name: None,
+                inner: crate::TypeInner::Scalar {
+                    width,
+                    kind: crate::ScalarKind::Float,
+                },
+            },
+            Span::UNDEFINED,
+        );
+
+        let handle = self.types.insert(
+            crate::Type {
+                name: Some("RayIntersection".to_string()),
+                inner: crate::TypeInner::Struct {
+                    members: vec![
+                        crate::StructMember {
+                            name: Some("kind".to_string()),
+                            ty: ty_flag,
+                            binding: None,
+                            offset: 0,
+                        },
+                        crate::StructMember {
+                            name: Some("t".to_string()),
+                            ty: ty_scalar,
+                            binding: None,
+                            offset: 4,
+                        },
+                        //TODO: the rest
+                    ],
+                    span: 8,
+                },
+            },
+            Span::UNDEFINED,
+        );
+
+        self.special_types.ray_intersection = Some(handle);
+        handle
+    }
+}

src/front/wgsl/lower/construction.rs

@@ -660,6 +660,10 @@ 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),
         };
 

src/front/wgsl/lower/mod.rs

@@ -234,14 +234,8 @@ impl<'a> ExpressionContext<'a, '_, '_> {
     /// [`self.resolved_inner(handle)`]: ExpressionContext::resolved_inner
     /// [`Typifier`]: Typifier
     fn grow_types(&mut self, handle: Handle<crate::Expression>) -> Result<&mut Self, Error<'a>> {
-        let resolve_ctx = ResolveContext {
-            constants: &self.module.constants,
-            types: &self.module.types,
-            global_vars: &self.module.global_variables,
-            local_vars: self.local_vars,
-            functions: &self.module.functions,
-            arguments: self.arguments,
-        };
+        let resolve_ctx =
+            ResolveContext::with_locals(&self.module, self.local_vars, self.arguments);
         self.typifier
             .grow(handle, self.naga_expressions, &resolve_ctx)
             .map_err(Error::InvalidResolve)?;
@@ -1919,6 +1913,54 @@ impl<'source, 'temp> Lowerer<'source, 'temp> {
                                 query: crate::ImageQuery::NumSamples,
                             }
                         }
+                        "rayQueryInitialize" => {
+                            let mut args = ctx.prepare_args(arguments, 3, span);
+                            let query = self.expression(args.next()?, ctx.reborrow())?;
+                            let acceleration_structure =
+                                self.expression(args.next()?, ctx.reborrow())?;
+                            let descriptor = self.expression(args.next()?, ctx.reborrow())?;
+                            args.finish()?;
+
+                            let _ = ctx.module.generate_ray_desc_type();
+                            let fun = crate::RayQueryFunction::Initialize {
+                                acceleration_structure,
+                                descriptor,
+                            };
+
+                            ctx.block.extend(ctx.emitter.finish(ctx.naga_expressions));
+                            ctx.emitter.start(ctx.naga_expressions);
+                            ctx.block
+                                .push(crate::Statement::RayQuery { query, fun }, span);
+                            return Ok(None);
+                        }
+                        "rayQueryProceed" => {
+                            let mut args = ctx.prepare_args(arguments, 1, span);
+                            let query = self.expression(args.next()?, ctx.reborrow())?;
+                            args.finish()?;
+
+                            let fun = crate::RayQueryFunction::Proceed;
+
+                            ctx.block.extend(ctx.emitter.finish(ctx.naga_expressions));
+                            let result = ctx
+                                .naga_expressions
+                                .append(crate::Expression::RayQueryProceedResult, span);
+                            ctx.emitter.start(ctx.naga_expressions);
+                            ctx.block
+                                .push(crate::Statement::RayQuery { query, fun }, span);
+                            return Ok(Some(result));
+                        }
+                        "rayQueryGetCommittedIntersection" => {
+                            let mut args = ctx.prepare_args(arguments, 1, span);
+                            let query = self.expression(args.next()?, ctx.reborrow())?;
+                            args.finish()?;
+
+                            let _ = ctx.module.generate_ray_intersection_type();
+
+                            crate::Expression::RayQueryGetIntersection {
+                                query,
+                                committed: true,
+                            }
+                        }
                         _ => return Err(Error::UnknownIdent(function.span, function.name)),
                     }
                 };

src/front/wgsl/parse/ast.rs

@@ -370,6 +370,9 @@ 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