Move shared WGSL code to imported module (#264)

Move most shared WGSL code into an import module `vfx_common.wgsl`, and
document it. This ensures the code remains in sync between the various passes
(init/update/render) without needing some error-prone copy/paste, while making
the code easier to read too.

This change moves `naga` and `naga_oil` from being dev-only dependencies to
being proper crate dependencies. This is required to use `naga_oil`'s import
system, like `bevy_render`. This makes the change a breaking one.

CHANGELOG.md

@@ -3,6 +3,12 @@
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [Unreleased]
+
+### Changed
+
+- Moved most shared WGSL code into an import module `vfx_common.wgsl`. This requires using `naga_oil` for import resolution, which in turns means `naga` and `naga_oil` are now dependencies of `bevy_hanabi` itself.
+
 ## [0.9.0] 2023-12-26
 
 ### Added

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "bevy_hanabi"
-version = "0.9.0"
+version = "0.10.0-dev"
 authors = ["Jerome Humbert <[email protected]>"]
 edition = "2021"
 description = "Hanabi GPU particle system for the Bevy game engine"
@@ -38,6 +38,9 @@ ron = "0.8"
 bitflags = "2.3"
 typetag = "0.2"
 thiserror = "1.0"
+# Same versions as Bevy 0.12 (bevy_render)
+naga = "0.13"
+naga_oil = "0.10"
 
 [dependencies.bevy]
 version = "0.12"
@@ -50,8 +53,6 @@ all-features = true
 [dev-dependencies]
 # Same versions as Bevy 0.12 (bevy_render)
 wgpu = "0.17.1"
-naga = "0.13"
-naga_oil = "0.10"
 
 # For procedural texture generation in examples
 noise = "0.8"

src/lib.rs

@@ -1773,6 +1773,9 @@ else { return c1; }
                 "PARTICLE_SCREEN_SPACE_SIZE".into(),
                 ShaderDefValue::Bool(true),
             );
+            if name == "Update" {
+                shader_defs.insert("RI_MAX_SPAWN_ATOMIC".into(), ShaderDefValue::Bool(true));
+            }
             let mut composer = Composer::default();
 
             // Import bevy_render::view for the render shader
@@ -1791,6 +1794,15 @@ else { return c1; }
                 assert!(res.is_ok());
             }
 
