more shader assignment work

src/app/app_builder.rs

@@ -16,6 +16,7 @@ use crate::{
 };
 
 use bevy_transform::{prelude::LocalToWorld, transform_system_bundle};
+use render_graph_2::CompiledShaderMap;
 use std::collections::HashMap;
 
 pub struct AppBuilder {
@@ -166,6 +167,7 @@ impl AppBuilder {
         resources.insert(AssetStorage::<Texture>::new());
         resources.insert(AssetStorage::<Shader>::new());
         resources.insert(ShaderAssignments::new());
+        resources.insert(CompiledShaderMap::new());
         self
     }
 

src/asset/mod.rs

@@ -8,7 +8,7 @@ pub use texture::*;
 
 use std::{collections::HashMap, marker::PhantomData};
 
-#[derive(Hash)]
+#[derive(Hash, Eq, PartialEq)]
 pub struct Handle<T> {
     pub id: usize,
     marker: PhantomData<T>,
@@ -75,7 +75,11 @@ impl<T> AssetStorage<T> {
         handle
     }
 
-    pub fn get(&mut self, id: usize) -> Option<&mut T> {
+    pub fn get_id(&mut self, id: usize) -> Option<&mut T> {
         self.assets.get_mut(&id)
     }
+
+    pub fn get(&mut self, handle: &Handle<T>) -> Option<&mut T> {
+        self.assets.get_mut(&handle.id)
+    }
 }

src/render/passes/forward/forward_pipeline.rs

@@ -157,14 +157,14 @@ impl Pipeline for ForwardPipeline {
             }
 
             if should_load_mesh {
-                if let Some(mesh_asset) = mesh_storage.get(mesh.id) {
+                if let Some(mesh_asset) = mesh_storage.get_id(mesh.id) {
                     mesh_asset.setup_buffers(&render_graph.device);
                     pass.set_index_buffer(mesh_asset.index_buffer.as_ref().unwrap(), 0);
                     pass.set_vertex_buffers(0, &[(&mesh_asset.vertex_buffer.as_ref().unwrap(), 0)]);
                 };
             }
 
-            if let Some(ref mesh_asset) = mesh_storage.get(mesh.id) {
+            if let Some(ref mesh_asset) = mesh_storage.get_id(mesh.id) {
                 pass.set_bind_group(1, material.bind_group.as_ref().unwrap(), &[]);
                 pass.draw_indexed(0..mesh_asset.indices.len() as u32, 0, 0..1);
             };

src/render/passes/forward_instanced/mod.rs

@@ -285,7 +285,7 @@ impl Pipeline for ForwardInstancedPipeline {
 
         let mut mesh_storage = world.resources.get_mut::<AssetStorage<Mesh>>().unwrap();
         for instance_buffer_info in self.instance_buffer_infos.as_ref().unwrap().iter() {
-            if let Some(mesh_asset) = mesh_storage.get(instance_buffer_info.mesh_id) {
+            if let Some(mesh_asset) = mesh_storage.get_id(instance_buffer_info.mesh_id) {
                 mesh_asset.setup_buffers(&render_graph.device);
                 pass.set_index_buffer(mesh_asset.index_buffer.as_ref().unwrap(), 0);
                 pass.set_vertex_buffers(0, &[(&mesh_asset.vertex_buffer.as_ref().unwrap(), 0)]);

src/render/passes/forward_shadow/mod.rs

@@ -176,7 +176,7 @@ impl Pipeline for ForwardShadowPassNew {
 
         let mut mesh_storage = world.resources.get_mut::<AssetStorage<Mesh>>().unwrap();
         for (material, mesh) in mesh_query.iter(world) {
-            if let Some(mesh_asset) = mesh_storage.get(mesh.id) {
+            if let Some(mesh_asset) = mesh_storage.get_id(mesh.id) {
                 mesh_asset.setup_buffers(&render_graph.device);
                 pass.set_bind_group(1, material.bind_group.as_ref().unwrap(), &[]);
                 pass.set_index_buffer(mesh_asset.index_buffer.as_ref().unwrap(), 0);

src/render/passes/shadow/shadow_pipeline.rs

@@ -166,7 +166,7 @@ impl Pipeline for ShadowPipeline {
 
         let mut mesh_storage = world.resources.get_mut::<AssetStorage<Mesh>>().unwrap();
         for (material, mesh) in mesh_query.iter(world) {
-            if let Some(mesh_asset) = mesh_storage.get(mesh.id) {
+            if let Some(mesh_asset) = mesh_storage.get_id(mesh.id) {
                 mesh_asset.setup_buffers(&render_graph.device);
 
                 pass.set_bind_group(1, material.bind_group.as_ref().unwrap(), &[]);

src/render/passes/ui/mod.rs

@@ -242,7 +242,7 @@ impl Pipeline for UiPipeline {
 
         let mut mesh_storage = world.resources.get_mut::<AssetStorage<Mesh>>().unwrap();
         for instance_buffer_info in instance_buffer_infos.as_ref().unwrap().iter() {
-            if let Some(mesh_asset) = mesh_storage.get(instance_buffer_info.mesh_id) {
+            if let Some(mesh_asset) = mesh_storage.get_id(instance_buffer_info.mesh_id) {
                 mesh_asset.setup_buffers(&render_graph.device);
                 pass.set_index_buffer(mesh_asset.index_buffer.as_ref().unwrap(), 0);
                 pass.set_vertex_buffers(0, &[(&mesh_asset.vertex_buffer.as_ref().unwrap(), 0)]);

src/render/render_graph_2/draw_targets/mesh_draw_target.rs

@@ -26,7 +26,7 @@ pub fn mesh_draw_target(world: &World, render_pass: &mut dyn RenderPass) {
         }
 
         if should_load_mesh {
-            if let Some(mesh_asset) = mesh_storage.get(mesh.id) {
+            if let Some(mesh_asset) = mesh_storage.get_id(mesh.id) {
                 let renderer = render_pass.get_renderer();
                 renderer.create_buffer_with_data(
                     resource_name::buffer::TEMP_MESH_VERTEX_BUFFER_NAME,

src/render/render_graph_2/draw_targets/ui_draw_target.rs

@@ -20,7 +20,7 @@ pub fn ui_draw_target(world: &World, render_pass: &mut dyn RenderPass) {
         };
 
         if let Some((instance_count, mesh_id)) = result {
-            if let Some(mesh_asset) = mesh_storage.get(mesh_id) {
+            if let Some(mesh_asset) = mesh_storage.get_id(mesh_id) {
                 renderer.create_buffer_with_data(
                     resource_name::buffer::TEMP_MESH_VERTEX_BUFFER_NAME,
                     mesh_asset.vertices.as_bytes(),

src/render/render_graph_2/renderable.rs

@@ -1,4 +1,7 @@
-use crate::{asset::{AssetStorage, Handle}, render::Shader, render::render_graph_2::RenderGraph};
+use crate::{
+    asset::{AssetStorage, Handle},
+    render::{render_graph_2::RenderGraph, Shader},
+};
 use legion::prelude::*;
 use std::collections::{HashMap, HashSet};
 
@@ -18,6 +21,19 @@ impl Default for Renderable {
     }
 }
 
+pub struct CompiledShaderMap {
+    // TODO: need macro hash lookup
+    pub source_to_compiled: HashMap<Handle<Shader>, Vec<(HashSet<String>, Handle<Shader>)>>,
+}
+
+impl CompiledShaderMap {
+    pub fn new() -> Self {
+        CompiledShaderMap {
+            source_to_compiled: HashMap::new(),
+        }
+    }
+}
+
 pub struct ShaderAssignments {
     pub assignments: HashMap<usize, Vec<Entity>>,
 }
@@ -35,10 +51,27 @@ pub fn update_shader_assignments(world: &mut World, render_graph: &mut RenderGra
     // lots of string + hashset allocations. sees uniform_resource_provider for more context
     {
         let shader_assignments = world.resources.get_mut::<ShaderAssignments>().unwrap();
-        let shader_storage = world.resources.get_mut::<AssetStorage<Shader>>().unwrap();
+        let mut compiled_shader_map = world.resources.get_mut::<CompiledShaderMap>().unwrap();
+        let mut shader_storage = world.resources.get_mut::<AssetStorage<Shader>>().unwrap();
         for (entity, renderable) in <Read<Renderable>>::query().iter_entities(world) {
             for shader in renderable.shaders.iter() {
-                
+                if let None = compiled_shader_map.source_to_compiled.get(shader) {
+                    compiled_shader_map
+                        .source_to_compiled
+                        .insert(shader.clone(), Vec::new());
+                }
+
+                let compiled_shaders = compiled_shader_map.source_to_compiled.get_mut(shader).unwrap();
+                if let None = compiled_shaders.iter().find(|(shader_defs, _shader)| *shader_defs == renderable.shader_defs) {
+                    let shader_resource = shader_storage.get(shader).unwrap();
+                    let shader_def_vec = renderable.shader_defs.iter().cloned().collect::<Vec<String>>();
+                    let compiled_shader = shader_resource.get_spirv_shader(Some(&shader_def_vec));
+                    compiled_shaders.push((renderable.shader_defs.clone(), shader.clone()));
+                    let compiled_shader_handle = shader_storage.add(compiled_shader);
+                    // TODO: collecting assigments in a map means they won't be removed when the macro changes
+                    // TODO: need to somehow grab base shader's pipeline, then copy it 
+                    // shader_assignments.assignments.insert()
+                }
             }
         }
     }
@@ -47,4 +80,4 @@ pub fn update_shader_assignments(world: &mut World, render_graph: &mut RenderGra
     for mut renderable in <Write<Renderable>>::query().iter_mut(world) {
         renderable.shader_defs = HashSet::new();
     }
-}
+}

src/render/shader.rs

@@ -1,6 +1,6 @@
 use std::marker::Copy;
 
-#[derive(Copy, Clone)]
+#[derive(Hash, Eq, PartialEq, Copy, Clone)]
 pub enum ShaderStage {
     Vertex,
     Fragment,
@@ -43,15 +43,18 @@ pub fn glsl_to_spirv(
     binary_result.as_binary().into()
 }
 
+#[derive(Hash, Eq, PartialEq)]
 pub enum ShaderSource {
     Spirv(Vec<u32>),
     Glsl(String),
 }
 
+#[derive(Hash, Eq, PartialEq)]
 pub struct Shader {
     pub source: ShaderSource,
     pub stage: ShaderStage,
     pub entry_point: String,
+    // TODO: add "precompile" flag?
 }
 
 impl Shader {