statically defined storage type

Author: Frizi

Cargo.toml

@@ -45,7 +45,7 @@ bevy_audio = ["bevy_internal/bevy_audio"]
 bevy_dynamic_plugin = ["bevy_internal/bevy_dynamic_plugin"]
 bevy_gilrs = ["bevy_internal/bevy_gilrs"]
 bevy_gltf = ["bevy_internal/bevy_gltf"]
-bevy_wgpu = ["bevy_internal/bevy_wgpu"]
+bevy_wgpu = ["bevy_internal/bevy_wgpu"] 
 bevy_winit = ["bevy_internal/bevy_winit"]
 
 trace_chrome = ["bevy_internal/trace_chrome"]

crates/bevy_ecs/macros/src/component.rs

@@ -1,11 +1,20 @@
 use proc_macro::TokenStream;
+use proc_macro2::{Span, TokenStream as TokenStream2};
 use quote::quote;
-use syn::{parse_macro_input, parse_quote, DeriveInput, Path};
+use syn::{parse_macro_input, parse_quote, Attribute, DeriveInput, Error, Ident, Path, Result};
 
 pub fn derive_component(input: TokenStream) -> TokenStream {
     let mut ast = parse_macro_input!(input as DeriveInput);
     let bevy_ecs_path: Path = crate::bevy_ecs_path();
 
+    let storage = ast
+        .attrs
+        .iter()
+        .find(|attr| attr.path.is_ident("storage"))
+        .map_or(Ok(StorageTy::Table), parse_storage_attribute)
+        .map(|ty| storage_path(bevy_ecs_path.clone(), ty))
+        .unwrap_or_else(|err| err.to_compile_error());
+
     ast.generics
         .make_where_clause()
         .predicates
@@ -15,6 +24,31 @@ pub fn derive_component(input: TokenStream) -> TokenStream {
     let (impl_generics, type_generics, where_clause) = &ast.generics.split_for_impl();
 
     TokenStream::from(quote! {
-        impl #impl_generics #bevy_ecs_path::component::Component for #struct_name #type_generics #where_clause {}
+        impl #impl_generics #bevy_ecs_path::component::Component for #struct_name #type_generics #where_clause {
+            type Storage = #storage;
+        }
     })
 }
