[naga wgsl] Impl const_assert (#6198)

sagudev · web-flow · commit 4e9a2a500328 · 2024-09-02T17:37:04.000Z
Signed-off-by: sagudev &lt;16504129+sagudev@users.noreply.github.com&gt;
diff --git a/naga/src/front/wgsl/error.rs b/naga/src/front/wgsl/error.rs
@@ -263,6 +263,8 @@ pub(crate) enum Error<'a> {
         limit: u8,
     },
     PipelineConstantIDValue(Span),
+    NotBool(Span),
+    ConstAssertFailed(Span),
 }
 
 #[derive(Clone, Debug)]
@@ -815,6 +817,22 @@ impl<'a> Error<'a> {
                 )],
                 notes: vec![],
             },
+            Error::NotBool(span) => ParseError {
+                message: "must be a const-expression that resolves to a bool".to_string(),
+                labels: vec![(
+                    span,
+                    "must resolve to bool".into(),
+                )],
+                notes: vec![],
+            },
+            Error::ConstAssertFailed(span) => ParseError {
+                message: "const_assert failure".to_string(),
+                labels: vec![(
+                    span,
+                    "evaluates to false".into(),
+                )],
+                notes: vec![],
+            },
         }
     }
 }
diff --git a/naga/src/front/wgsl/index.rs b/naga/src/front/wgsl/index.rs
@@ -20,13 +20,16 @@ impl<'a> Index<'a> {
         // While doing so, reject conflicting definitions.
         let mut globals = FastHashMap::with_capacity_and_hasher(tu.decls.len(), Default::default());
         for (handle, decl) in tu.decls.iter() {
-            let ident = decl_ident(decl);
-            let name = ident.name;
-            if let Some(old) = globals.insert(name, handle) {
-                return Err(Error::Redefinition {
-                    previous: decl_ident(&tu.decls[old]).span,
-                    current: ident.span,
-                });
+            if let Some(ident) = decl_ident(decl) {
+                let name = ident.name;
+                if let Some(old) = globals.insert(name, handle) {
+                    return Err(Error::Redefinition {
+                        previous: decl_ident(&tu.decls[old])
+                            .expect("decl should have ident for redefinition")
+                            .span,
+                        current: ident.span,
+                    });
+                }
             }
         }
 
@@ -130,7 +133,7 @@ impl<'a> DependencySolver<'a, '_> {
                     return if dep_id == id {
                         // A declaration refers to itself directly.
                         Err(Error::RecursiveDeclaration {
-                            ident: decl_ident(decl).span,
+                            ident: decl_ident(decl).expect("decl should have ident").span,
                             usage: dep.usage,
                         })
                     } else {
@@ -146,14 +149,19 @@ impl<'a> DependencySolver<'a, '_> {
                             .unwrap_or(0);
 
                         Err(Error::CyclicDeclaration {
-                            ident: decl_ident(&self.module.decls[dep_id]).span,
+                            ident: decl_ident(&self.module.decls[dep_id])
+                                .expect("decl should have ident")
+                                .span,
                             path: self.path[start_at..]
                                 .iter()
                                 .map(|curr_dep| {
                                     let curr_id = curr_dep.decl;
                                     let curr_decl = &self.module.decls[curr_id];
 
-                                    (decl_ident(curr_decl).span, curr_dep.usage)
+                                    (
+                                        decl_ident(curr_decl).expect("decl should have ident").span,
+                                        curr_dep.usage,
+                                    )
                                 })
                                 .collect(),
                         })
@@ -182,13 +190,14 @@ impl<'a> DependencySolver<'a, '_> {
     }
 }
 
-const fn decl_ident<'a>(decl: &ast::GlobalDecl<'a>) -> ast::Ident<'a> {
+const fn decl_ident<'a>(decl: &ast::GlobalDecl<'a>) -> Option<ast::Ident<'a>> {
     match decl.kind {
-        ast::GlobalDeclKind::Fn(ref f) => f.name,
-        ast::GlobalDeclKind::Var(ref v) => v.name,
-        ast::GlobalDeclKind::Const(ref c) => c.name,
-        ast::GlobalDeclKind::Override(ref o) => o.name,
-        ast::GlobalDeclKind::Struct(ref s) => s.name,
-        ast::GlobalDeclKind::Type(ref t) => t.name,
+        ast::GlobalDeclKind::Fn(ref f) => Some(f.name),
+        ast::GlobalDeclKind::Var(ref v) => Some(v.name),
+        ast::GlobalDeclKind::Const(ref c) => Some(c.name),
+        ast::GlobalDeclKind::Override(ref o) => Some(o.name),
+        ast::GlobalDeclKind::Struct(ref s) => Some(s.name),
+        ast::GlobalDeclKind::Type(ref t) => Some(t.name),
+        ast::GlobalDeclKind::ConstAssert(_) => None,
     }
 }
