AtomicU32 for ids instead of AtomicRefCell poop

darthdeus · Dec 16, 2023 · feed94a · feed94a
1 parent 958e4c1
commit feed94a
Show file tree

Hide file tree

Showing 5 changed files with 34 additions and 26 deletions.
diff --git a/comfy-core/src/render_queues.rs b/comfy-core/src/render_queues.rs
@@ -21,16 +21,19 @@ pub fn get_shader_instance(
 }
 
 pub fn set_uniform(name: impl Into<String>, value: Uniform) {
-    if let Some(instance_id) = &mut *CURRENT_SHADER.borrow_mut() {
+    let instance_id = CURRENT_SHADER_INSTANCE_ID.load(Ordering::SeqCst);
+
+    if instance_id > 0 {
         let mut table = SHADER_UNIFORM_TABLE.borrow_mut();
 
-        if let Some(instance) = table.instances.get(instance_id.0 as usize) {
+        if let Some(instance) = table.instances.get(instance_id 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);
+            CURRENT_SHADER_INSTANCE_ID
+                .store(table.instances.len() as u32 - 1, Ordering::SeqCst);
         } else {
             panic!(
     "Current shader instance id is invalid.
@@ -45,8 +48,7 @@ pub fn set_uniform(name: impl Into<String>, value: Uniform) {
     }
 }
 
-static CURRENT_SHADER: Lazy<AtomicRefCell<Option<ShaderInstanceId>>> =
-    Lazy::new(|| AtomicRefCell::new(None));
+static CURRENT_SHADER_INSTANCE_ID: AtomicU32 = AtomicU32::new(0);
 
 /// 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()`.
@@ -57,21 +59,25 @@ pub fn use_shader(shader_id: ShaderId) {
         .instances
         .push(ShaderInstance { id: shader_id, uniforms: Default::default() });
 
-    *CURRENT_SHADER.borrow_mut() =
-        Some(ShaderInstanceId(table.instances.len() as u32 - 1));
+    CURRENT_SHADER_INSTANCE_ID
+        .store(table.instances.len() as u32 - 1, Ordering::SeqCst);
 }
 
 /// Switches back to the default shader.
 pub fn use_default_shader() {
-    *CURRENT_SHADER.borrow_mut() = None;
+    CURRENT_SHADER_INSTANCE_ID.store(0, Ordering::SeqCst);
 }
 
 /// Returns the current `ShaderInstance` if any. Currently intended only for internal use.
-pub fn get_current_shader() -> Option<ShaderInstanceId> {
-    *CURRENT_SHADER.borrow()
+pub fn get_current_shader() -> ShaderInstanceId {
+    // TODO: we probably don't need SeqCst for any of these, but that can be fixed later
+    ShaderInstanceId(CURRENT_SHADER_INSTANCE_ID.load(Ordering::SeqCst))
 }
 
-use std::{collections::BTreeMap, sync::atomic::AtomicU64};
+use std::{
+    collections::BTreeMap,
+    sync::atomic::{AtomicU32, AtomicU64, Ordering},
+};
 
 static SHADER_IDS: AtomicU64 = AtomicU64::new(0);
 
@@ -87,7 +93,7 @@ pub fn gen_shader_id() -> ShaderId {
 /// Represents a set of shader uniform parameters.
 ///
 /// u32 ID is exposed for debugging purposes only, do not modify by hand.
-#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)]
+#[derive(Copy, Clone, Debug, Default, PartialEq, Eq, Hash, PartialOrd, Ord)]
 pub struct ShaderInstanceId(pub u32);
 
 static TEXT_QUEUE: Lazy<AtomicRefCell<Vec<DrawText>>> =
@@ -108,15 +114,14 @@ pub type RenderQueue = Vec<Mesh>;
 #[derive(Default)]
 struct RenderQueues {
     data: BTreeMap<MeshGroupKey, RenderQueue>,
-    // data: FxHashMap<MeshGroupKey, RenderQueue>,
 }
 
 #[derive(Copy, Clone, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)]
 pub struct MeshGroupKey {
     pub z_index: i32,
     pub blend_mode: BlendMode,
     pub texture_id: TextureHandle,
-    pub shader: Option<ShaderInstanceId>,
+    pub shader: ShaderInstanceId,
     pub render_target: RenderTargetId,
 }
 

diff --git a/comfy-wgpu/src/batching.rs b/comfy-wgpu/src/batching.rs
@@ -16,7 +16,7 @@ pub fn run_batched_render_passes(
             z_index: 0,
             blend_mode: BlendMode::Alpha,
             texture_id: TextureHandle::from_path("1px"),
-            shader: None,
+            shader: ShaderInstanceId::default(),
             render_target: RenderTargetId::default(),
         },
         RenderPassData {
@@ -150,7 +150,7 @@ pub fn run_batched_render_passes(
             MeshDrawData {
                 blend_mode: BlendMode::Alpha,
                 texture: TextureHandle::from_path("1px"),
-                shader: None,
+                shader: ShaderInstanceId::default(),
                 render_target: RenderTargetId::default(),
                 data: SmallVec::new(),
             },

diff --git a/comfy-wgpu/src/pipelines.rs b/comfy-wgpu/src/pipelines.rs
@@ -112,15 +112,18 @@ pub fn ensure_pipeline_exists(
 
     let maybe_shader_instance_id = pass_data.shader;
 
-    let maybe_shader =
-        maybe_shader_instance_id.as_ref().and_then(|instance_id| {
-            let instance = get_shader_instance(*instance_id);
+    let maybe_shader = {
+        if maybe_shader_instance_id.0 > 0 {
+            let instance = get_shader_instance(maybe_shader_instance_id);
             shaders.get(instance.id)
-        });
+        } else {
+            None
+        }
+    };
 
     let name = format!(
         "{} {:?} {:?} {:?}",
-        if maybe_shader_instance_id.is_some() {
+        if maybe_shader_instance_id.0 > 0 {
             "USER(Mesh)"
         } else {
             "BUILTIN(Mesh)"
@@ -163,8 +166,8 @@ pub fn ensure_pipeline_exists(
     };
 
     if let RenderPipeline::User(user_pipeline) = mesh_pipeline {
-        if let Some(shader_instance_id) = maybe_shader_instance_id {
-            let shader_instance = get_shader_instance(shader_instance_id);
+        if maybe_shader_instance_id.0 > 0 {
+            let shader_instance = get_shader_instance(maybe_shader_instance_id);
             let shader = shaders.get(shader_instance.id).unwrap();
 
             for (buffer_name, buffer) in

diff --git a/comfy-wgpu/src/render_pass.rs b/comfy-wgpu/src/render_pass.rs
@@ -3,7 +3,7 @@ use crate::*;
 pub struct MeshDrawData {
     pub blend_mode: BlendMode,
     pub texture: TextureHandle,
-    pub shader: Option<ShaderInstanceId>,
+    pub shader: ShaderInstanceId,
     pub render_target: RenderTargetId,
     pub data: smallvec::SmallVec<[Mesh; 1]>,
 }

diff --git a/comfy/src/update_stages.rs b/comfy/src/update_stages.rs
@@ -1032,7 +1032,7 @@ fn renderer_update(c: &mut EngineContext) {
             z_index: particle_system.z_index,
             blend_mode: p.blend_mode,
             texture_id: p.texture,
-            shader: None,
+            shader: ShaderInstanceId::default(),
             render_target: RenderTargetId::default(),
         };
 
@@ -1065,7 +1065,7 @@ fn renderer_update(c: &mut EngineContext) {
                 z_index: p.z_index,
                 blend_mode: p.blend_mode,
                 texture_id: p.texture,
-                shader: None,
+                shader: ShaderInstanceId::default(),
                 render_target: RenderTargetId::default(),
             })
             .or_default()