[naga] parse and validate @must_use attribute (#6801)

* feat: parse and enforce `@must_use` attribute

* feat: refuse `@must_use` of void function

* generated naga/tests/out/ir

* revert extraneous changes

* fix clippy

* update repeated attribute error to trunk format

* struct members can't have `@must_use`

* review fixes

* Remove must_use from FunctionResult and revert the related changes

---------

Co-authored-by: turbocrime <[email protected]>
Co-authored-by: Jamie Nicol <[email protected]>

CHANGELOG.md

@@ -40,6 +40,12 @@ Bottom level categories:
 
 ## Unreleased
 
+### New Features
+
+#### Naga
+
+- Support @must_use attribute on function declarations. By @turbocrime in [#6801](https://github.com/gfx-rs/wgpu/pull/6801).
+
 ### Changes
 
 #### General

naga/src/front/wgsl/error.rs

@@ -262,6 +262,8 @@ pub(crate) enum Error<'a> {
         found: u32,
     },
     FunctionReturnsVoid(Span),
+    FunctionMustUseUnused(Span),
+    FunctionMustUseReturnsVoid(Span, Span),
     InvalidWorkGroupUniformLoad(Span),
     Internal(&'static str),
     ExpectedConstExprConcreteIntegerScalar(Span),
@@ -820,6 +822,27 @@ impl<'a> Error<'a> {
                     "perhaps you meant to call the function in a separate statement?".into(),
                 ],
             },
+            Error::FunctionMustUseUnused(call) => ParseError {
+                message: "unused return value from function annotated with @must_use".into(),
+                labels: vec![(call, "".into())],
+                notes: vec![
+                    format!(
+                        "function '{}' is declared with `@must_use` attribute",
+                        &source[call],
+                    ),
+                    "use a phony assignment or declare a value using the function call as the initializer".into(),
+                ],
+            },
+            Error::FunctionMustUseReturnsVoid(attr, signature) => ParseError {
+                message: "function annotated with @must_use but does not return any value".into(),
+                labels: vec![
+                    (attr, "".into()),
+                    (signature, "".into()),
+                ],
+                notes: vec![
+                    "declare a return type or remove the attribute".into(),
+                ],
+            },
             Error::InvalidWorkGroupUniformLoad(span) => ParseError {
                 message: "incorrect type passed to workgroupUniformLoad".into(),
                 labels: vec![(span, "".into())],

naga/src/front/wgsl/lower/mod.rs

@@ -934,7 +934,10 @@ impl Components {
 
 /// An `ast::GlobalDecl` for which we have built the Naga IR equivalent.
 enum LoweredGlobalDecl {
-    Function(Handle<crate::Function>),
+    Function {
+        handle: Handle<crate::Function>,
+        must_use: bool,
+    },
     Var(Handle<crate::GlobalVariable>),
     Const(Handle<crate::Constant>),
     Override(Handle<crate::Override>),
@@ -1365,7 +1368,10 @@ impl<'source, 'temp> Lowerer<'source, 'temp> {
             Ok(LoweredGlobalDecl::EntryPoint)
         } else {
             let handle = ctx.module.functions.append(function, span);
-            Ok(LoweredGlobalDecl::Function(handle))
+            Ok(LoweredGlobalDecl::Function {
+                handle,
+                must_use: f.result.as_ref().is_some_and(|res| res.must_use),
+            })
         }
     }
 
@@ -1934,7 +1940,7 @@ impl<'source, 'temp> Lowerer<'source, 'temp> {
                     LoweredGlobalDecl::Override(handle) => {
                         Typed::Plain(crate::Expression::Override(handle))
                     }
-                    LoweredGlobalDecl::Function(_)
+                    LoweredGlobalDecl::Function { .. }
                     | LoweredGlobalDecl::Type(_)
                     | LoweredGlobalDecl::EntryPoint => {
                         return Err(Error::Unexpected(span, ExpectedToken::Variable));
@@ -2181,12 +2187,13 @@ impl<'source, 'temp> Lowerer<'source, 'temp> {
         ctx: &mut ExpressionContext<'source, '_, '_>,
         is_statement: bool,
     ) -> Result<Option<Handle<crate::Expression>>, Error<'source>> {
+        let function_span = function.span;
         match ctx.globals.get(function.name) {
             Some(&LoweredGlobalDecl::Type(ty)) => {
                 let handle = self.construct(
                     span,
                     &ast::ConstructorType::Type(ty),
-                    function.span,
+                    function_span,
                     arguments,
                     ctx,
                 )?;
@@ -2196,9 +2203,12 @@ impl<'source, 'temp> Lowerer<'source, 'temp> {
                 &LoweredGlobalDecl::Const(_)
                 | &LoweredGlobalDecl::Override(_)
                 | &LoweredGlobalDecl::Var(_),
-            ) => Err(Error::Unexpected(function.span, ExpectedToken::Function)),
-            Some(&LoweredGlobalDecl::EntryPoint) => Err(Error::CalledEntryPoint(function.span)),
-            Some(&LoweredGlobalDecl::Function(function)) => {
+            ) => Err(Error::Unexpected(function_span, ExpectedToken::Function)),
+            Some(&LoweredGlobalDecl::EntryPoint) => Err(Error::CalledEntryPoint(function_span)),
+            Some(&LoweredGlobalDecl::Function {
+                handle: function,
+                must_use,
+            }) => {
                 let arguments = arguments
                     .iter()
                     .enumerate()
@@ -2223,6 +2233,11 @@ impl<'source, 'temp> Lowerer<'source, 'temp> {
                     .collect::<Result<Vec<_>, _>>()?;
 
                 let has_result = ctx.module.functions[function].result.is_some();
+
+                if must_use && is_statement {
+                    return Err(Error::FunctionMustUseUnused(function_span));
+                }
+
                 let rctx = ctx.runtime_expression_ctx(span)?;
                 // we need to always do this before a fn call since all arguments need to be emitted before the fn call
                 rctx.block
@@ -2249,7 +2264,7 @@ impl<'source, 'temp> Lowerer<'source, 'temp> {
                 Ok(result)
             }
             None => {
-                let span = function.span;
+                let span = function_span;
                 let expr = if let Some(fun) = conv::map_relational_fun(function.name) {
                     let mut args = ctx.prepare_args(arguments, 1, span);
                     let argument = self.expression(args.next()?, ctx)?;

naga/src/front/wgsl/parse/ast.rs

@@ -114,6 +114,7 @@ pub struct FunctionArgument<'a> {
 pub struct FunctionResult<'a> {
     pub ty: Handle<Type<'a>>,
     pub binding: Option<Binding<'a>>,
+    pub must_use: bool,
 }
 
 #[derive(Debug)]

naga/src/front/wgsl/parse/mod.rs

@@ -2332,6 +2332,7 @@ impl Parser {
         &mut self,
         lexer: &mut Lexer<'a>,
         diagnostic_filter_leaf: Option<Handle<DiagnosticFilterNode>>,
+        must_use: Option<Span>,
         out: &mut ast::TranslationUnit<'a>,
         dependencies: &mut FastIndexSet<ast::Dependency<'a>>,
     ) -> Result<ast::Function<'a>, Error<'a>> {
@@ -2383,7 +2384,17 @@ impl Parser {
         let result = if lexer.skip(Token::Arrow) {
             let binding = self.varying_binding(lexer, &mut ctx)?;
             let ty = self.type_decl(lexer, &mut ctx)?;
-            Some(ast::FunctionResult { ty, binding })
+            let must_use = must_use.is_some();
+            Some(ast::FunctionResult {
+                ty,
+                binding,
+                must_use,
+            })
+        } else if let Some(must_use) = must_use {
+            return Err(Error::FunctionMustUseReturnsVoid(
+                must_use,
+                self.peek_rule_span(lexer),
+            ));
         } else {
             None
         };
@@ -2457,6 +2468,8 @@ impl Parser {
             (ParsedAttribute::default(), ParsedAttribute::default());
         let mut id = ParsedAttribute::default();
 
+        let mut must_use: ParsedAttribute<Span> = ParsedAttribute::default();
+
         let mut dependencies = FastIndexSet::default();
         let mut ctx = ExpressionContext {
             expressions: &mut out.expressions,
@@ -2541,6 +2554,9 @@ impl Parser {
                     };
                     early_depth_test.set(crate::EarlyDepthTest { conservative }, name_span)?;
                 }
+                "must_use" => {
+                    must_use.set(name_span, name_span)?;
+                }
                 _ => return Err(Error::UnknownAttribute(name_span)),
             }
         }
@@ -2646,8 +2662,14 @@ impl Parser {
                     diagnostic_filters,
                     out.diagnostic_filter_leaf,
                 );
-                let function =
-                    self.function_decl(lexer, diagnostic_filter_leaf, out, &mut dependencies)?;
+
+                let function = self.function_decl(
+                    lexer,
+                    diagnostic_filter_leaf,
+                    must_use.value,
+                    out,
+                    &mut dependencies,
+                )?;
                 Some(ast::GlobalDeclKind::Fn(ast::Function {
                     entry_point: if let Some(stage) = stage.value {
                         if stage == ShaderStage::Compute && workgroup_size.value.is_none() {

naga/tests/in/must-use.wgsl

@@ -0,0 +1,23 @@
+@compute @workgroup_size(1)
+fn main() {}
+
+@must_use
+fn use_me() -> i32 { return 10; }
+
+fn use_return() -> i32 {
+    return use_me();
+}
+
+fn use_assign_var() -> i32 {
+    var q = use_me();
+    return q;
+}
+
+fn use_assign_let() -> i32 {
+    let q = use_me();
+    return q;
+}
+
+fn use_phony_assign() {
+    _ = use_me();
+}