From 0d98327ce76f2658e64e50762da8ef027fcc1074 Mon Sep 17 00:00:00 2001
From: JoJoJet <21144246+JoJoJet@users.noreply.github.com>
Date: Sun, 25 Dec 2022 00:06:23 +0000
Subject: [PATCH] Support `SystemParam` types with const generics (#7001)

# Objective

* Currently, the `SystemParam` derive does not support types with const generic parameters.
  * If you try to use const generics, the error message is cryptic and unhelpful.
* Continuation of the work started in #6867 and #6957.

## Solution

Allow const generic parameters to be used with `#[derive(SystemParam)]`.
---
 crates/bevy_ecs/macros/src/lib.rs          | 32 ++++++++++++++++++----
 crates/bevy_ecs/src/system/system_param.rs |  3 ++
 2 files changed, 29 insertions(+), 6 deletions(-)

diff --git a/crates/bevy_ecs/macros/src/lib.rs b/crates/bevy_ecs/macros/src/lib.rs
index 1881ba2683bcf..1a0d2f0c4770d 100644
--- a/crates/bevy_ecs/macros/src/lib.rs
+++ b/crates/bevy_ecs/macros/src/lib.rs
@@ -14,8 +14,8 @@ use syn::{
     punctuated::Punctuated,
     spanned::Spanned,
     token::Comma,
-    DeriveInput, Field, GenericParam, Ident, Index, LitInt, Meta, MetaList, NestedMeta, Result,
-    Token, TypeParam,
+    ConstParam, DeriveInput, Field, GenericParam, Ident, Index, LitInt, Meta, MetaList, NestedMeta,
+    Result, Token, TypeParam,
 };
 
 struct AllTuples {
@@ -416,7 +416,7 @@ pub fn derive_system_param(input: TokenStream) -> TokenStream {
     let lifetimeless_generics: Vec<_> = generics
         .params
         .iter()
-        .filter(|g| matches!(g, GenericParam::Type(_)))
+        .filter(|g| !matches!(g, GenericParam::Lifetime(_)))
         .collect();
 
     let mut punctuated_generics = Punctuated::<_, Token![,]>::new();
@@ -425,12 +425,32 @@ pub fn derive_system_param(input: TokenStream) -> TokenStream {
             default: None,
             ..g.clone()
         }),
+        GenericParam::Const(g) => GenericParam::Const(ConstParam {
+            default: None,
+            ..g.clone()
+        }),
         _ => unreachable!(),
     }));
 
+    let mut punctuated_generics_no_bounds = punctuated_generics.clone();
+    for g in &mut punctuated_generics_no_bounds {
+        match g {
+            GenericParam::Type(g) => g.bounds.clear(),
+            GenericParam::Lifetime(g) => g.bounds.clear(),
+            GenericParam::Const(_) => {}
+        }
+    }
+
+    let mut punctuated_type_generic_idents = Punctuated::<_, Token![,]>::new();
+    punctuated_type_generic_idents.extend(lifetimeless_generics.iter().filter_map(|g| match g {
+        GenericParam::Type(g) => Some(&g.ident),
+        _ => None,
+    }));
+
     let mut punctuated_generic_idents = Punctuated::<_, Token![,]>::new();
     punctuated_generic_idents.extend(lifetimeless_generics.iter().map(|g| match g {
         GenericParam::Type(g) => &g.ident,
+        GenericParam::Const(g) => &g.ident,
         _ => unreachable!(),
     }));
 
@@ -470,15 +490,15 @@ pub fn derive_system_param(input: TokenStream) -> TokenStream {
             }
 
             #[doc(hidden)]
-            type State<'w, 's, #punctuated_generic_idents> = FetchState<
+            type State<'w, 's, #punctuated_generics_no_bounds> = FetchState<
                 (#(<#tuple_types as #path::system::SystemParam>::State,)*),
                 #punctuated_generic_idents
             >;
 
             #[doc(hidden)]
-            #state_struct_visibility struct FetchState <TSystemParamState, #punctuated_generic_idents> {
+            #state_struct_visibility struct FetchState <TSystemParamState, #punctuated_generics> {
                 state: TSystemParamState,
-                marker: std::marker::PhantomData<fn()->(#punctuated_generic_idents)>
+                marker: std::marker::PhantomData<fn()->(#punctuated_type_generic_idents)>
             }
 
             unsafe impl<'__w, '__s, #punctuated_generics> #path::system::SystemParamState for
diff --git a/crates/bevy_ecs/src/system/system_param.rs b/crates/bevy_ecs/src/system/system_param.rs
index 891a2d6768af8..8c9ff2ddfa01b 100644
--- a/crates/bevy_ecs/src/system/system_param.rs
+++ b/crates/bevy_ecs/src/system/system_param.rs
@@ -1772,6 +1772,9 @@ mod tests {
     #[derive(Resource)]
     pub struct R<const I: usize>;
 
+    #[derive(SystemParam)]
+    pub struct ConstGenericParam<'w, const I: usize>(Res<'w, R<I>>);
+
     #[derive(SystemParam)]
     pub struct LongParam<'w> {
         _r0: Res<'w, R<0>>,