+            // Import bevy_hanabi::vfx_common
+            {
+                let min_storage_buffer_offset_alignment = 256usize;
+                let common_shader =
+                    HanabiPlugin::make_common_shader(min_storage_buffer_offset_alignment);
+                let res = composer.add_composable_module((&common_shader).into());
+                assert!(res.is_ok());
+            }
+
             match composer.make_naga_module(NagaModuleDescriptor {
                 source: code,
                 file_path: "init.wgsl",

src/plugin.rs

@@ -21,8 +21,9 @@ use crate::{
     render::{
         extract_effect_events, extract_effects, prepare_effects, prepare_resources, queue_effects,
         DispatchIndirectPipeline, DrawEffects, EffectAssetEvents, EffectBindGroups, EffectSystems,
-        EffectsMeta, ExtractedEffects, ParticlesInitPipeline, ParticlesRenderPipeline,
-        ParticlesUpdatePipeline, ShaderCache, SimParams, VfxSimulateDriverNode, VfxSimulateNode,
+        EffectsMeta, ExtractedEffects, GpuSpawnerParams, ParticlesInitPipeline,
+        ParticlesRenderPipeline, ParticlesUpdatePipeline, ShaderCache, SimParams,
+        VfxSimulateDriverNode, VfxSimulateNode,
     },
     spawn::{self, Random},
     tick_spawners, update_properties_from_asset, ParticleEffect, RemovedEffectsEvent, Spawner,
@@ -46,10 +47,39 @@ pub mod simulate_graph {
     }
 }
 
+// {626E7AD3-4E54-487E-B796-9A90E34CC1EC}
+const HANABI_COMMON_TEMPLATE_HANDLE: Handle<Shader> =
+    Handle::weak_from_u128(0x626E7AD34E54487EB7969A90E34CC1ECu128);
+
 /// Plugin to add systems related to Hanabi.
 #[derive(Debug, Clone, Copy)]
 pub struct HanabiPlugin;
 
+impl HanabiPlugin {
+    /// Create the `vfx_common.wgsl` shader with proper alignment.
+    ///
+    /// This creates a new [`Shader`] from the `vfx_common.wgsl` code, by
+    /// applying the given alignment for storage buffers. This produces a shader
+    /// ready for the specific GPU device associated with that alignment.
+    pub(crate) fn make_common_shader(min_storage_buffer_offset_alignment: usize) -> Shader {
+        let spawner_padding_code =
+            GpuSpawnerParams::padding_code(min_storage_buffer_offset_alignment);
+        let common_code = include_str!("render/vfx_common.wgsl")
+            .replace("{{SPAWNER_PADDING}}", &spawner_padding_code);
+        Shader::from_wgsl(
+            common_code,
+            std::path::Path::new(file!())
+                .parent()
+                .unwrap()
+                .join(format!(
+                    "render/vfx_common_{}.wgsl",
+                    min_storage_buffer_offset_alignment
+                ))
+                .to_string_lossy(),
+        )
+    }
+}
+
 impl Plugin for HanabiPlugin {
     fn build(&self, app: &mut App) {
         // Register asset
@@ -94,8 +124,7 @@ impl Plugin for HanabiPlugin {
         let render_device = app
             .sub_app(RenderApp)
             .world
-            .get_resource::<RenderDevice>()
-            .unwrap()
+            .resource::<RenderDevice>()
             .clone();
 
         let adapter_name = app
@@ -113,6 +142,17 @@ impl Plugin for HanabiPlugin {
             info!("Initializing Hanabi for GPU adapter {}", adapter_name);
         }
 
+        // Insert the properly aligned `vfx_common.wgsl` shader into Assets<Shader>, so
+        // that the automated Bevy shader processing finds it as an import. This is used
+        // for init/update/render shaders (but not the indirect one).
+        {
+            let common_shader = HanabiPlugin::make_common_shader(
+                render_device.limits().min_storage_buffer_offset_alignment as usize,
+            );
+            let mut assets = app.world.resource_mut::<Assets<Shader>>();
+            assets.insert(HANABI_COMMON_TEMPLATE_HANDLE, common_shader);
+        }
+
         let effects_meta = EffectsMeta::new(render_device);
 
         // Register the custom render pipeline

src/render/mod.rs

@@ -28,6 +28,7 @@ use bevy::{
     utils::HashMap,
 };
 use bitflags::bitflags;
+use naga_oil::compose::{Composer, NagaModuleDescriptor};
 use rand::random;
 use std::marker::PhantomData;
 use std::{
@@ -46,8 +47,8 @@ use crate::{
     next_multiple_of,
     render::batch::{BatchInput, BatchState, Batcher, EffectBatch},
     spawn::EffectSpawner,
-    CompiledParticleEffect, EffectProperties, EffectShader, ParticleLayout, PropertyLayout,
-    RemovedEffectsEvent, SimulationCondition, SimulationSpace,
+    CompiledParticleEffect, EffectProperties, EffectShader, HanabiPlugin, ParticleLayout,
+    PropertyLayout, RemovedEffectsEvent, SimulationCondition, SimulationSpace,
 };
 
 mod aligned_buffer_vec;
@@ -246,7 +247,7 @@ impl From<&Mat4> for GpuCompressedTransform {
 /// together form the spawner parameter buffer.
 #[repr(C)]
 #[derive(Debug, Default, Clone, Copy, Pod, Zeroable, ShaderType)]
-struct GpuSpawnerParams {
+pub(crate) struct GpuSpawnerParams {
     /// Transform of the effect, as a Mat4 without the last row (which is always
     /// (0,0,0,1) for an affine transform), stored transposed as a mat3x4 to
     /// avoid padding in WGSL. This is either added to emitted particles at
@@ -269,6 +270,33 @@ struct GpuSpawnerParams {
     force_field: [GpuForceFieldSource; ForceFieldSource::MAX_SOURCES],
 }
 
+impl GpuSpawnerParams {
+    /// Get the aligned size of this type based on the given alignment in bytes.
+    pub fn aligned_size(align_size: usize) -> usize {
+        next_multiple_of(GpuSpawnerParams::min_size().get() as usize, align_size)
+    }
+
+    /// Get the WGSL padding code to append to the GPU struct to align it.
+    pub fn padding_code(align_size: usize) -> String {
+        let aligned_size = GpuSpawnerParams::aligned_size(align_size);
+        trace!(
+            "Aligning spawner params to {} bytes as device limits requires. Aligned size: {} bytes.",
+            align_size,
+            aligned_size
+        );
+
+        // We need to pad the Spawner WGSL struct based on the device padding so that we
+        // can use it as an array element but also has a direct struct binding.
+        if GpuSpawnerParams::min_size().get() as usize != aligned_size {
+            let padding_size = aligned_size - GpuSpawnerParams::min_size().get() as usize;
+            assert!(padding_size % 4 == 0);
+            format!("padding: array<u32, {}>", padding_size / 4)
+        } else {
+            "".to_string()
+        }
+    }
+}
+
 // FIXME - min_storage_buffer_offset_alignment
 #[repr(C)]
 #[derive(Debug, Clone, Copy, Pod, Zeroable, ShaderType)]
@@ -327,8 +355,7 @@ impl FromWorld for DispatchIndirectPipeline {
         // so needs the proper align. Because WGSL removed the @stride attribute, we pad
         // the WGSL type manually, so need to enforce min_binding_size everywhere.
         let item_align = render_device.limits().min_storage_buffer_offset_alignment as usize;
-        let spawner_aligned_size =
-            next_multiple_of(GpuSpawnerParams::min_size().get() as usize, item_align);
+        let spawner_aligned_size = GpuSpawnerParams::aligned_size(item_align);
         trace!(
             "Aligning spawner params to {} bytes as device limits requires. Size: {} bytes.",
             item_align,
@@ -403,22 +430,49 @@ impl FromWorld for DispatchIndirectPipeline {
 
         // We need to pad the Spawner WGSL struct based on the device padding so that we
         // can use it as an array element but also has a direct struct binding.
-        let spawner_padding_code = if GpuSpawnerParams::min_size().get() as usize
-            != spawner_aligned_size
-        {
-            let padding_size = spawner_aligned_size - GpuSpawnerParams::min_size().get() as usize;
-            assert!(padding_size % 4 == 0);
-            format!("padding: array<u32, {}>", padding_size / 4)
-        } else {
-            "".to_string()
+        let spawner_padding_code = GpuSpawnerParams::padding_code(item_align);
+        let indirect_code =
+            include_str!("vfx_indirect.wgsl").replace("{{SPAWNER_PADDING}}", &spawner_padding_code);
+
+        // Resolve imports. Because we don't insert this shader into Bevy' pipeline
+        // cache, we don't get that part "for free", so we have to do it manually here.
+        let indirect_naga_module = {
+            let mut composer = Composer::default();
+
+            // Import bevy_hanabi::vfx_common
+            {
+                let common_shader = HanabiPlugin::make_common_shader(item_align);
+                let mut desc: naga_oil::compose::ComposableModuleDescriptor<'_> =
+                    (&common_shader).into();
+                desc.shader_defs.insert(
+                    "SPAWNER_PADDING".to_string(),
+                    naga_oil::compose::ShaderDefValue::Bool(true),
+                );
+                let res = composer.add_composable_module(desc);
+                assert!(res.is_ok());
+            }
+
+            let shader_defs = default();
+
+            match composer.make_naga_module(NagaModuleDescriptor {
+                source: &indirect_code,
+                file_path: "vfx_indirect.wgsl",
+                shader_defs,
+                ..Default::default()
+            }) {
+                Ok(naga_module) => ShaderSource::Naga(Cow::Owned(naga_module)),
+                Err(compose_error) => panic!(
+                    "Failed to compose vfx_indirect.wgsl, naga_oil returned: {}",
+                    compose_error.emit_to_string(&composer)
+                ),
+            }
         };
-        let indirect_code = include_str!("vfx_indirect.wgsl")
-            .to_string()
-            .replace("{{SPAWNER_PADDING}}", &spawner_padding_code);
+
         debug!("Create indirect dispatch shader:\n{}", indirect_code);
+
         let shader_module = render_device.create_shader_module(ShaderModuleDescriptor {
             label: Some("hanabi:vfx_indirect_shader"),
-            source: ShaderSource::Wgsl(Cow::Owned(indirect_code)),
+            source: indirect_naga_module,
         });
 
         let pipeline = render_device.create_compute_pipeline(&RawComputePipelineDescriptor {
@@ -772,7 +826,7 @@ impl SpecializedComputePipeline for ParticlesUpdatePipeline {
                 self.render_indirect_layout.clone(),
             ],
             shader: key.shader,
-            shader_defs: vec![],
+            shader_defs: vec!["RI_MAX_SPAWN_ATOMIC".into()],
             entry_point: "main".into(),
             push_constant_ranges: Vec::new(),
         }
@@ -2049,7 +2103,7 @@ pub(crate) struct BufferBindGroups {
     /// Bind group for the render graphic shader.
     ///
     /// ```wgsl
-    /// @binding(0) var<storage, read> particle_buffer : ParticlesBuffer;
+    /// @binding(0) var<storage, read> particle_buffer : ParticleBuffer;
     /// @binding(1) var<storage, read> indirect_buffer : IndirectBuffer;
     /// @binding(2) var<storage, read> dispatch_indirect : DispatchIndirect;
     /// ```

src/render/shader_cache.rs

@@ -3,7 +3,7 @@ use std::hash::{Hash, Hasher};
 use bevy::{
     asset::{Assets, Handle},
     ecs::{change_detection::ResMut, system::Resource},
-    log::debug,
+    log::{debug, trace},
     render::render_resource::Shader,
     utils::HashMap,
 };
@@ -40,10 +40,17 @@ impl ShaderCache {
             let mut hasher = bevy::utils::AHasher::default();
             source.hash(&mut hasher);
             let hash = hasher.finish();
-            let handle = shaders.add(Shader::from_wgsl(
+            let shader = Shader::from_wgsl(
                 source.to_string(),
                 format!("hanabi/{}_{}.wgsl", filename, hash),
-            ));
+            );
+            trace!(
+                "Shader path={} import_path={:?} imports={:?}",
+                shader.path,
+                shader.import_path,
+                shader.imports
+            );
+            let handle = shaders.add(shader);
             debug!("Inserted new configured shader: {:?}\n{}", handle, source);
             self.cache.insert(source.to_string(), handle.clone());
             handle