fix(engine): make shader bindings work for multiple windows

3 files changed, 51 insertions, 30 deletions

diff --git a/engine/src/renderer.rs b/engine/src/renderer.rs
index 9f4ebda..1a27c62 100644
--- a/engine/src/renderer.rs
+++ b/engine/src/renderer.rs
@@ -520,6 +520,8 @@ fn handle_window_removed(
 pub struct PendingShaderBindings
 {
     pub bindings: Vec<(ShaderBindingLocation, ShaderBindingValue)>,
+    pub surface_specific_bindings:
+        Vec<(SurfaceId, ShaderBindingLocation, ShaderBindingValue)>,
 }
 
 impl<'a> Extend<(ShaderCursor<'a>, ShaderBindingValue)> for PendingShaderBindings
diff --git a/engine/src/renderer/main_render_pass.rs b/engine/src/renderer/main_render_pass.rs
index e7964dd..7492379 100644
--- a/engine/src/renderer/main_render_pass.rs
+++ b/engine/src/renderer/main_render_pass.rs
@@ -84,10 +84,13 @@ pub fn add_main_render_passes(
                 None => &default_shader_asset,
             };
 
-            if pending_shader_bindings.as_ref().map_or_else(
-                || true,
-                |pending_shader_bindings| pending_shader_bindings.bindings.is_empty(),
-            ) {
+            let Some(pending_shader_bindings) = pending_shader_bindings.as_mut() else {
+                continue;
+            };
+
+            if pending_shader_bindings.bindings.is_empty()
+                && pending_shader_bindings.surface_specific_bindings.is_empty()
+            {
                 continue;
             }
 
@@ -167,15 +170,26 @@ pub fn add_main_render_passes(
                 }
             }
 
-            if let Some(pending_shader_bindings) = &mut pending_shader_bindings {
-                for (shader_binding_loc, shader_binding_val) in
-                    pending_shader_bindings.bindings.drain(..)
-                {
-                    render_pass.commands.push(RendererCommand::SetShaderBinding(
-                        shader_binding_loc,
-                        shader_binding_val,
-                    ));
+            for (shader_binding_loc, shader_binding_val) in
+                &pending_shader_bindings.bindings
+            {
+                render_pass.commands.push(RendererCommand::SetShaderBinding(
+                    shader_binding_loc.clone(),
+                    shader_binding_val.clone(),
+                ));
+            }
+
+            for (shader_binding_surface_id, shader_binding_loc, shader_binding_val) in
+                &pending_shader_bindings.surface_specific_bindings
+            {
+                if *shader_binding_surface_id != surface_spec.id {
+                    continue;
                 }
+
+                render_pass.commands.push(RendererCommand::SetShaderBinding(
+                    shader_binding_loc.clone(),
+                    shader_binding_val.clone(),
+                ));
             }
 
             if let Some(draw_flags) = draw_flags.as_deref()
diff --git a/engine/src/shader/default.rs b/engine/src/shader/default.rs
index 8690eb9..7f0d1bb 100644
--- a/engine/src/shader/default.rs
+++ b/engine/src/shader/default.rs
@@ -15,7 +15,7 @@ use crate::material::{Flags as MaterialFlags, Material};
 use crate::matrix::Matrix;
 use crate::model::{MaterialSearchResult, Model};
 use crate::projection::{ClipVolume as ProjectionClipVolume, Projection};
-use crate::renderer::PendingShaderBindings;
+use crate::renderer::{PendingShaderBindings, SurfaceSpec};
 use crate::shader::cursor::{BindingValue as ShaderBindingValue, Cursor as ShaderCursor};
 use crate::shader::{
     Context as ShaderContext,
@@ -35,7 +35,7 @@ pub static ASSET_LABEL: LazyLock<AssetLabel> = LazyLock::new(|| AssetLabel {
 pub fn enqueue_set_shader_bindings(
     renderable_query: Query<RenderableEntity<'_>, (Without<NoDraw>,)>,
     camera_query: Query<(&Camera, &WorldPosition, &ActiveCamera)>,
-    window_query: Query<(&Window,)>,
+    window_query: Query<(&Window, &SurfaceSpec)>,
     point_light_query: Query<(&PointLight, &WorldPosition)>,
     directional_light_query: Query<(&DirectionalLight,)>,
     assets: Single<Assets>,
@@ -49,11 +49,6 @@ pub fn enqueue_set_shader_bindings(
         return;
     };
 
-    let Some((window,)) = window_query.iter().next() else {
-        // tracing::warn!("No window");
-        return;
-    };
-
     let default_shader_asset = assets
         .get_handle_to_loaded::<ShaderModuleSource>(ASSET_LABEL.clone())
         .unwrap();
@@ -72,6 +67,17 @@ pub fn enqueue_set_shader_bindings(
             continue;
         }
 
+        let has_pending_shader_bindings_comp = pending_shader_bindings.is_some();
+
+        let shader_bindings = match pending_shader_bindings.as_deref_mut() {
+            Some(pending_shader_bindings) => pending_shader_bindings,
+            None => &mut PendingShaderBindings::default(),
+        };
+
+        // Bindings are cleared to prevent them from growing infinitely
+        shader_bindings.bindings.clear();
+        shader_bindings.surface_specific_bindings.clear();
+
         let Some(shader_program) = shader_context.get_program(&shader_asset.id()) else {
             continue;
         };
@@ -80,13 +86,6 @@ pub fn enqueue_set_shader_bindings(
             continue;
         };
 
-        let has_pending_shader_bindings_comp = pending_shader_bindings.is_some();
-
-        let shader_bindings = match pending_shader_bindings.as_deref_mut() {
-            Some(pending_shader_bindings) => pending_shader_bindings,
-            None => &mut PendingShaderBindings::default(),
-        };
-
         let shader_cursor = ShaderCursor::new(
             shader_program
                 .reflection(0)
@@ -147,10 +146,6 @@ pub fn enqueue_set_shader_bindings(
                 create_view_matrix(&camera, &camera_world_pos.position).into(),
             ),
             (
-                model_3d_shader_cursor.field("projection"),
-                create_projection_matrix(&camera, window.inner_size()).into(),
-            ),
-            (
                 lighting_shader_cursor.field("view_pos"),
                 camera_world_pos.position.into(),
             ),
@@ -194,6 +189,16 @@ pub fn enqueue_set_shader_bindings(
             ),
         ]);
 
+        for (window, window_surface_spec) in &window_query {
+            shader_bindings.surface_specific_bindings.push((
+                window_surface_spec.id,
+                model_3d_shader_cursor
+                    .field("projection")
+                    .into_binding_location(),
+                create_projection_matrix(&camera, window.inner_size()).into(),
+            ));
+        }
+
         shader_bindings.extend(point_light_query.iter().enumerate().flat_map(
             |(point_light_index, (point_light, point_light_world_pos))| {
                 let point_light_shader_cursor = lighting_shader_cursor