feat(engine): improve flexibility of binding shader bindings

4 files changed, 266 insertions, 122 deletions

diff --git a/engine/src/rendering.rs b/engine/src/rendering.rs
index 140d99e..1bb11cc 100644
--- a/engine/src/rendering.rs
+++ b/engine/src/rendering.rs
@@ -219,7 +219,7 @@ pub enum Command
     SwapBuffers(SurfaceId),
     CreateShaderProgram(ObjectId, ShaderProgram),
     ActivateShader(ObjectId),
-    SetShaderBinding(ShaderBinding),
+    SetShaderBinding(ObjectId, ShaderBinding),
     CreateTexture
     {
         obj_id: ObjectId,
diff --git a/engine/src/rendering/backend/opengl.rs b/engine/src/rendering/backend/opengl.rs
index 2f373fc..cd53bc5 100644
--- a/engine/src/rendering/backend/opengl.rs
+++ b/engine/src/rendering/backend/opengl.rs
@@ -23,6 +23,10 @@ use opengl_bindings::blending::{
     Equation as GlBlendingEquation,
     Factor as GlBlendingFactor,
 };
+use opengl_bindings::buffer::{
+    BindingTarget as GlBufferBindingTarget,
+    Buffer as GlBuffer,
+};
 use opengl_bindings::debug::{
     set_debug_message_callback,
     set_debug_message_control,
@@ -104,6 +108,7 @@ use crate::rendering::object::{
 };
 use crate::rendering::shader::cursor::{
     Binding as ShaderBinding,
+    BindingLocation as ShaderBindingLocation,
     BindingValue as ShaderBindingValue,
 };
 use crate::rendering::shader::{
@@ -164,8 +169,6 @@ struct GraphicsContext
 {
     gl_context: Option<MaybeCurrentContextWithFns>,
     surfaces: HashMap<SurfaceId, GraphicsContextSurface>,
-    shader_uniform_buffer_objs:
-        HashMap<ObjectId, HashMap<u32, opengl_bindings::buffer::Buffer<u8>>>,
     backend_resources: BackendResourceStore,
 }
 
@@ -368,6 +371,51 @@ struct GraphicsContextSurface
     size: PhysicalSize<u32>,
 }
 
+#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
+struct ShaderBindingIndex(u32);
+
+impl intmap::IntKey for ShaderBindingIndex
+{
+    type Int = u32;
+
+    const PRIME: Self::Int = <u32 as intmap::IntKey>::PRIME;
+
+    fn into_int(self) -> Self::Int
+    {
+        self.0
+    }
+}
+
+#[derive(Debug)]
+enum BackendShaderBinding
+{
+    Uniform(GlBuffer<u8>),
+    Texture(GlTexture),
+}
+
+impl BackendShaderBinding
+{
+    fn bind(
+        &self,
+        gl_context: &MaybeCurrentContextWithFns,
+        binding_index: ShaderBindingIndex,
+    )
+    {
+        match self {
+            BackendShaderBinding::Uniform(uniform_buf) => {
+                uniform_buf.bind_to_indexed_target(
+                    gl_context,
+                    GlBufferBindingTarget::UniformBuffer,
+                    binding_index.0,
+                );
+            }
+            BackendShaderBinding::Texture(texture) => {
+                texture.bind_to_texture_unit(gl_context, binding_index.0)
+            }
+        }
+    }
+}
+
 #[derive(Debug)]
 enum BackendResource
 {
@@ -380,6 +428,7 @@ enum BackendResource
     {
         program: GlShaderProgram,
         program_metadata: ShaderProgramMetadata,
+        bindings: IntMap<ShaderBindingIndex, BackendShaderBinding>,
     },
     Texture
     {
@@ -716,7 +765,6 @@ fn handle_commands(
     let GraphicsContext {
         ref gl_context,
         ref mut surfaces,
-        ref mut shader_uniform_buffer_objs,
         ref mut backend_resources,
     } = *graphics_ctx;
 
@@ -845,6 +893,7 @@ fn handle_commands(
                                     &shader_program,
                                 )?,
                                 program_metadata: shader_program.metadata().clone(),
+                                bindings: IntMap::with_capacity(4),
                             })
                         },
                     )
@@ -852,135 +901,58 @@ fn handle_commands(
                     tracing::error!("Failed to create shader program: {err}");
                 }
             }