diff --git a/naga/src/front/wgsl/lower/mod.rs b/naga/src/front/wgsl/lower/mod.rs
@@ -1204,6 +1204,20 @@ impl<'source, 'temp> Lowerer<'source, 'temp> {
                     ctx.globals
                         .insert(alias.name.name, LoweredGlobalDecl::Type(ty));
                 }
+                ast::GlobalDeclKind::ConstAssert(condition) => {
+                    let condition = self.expression(condition, &mut ctx.as_const())?;
+
+                    let span = ctx.module.global_expressions.get_span(condition);
+                    match ctx
+                        .module
+                        .to_ctx()
+                        .eval_expr_to_bool_from(condition, &ctx.module.global_expressions)
+                    {
+                        Some(true) => Ok(()),
+                        Some(false) => Err(Error::ConstAssertFailed(span)),
+                        _ => Err(Error::NotBool(span)),
+                    }?;
+                }
             }
         }
 
@@ -1742,6 +1756,28 @@ impl<'source, 'temp> Lowerer<'source, 'temp> {
                     value,
                 }
             }
+            ast::StatementKind::ConstAssert(condition) => {
+                let mut emitter = Emitter::default();
+                emitter.start(&ctx.function.expressions);
+
+                let condition =
+                    self.expression(condition, &mut ctx.as_const(block, &mut emitter))?;
+
+                let span = ctx.function.expressions.get_span(condition);
+                match ctx
+                    .module
+                    .to_ctx()
+                    .eval_expr_to_bool_from(condition, &ctx.function.expressions)
+                {
+                    Some(true) => Ok(()),
+                    Some(false) => Err(Error::ConstAssertFailed(span)),
+                    _ => Err(Error::NotBool(span)),
+                }?;
+
+                block.extend(emitter.finish(&ctx.function.expressions));
+
+                return Ok(());
+            }
             ast::StatementKind::Ignore(expr) => {
                 let mut emitter = Emitter::default();
                 emitter.start(&ctx.function.expressions);
diff --git a/naga/src/front/wgsl/parse/ast.rs b/naga/src/front/wgsl/parse/ast.rs
@@ -85,6 +85,7 @@ pub enum GlobalDeclKind<'a> {
     Override(Override<'a>),
     Struct(Struct<'a>),
     Type(TypeAlias<'a>),
+    ConstAssert(Handle<Expression<'a>>),
 }
 
 #[derive(Debug)]
@@ -284,6 +285,7 @@ pub enum StatementKind<'a> {
     Increment(Handle<Expression<'a>>),
     Decrement(Handle<Expression<'a>>),
     Ignore(Handle<Expression<'a>>),
+    ConstAssert(Handle<Expression<'a>>),
 }
 
 #[derive(Debug)]
diff --git a/naga/src/front/wgsl/parse/mod.rs b/naga/src/front/wgsl/parse/mod.rs
@@ -2016,6 +2016,20 @@ impl Parser {
                         lexer.expect(Token::Separator(';'))?;
                         ast::StatementKind::Kill
                     }
+                    // https://www.w3.org/TR/WGSL/#const-assert-statement
+                    "const_assert" => {
+                        let _ = lexer.next();
+                        // parentheses are optional
+                        let paren = lexer.skip(Token::Paren('('));
+
+                        let condition = self.general_expression(lexer, ctx)?;
+
+                        if paren {
+                            lexer.expect(Token::Paren(')'))?;
+                        }
+                        lexer.expect(Token::Separator(';'))?;
+                        ast::StatementKind::ConstAssert(condition)
+                    }
                     // assignment or a function call
                     _ => {
                         self.function_call_or_assignment_statement(lexer, ctx, block)?;
@@ -2419,6 +2433,18 @@ impl Parser {
                     ..function
                 }))
             }
+            (Token::Word("const_assert"), _) => {
+                // parentheses are optional
+                let paren = lexer.skip(Token::Paren('('));
+
+                let condition = self.general_expression(lexer, &mut ctx)?;
+
+                if paren {
+                    lexer.expect(Token::Paren(')'))?;
+                }
+                lexer.expect(Token::Separator(';'))?;
+                Some(ast::GlobalDeclKind::ConstAssert(condition))
+            }
             (Token::End, _) => return Ok(()),
             other => return Err(Error::Unexpected(other.1, ExpectedToken::GlobalItem)),
         };
diff --git a/naga/src/proc/mod.rs b/naga/src/proc/mod.rs
@@ -674,6 +674,19 @@ impl GlobalCtx<'_> {
         }
     }
 
+    /// Try to evaluate the expression in the `arena` using its `handle` and return it as a `bool`.
+    #[allow(dead_code)]
+    pub(super) fn eval_expr_to_bool_from(
+        &self,
+        handle: crate::Handle<crate::Expression>,
+        arena: &crate::Arena<crate::Expression>,
+    ) -> Option<bool> {
+        match self.eval_expr_to_literal_from(handle, arena) {
+            Some(crate::Literal::Bool(value)) => Some(value),
+            _ => None,
+        }
+    }
+
     #[allow(dead_code)]
     pub(crate) fn eval_expr_to_literal(
         &self,
diff --git a/naga/tests/in/const_assert.wgsl b/naga/tests/in/const_assert.wgsl
@@ -0,0 +1,11 @@
+// Sourced from https://www.w3.org/TR/WGSL/#const-assert-statement
+const x = 1;
+const y = 2;
+const_assert x < y; // valid at module-scope.
+const_assert(y != 0); // parentheses are optional.
+
+fn foo() {
+  const z = x + y - 2;
+  const_assert z > 0; // valid in functions.
+  const_assert(z > 0);
+}
diff --git a/naga/tests/out/ir/const_assert.compact.ron b/naga/tests/out/ir/const_assert.compact.ron
@@ -0,0 +1,54 @@
+(
+    types: [
+        (
+            name: None,
+            inner: Scalar((
+                kind: Sint,
+                width: 4,
+            )),
+        ),
+    ],
+    special_types: (
+        ray_desc: None,
+        ray_intersection: None,
+        predeclared_types: {},
+    ),
+    constants: [
+        (
+            name: Some("x"),
+            ty: 0,
+            init: 0,
+        ),
+        (
+            name: Some("y"),
+            ty: 0,
+            init: 1,
+        ),
+    ],
+    overrides: [],
+    global_variables: [],
+    global_expressions: [
+        Literal(I32(1)),
+        Literal(I32(2)),
+    ],
+    functions: [
+        (
+            name: Some("foo"),
+            arguments: [],
+            result: None,
+            local_variables: [],
+            expressions: [
+                Literal(I32(1)),
+            ],
+            named_expressions: {
+                0: "z",
+            },
+            body: [
+                Return(
+                    value: None,
+                ),
+            ],
+        ),
+    ],
+    entry_points: [],
+)
diff --git a/naga/tests/out/ir/const_assert.ron b/naga/tests/out/ir/const_assert.ron
@@ -0,0 +1,54 @@
+(
+    types: [
+        (
+            name: None,
+            inner: Scalar((
+                kind: Sint,
+                width: 4,
+            )),
+        ),
+    ],
+    special_types: (
+        ray_desc: None,
+        ray_intersection: None,
+        predeclared_types: {},
+    ),
+    constants: [
+        (
+            name: Some("x"),
+            ty: 0,
+            init: 0,
+        ),
+        (
+            name: Some("y"),
+            ty: 0,
+            init: 1,
+        ),
+    ],
+    overrides: [],
+    global_variables: [],
+    global_expressions: [
+        Literal(I32(1)),
+        Literal(I32(2)),
+    ],
+    functions: [
+        (
+            name: Some("foo"),
+            arguments: [],
+            result: None,
+            local_variables: [],
+            expressions: [
+                Literal(I32(1)),
+            ],
+            named_expressions: {
+                0: "z",
+            },
+            body: [
+                Return(
+                    value: None,
+                ),
+            ],
+        ),
+    ],
+    entry_points: [],
+)
diff --git a/naga/tests/out/wgsl/const_assert.wgsl b/naga/tests/out/wgsl/const_assert.wgsl
@@ -0,0 +1,7 @@
+const x: i32 = 1i;
+const y: i32 = 2i;
+
+fn foo() {
+    return;
+}
+
diff --git a/naga/tests/snapshots.rs b/naga/tests/snapshots.rs
@@ -868,6 +868,7 @@ fn convert_wgsl() {
             "const-exprs",
             Targets::SPIRV | Targets::METAL | Targets::GLSL | Targets::HLSL | Targets::WGSL,
         ),
+        ("const_assert", Targets::WGSL | Targets::IR),
         ("separate-entry-points", Targets::SPIRV | Targets::GLSL),
         (
             "struct-layout",
diff --git a/naga/tests/wgsl_errors.rs b/naga/tests/wgsl_errors.rs

Original file line number	Diff line number	Diff line change
`@@ -85,6 +85,7 @@ pub enum GlobalDeclKind<'a> {`
`85`	`85`	`Override(Override<'a>),`
`86`	`86`	`Struct(Struct<'a>),`
`87`	`87`	`Type(TypeAlias<'a>),`
	`88`	`+ ConstAssert(Handle<Expression<'a>>),`
`88`	`89`	`}`
`89`	`90`
`90`	`91`	`#[derive(Debug)]`
`@@ -284,6 +285,7 @@ pub enum StatementKind<'a> {`
`284`	`285`	`Increment(Handle<Expression<'a>>),`
`285`	`286`	`Decrement(Handle<Expression<'a>>),`
`286`	`287`	`Ignore(Handle<Expression<'a>>),`
	`288`	`+ ConstAssert(Handle<Expression<'a>>),`
`287`	`289`	`}`
`288`	`290`
`289`	`291`	`#[derive(Debug)]`