render: implement simple texture atlas for the shadow maps

lyra-game/src/render/graph/node.rs

@@ -56,7 +56,7 @@ pub enum SlotValue {
     Lazy,
     TextureView(Arc<wgpu::TextureView>),
     Sampler(Rc<wgpu::Sampler>),
-    Texture(Rc<wgpu::Texture>),
+    Texture(Arc<wgpu::Texture>),
     Buffer(Arc<wgpu::Buffer>),
     RenderTarget(Rc<RefCell<RenderTarget>>),
     Frame(Rc<RefCell<Option<Frame>>>),
@@ -71,7 +71,7 @@ impl SlotValue {
         bind_match!(self, Self::Sampler(v) => v)
     }
 
-    pub fn as_texture(&self) -> Option<&Rc<wgpu::Texture>> {
+    pub fn as_texture(&self) -> Option<&Arc<wgpu::Texture>> {
         bind_match!(self, Self::Texture(v) => v)
     }
 

lyra-game/src/render/graph/passes/shadows.rs

@@ -1,6 +1,10 @@
 use std::{mem, num::NonZeroU64, rc::Rc, sync::Arc};
 
-use lyra_ecs::{query::{filter::Has, Entities}, AtomicRef, Entity, ResourceData};
+use glam::UVec2;
+use lyra_ecs::{
+    query::{filter::Has, Entities},
+    AtomicRef, Entity, ResourceData,
+};
 use lyra_game_derive::RenderGraphLabel;
 use lyra_math::Transform;
 use rustc_hash::FxHashMap;
@@ -10,12 +14,10 @@ use wgpu::util::DeviceExt;
 use crate::render::{
     graph::{Node, NodeDesc, NodeType, SlotAttribute, SlotValue},
     light::directional::DirectionalLight,
-    resource::{
+    resource::{RenderPipeline, RenderPipelineDescriptor, Shader, VertexState},
-        RenderPipeline, RenderPipelineDescriptor, Shader,
-        VertexState,
-    },
     transform_buffer_storage::TransformBuffers,
     vertex::Vertex,
+    TextureAtlas,
 };
 
 use super::{MeshBufferStorage, RenderAssets, RenderMeshes};
@@ -50,10 +52,7 @@ pub struct ShadowMapsPass {
     mesh_buffers: Option<ResourceData>,
     pipeline: Option<RenderPipeline>,
 
-    /// The depth map atlas texture
+    atlas: Arc<TextureAtlas>,
-    atlas_texture: Rc<wgpu::Texture>,
-    /// The depth map atlas texture view
-    atlas_view: Arc<wgpu::TextureView>,
     /// The depth map atlas sampler
     atlas_sampler: Rc<wgpu::Sampler>,
 }
@@ -78,7 +77,7 @@ impl ShadowMapsPass {
             }),
         );
 