-            Command::ActivateShader(object_id) => {
+            Command::ActivateShader(shader_object_id) => {
                 let Some(resource) = backend_resources.get_resource(
                     object_store,
-                    object_id,
+                    shader_object_id,
                     ObjectKind::ShaderProgram,
                 ) else {
                     continue;
                 };
 
-                let BackendResource::Shader { program, program_metadata } = resource
+                let BackendResource::Shader { program, program_metadata, bindings } =
+                    resource
                 else {
                     unreachable!();
                 };
 
                 program.activate(gl_context);
 
-                activated_gl_shader_program =
-                    Some((object_id, program.clone(), program_metadata.clone()));
-            }
-            Command::SetShaderBinding(ShaderBinding {
-                location: binding_location,
-                value: binding_value,
-            }) => {
-                let Some((activated_gl_shader_program_obj_id, _, _)) =
-                    &activated_gl_shader_program
-                else {
-                    tracing::error!("No shader program is activated");
-                    continue;
-                };
-
-                if let ShaderBindingValue::Texture(texture_object_id, _) = &binding_value
-                {
-                    let Some(texture_resource) = backend_resources.get_resource(
-                        object_store,
-                        *texture_object_id,
-                        ObjectKind::Texture,
-                    ) else {
-                        continue;
-                    };
-
-                    let BackendResource::Texture { texture: gl_texture, .. } =
-                        texture_resource
-                    else {
-                        unreachable!();
-                    };
-
-                    gl_texture
-                        .bind_to_texture_unit(gl_context, binding_location.binding_index);
-
-                    // gl_shader_program.set_uniform_at_location(
-                    //     curr_gl_ctx,
-                    //     GlUniformLocation::from_number(
-                    //         binding_location.binding_index as i32,
-                    //     ),
-                    //     &binding_location.binding_index,
-                    // );
-
-                    continue;
+                for (binding_index, binding) in bindings.iter() {
+                    binding.bind(gl_context, binding_index);
                 }
 
-                let binding_index = binding_location.binding_index;
-
-                let uniform_buffer_objs = shader_uniform_buffer_objs
-                    .entry(*activated_gl_shader_program_obj_id)
-                    .or_default();
-
-                let uniform_buffer =
-                    uniform_buffer_objs.entry(binding_index).or_insert_with(|| {
-                        let uniform_buf =
-                            opengl_bindings::buffer::Buffer::<u8>::new(gl_context);
-
-                        uniform_buf
-                            .init(
-                                gl_context,
-                                binding_location.binding_size,
-                                opengl_bindings::buffer::Usage::Dynamic,
-                            )
-                            .unwrap();
-
-                        uniform_buf
-                    });
-
-                // The index into the uniform buffer binding target is for whatever
-                // shader program is currently bound so the uniform buffer object has
-                // to be re-bound so that a uniform buffer from another shader isn't
-                // used
-                uniform_buffer.bind_to_indexed_target(
+                activated_gl_shader_program =
+                    Some((shader_object_id, program.clone(), program_metadata.clone()));
+            }
+            Command::SetShaderBinding(
+                shader_object_id,
+                ShaderBinding {
+                    location: binding_location,
+                    value: binding_value,
+                },
+            ) => {
+                let binding = set_shader_binding(
+                    backend_resources,
+                    object_store,
                     gl_context,
-                    opengl_bindings::buffer::BindingTarget::UniformBuffer,
-                    binding_index as u32,
+                    shader_object_id.clone(),
+                    &binding_location,
+                    binding_value,
                 );
 
-                let fvec3_value;
-                let mut color_value = [0.0; 4];
-
-                uniform_buffer
-                    .store_at_byte_offset(
-                        gl_context,
-                        binding_location.byte_offset,
-                        match binding_value {
-                            ShaderBindingValue::Uint(ref value) => value.as_bytes(),
-                            ShaderBindingValue::Int(ref value) => value.as_bytes(),
-                            ShaderBindingValue::Float(ref value) => value.as_bytes(),
-                            ShaderBindingValue::FVec3(value) => {
-                                fvec3_value = [value.x, value.y, value.z];
-                                fvec3_value.as_bytes()
-                            }
-                            ShaderBindingValue::Color(value) => match value {
-                                Color::Rgb(value) => {
-                                    color_value[..3]
-                                        .copy_from_slice(&[value.r, value.g, value.b]);
-
-                                    color_value[..3].as_bytes()
-                                }
-                                Color::Rgba(value) => {
-                                    color_value[..4].copy_from_slice(&[
-                                        value.r, value.g, value.b, value.a,
-                                    ]);
-
-                                    color_value[..4].as_bytes()
-                                }
-                            },
-                            ShaderBindingValue::FMat4x4(ref value) => {
-                                value.items().as_bytes()
-                            }
-                            ShaderBindingValue::Texture(_, _) => unreachable!(),
-                        },
-                    )
-                    .unwrap();
+                if let Some(binding) = binding {
+                    if Some(shader_object_id)
+                        == activated_gl_shader_program
+                            .as_ref()
+                            .map(|(obj_id, ..)| *obj_id)
+                    {
+                        binding.bind(
+                            gl_context,
+                            ShaderBindingIndex(binding_location.binding_index),
+                        );
+                    }
+                }
             }
             Command::CreateTexture {
                 obj_id,
@@ -1527,6 +1499,146 @@ fn update_framebuffer_properties(
     }
 }
 
+fn set_shader_binding<'backend_resources>(
+    backend_resources: &'backend_resources mut BackendResourceStore,
+    object_store: &ObjectStore,
+    gl_context: &MaybeCurrentContextWithFns,
+    shader_object_id: ObjectId,
+    binding_location: &ShaderBindingLocation,
+    binding_value: ShaderBindingValue,
+) -> Option<&'backend_resources BackendShaderBinding>
+{
+    if let ShaderBindingValue::Texture(texture_object_id, _) = &binding_value {
+        let gl_texture = match backend_resources.get_resource(
+            object_store,
+            *texture_object_id,
+            ObjectKind::Texture,
+        ) {
+            Some(BackendResource::Texture { texture, .. }) => texture.clone(),
+            Some(_) => unreachable!(),
+            None => return None,
+        };
+
+        let shader_bindings = match backend_resources.get_resource_mut(
+            object_store,
+            shader_object_id,
+            ObjectKind::ShaderProgram,
+        ) {
+            Some(BackendResource::Shader { bindings, .. }) => bindings,
+            Some(_) => unreachable!(),
+            None => return None,
+        };
+
+        if let Some(prev_binding) =
+            shader_bindings.remove(ShaderBindingIndex(binding_location.binding_index))
+        {
+            // TODO: Textures should probably also be handled somehow here
+            if let BackendShaderBinding::Uniform(prev_binding_uniform_buf) = prev_binding
+            {
+                prev_binding_uniform_buf.delete(gl_context);
+            }
+        }
+
+        return Some(
+            shader_bindings
+                .entry(ShaderBindingIndex(binding_location.binding_index))
+                .set_or_insert_with(|| BackendShaderBinding::Texture(gl_texture.clone())),
+        );
+
+        // if Some(shader_object_id)
+        //     == activated_gl_shader_program
+        //         .as_ref()
+        //         .map(|(obj_id, ..)| *obj_id)
+        // {
+        //     BackendShaderBinding::Texture(gl_texture.clone()).bind(
+        //         gl_context,
+        //         ShaderBindingIndex(binding_location.binding_index),
+        //     );
+        // }
+    }
+
+    let shader_bindings = match backend_resources.get_resource_mut(
+        object_store,
+        shader_object_id,
+        ObjectKind::ShaderProgram,
+    ) {
+        Some(BackendResource::Shader { bindings, .. }) => bindings,
+        Some(_) => unreachable!(),
+        None => return None,
+    };
+
+    let binding =
+        match shader_bindings.get(ShaderBindingIndex(binding_location.binding_index)) {
+            Some(binding @ BackendShaderBinding::Uniform(_)) => binding,
+            Some(_) | None => {
+                let uniform_buf = GlBuffer::<u8>::new(gl_context);
+
+                uniform_buf
+                    .init(
+                        gl_context,
+                        binding_location.binding_size,
+                        opengl_bindings::buffer::Usage::Dynamic,
+                    )
+                    .unwrap();
+
+                shader_bindings
+                    .entry(ShaderBindingIndex(binding_location.binding_index))
+                    .set_or_insert_with(|| BackendShaderBinding::Uniform(uniform_buf))
+            }
+        };
+
+    let BackendShaderBinding::Uniform(binding_uniform_buf) = binding else {
+        unreachable!();
+    };
+
+    let fvec3_value;
+    let mut color_value = [0.0; 4];
+
+    binding_uniform_buf
+        .store_at_byte_offset(
+            gl_context,
+            binding_location.byte_offset,
+            match binding_value {
+                ShaderBindingValue::Uint(ref value) => value.as_bytes(),
+                ShaderBindingValue::Int(ref value) => value.as_bytes(),
+                ShaderBindingValue::Float(ref value) => value.as_bytes(),
+                ShaderBindingValue::FVec3(value) => {
+                    fvec3_value = [value.x, value.y, value.z];
+                    fvec3_value.as_bytes()
+                }
+                ShaderBindingValue::Color(value) => match value {
+                    Color::Rgb(value) => {
+                        color_value[..3].copy_from_slice(&[value.r, value.g, value.b]);
+
+                        color_value[..3].as_bytes()
+                    }
+                    Color::Rgba(value) => {
+                        color_value[..4]
+                            .copy_from_slice(&[value.r, value.g, value.b, value.a]);
+
+                        color_value[..4].as_bytes()
+                    }
+                },
+                ShaderBindingValue::FMat4x4(ref value) => value.items().as_bytes(),
+                ShaderBindingValue::Texture(_, _) => unreachable!(),
+            },
+        )
+        .unwrap();
+
+    Some(binding)
+
+    // if Some(shader_object_id)
+    //     == activated_gl_shader_program
+    //         .as_ref()
+    //         .map(|(obj_id, ..)| *obj_id)
+    // {
+    //     BackendShaderBinding::Uniform(binding_uniform_buf.clone()).bind(
+    //         gl_context,
+    //         ShaderBindingIndex(binding_location.binding_index),
+    //     );
+    // }
+}
+
 fn draw_mesh(
     current_context: &MaybeCurrentContextWithFns,
     graphics_mesh: &GraphicsMesh,
@@ -1901,3 +2013,28 @@ fn blending_equation_to_gl(blending_equation: BlendingEquation) -> GlBlendingEqu
         BlendingEquation::Max => GlBlendingEquation::Max,
     }
 }
+
+trait IntMapEntryExt<'map, Value>
+{
+    fn set_or_insert_with(self, func: impl FnOnce() -> Value) -> &'map mut Value;
+}
+
+impl<'map, Key, Value> IntMapEntryExt<'map, Value> for intmap::Entry<'map, Key, Value>
+where
+    Key: intmap::IntKey,
+    Value: 'map,
+{
+    fn set_or_insert_with(self, func: impl FnOnce() -> Value) -> &'map mut Value
+    {
+        match self {
+            Self::Occupied(occupied) => {
+                let value = occupied.into_mut();
+
+                *value = func();
+
+                value
+            }
+            Self::Vacant(vacant) => vacant.insert(func()),
+        }
+    }
+}
diff --git a/engine/src/rendering/main_render_pass.rs b/engine/src/rendering/main_render_pass.rs
index 895517b..3748d2e 100644
--- a/engine/src/rendering/main_render_pass.rs
+++ b/engine/src/rendering/main_render_pass.rs
@@ -357,6 +357,7 @@ pub fn add_main_render_pass(
                 &point_light_query,
                 &directional_light_query,
                 shader_program,
+                ObjectId::Asset(main_3d_shader_asset.id()),
                 white_1x1_tex_obj_id,
             )?;
 
@@ -463,9 +464,10 @@ pub fn add_main_render_pass(
             ];
 
             for binding in sky_box_shader_bindings {
-                render_pass
-                    .commands
-                    .push(Command::SetShaderBinding(binding));
+                render_pass.commands.push(Command::SetShaderBinding(
+                    sky_box_ids.shader_object,
+                    binding,
+                ));
             }
 
             render_pass.commands.push(Command::UpdateDrawProperties(
@@ -502,6 +504,7 @@ fn add_set_3d_shader_bindings(
         Traverse<(With<Scene>, With<ActiveScene>), TraverseUp, ChildOf>,
     )>,
     shader_program: &ShaderProgram,
+    shader_object_id: ObjectId,
     white_1x1_tex_obj_id: ObjectId,
 ) -> Result<(), crate::Error>
 {
@@ -620,7 +623,7 @@ fn add_set_3d_shader_bindings(
             )?,
         ]
         .into_iter()
-        .map(Command::SetShaderBinding),
+        .map(|binding| Command::SetShaderBinding(shader_object_id, binding)),
     );
 
     render_pass
@@ -660,7 +663,7 @@ fn add_set_3d_shader_bindings(
                     .field("quadratic")
                     .binding(point_light.attenuation_params.quadratic.into())?,
             ]
-            .map(Command::SetShaderBinding),
+            .map(|binding| Command::SetShaderBinding(shader_object_id, binding)),
         );
     }
 
