1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96

use std::{collections::BTreeMap, sync::RwLock};

use bevy::prelude::{
    default, AssetEvent, Assets, CoreSet, Deref, DerefMut, EventReader, Handle, IntoSystemConfig,
    Plugin, Res, ResMut, Shader,
};
use once_cell::sync::Lazy;
use rust_gpu_builder_shared::RustGpuBuilderOutput;

/// Static container for `RustGpuArtifacts` to allow access from `Material::specialize`
pub static RUST_GPU_ARTIFACTS: Lazy<RwLock<RustGpuArtifacts>> = Lazy::new(default);

pub struct BuilderOutputPlugin;

impl Plugin for BuilderOutputPlugin {
    fn build(&self, app: &mut bevy::prelude::App) {
        #[cfg(feature = "json")]
        app.add_plugin(bevy_common_assets::json::JsonAssetPlugin::<
            RustGpuBuilderOutput,
        >::new(&["rust-gpu.json"]));

        #[cfg(feature = "msgpack")]
        app.add_plugin(bevy_common_assets::msgpack::MsgPackAssetPlugin::<
            RustGpuBuilderOutput,
        >::new(&["rust-gpu.msgpack"]));

        app.add_system(builder_output_events.in_base_set(CoreSet::PreUpdate));
    }
}

/// Module shader handle container.
#[derive(Debug, Clone)]
pub enum RustGpuModules {
    /// Contains a single unnamed shader.
    Single(Handle<Shader>),
    /// Contains multiple named shaders.
    Multi(BTreeMap<String, Handle<Shader>>),
}

/// Asset containing loaded rust-gpu shaders and entry point metadata.
#[derive(Debug, Clone)]
pub struct RustGpuArtifact {
    pub entry_points: Vec<String>,
    pub modules: RustGpuModules,
}

#[derive(Debug, Default, Clone, Deref, DerefMut)]
pub struct RustGpuArtifacts {
    pub artifacts: BTreeMap<Handle<RustGpuBuilderOutput>, RustGpuArtifact>,
}

/// [`RustGpuBuilderOutput`] asset event handler.
///
/// Handles loading shader assets, maintaining static material data, and respecializing materials on reload.
pub fn builder_output_events(
    mut builder_output_events: EventReader<AssetEvent<RustGpuBuilderOutput>>,
    builder_outputs: Res<Assets<RustGpuBuilderOutput>>,
    mut shaders: ResMut<Assets<Shader>>,
) {
    for event in builder_output_events.iter() {
        if let AssetEvent::Created { handle } | AssetEvent::Modified { handle } = event {
            let asset = builder_outputs.get(handle).unwrap().clone();

            // Create a `RustGpuArtifact` from the affected asset
            let artifact = match asset.modules {
                rust_gpu_builder_shared::RustGpuBuilderModules::Single(ref single) => {
                    let shader = shaders.add(Shader::from_spirv(single.clone()));
                    RustGpuArtifact {
                        entry_points: asset.entry_points,
                        modules: RustGpuModules::Single(shader),
                    }
                }
                rust_gpu_builder_shared::RustGpuBuilderModules::Multi(multi) => RustGpuArtifact {
                    entry_points: asset.entry_points,
                    modules: RustGpuModules::Multi(
                        multi
                            .into_iter()
                            .map(|(k, module)| (k.clone(), shaders.add(Shader::from_spirv(module))))
                            .collect(),
                    ),
                },
            };

            // Emplace it in static storage
            RUST_GPU_ARTIFACTS
                .write()
                .unwrap()
                .insert(handle.clone_weak(), artifact);
        }

        // On remove, remove the corresponding artifact from static storage
        if let AssetEvent::Removed { handle } = event {
            RUST_GPU_ARTIFACTS.write().unwrap().remove(handle);
        }
    }
}