Fix new uniform instance per set_uniform

comfy-core/src/render_queues.rs

@@ -5,14 +5,83 @@ static SHADER_UNIFORM_TABLE: Lazy<AtomicRefCell<ShaderUniformTable>> =
 
 #[derive(Default)]
 pub struct ShaderUniformTable {
-    next_id: u32,
-    table: Vec<ShaderInstance>,
+    instances: Vec<ShaderInstance>,
+}
+
+pub fn clear_shader_uniform_table() {
+    SHADER_UNIFORM_TABLE.borrow_mut().instances.clear();
 }
 
 pub fn get_shader_instance(
     id: ShaderInstanceId,
 ) -> AtomicRef<'static, ShaderInstance> {
-    AtomicRef::map(SHADER_UNIFORM_TABLE.borrow(), |x| &x.table[id.0 as usize])
+    AtomicRef::map(SHADER_UNIFORM_TABLE.borrow(), |x| {
+        &x.instances[id.0 as usize]
+    })
+}
+
+pub fn set_uniform(name: impl Into<String>, value: Uniform) {
+    if let Some(instance_id) = &mut *CURRENT_SHADER.borrow_mut() {
+        let mut table = SHADER_UNIFORM_TABLE.borrow_mut();
+
+        if let Some(instance) = table.instances.get(instance_id.0 as usize) {
+            let mut new_instance = instance.clone();
+            new_instance.uniforms.insert(name.into(), value);
+
+            table.instances.push(new_instance);
+
+            *instance_id = ShaderInstanceId(table.instances.len() as u32 - 1);
+        } else {
+            panic!(
+    "Current shader instance id is invalid.
+
+                This is likely a bug, \
+     please report an issue on https://github.com/darthdeus/comfy/issues with \
+     some information on what you did."
+);
+        }
+    } else {
+        panic!("Trying to set a uniform with no shader active");
+    }
+}
+
+static CURRENT_SHADER: Lazy<AtomicRefCell<Option<ShaderInstanceId>>> =
+    Lazy::new(|| AtomicRefCell::new(None));
+
+/// Switches to the shader with the given ID. The shader must already exist. To revert back to the
+/// default shader simply call `use_default_shader()`.
+pub fn use_shader(shader_id: ShaderId) {
+    let mut table = SHADER_UNIFORM_TABLE.borrow_mut();
+
+    table
+        .instances
+        .push(ShaderInstance { id: shader_id, uniforms: Default::default() });
+
+    *CURRENT_SHADER.borrow_mut() =
+        Some(ShaderInstanceId(table.instances.len() as u32 - 1));
+}
+
+/// Switches back to the default shader.
+pub fn use_default_shader() {
+    *CURRENT_SHADER.borrow_mut() = None;
+}
+
+/// Returns the current `ShaderInstance` if any. Currently intended only for internal use.
+pub fn get_current_shader() -> Option<ShaderInstanceId> {
+    *CURRENT_SHADER.borrow()
+}
+
+use std::sync::atomic::AtomicU64;
+
+static SHADER_IDS: AtomicU64 = AtomicU64::new(0);
+
+/// Generates a new ShaderId. This is intended for internal use only.
+pub fn gen_shader_id() -> ShaderId {
+    let id = SHADER_IDS.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
+
+    info!("Generated ShaderId: {}", id);
+
+    ShaderId(id)
 }
 
 /// Represents a set of shader uniform parameters.
@@ -59,12 +128,10 @@ pub fn consume_render_queues() -> HashMap<MeshGroupKey, RenderQueue> {
 }
 
 pub fn queue_mesh_draw(mesh: Mesh, blend_mode: BlendMode) {
-    // let shader = get_current_shader();
+    let shader = get_current_shader();
     let render_target = get_current_render_target();
     let white_px = TextureHandle::from_path("1px");
 
-    let shader = None;
-
     RENDER_QUEUES
         .borrow_mut()
         .data

comfy-core/src/shaders.rs

@@ -1,7 +1,4 @@
 use crate::*;
-use std::sync::atomic::AtomicU64;
-
-static SHADER_IDS: AtomicU64 = AtomicU64::new(0);
 
 #[derive(Debug)]
 pub struct ShaderMap {
@@ -65,7 +62,7 @@ impl UniformDef {
     pub fn to_wgsl(&self) -> &str {
         match self {
             UniformDef::F32(_) => "f32",
-            UniformDef::Custom { wgsl_decl, .. } => &wgsl_decl,
+            UniformDef::Custom { wgsl_decl, .. } => wgsl_decl,
         }
     }
 }
@@ -91,46 +88,11 @@ pub fn get_current_render_target() -> Option<RenderTargetId> {
     *CURRENT_RENDER_TARGET.borrow()
 }
 
-static CURRENT_SHADER: Lazy<AtomicRefCell<Option<ShaderInstance>>> =
-    Lazy::new(|| AtomicRefCell::new(None));
-
-/// Switches to the shader with the given ID. The shader must already exist. To revert back to the
-/// default shader simply call `use_default_shader()`.
-pub fn use_shader(shader_id: ShaderId) {
-    *CURRENT_SHADER.borrow_mut() =
-        Some(ShaderInstance { id: shader_id, uniforms: Default::default() });
-}
-
-/// Switches back to the default shader.
-pub fn use_default_shader() {
-    *CURRENT_SHADER.borrow_mut() = None;
-}
-
-/// Returns the current `ShaderInstance` if any. Currently intended only for internal use.
-pub fn get_current_shader() -> Option<ShaderInstance> {
-    CURRENT_SHADER.borrow().clone()
-}
-
-/// Generates a new ShaderId. This is intended for internal use only.
-pub fn gen_shader_id() -> ShaderId {
-    let id = SHADER_IDS.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
-
-    info!("Generated ShaderId: {}", id);
-
-    ShaderId(id)
-}
-
 /// Sets a `f32` uniform value by name. The uniform must exist in the shader.
 pub fn set_uniform_f32(name: impl Into<String>, value: f32) {
     set_uniform(name, Uniform::F32(OrderedFloat(value)));
 }
 
-pub fn set_uniform(name: impl Into<String>, value: Uniform) {
-    if let Some(shader) = &mut *CURRENT_SHADER.borrow_mut() {
-        shader.uniforms.insert(name.into(), value);
-    }
-}
-
 /// Creates a new shader and returns its ID. The `source` parameter should only contain the
 /// fragment function, as the rest of the shader is automatically generated.
 ///

comfy/src/update_stages.rs

@@ -101,6 +101,7 @@ pub(crate) fn run_late_update_stages(c: &mut EngineContext, delta: f32) {
     show_errors(c);
     commands().run_on(&mut world_mut());
     world_mut().flush();
+    clear_shader_uniform_table();
 }
 
 fn dev_hotkeys(_c: &EngineContext) {