Preventing wgsl helper usage in places they shouldn't be, also slightly cleaned up recently added feature related to "on_cpu" version of shader functions.

bevy_gpu_compute_core/src/wgsl/shader_module/derived_portion.rs

@@ -82,9 +82,11 @@ mod tests {
 
     use crate::{
         IterSpaceDimmension,
-        wgsl::shader_custom_type_name::ShaderCustomTypeName,
-        wgsl::shader_module::complete_shader_module::WgslShaderModule,
-        wgsl::shader_sections::{WgslInputArray, WgslOutputArray, WgslShaderModuleSectionCode},
+        wgsl::{
+            shader_custom_type_name::ShaderCustomTypeName,
+            shader_module::complete_shader_module::WgslShaderModule,
+            shader_sections::{WgslInputArray, WgslOutputArray, WgslShaderModuleSectionCode},
+        },
     };
 
     use super::*;

bevy_gpu_compute_macro/src/lib.rs

@@ -8,7 +8,6 @@ use transformer::{
     custom_types::get_all_custom_types::get_custom_types, module_parser::lib::parse_shader_module,
     output::produce_expanded_output, remove_doc_comments::DocCommentRemover,
     transform_wgsl_helper_methods::run::transform_wgsl_helper_methods,
-    transform_wgsl_helper_methods_for_cpu::run::transform_wgsl_helper_methods_for_cpu,
 };
 mod state;
 mod transformer;
@@ -71,8 +70,8 @@ pub fn wgsl_shader_module(_attr: TokenStream, item: TokenStream) -> TokenStream
     DocCommentRemover {}.visit_item_mod(&module);
     let mut state = ModuleTransformState::empty(module, content);
     get_custom_types(&mut state);
-    transform_wgsl_helper_methods(&mut state);
-    transform_wgsl_helper_methods_for_cpu(&mut state);
+    transform_wgsl_helper_methods(&state.custom_types, &mut state.rust_module, false);
+    transform_wgsl_helper_methods(&state.custom_types, &mut state.rust_module_for_cpu, true);
     parse_shader_module(&mut state);
     let output = produce_expanded_output(&mut state);
     output.into()

bevy_gpu_compute_macro/src/transformer/mod.rs

@@ -5,4 +5,3 @@ pub mod output;
 pub mod remove_doc_comments;
 pub mod to_wgsl_syntax;
 pub mod transform_wgsl_helper_methods;
-pub mod transform_wgsl_helper_methods_for_cpu;

bevy_gpu_compute_macro/src/transformer/module_parser/lib.rs

@@ -6,7 +6,7 @@ use crate::state::ModuleTransformState;
 use super::constants::find_constants;
 use super::divide_custom_types::divide_custom_types_by_category;
 use super::helper_functions::find_helper_functions;
-use super::main_function::find_main_function;
+use super::main_function::parse_main_function;
 use super::use_statements::handle_use_statements;
 use super::validate_no_global_id_assignments::check_module_for_global_id_assignment;
 
@@ -17,7 +17,7 @@ pub fn parse_shader_module(state: &mut ModuleTransformState) {
             "Shader module must have a body"
         );
     }
-    find_main_function(state);
+    parse_main_function(state);
     handle_use_statements(state);
     state.module_ident = Some(state.rust_module.ident.to_string());
     state.module_visibility = Some(state.rust_module.vis.to_token_stream().to_string());

bevy_gpu_compute_macro/src/transformer/module_parser/main_function.rs

@@ -5,7 +5,7 @@ use proc_macro_error::abort;
 use quote::ToTokens;
 use syn::{ItemFn, spanned::Spanned, visit::Visit};
 