@@ -689,7 +692,7 @@ fn add_set_3d_shader_bindings(
                     .field("direction")
                     .binding(directional_light.direction.into())?,
             ]
-            .map(Command::SetShaderBinding),
+            .map(|binding| Command::SetShaderBinding(shader_object_id, binding)),
         );
     }
 
diff --git a/engine/src/ui/dear_imgui.rs b/engine/src/ui/dear_imgui.rs
index b589696..4da5c5f 100644
--- a/engine/src/ui/dear_imgui.rs
+++ b/engine/src/ui/dear_imgui.rs
@@ -674,10 +674,13 @@ fn add_drawing_render_pass(
 
     // let [scale_width, scale_height] = draw_data.framebuffer_scale;
 
+    let shader_object_id = RenderingObjectId::Asset(shader_asset.id());
+
     render_pass.commands.extend([
         RenderingCommand::MakeCurrent(window_surface_id),
-        RenderingCommand::ActivateShader(RenderingObjectId::Asset(shader_asset.id())),
+        RenderingCommand::ActivateShader(shader_object_id),
         RenderingCommand::SetShaderBinding(
+            shader_object_id,
             shader_cursor
                 .field("Uniforms")
                 .field("projection")
@@ -756,6 +759,7 @@ fn add_drawing_render_pass(
                             DrawPropertiesUpdateFlags::SCISSOR_BOX,
                         ),
                         RenderingCommand::SetShaderBinding(
+                            shader_object_id,
                             shader_cursor.field("main_texture").binding(
                                 ShaderBindingValue::Texture(
                                     *texture_object_id,