From 8039f34b0d4254ae3ca62041052f514f6fa34e7c Mon Sep 17 00:00:00 2001
From: Gino Valente <49806985+MrGVSV@users.noreply.github.com>
Date: Mon, 7 Oct 2024 09:30:34 -0700
Subject: [PATCH] bevy_ecs: Replace panics in `QueryData` derive compile errors
 (#15691)

# Objective

The current `QueryData` derive panics when it encounters an error.
Additionally, it doesn't provide the clearest error message:

```rust
#[derive(QueryData)]
#[query_data(mut)]
struct Foo {
    // ...
}
```

```
error: proc-macro derive panicked
  --> src/foo.rs:16:10
   |
16 | #[derive(QueryData)]
   |          ^^^^^^^^^
   |
   = help: message: Invalid `query_data` attribute format
```

## Solution

Updated the derive logic to not panic and gave a bit more detail in the
error message.

This is makes the error message just a bit clearer and maintains the
correct span:

```
error: invalid attribute, expected `mutable` or `derive`
  --> src/foo.rs:17:14
   |
17 | #[query_data(mut)]
   |              ^^^
```

## Testing

You can test locally by running the following in
`crates/bevy_ecs/compile_fail`:

```
cargo test --target-dir ../../../target
```
---
 .../tests/ui/world_query_derive.rs            | 21 ++++++
 crates/bevy_ecs/macros/src/query_data.rs      | 65 +++++++------------
 2 files changed, 44 insertions(+), 42 deletions(-)

diff --git a/crates/bevy_ecs/compile_fail/tests/ui/world_query_derive.rs b/crates/bevy_ecs/compile_fail/tests/ui/world_query_derive.rs
index d990f59ecc..44e4343047 100644
--- a/crates/bevy_ecs/compile_fail/tests/ui/world_query_derive.rs
+++ b/crates/bevy_ecs/compile_fail/tests/ui/world_query_derive.rs
@@ -10,6 +10,27 @@ struct MutableUnmarked {
     a: &'static mut Foo,
 }
 
+#[derive(QueryData)]
+#[query_data(mut)]
+//~^ ERROR: invalid attribute, expected `mutable` or `derive`
+struct MutableInvalidAttribute {
+    a: &'static mut Foo,
+}
+
+#[derive(QueryData)]
+#[query_data(mutable(foo))]
+//~^ ERROR: `mutable` does not take any arguments
+struct MutableInvalidAttributeParameters {
+    a: &'static mut Foo,
+}
+
+#[derive(QueryData)]
+#[query_data(derive)]
+//~^ ERROR: `derive` requires at least one argument
+struct MutableMissingAttributeParameters {
+    a: &'static mut Foo,
+}
+
 #[derive(QueryData)]
 #[query_data(mutable)]
 struct MutableMarked {
diff --git a/crates/bevy_ecs/macros/src/query_data.rs b/crates/bevy_ecs/macros/src/query_data.rs
index 0a78b705a0..3f198b1ad1 100644
--- a/crates/bevy_ecs/macros/src/query_data.rs
+++ b/crates/bevy_ecs/macros/src/query_data.rs
@@ -1,13 +1,10 @@
 use bevy_macro_utils::ensure_no_collision;
 use proc_macro::TokenStream;
 use proc_macro2::{Ident, Span};
-use quote::{format_ident, quote, ToTokens};
+use quote::{format_ident, quote};
 use syn::{
-    parse::{Parse, ParseStream},
-    parse_macro_input, parse_quote,
-    punctuated::Punctuated,
-    token::Comma,
-    Attribute, Data, DataStruct, DeriveInput, Field, Index, Meta,
+    parse_macro_input, parse_quote, punctuated::Punctuated, token, token::Comma, Attribute, Data,
+    DataStruct, DeriveInput, Field, Index, Meta,
 };
 
 use crate::{
@@ -47,45 +44,29 @@ pub fn derive_query_data_impl(input: TokenStream) -> TokenStream {
             continue;
         }
 
-        attr.parse_args_with(|input: ParseStream| {
-            let meta = input.parse_terminated(Meta::parse, Comma)?;
-            for meta in meta {
-                let ident = meta.path().get_ident().unwrap_or_else(|| {
-                    panic!(
-                        "Unrecognized attribute: `{}`",
-                        meta.path().to_token_stream()
-                    )
-                });
-                if ident == MUTABLE_ATTRIBUTE_NAME {
-                    if let Meta::Path(_) = meta {
-                        attributes.is_mutable = true;
-                    } else {
-                        panic!(
-                            "The `{MUTABLE_ATTRIBUTE_NAME}` attribute is expected to have no value or arguments",
-                        );
-                    }
-                }
-                else if ident == DERIVE_ATTRIBUTE_NAME {
-                    if let Meta::List(meta_list) = meta {
-                        meta_list.parse_nested_meta(|meta| {
-                            attributes.derive_args.push(Meta::Path(meta.path));
-                            Ok(())
-                        })?;
-                    } else {
-                        panic!(
-                            "Expected a structured list within the `{DERIVE_ATTRIBUTE_NAME}` attribute",
-                        );
-                    }
+        let result = attr.parse_nested_meta(|meta| {
+            if meta.path.is_ident(MUTABLE_ATTRIBUTE_NAME) {
+                attributes.is_mutable = true;
+                if meta.input.peek(token::Paren) {
+                    Err(meta.error(format_args!("`{MUTABLE_ATTRIBUTE_NAME}` does not take any arguments")))
                 } else {
-                    panic!(
-                        "Unrecognized attribute: `{}`",
-                        meta.path().to_token_stream()
-                    );
+                    Ok(())
                 }
+            } else if meta.path.is_ident(DERIVE_ATTRIBUTE_NAME) {
+                meta.parse_nested_meta(|meta| {
+                    attributes.derive_args.push(Meta::Path(meta.path));
+                    Ok(())
+                }).map_err(|_| {
+                    meta.error(format_args!("`{DERIVE_ATTRIBUTE_NAME}` requires at least one argument"))
+                })
+            } else {
+                Err(meta.error(format_args!("invalid attribute, expected `{MUTABLE_ATTRIBUTE_NAME}` or `{DERIVE_ATTRIBUTE_NAME}`")))
             }
-            Ok(())
-        })
-        .unwrap_or_else(|_| panic!("Invalid `{QUERY_DATA_ATTRIBUTE_NAME}` attribute format"));
+        });
+
+        if let Err(err) = result {
+            return err.to_compile_error().into();
+        }
     }
 
     let path = bevy_ecs_path();