+
+enum StorageTy {
+    Table,
+    Sparse,
+}
+
+fn parse_storage_attribute(attr: &Attribute) -> Result<StorageTy> {
+    let ident = attr.parse_args::<Ident>()?;
+    match ident.to_string().as_str() {
+        "table" => Ok(StorageTy::Table),
+        "sparse" => Ok(StorageTy::Sparse),
+        _ => Err(Error::new(ident.span(), "Invalid storage type, expected 'table' or 'sparse'.")),
+    }
+}
+
+fn storage_path(bevy_ecs_path: Path, ty: StorageTy) -> TokenStream2 {
+    let typename = match ty {
+        StorageTy::Table => Ident::new("TableStorage", Span::call_site()),
+        StorageTy::Sparse => Ident::new("SparseStorage", Span::call_site()),
+    };
+
+    quote! { #bevy_ecs_path::component::#typename }
+}

crates/bevy_ecs/macros/src/lib.rs

@@ -115,6 +115,8 @@ pub fn derive_bundle(input: TokenStream) -> TokenStream {
     let mut field_type_infos = Vec::new();
     let mut field_get_components = Vec::new();
     let mut field_from_components = Vec::new();
+    let mut is_dense_const_exprs = Vec::new();
+    let mut is_dense_fn_exprs = Vec::new();
     for ((field_type, is_bundle), field) in
         field_type.iter().zip(is_bundle.iter()).zip(field.iter())
     {
@@ -128,6 +130,9 @@ pub fn derive_bundle(input: TokenStream) -> TokenStream {
             field_from_components.push(quote! {
                 #field: <#field_type as #ecs_path::bundle::Bundle>::from_components(&mut func),
             });
+            is_dense_fn_exprs.push(quote! {
+                <#field_type as #ecs_path::bundle::Bundle>::is_dense()
+            });
         } else {
             field_type_infos.push(quote! {
                 type_info.push(#ecs_path::component::TypeInfo::of::<#field_type>());
@@ -139,6 +144,10 @@ pub fn derive_bundle(input: TokenStream) -> TokenStream {
             field_from_components.push(quote! {
                 #field: func().cast::<#field_type>().read(),
             });
+            is_dense_const_exprs.push(quote! {
+                <<#field_type as #ecs_path::component::Component>::Storage as #ecs_path::component::ComponentStorage>::STORAGE_TYPE
+                    == #ecs_path::component::StorageType::Table
+            });
         }
     }
     let field_len = field.len();
@@ -155,6 +164,12 @@ pub fn derive_bundle(input: TokenStream) -> TokenStream {
                 type_info
             }
 
+            #[inline(always)]
+            fn is_dense() -> bool {
+                // insert const expressions in front to avoid evaluating non-const functions if possible
+                true #(&& (#is_dense_const_exprs))* #(&& #is_dense_fn_exprs)*
+            }
+
             #[allow(unused_variables, unused_mut, non_snake_case)]
             unsafe fn from_components(mut func: impl FnMut() -> *mut u8) -> Self {
                 Self {
@@ -481,7 +496,7 @@ pub(crate) fn bevy_ecs_path() -> syn::Path {
     BevyManifest::default().get_path("bevy_ecs")
 }
 
-#[proc_macro_derive(Component)]
+#[proc_macro_derive(Component, attributes(storage))]
 pub fn derive_component(input: TokenStream) -> TokenStream {
     component::derive_component(input)
 }

crates/bevy_ecs/src/bundle.rs

@@ -2,7 +2,9 @@ pub use bevy_ecs_macros::Bundle;
 
 use crate::{
     archetype::ComponentStatus,
-    component::{Component, ComponentId, ComponentTicks, Components, StorageType, TypeInfo},
+    component::{
+        Component, ComponentId, ComponentStorage, ComponentTicks, Components, StorageType, TypeInfo,
+    },
     entity::Entity,
     storage::{SparseSetIndex, SparseSets, Table},
 };
@@ -60,6 +62,8 @@ pub unsafe trait Bundle: Send + Sync + 'static {
     /// "mem::forget" the bundle fields, so callers are responsible for dropping the fields if
     /// that is desirable.
     fn get_components(self, func: impl FnMut(*mut u8));
+
+    fn is_dense() -> bool;
 }
 
 macro_rules! tuple_impl {
@@ -70,6 +74,11 @@ macro_rules! tuple_impl {
                 vec![$(TypeInfo::of::<$name>()),*]
             }
 
+            #[inline(always)]
+            fn is_dense() -> bool {
+                true $(&& $name::Storage::STORAGE_TYPE == StorageType::Table)*
+            }
+
             #[allow(unused_variables, unused_mut)]
             unsafe fn from_components(mut func: impl FnMut() -> *mut u8) -> Self {
                 #[allow(non_snake_case)]

crates/bevy_ecs/src/component/mod.rs

@@ -31,26 +31,72 @@ use thiserror::Error;
 /// as one of the arguments.
 ///
 /// Components can be grouped together into a [`Bundle`](crate::bundle::Bundle).
-pub trait Component: Send + Sync + 'static {}
+pub trait Component: Send + Sync + 'static {
+    type Storage: ComponentStorage;
+}
+
+pub struct TableStorage;
+pub struct SparseStorage;
+
+pub trait ComponentStorage: sealed::Sealed {
+    // because the trait is selaed, those items are private API.
+    const STORAGE_TYPE: StorageType;
+}
+
+impl ComponentStorage for TableStorage {
+    const STORAGE_TYPE: StorageType = StorageType::Table;
+}
+impl ComponentStorage for SparseStorage {
+    const STORAGE_TYPE: StorageType = StorageType::SparseSet;
+}
+
+mod sealed {
+    pub trait Sealed {}
+    impl Sealed for super::TableStorage {}
+    impl Sealed for super::SparseStorage {}
+}
 
 // ECS dependencies cannot derive Component, so we must implement it manually for relevant structs.
-impl<T> Component for bevy_tasks::Task<T> where Self: Send + Sync + 'static {}
+impl<T> Component for bevy_tasks::Task<T>
+where
+    Self: Send + Sync + 'static,
+{
+    type Storage = TableStorage;
+}
 
 // For our own convinience, let's implement Component for primitives in tests.
 // It will eventually be removed, once tests are not using them anymore.
 // Consider those impls deprecated.
 #[cfg(test)]
 mod private_test_component_impls {
-    use super::Component;
-    impl Component for &'static str {}
-    impl Component for usize {}
-    impl Component for i32 {}
-    impl Component for u32 {}
-    impl Component for u64 {}
-    impl Component for f32 {}
-    impl Component for f64 {}
-    impl Component for u8 {}
-    impl Component for bool {}
+    use super::{Component, TableStorage};
+    impl Component for &'static str {
+        type Storage = TableStorage;
+    }
+    impl Component for usize {
+        type Storage = TableStorage;
+    }
+    impl Component for i32 {
+        type Storage = TableStorage;
+    }
+    impl Component for u32 {
+        type Storage = TableStorage;
+    }
+    impl Component for u64 {
+        type Storage = TableStorage;
+    }
+    impl Component for f32 {
+        type Storage = TableStorage;
+    }
+    impl Component for f64 {
+        type Storage = TableStorage;
+    }
+    impl Component for u8 {
+        type Storage = TableStorage;
+    }
+    impl Component for bool {
+        type Storage = TableStorage;
+    }
 }
 
 #[derive(Debug, Copy, Clone, Eq, PartialEq)]
@@ -165,10 +211,10 @@ pub struct ComponentDescriptor {
 }
 
 impl ComponentDescriptor {
-    pub fn new<T: Component>(storage_type: StorageType) -> Self {
+    pub fn new<T: Component>(_storage_type: StorageType) -> Self {
         Self {
             name: std::any::type_name::<T>().to_string(),
-            storage_type,
+            storage_type: T::Storage::STORAGE_TYPE,
             is_send_and_sync: true,
             type_id: Some(TypeId::of::<T>()),
             layout: Layout::new::<T>(),