-        let tex = device.create_texture(&wgpu::TextureDescriptor {
+        /* let tex = device.create_texture(&wgpu::TextureDescriptor {
             label: Some("texture_shadow_map_atlas"),
             size: wgpu::Extent3d {
                 width: SHADOW_SIZE.x,
@@ -96,7 +95,15 @@ impl ShadowMapsPass {
         let view = tex.create_view(&wgpu::TextureViewDescriptor {
             label: Some("shadows_map_view"),
             ..Default::default()
-        });
+        }); */
+
+        let atlas = TextureAtlas::new(
+            device,
+            wgpu::TextureFormat::Depth32Float,
+            wgpu::TextureUsages::RENDER_ATTACHMENT | wgpu::TextureUsages::TEXTURE_BINDING,
+            SHADOW_SIZE,
+            UVec2::new(4, 4),
+        );
 
         let sampler = device.create_sampler(&wgpu::SamplerDescriptor {
             label: Some("sampler_shadow_map_atlas"),
@@ -119,14 +126,13 @@ impl ShadowMapsPass {
             pipeline: None,
 
             atlas_sampler: Rc::new(sampler),
-            atlas_texture: Rc::new(tex),
+            atlas: Arc::new(atlas),
-            atlas_view: Arc::new(view),
         }
     }
 
     fn create_depth_map(&mut self, device: &wgpu::Device, entity: Entity, light_pos: Transform) {
         const NEAR_PLANE: f32 = 0.1;
-        const FAR_PLANE: f32 = 25.5;
+        const FAR_PLANE: f32 = 45.0;
 
         let ortho_proj =
             glam::Mat4::orthographic_rh(-10.0, 10.0, -10.0, 10.0, NEAR_PLANE, FAR_PLANE);
@@ -188,13 +194,13 @@ impl Node for ShadowMapsPass {
         node.add_texture_slot(
             ShadowMapsPassSlots::ShadowAtlasTexture,
             SlotAttribute::Output,
-            Some(SlotValue::Texture(self.atlas_texture.clone())),
+            Some(SlotValue::Texture(self.atlas.texture().clone())),
         );
 
         node.add_texture_view_slot(
             ShadowMapsPassSlots::ShadowAtlasTextureView,
             SlotAttribute::Output,
-            Some(SlotValue::TextureView(self.atlas_view.clone())),
+            Some(SlotValue::TextureView(self.atlas.view().clone())),
         );
 
         node.add_sampler_slot(
@@ -222,6 +228,8 @@ impl Node for ShadowMapsPass {
         self.transform_buffers = world.try_get_resource_data::<TransformBuffers>();
         self.mesh_buffers = world.try_get_resource_data::<RenderAssets<MeshBufferStorage>>();
 
+        world.add_resource(self.atlas.clone());
+
         for (entity, pos, _) in world.view_iter::<(Entities, &Transform, Has<DirectionalLight>)>() {
             if !self.depth_maps.contains_key(&entity) {
                 self.create_depth_map(graph.device(), entity, *pos);
@@ -231,7 +239,8 @@ impl Node for ShadowMapsPass {
 
         // update the light projection buffer slot
         let (_, dir_depth_map) = self.depth_maps.iter().next().unwrap();
-        let val = graph.slot_value_mut(ShadowMapsPassSlots::DirLightProjectionBuffer)
+        let val = graph
+            .slot_value_mut(ShadowMapsPassSlots::DirLightProjectionBuffer)
             .unwrap();
         *val = SlotValue::Buffer(dir_depth_map.light_projection_buffer.clone());
 
@@ -308,7 +317,7 @@ impl Node for ShadowMapsPass {
                 label: Some("pass_shadow_map"),
                 color_attachments: &[],
                 depth_stencil_attachment: Some(wgpu::RenderPassDepthStencilAttachment {
-                    view: &self.atlas_view,
+                    view: self.atlas.view(),
                     depth_ops: Some(wgpu::Operations {
                         load: wgpu::LoadOp::Clear(1.0),
                         store: true,
@@ -317,6 +326,11 @@ impl Node for ShadowMapsPass {
                 }),
             });
             pass.set_pipeline(&pipeline);
+            let viewport = self.atlas.texture_viewport(0);
+            // only render to the light's map in the atlas
+            pass.set_viewport(viewport.offset.x as _, viewport.offset.y as _, viewport.size.x as _, viewport.size.y as _, 0.0, 1.0);
+            // only clear the light map in the atlas
+            pass.set_scissor_rect(viewport.offset.x, viewport.offset.y, viewport.size.x, viewport.size.y);
 
             for job in render_meshes.iter() {
                 // get the mesh (containing vertices) and the buffers from storage

lyra-game/src/render/mod.rs

@@ -15,3 +15,6 @@ pub mod light;
 //pub mod light_cull_compute;
 pub mod avec;
 pub mod graph;
+
+mod texture_atlas;
+pub use texture_atlas::*;

lyra-game/src/render/shaders/base.wgsl

@@ -149,6 +149,34 @@ fn fs_main(in: VertexOutput) -> @location(0) vec4<f32> {
         let light: Light = u_lights.data[light_index];
 
         if (light.light_ty == LIGHT_TY_DIRECTIONAL) {
+            /*var proj_coords = in.frag_pos_light_space.xyz / in.frag_pos_light_space.w;
+            // for some reason the y component is clipped after transforming
+            proj_coords.y = -proj_coords.y;
+
+            // Remap xy to [0.0, 1.0]
+            let xy_remapped = proj_coords.xy * 0.5 + 0.5;
+            proj_coords.x = mix(0.0, 1024.0 / 4096.0, xy_remapped.x);
+            proj_coords.y = mix(0.0, 1024.0 / 4096.0, xy_remapped.y);
+
+            let closest_depth = textureSampleLevel(t_shadow_maps_atlas, s_shadow_maps_atlas, proj_coords.xy, 0.0, vec2<i32>(0, 0));
+            let current_depth = proj_coords.z;
+
+            // use a bias to avoid shadow acne
+            let light_dir = normalize(-light.direction);
+            let bias = max(0.05 * (1.0 - dot(in.world_normal, light_dir)), 0.005);  
+            var shadow = 0.0;
+            if current_depth - bias > closest_depth {
+                shadow = 1.0;
+            }
+
+            // dont cast shadows outside the light's far plane
+            if (proj_coords.z > 1.0) {
+                shadow = 0.0;
+            }
+
+            return vec4<f32>(vec3<f32>(closest_depth), 1.0);*/
+
+
             let light_dir = normalize(-light.direction);
             let shadow = calc_shadow(in.world_normal, light_dir, in.frag_pos_light_space);
             light_res += blinn_phong_dir_light(in.world_position, in.world_normal, light, u_material, specular_color, shadow);
@@ -168,12 +196,27 @@ fn calc_shadow(normal: vec3<f32>, light_dir: vec3<f32>, frag_pos_light_space: ve
     // for some reason the y component is clipped after transforming
     proj_coords.y = -proj_coords.y;
 
+    // dont cast shadows outside the light's far plane
+    if (proj_coords.z > 1.0) {
+        return 0.0;
+    }
+
     // Remap xy to [0.0, 1.0]
     let xy_remapped = proj_coords.xy * 0.5 + 0.5;
-    proj_coords.x = xy_remapped.x;
+     // TODO: when more lights are added, change the index, and the atlas sizes
-    proj_coords.y = xy_remapped.y;
+    let shadow_map_index = 0; 
+    let shadow_map_region = vec2<f32>( (f32(shadow_map_index) * 1024.0) / 4096.0, (f32(shadow_map_index + 1) * 1024.0) / 4096.0);
+    // lerp the tex coords to the shadow map for this light.
+    proj_coords.x = mix(shadow_map_region.x, shadow_map_region.y, xy_remapped.x);
+    proj_coords.y = mix(shadow_map_region.x, shadow_map_region.y, xy_remapped.y);
 
-    let closest_depth = textureSampleLevel(t_shadow_maps_atlas, s_shadow_maps_atlas, proj_coords.xy, 0.0);
+    // simulate `ClampToBorder`, not creating shadows past the shadow map regions
+    if (proj_coords.x > shadow_map_region.y && proj_coords.y > shadow_map_region.y)
+            || (proj_coords.x < shadow_map_region.x && proj_coords.y < shadow_map_region.x) {
+        return 0.0;
+    }
+
+    let closest_depth = textureSampleLevel(t_shadow_maps_atlas, s_shadow_maps_atlas, proj_coords.xy, 0.0, vec2<i32>(0, 0));
     let current_depth = proj_coords.z;
 
     // use a bias to avoid shadow acne
@@ -183,11 +226,6 @@ fn calc_shadow(normal: vec3<f32>, light_dir: vec3<f32>, frag_pos_light_space: ve
         shadow = 1.0;
     }
 
-    // dont cast shadows outside the light's far plane
-    if (proj_coords.z > 1.0) {
-        shadow = 0.0;
-    }
-
     return shadow;
 }
 

lyra-game/src/render/texture_atlas.rs

@@ -0,0 +1,78 @@
+use std::sync::Arc;
+
+use glam::UVec2;
+
+#[derive(Debug, Clone, Copy)]
+pub struct AtlasViewport {
+    pub offset: UVec2,
+    pub size: UVec2,
+}
+
+pub struct TextureAtlas {
+    /// The size of each texture in the atlas.
+    texture_size: UVec2,
+    /// The amount of textures in the atlas.
+    texture_count: UVec2,
+
+    texture_format: wgpu::TextureFormat,
+    texture: Arc<wgpu::Texture>,
+    view: Arc<wgpu::TextureView>,
+}
+
+impl TextureAtlas {
+    pub fn new(device: &wgpu::Device, format: wgpu::TextureFormat, usages: wgpu::TextureUsages, texture_size: UVec2, texture_count: UVec2) -> Self {
+        let total_size = texture_size * texture_count;
+        
+        let texture = device.create_texture(&wgpu::TextureDescriptor {
+            label: Some("texture_atlas"),
+            size: wgpu::Extent3d { width: total_size.x, height: total_size.y, depth_or_array_layers: 1 },
+            mip_level_count: 1,
+            sample_count: 1,
+            dimension: wgpu::TextureDimension::D2,
+            format,
+            usage: usages,
+            view_formats: &[],
+        });
+        let view = texture.create_view(&wgpu::TextureViewDescriptor::default());
+
+        Self {
+            texture_size,
+            texture_count,
+            texture_format: format,
+            texture: Arc::new(texture),
+            view: Arc::new(view),
+        }
+    }
+
+    /// Get the viewport of a texture index in the atlas.
+    pub fn texture_viewport(&self, atlas_index: u32) -> AtlasViewport {
+        let x = (atlas_index % self.texture_count.x) * self.texture_size.x;
+        let y = (atlas_index / self.texture_count.y) * self.texture_size.y;
+
+        AtlasViewport { offset: UVec2::new(x, y), size: self.texture_size }
+    }
+
+    pub fn view(&self) -> &Arc<wgpu::TextureView> {
+        &self.view
+    }
+
+    pub fn texture(&self) -> &Arc<wgpu::Texture> {
+        &self.texture
+    }
+
+    pub fn texture_format(&self) -> &wgpu::TextureFormat {
+        &self.texture_format
+    }
+
+    pub fn texture_size(&self) -> UVec2 {
+        self.texture_size
+    }
+
+    pub fn texture_count(&self) -> UVec2 {
+        self.texture_count
+    }
+
+    pub fn total_texture_count(&self) -> u32 {
+        self.texture_count.x * self.texture_count.y
+    }
+}