-pub fn find_main_function(state: &mut ModuleTransformState) {
+pub fn parse_main_function(state: &mut ModuleTransformState) {
     let module = state.rust_module.clone();
     let mut extractor = MainFunctionsExtractor::new(state, false);
     extractor.visit_item_mod(&module);

bevy_gpu_compute_macro/src/transformer/transform_wgsl_helper_methods/erroneous_usage_finder.rs

@@ -0,0 +1,39 @@
+use syn::{
+    Expr,
+    visit::{self, Visit},
+};
+
+use crate::transformer::custom_types::custom_type::CustomType;
+
+use super::parse::parse_possible_wgsl_helper;
+
+pub struct ErroneousUsageFinder {
+    custom_types: Vec<CustomType>,
+}
+impl ErroneousUsageFinder {
+    pub fn new(custom_types: &[CustomType]) -> Self {
+        Self {
+            custom_types: custom_types.to_vec(),
+        }
+    }
+}
+impl Visit<'_> for ErroneousUsageFinder {
+    /// This error message relies on `WgslVecInput` being in `bevy_gpu_compute_core::wgsl_helpers`
+    /**
+    ```rust
+    // ensure that the crate structure is what we expect, otherwise the error message will be incorrect
+       use bevy_gpu_compute_core::wgsl_helpers::WgslVecInput;
+    ```
+    */
+    fn visit_expr(&mut self, expr: &Expr) {
+        visit::visit_expr(self, expr);
+        if let Expr::Call(call) = expr {
+            let helper_method = parse_possible_wgsl_helper(call, &self.custom_types);
+            if helper_method.is_some() {
+                panic!(
+                    "WGSL Helpers (`bevy_gpu_compute_core::wgsl_helpers`) not allowed outside of functions."
+                );
+            }
+        }
+    }
+}

bevy_gpu_compute_macro/src/transformer/transform_wgsl_helper_methods/helper_method.rs

@@ -7,11 +7,11 @@ use super::{
     to_expanded_format::ToExpandedFormatMethodKind,
 };
 
-pub struct WgslHelperMethod<'a> {
+pub struct WgslHelperMethod {
     pub category: WgslHelperCategory,
     pub method: WgslHelperMethodName,
-    pub t_def: &'a CustomType,
-    pub arg1: Option<&'a Expr>,
-    pub arg2: Option<&'a Expr>,
+    pub t_def: CustomType,
+    pub arg1: Option<Expr>,
+    pub arg2: Option<Expr>,
     pub method_expander_kind: Option<ToExpandedFormatMethodKind>,
 }

bevy_gpu_compute_macro/src/transformer/transform_wgsl_helper_methods/mod.rs

@@ -1,7 +1,11 @@
 pub mod category;
+mod erroneous_usage_finder;
 pub mod helper_method;
 pub mod matcher;
 pub mod method_name;
+mod parse;
 pub mod run;
 pub mod test;
+mod test_for_cpu;
 pub mod to_expanded_format;
+mod to_expanded_format_for_cpu;

bevy_gpu_compute_macro/src/transformer/transform_wgsl_helper_methods/parse.rs

@@ -0,0 +1,73 @@
+use syn::{Expr, ExprCall, GenericArgument, PathArguments, Type};
+
+use crate::transformer::{
+    custom_types::custom_type::CustomType,
+    transform_wgsl_helper_methods::helper_method::WgslHelperMethod,
+};
+
+use super::{
+    category::WgslHelperCategory, matcher::WgslHelperMethodMatcher,
+    method_name::WgslHelperMethodName,
+};
+fn get_special_function_category(call: &ExprCall) -> Option<WgslHelperCategory> {
+    if let Expr::Path(path) = &*call.func {
+        if let Some(first_seg) = path.path.segments.first() {
+            return WgslHelperCategory::from_ident(first_seg.ident.clone());
+        }
+    }
+    None
+}
+fn get_special_function_method(call: &ExprCall) -> Option<WgslHelperMethodName> {
+    if let Expr::Path(path) = &*call.func {
+        if let Some(last_seg) = path.path.segments.last() {
+            return WgslHelperMethodName::from_ident(last_seg.ident.clone());
+        }
+    }
+    None
+}
+fn get_special_function_generic_type<'a>(
+    call: &'a ExprCall,
+    custom_types: &'a [CustomType],
+) -> Option<&'a CustomType> {
+    if let Expr::Path(path) = &*call.func {
+        if let Some(last_seg) = path.path.segments.last() {
+            if let PathArguments::AngleBracketed(args) = &last_seg.arguments {
+                if let Some(GenericArgument::Type(Type::Path(type_path))) = args.args.first() {
+                    if let Some(last_seg) = type_path.path.segments.last() {
+                        return custom_types.iter().find(|t| t.name.eq(&last_seg.ident));
+                    }
+                }
+            }
+        }
+    }
+    None
+}
+
+pub fn parse_possible_wgsl_helper<'a>(
+    call: &'a ExprCall,
+    custom_types: &'a [CustomType],
+) -> Option<WgslHelperMethod> {
+    let category = get_special_function_category(call);
+    let method = get_special_function_method(call);
+    let type_name = get_special_function_generic_type(call, custom_types);
+    if let Some(cat) = category {
+        if let Some(met) = method {
+            if let Some(ty) = type_name {
+                let args = call.args.clone();
+                let mut method = WgslHelperMethod {
+                    category: cat,
+                    method: met,
+                    t_def: ty.clone(),
+                    arg1: args.first().cloned(),
+                    arg2: args.get(1).cloned(),
+                    method_expander_kind: None,
+                };
+                WgslHelperMethodMatcher::choose_expand_format(&mut method);
+                if method.method_expander_kind.is_some() {
+                    return Some(method);
+                }
+            }
+        }
+    }
+    None
+}