render: create a GpuSlotBuffer for stable indices in a gpu buffer

Cargo.lock

@@ -1881,6 +1881,8 @@ dependencies = [
  "lyra-scene",
  "petgraph",
  "quote",
+ "rectangle-pack",
+ "round_mult",
  "rustc-hash",
  "syn 2.0.51",
  "thiserror",
@@ -2767,6 +2769,12 @@ dependencies = [
  "rand_core 0.3.1",
 ]
 
+[[package]]
+name = "rectangle-pack"
+version = "0.4.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a0d463f2884048e7153449a55166f91028d5b0ea53c79377099ce4e8cf0cf9bb"
+
 [[package]]
 name = "redox_syscall"
 version = "0.3.5"
@@ -2884,6 +2892,15 @@ dependencies = [
  "winreg",
 ]
 
+[[package]]
+name = "round_mult"
+version = "0.1.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "74bc7d5286c4d36f09aa6ae93f76acf6aa068cd62bc02970a9deb24763655dee"
+dependencies = [
+ "rustc_version",
+]
+
 [[package]]
 name = "rustc-demangle"
 version = "0.1.23"
@@ -2896,6 +2913,15 @@ version = "1.1.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "08d43f7aa6b08d49f382cde6a7982047c3426db949b1424bc4b7ec9ae12c6ce2"
 
+[[package]]
+name = "rustc_version"
+version = "0.4.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "bfa0f585226d2e68097d4f95d113b15b83a82e819ab25717ec0590d9584ef366"
+dependencies = [
+ "semver",
+]
+
 [[package]]
 name = "rustix"
 version = "0.37.27"
@@ -3010,6 +3036,12 @@ dependencies = [
  "libc",
 ]
 
+[[package]]
+name = "semver"
+version = "1.0.23"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "61697e0a1c7e512e84a621326239844a24d8207b4669b41bc18b32ea5cbf988b"
+
 [[package]]
 name = "serde"
 version = "1.0.194"

lyra-game/Cargo.toml

@@ -38,6 +38,8 @@ unique = "0.9.1"
 rustc-hash = "1.1.0"
 petgraph = { version = "0.6.5", features = ["matrix_graph"] }
 bind_match = "0.1.2"
+rectangle-pack = "0.4.2"
+round_mult = "0.1.3"
 
 [features]
 tracy = ["dep:tracing-tracy"]

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

@@ -14,8 +14,7 @@ use crate::render::{
 };
 
 use super::{
-    BasePassSlots, LightBasePassSlots, LightCullComputePassSlots, MeshBufferStorage, RenderAssets,
-    RenderMeshes, ShadowMapsPassSlots,
+    BasePassSlots, LightBasePassSlots, LightCullComputePassSlots, MeshBufferStorage, RenderAssets, RenderMeshes, ShadowMapsPassSlots
 };
 
 #[derive(Debug, Hash, Clone, Default, PartialEq, RenderGraphLabel)]
@@ -116,9 +115,14 @@ impl Node for MeshPass {
                 .expect("missing ShadowMapsPassSlots::ShadowAtlasSampler")
                 .as_sampler()
                 .unwrap();
-            let dir_light_projection_buf = graph
-                .slot_value(ShadowMapsPassSlots::DirLightProjectionBuffer)
-                .expect("missing ShadowMapsPassSlots::DirLightProjectionBuffer")
+            let atlas_size_buf = graph
+                .slot_value(ShadowMapsPassSlots::ShadowAtlasSizeBuffer)
+                .expect("missing ShadowMapsPassSlots::ShadowAtlasSizeBuffer")
+                .as_buffer()
+                .unwrap();
+            let light_uniform_buf = graph
+                .slot_value(ShadowMapsPassSlots::ShadowLightUniformsBuffer)
+                .expect("missing ShadowMapsPassSlots::ShadowLightUniformsBuffer")
                 .as_buffer()
                 .unwrap();
 
@@ -151,6 +155,16 @@ impl Node for MeshPass {
                         },
                         count: None,
                     },
+                    wgpu::BindGroupLayoutEntry {
+                        binding: 3,
+                        visibility: wgpu::ShaderStages::VERTEX_FRAGMENT,
+                        ty: wgpu::BindingType::Buffer {
+                            ty: wgpu::BufferBindingType::Storage { read_only: true },
+                            has_dynamic_offset: false,
+                            min_binding_size: None,
+                        },
+                        count: None,
+                    },
                 ],
             });
 
@@ -169,7 +183,15 @@ impl Node for MeshPass {
                     wgpu::BindGroupEntry {
                         binding: 2,
                         resource: wgpu::BindingResource::Buffer(wgpu::BufferBinding {
-                            buffer: dir_light_projection_buf,
+                            buffer: atlas_size_buf,
+                            offset: 0,
+                            size: None,
+                        }),
+                    },
+                    wgpu::BindGroupEntry {
+                        binding: 3,
+                        resource: wgpu::BindingResource::Buffer(wgpu::BufferBinding {
+                            buffer: light_uniform_buf,
                             offset: 0,
                             size: None,
                         }),

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

@@ -1,9 +1,14 @@
-use std::{mem, num::NonZeroU64, rc::Rc, sync::Arc};
+use std::{
+    collections::VecDeque,
+    mem,
+    num::NonZeroU64,
+    rc::Rc,
+    sync::{Arc, RwLock, RwLockReadGuard, RwLockWriteGuard},
+};
 
-use glam::UVec2;
 use lyra_ecs::{
     query::{filter::Has, Entities},
-    AtomicRef, Entity, ResourceData,
+    AtomicRef, Component, Entity, ResourceData,
 };
 use lyra_game_derive::RenderGraphLabel;
 use lyra_math::Transform;
@@ -17,31 +22,38 @@ use crate::render::{
     resource::{RenderPipeline, RenderPipelineDescriptor, Shader, VertexState},
     transform_buffer_storage::TransformBuffers,
     vertex::Vertex,
-    TextureAtlas,
+    AtlasViewport, GpuSlotBuffer, TextureAtlas,
 };
 
 use super::{MeshBufferStorage, RenderAssets, RenderMeshes};
 
-const SHADOW_SIZE: glam::UVec2 = glam::UVec2::new(1024, 1024);
+const SHADOW_SIZE: glam::UVec2 = glam::uvec2(1024, 1024);
 
 #[derive(Debug, Clone, Hash, PartialEq, RenderGraphLabel)]
 pub enum ShadowMapsPassSlots {
     ShadowAtlasTexture,
     ShadowAtlasTextureView,
     ShadowAtlasSampler,
-    DirLightProjectionBuffer,
+    ShadowAtlasSizeBuffer,
+    ShadowLightUniformsBuffer,
 }
 
 #[derive(Debug, Clone, Hash, PartialEq, RenderGraphLabel)]
 pub struct ShadowMapsPassLabel;
 
+#[derive(Clone, Copy)]
 struct LightDepthMap {
-    light_projection_buffer: Arc<wgpu::Buffer>,
-    bindgroup: wgpu::BindGroup,
+    //light_projection_buffer: Arc<wgpu::Buffer>,
+    //bindgroup: wgpu::BindGroup,
+    atlas_index: u64,
+    uniform_index: u64,
 }
 
 pub struct ShadowMapsPass {
     bgl: Arc<wgpu::BindGroupLayout>,
+    atlas_size_buffer: Arc<wgpu::Buffer>,
+    light_uniforms_buffer: GpuSlotBuffer<LightShadowUniform>,
+    uniforms_bg: Arc<wgpu::BindGroup>,
     /// depth maps for a light owned by an entity.
     depth_maps: FxHashMap<Entity, LightDepthMap>,
 
@@ -52,7 +64,7 @@ pub struct ShadowMapsPass {
     mesh_buffers: Option<ResourceData>,
     pipeline: Option<RenderPipeline>,
 
-    atlas: Arc<TextureAtlas>,
+    atlas: LightShadowMapAtlas,
     /// The depth map atlas sampler
     atlas_sampler: Rc<wgpu::Sampler>,
 }
@@ -61,15 +73,15 @@ impl ShadowMapsPass {
     pub fn new(device: &wgpu::Device) -> Self {
         let bgl = Arc::new(
             device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
-                label: Some("bgl_shadows_light_projection"),
+                label: Some("bgl_shadow_maps_lights"),
                 entries: &[wgpu::BindGroupLayoutEntry {
                     binding: 0,
-                    visibility: wgpu::ShaderStages::VERTEX,
+                    visibility: wgpu::ShaderStages::VERTEX_FRAGMENT,
                     ty: wgpu::BindingType::Buffer {
-                        ty: wgpu::BufferBindingType::Uniform,
-                        has_dynamic_offset: false,
+                        ty: wgpu::BufferBindingType::Storage { read_only: true },
+                        has_dynamic_offset: true,
                         min_binding_size: Some(
-                            NonZeroU64::new(mem::size_of::<glam::Mat4>() as _).unwrap(),
+                            NonZeroU64::new(mem::size_of::<LightShadowUniform>() as _).unwrap(),
                         ),
                     },
                     count: None,
@@ -77,34 +89,19 @@ impl ShadowMapsPass {
             }),
         );
 
-        /* let tex = device.create_texture(&wgpu::TextureDescriptor {
-            label: Some("texture_shadow_map_atlas"),
-            size: wgpu::Extent3d {
-                width: SHADOW_SIZE.x,
-                height: SHADOW_SIZE.y,
-                depth_or_array_layers: 1,
-            },
-            mip_level_count: 1,
-            sample_count: 1,
-            dimension: wgpu::TextureDimension::D2,
-            format: wgpu::TextureFormat::Depth32Float,
-            usage: wgpu::TextureUsages::RENDER_ATTACHMENT | wgpu::TextureUsages::TEXTURE_BINDING,
-            view_formats: &[],
-        });
-
-        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),
+            SHADOW_SIZE * 4,
         );
 
+        let atlas_size_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
+            label: Some("buffer_shadow_maps_atlas_size"),
+            usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
+            contents: bytemuck::bytes_of(&atlas.atlas_size()),
+        });
+
         let sampler = device.create_sampler(&wgpu::SamplerDescriptor {
             label: Some("sampler_shadow_map_atlas"),
             address_mode_u: wgpu::AddressMode::ClampToBorder,
@@ -117,8 +114,34 @@ impl ShadowMapsPass {
             ..Default::default()
         });
 
+        let cap = device.limits().max_storage_buffer_binding_size as u64
+            / mem::size_of::<LightShadowUniform>() as u64;
+        let uniforms_buffer = GpuSlotBuffer::new_aligned(
+            device,
+            Some("buffer_shadow_maps_light"),
+            wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
+            cap,
+            256,
+        );
+
+        let uniforms_bg = device.create_bind_group(&wgpu::BindGroupDescriptor {
+            label: Some("bind_group_shadows"),
+            layout: &bgl,
+            entries: &[wgpu::BindGroupEntry {
+                binding: 0,
+                resource: wgpu::BindingResource::Buffer(wgpu::BufferBinding {
+                    buffer: uniforms_buffer.buffer(),
+                    offset: 0,
+                    size: Some(NonZeroU64::new(mem::size_of::<LightShadowUniform>() as _).unwrap()),
+                }),
+            }],
+        });
+
         Self {
             bgl,
+            light_uniforms_buffer: uniforms_buffer,
+            uniforms_bg: Arc::new(uniforms_bg),
+            atlas_size_buffer: Arc::new(atlas_size_buffer),
             depth_maps: Default::default(),
             transform_buffers: None,
             render_meshes: None,
@@ -126,14 +149,26 @@ impl ShadowMapsPass {
             pipeline: None,
 
             atlas_sampler: Rc::new(sampler),
-            atlas: Arc::new(atlas),
+            atlas: LightShadowMapAtlas(Arc::new(RwLock::new(atlas))),
         }
     }
 
-    fn create_depth_map(&mut self, device: &wgpu::Device, entity: Entity, light_pos: Transform) {
+    /// Create a depth map and return the id of the depth map in the texture atlas.
+    fn create_depth_map(
+        &mut self,
+        queue: &wgpu::Queue,
+        entity: Entity,
+        light_pos: Transform,
+    ) -> LightDepthMap {
         const NEAR_PLANE: f32 = 0.1;
         const FAR_PLANE: f32 = 45.0;
 
+        let mut atlas = self.atlas.get_mut();
+        let atlas_index = atlas
+            .pack_new_texture(SHADOW_SIZE.x as _, SHADOW_SIZE.y as _)
+            .expect("failed to pack new shadow map into texture atlas");
+        let atlas_frame = atlas.texture_viewport(atlas_index);
+
         let ortho_proj =
             glam::Mat4::orthographic_rh(-10.0, 10.0, -10.0, 10.0, NEAR_PLANE, FAR_PLANE);
 
@@ -141,34 +176,26 @@ impl ShadowMapsPass {
             glam::Mat4::look_to_rh(light_pos.translation, light_pos.forward(), light_pos.up());
 
         let light_proj = ortho_proj * look_view;
+        let uniform = LightShadowUniform {
+            space_mat: light_proj,
+            atlas_frame,
+        };
 
-        let light_projection_buffer =
-            device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
-                label: Some("shadows_light_view_mat_buffer"),
-                usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
-                contents: bytemuck::bytes_of(&light_proj),
-            });
+        /* let uniform_index = self.light_uniforms_index;
+        self.light_uniforms_index += 1;
 
-        let bg = device.create_bind_group(&wgpu::BindGroupDescriptor {
-            label: Some("shadows_bind_group"),
-            layout: &self.bgl,
-            entries: &[wgpu::BindGroupEntry {
-                binding: 0,
-                resource: wgpu::BindingResource::Buffer(wgpu::BufferBinding {
-                    buffer: &light_projection_buffer,
-                    offset: 0,
-                    size: None,
-                }),
-            }],
-        });
+        //self.light_uniforms_buffer
+        let offset = uniform_index_offset(&device.limits(), uniform_index);
+        queue.write_buffer(&self.light_uniforms_buffer, offset as u64, bytemuck::bytes_of(&uniform)); */
+        let uniform_index = self.light_uniforms_buffer.insert(queue, &uniform);
 
-        self.depth_maps.insert(
-            entity,
-            LightDepthMap {
-                light_projection_buffer: Arc::new(light_projection_buffer),
-                bindgroup: bg,
-            },
-        );
+        let v = LightDepthMap {
+            atlas_index,
+            uniform_index,
+        };
+        self.depth_maps.insert(entity, v);
+
+        v
     }
 
     fn transform_buffers(&self) -> AtomicRef<TransformBuffers> {
@@ -191,16 +218,18 @@ impl Node for ShadowMapsPass {
     ) -> crate::render::graph::NodeDesc {
         let mut node = NodeDesc::new(NodeType::Render, None, vec![]);
 
+        let atlas = self.atlas.get();
+
         node.add_texture_slot(
             ShadowMapsPassSlots::ShadowAtlasTexture,
             SlotAttribute::Output,
-            Some(SlotValue::Texture(self.atlas.texture().clone())),
+            Some(SlotValue::Texture(atlas.texture().clone())),
         );
 
         node.add_texture_view_slot(
             ShadowMapsPassSlots::ShadowAtlasTextureView,
             SlotAttribute::Output,
-            Some(SlotValue::TextureView(self.atlas.view().clone())),
+            Some(SlotValue::TextureView(atlas.view().clone())),
         );
 
         node.add_sampler_slot(
@@ -209,10 +238,18 @@ impl Node for ShadowMapsPass {
             Some(SlotValue::Sampler(self.atlas_sampler.clone())),
         );
 
-        node.add_sampler_slot(
-            ShadowMapsPassSlots::DirLightProjectionBuffer,
+        node.add_buffer_slot(
+            ShadowMapsPassSlots::ShadowLightUniformsBuffer,
             SlotAttribute::Output,
-            Some(SlotValue::Lazy),
+            Some(SlotValue::Buffer(
+                self.light_uniforms_buffer.buffer().clone(),
+            )),
+        );
+
+        node.add_buffer_slot(
+            ShadowMapsPassSlots::ShadowAtlasSizeBuffer,
+            SlotAttribute::Output,
+            Some(SlotValue::Buffer(self.atlas_size_buffer.clone())),
         );
 
         node
@@ -222,7 +259,7 @@ impl Node for ShadowMapsPass {
         &mut self,
         graph: &mut crate::render::graph::RenderGraph,
         world: &mut lyra_ecs::World,
-        _: &mut crate::render::graph::RenderGraphContext,
+        context: &mut crate::render::graph::RenderGraphContext,
     ) {
         self.render_meshes = world.try_get_resource_data::<RenderMeshes>();
         self.transform_buffers = world.try_get_resource_data::<TransformBuffers>();
@@ -230,19 +267,30 @@ impl Node for ShadowMapsPass {
 
         world.add_resource(self.atlas.clone());
 
+        // use a queue for storing atlas ids to add to entities after the entities are iterated
+        let mut index_components_queue = VecDeque::new();
+
         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);
+                // TODO: dont pack the textures as they're added
+                let atlas_index =
+                    self.create_depth_map(&context.queue, entity, *pos);
+                index_components_queue.push_back((entity, atlas_index));
+
                 debug!("Created depth map for {:?} light entity", entity);
             }
         }
 
-        // update the light projection buffer slot
-        let (_, dir_depth_map) = self.depth_maps.iter().next().unwrap();
-        let val = graph
-            .slot_value_mut(ShadowMapsPassSlots::DirLightProjectionBuffer)
-            .unwrap();
-        *val = SlotValue::Buffer(dir_depth_map.light_projection_buffer.clone());
+        // now consume from the queue adding the components to the entities
+        while let Some((entity, depth)) = index_components_queue.pop_front() {
+            world.insert(
+                entity,
+                LightShadowMapId {
+                    atlas_index: depth.atlas_index,
+                    uniform_index: depth.uniform_index,
+                },
+            );
+        }
 
         if self.pipeline.is_none() {
             let shader = Rc::new(Shader {
@@ -313,11 +361,12 @@ impl Node for ShadowMapsPass {
             .expect("missing directional light in scene");
 
         {
+            let atlas = self.atlas.get();
             let mut pass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
                 label: Some("pass_shadow_map"),
                 color_attachments: &[],
                 depth_stencil_attachment: Some(wgpu::RenderPassDepthStencilAttachment {
-                    view: self.atlas.view(),
+                    view: atlas.view(),
                     depth_ops: Some(wgpu::Operations {
                         load: wgpu::LoadOp::Clear(1.0),
                         store: true,
@@ -326,11 +375,27 @@ impl Node for ShadowMapsPass {
                 }),
             });
             pass.set_pipeline(&pipeline);
-            let viewport = self.atlas.texture_viewport(0);
+            let viewport = atlas.texture_viewport(dir_depth_map.atlas_index);
+            debug!(
+                "Rendering shadow map to viewport: {viewport:?}, uniform index: {}",
+                dir_depth_map.uniform_index
+            );
             // 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);
+            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);
+            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
@@ -341,7 +406,10 @@ impl Node for ShadowMapsPass {
                 }
                 let buffers = buffers.unwrap();
 
-                pass.set_bind_group(0, &dir_depth_map.bindgroup, &[]);
+                let uniform_index =
+                    self.light_uniforms_buffer
+                        .offset_of(dir_depth_map.uniform_index) as u32;
+                pass.set_bind_group(0, &self.uniforms_bg, &[uniform_index]);
 
                 // Get the bindgroup for job's transform and bind to it using an offset.
                 let bindgroup = transforms.bind_group(job.transform_id);
@@ -371,3 +439,49 @@ impl Node for ShadowMapsPass {
         }
     }
 }
+
+#[repr(C)]
+#[derive(Debug, Copy, Clone, bytemuck::Pod, bytemuck::Zeroable)]
+pub struct LightShadowUniform {
+    space_mat: glam::Mat4,
+    atlas_frame: AtlasViewport, // 2xUVec2 (4xf32), so no padding needed
+}
+
+/// A component that stores the ID of a shadow map in the shadow map atlas for the entities.
+///
+/// An entity owns a light. If that light casts shadows, this will contain the ID of the shadow
+/// map inside of the [`TextureAtlas`].
+#[derive(Debug, Default, Copy, Clone, Component)]
+pub struct LightShadowMapId {
+    atlas_index: u64,
+    uniform_index: u64,
+}
+
+impl LightShadowMapId {
+    pub fn atlas_index(&self) -> u64 {
+        self.atlas_index
+    }
+
+    pub fn uniform_index(&self) -> u64 {
+        self.uniform_index
+    }
+}
+
+/// An ecs resource storing the [`TextureAtlas`] of shadow maps.
+#[derive(Clone)]
+pub struct LightShadowMapAtlas(Arc<RwLock<TextureAtlas>>);
+
+impl LightShadowMapAtlas {
+    pub fn get(&self) -> RwLockReadGuard<TextureAtlas> {
+        self.0.read().unwrap()
+    }
+
+    pub fn get_mut(&self) -> RwLockWriteGuard<TextureAtlas> {
+        self.0.write().unwrap()
+    }
+}
+
+/* fn uniform_index_offset(limits: &wgpu::Limits, uniform_idx: u64) -> u32 {
+    let t = uniform_idx as u32 % (limits.max_storage_buffer_binding_size / mem::size_of::<LightShadowUniform>() as u32);
+    t * limits.min_uniform_buffer_offset_alignment
+} */

lyra-game/src/render/light/mod.rs

@@ -12,6 +12,8 @@ use crate::math::Transform;
 
 use self::directional::DirectionalLight;
 
+use super::graph::LightShadowMapId;
+
 const MAX_LIGHT_COUNT: usize = 16;
 
 /// A struct that stores a list of lights in a wgpu::Buffer.
@@ -166,18 +168,21 @@ impl LightUniformBuffers {
         let _ = world_tick;
         let mut lights = vec![];
 
-        for (point_light, transform) in world.view_iter::<(&PointLight, &Transform)>() {
-            let uniform = LightUniform::from_point_light_bundle(&point_light, &transform);
+        for (point_light, transform, shadow_map_id) in world.view_iter::<(&PointLight, &Transform, Option<&LightShadowMapId>)>() {
+            let shadow_map_id = shadow_map_id.map(|m| m.clone());
+            let uniform = LightUniform::from_point_light_bundle(&point_light, &transform, shadow_map_id);
             lights.push(uniform);
         }
 
-        for (spot_light, transform) in world.view_iter::<(&SpotLight, &Transform)>() {
-            let uniform = LightUniform::from_spot_light_bundle(&spot_light, &transform);
+        for (spot_light, transform, shadow_map_id) in world.view_iter::<(&SpotLight, &Transform, Option<&LightShadowMapId>)>() {
+            let shadow_map_id = shadow_map_id.map(|m| m.clone());
+            let uniform = LightUniform::from_spot_light_bundle(&spot_light, &transform, shadow_map_id);
             lights.push(uniform);
         }
 
-        for (dir_light, transform) in world.view_iter::<(&DirectionalLight, &Transform)>() {
-            let uniform = LightUniform::from_directional_bundle(&dir_light, &transform);
+        for (dir_light, transform, shadow_map_id) in world.view_iter::<(&DirectionalLight, &Transform, Option<&LightShadowMapId>)>() {
+            let shadow_map_id = shadow_map_id.map(|m| m.clone());
+            let uniform = LightUniform::from_directional_bundle(&dir_light, &transform, shadow_map_id);
             lights.push(uniform);
         }
 
@@ -216,10 +221,11 @@ pub(crate) struct LightUniform {
 
     pub spot_cutoff_rad: f32,
     pub spot_outer_cutoff_rad: f32,
+    pub light_shadow_uniform_index: i32,
 }
 
 impl LightUniform {
-    pub fn from_point_light_bundle(light: &PointLight, transform: &Transform) -> Self {
+    pub fn from_point_light_bundle(light: &PointLight, transform: &Transform, map_id: Option<LightShadowMapId>) -> Self {
         Self {
             light_type: LightType::Point as u32,
             enabled: light.enabled as u32,
@@ -233,11 +239,12 @@ impl LightUniform {
 
             spot_cutoff_rad: 0.0,
             spot_outer_cutoff_rad: 0.0,
+            light_shadow_uniform_index:  map_id.map(|m| m.uniform_index() as i32).unwrap_or(-1),
 
         }
     }
 
-    pub fn from_directional_bundle(light: &DirectionalLight, transform: &Transform) -> Self {
+    pub fn from_directional_bundle(light: &DirectionalLight, transform: &Transform, map_id: Option<LightShadowMapId>) -> Self {
         Self {
             light_type: LightType::Directional as u32,
             enabled: light.enabled as u32,
@@ -251,11 +258,12 @@ impl LightUniform {
 
             spot_cutoff_rad: 0.0,
             spot_outer_cutoff_rad: 0.0,
+            light_shadow_uniform_index:  map_id.map(|m| m.uniform_index() as i32).unwrap_or(-1),
         }
     }
 
     // Create the SpotLightUniform from an ECS bundle
-    pub fn from_spot_light_bundle(light: &SpotLight, transform: &Transform) -> Self {
+    pub fn from_spot_light_bundle(light: &SpotLight, transform: &Transform, map_id: Option<LightShadowMapId>) -> Self {
         Self {
             light_type: LightType::Spotlight as u32,
             enabled: light.enabled as u32,
@@ -269,6 +277,7 @@ impl LightUniform {
 
             spot_cutoff_rad: light.cutoff.to_radians(),
             spot_outer_cutoff_rad: light.outer_cutoff.to_radians(),
+            light_shadow_uniform_index:  map_id.map(|m| m.uniform_index() as i32).unwrap_or(-1),
         }
     }
 }

lyra-game/src/render/mod.rs

@@ -18,3 +18,6 @@ pub mod graph;
 
 mod texture_atlas;
 pub use texture_atlas::*;
+
+mod slot_buffer;
+pub use slot_buffer::*;

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

@@ -22,6 +22,11 @@ struct VertexOutput {
     @location(3) frag_pos_light_space: vec4<f32>,
 }
 
+struct TextureAtlasFrame {
+    offset: vec2<u32>,
+    size: vec2<u32>,
+}
+
 struct TransformData {
     transform: mat4x4<f32>,
     normal_matrix: mat4x4<f32>,
@@ -49,6 +54,7 @@ struct Light {
 
     spot_cutoff: f32,
     spot_outer_cutoff: f32,
+    light_shadow_uniform_index: i32,
 };
 
 struct Lights {
@@ -80,9 +86,7 @@ fn vs_main(
     // the normal mat is actually only a mat3x3, but there's a bug in wgpu: https://github.com/gfx-rs/wgpu-rs/issues/36
     let normal_mat4 = u_model_transform_data.normal_matrix;
     let normal_mat = mat3x3(normal_mat4[0].xyz, normal_mat4[1].xyz, normal_mat4[2].xyz);
-    out.world_normal = normalize(normal_mat * model.normal, );
-    
-    out.frag_pos_light_space = u_light_space_matrix * world_position;
+    out.world_normal = normalize(normal_mat * model.normal);
 
     return out;
 }
@@ -104,25 +108,29 @@ var t_diffuse: texture_2d<f32>;
 @group(0) @binding(2)
 var s_diffuse: sampler;
 
-/*@group(4) @binding(0)
-var<uniform> u_material: Material;
+struct LightShadowMapUniform {
+    light_space_matrix: mat4x4<f32>,
+    atlas_frame: TextureAtlasFrame,
+}
 
-@group(5) @binding(0)
-var t_specular: texture_2d<f32>;
-@group(5) @binding(1)
-var s_specular: sampler;*/
+struct LightShadowMapUniformAligned {
+    @size(256)
+    inner: LightShadowMapUniform
+}
 
 @group(4) @binding(0)
 var<storage, read_write> u_light_indices: array<u32>;
 @group(4) @binding(1)
-var t_light_grid: texture_storage_2d<rg32uint, read_write>; // vec2<u32>
+var t_light_grid: texture_storage_2d<rg32uint, read_write>; // rg32uint = vec2<u32>
 
 @group(5) @binding(0)
 var t_shadow_maps_atlas: texture_depth_2d;
 @group(5) @binding(1)
 var s_shadow_maps_atlas: sampler;
 @group(5) @binding(2)
-var<uniform> u_light_space_matrix: mat4x4<f32>;
+var<uniform> u_shadow_maps_atlas_size: vec2<u32>;
+@group(5) @binding(3)
+var<storage, read> u_light_shadow: array<LightShadowMapUniformAligned>;
 
 @fragment
 fn fs_main(in: VertexOutput) -> @location(0) vec4<f32> {
@@ -144,41 +152,18 @@ fn fs_main(in: VertexOutput) -> @location(0) vec4<f32> {
     let light_offset = tile.x;
     let light_count = tile.y;
 
+    let atlas_dimensions: vec2<i32> = textureDimensions(t_shadow_maps_atlas);
+
     for (var i = 0u; i < light_count; i++) {
         let light_index = u_light_indices[light_offset + i];
         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;
-            }
+            let shadow_u: LightShadowMapUniform = u_light_shadow[light.light_shadow_uniform_index].inner;
+            let frag_pos_light_space = shadow_u.light_space_matrix * vec4<f32>(in.world_position, 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);
+            let shadow = calc_shadow(in.world_normal, light_dir, frag_pos_light_space, atlas_dimensions, shadow_u.atlas_frame);
             light_res += blinn_phong_dir_light(in.world_position, in.world_normal, light, u_material, specular_color, shadow);
         } else if (light.light_ty == LIGHT_TY_POINT) {
             light_res += blinn_phong_point_light(in.world_position, in.world_normal, light, u_material, specular_color);
@@ -191,7 +176,7 @@ fn fs_main(in: VertexOutput) -> @location(0) vec4<f32> {
     return vec4<f32>(light_object_res, object_color.a);
 }
 
-fn calc_shadow(normal: vec3<f32>, light_dir: vec3<f32>, frag_pos_light_space: vec4<f32>) -> f32 {
+fn calc_shadow(normal: vec3<f32>, light_dir: vec3<f32>, frag_pos_light_space: vec4<f32>, atlas_dimensions: vec2<i32>, atlas_region: TextureAtlasFrame) -> f32 {
     var proj_coords = frag_pos_light_space.xyz / frag_pos_light_space.w;
     // for some reason the y component is clipped after transforming
     proj_coords.y = -proj_coords.y;
@@ -203,20 +188,24 @@ fn calc_shadow(normal: vec3<f32>, light_dir: vec3<f32>, frag_pos_light_space: ve
 
     // Remap xy to [0.0, 1.0]
     let xy_remapped = proj_coords.xy * 0.5 + 0.5;
-     // TODO: when more lights are added, change the index, and the atlas sizes
-    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);
+
+    // no need to get the y since the maps are square
+    let atlas_start = f32(atlas_region.offset.x) / f32(atlas_dimensions.x);
+    let atlas_end = f32(atlas_region.offset.x + atlas_region.size.x) / f32(atlas_dimensions.x);
     // 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);
+    proj_coords.x = mix(atlas_start, atlas_end, xy_remapped.x);
+    proj_coords.y = mix(atlas_start, atlas_end, xy_remapped.y);
 
     // 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) {
+    if (proj_coords.x > atlas_end && proj_coords.y > atlas_end)
+            || (proj_coords.x < atlas_start && proj_coords.y < atlas_start) {
         return 0.0;
     }
 
-    let closest_depth = textureSampleLevel(t_shadow_maps_atlas, s_shadow_maps_atlas, proj_coords.xy, 0.0, vec2<i32>(0, 0));
+    // must manually apply offset to the texture coords since `textureSampleLevel` requires a
+    // const value.
+    let offset_coords = proj_coords.xy + (vec2<f32>(atlas_region.offset) / vec2<f32>(atlas_dimensions));
+    let closest_depth = textureSampleLevel(t_shadow_maps_atlas, s_shadow_maps_atlas, offset_coords, 0.0);
     let current_depth = proj_coords.z;
 
     // use a bias to avoid shadow acne

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

@@ -3,8 +3,18 @@ struct TransformData {
     normal_matrix: mat4x4<f32>,
 }
 
+struct TextureAtlasFrame {
+    offset: vec2<u32>,
+    size: vec2<u32>,
+}
+
+struct LightShadowMapUniform {
+    light_space_matrix: mat4x4<f32>,
+    atlas_frame: TextureAtlasFrame,
+}
+
 @group(0) @binding(0)
-var<uniform> u_light_space_matrix: mat4x4<f32>;
+var<storage, read> u_light_shadow: LightShadowMapUniform;
 
 @group(1) @binding(0)
 var<uniform> u_model_transform_data: TransformData;
@@ -18,6 +28,6 @@ struct VertexOutput {
 fn vs_main(
     @location(0) position: vec3<f32>
 ) -> VertexOutput {
-    let pos = u_light_space_matrix * u_model_transform_data.transform * vec4<f32>(position, 1.0);
+    let pos = u_light_shadow.light_space_matrix * u_model_transform_data.transform * vec4<f32>(position, 1.0);
     return VertexOutput(pos);
 }

lyra-game/src/render/slot_buffer.rs

@@ -0,0 +1,150 @@
+use std::{collections::VecDeque, marker::PhantomData, mem, num::NonZeroU64, sync::Arc};
+
+/// A buffer on the GPU that has persistent indices.
+/// 
+/// `GpuSlotBuffer` allocates a buffer on the GPU and keeps stable indices of elements and
+/// reuses ones that were removed. It supports aligned buffers with [`GpuSlotBuffer::new_aligned`],
+/// as well as unaligned buffers with [`GpuSlotBuffer::new`].
+pub struct GpuSlotBuffer<T: bytemuck::Pod + bytemuck::Zeroable> {
+    /// The amount of elements that can fit in the buffer.
+    capacity: u64,
+    /// The ending point of the buffer elements.
+    len: u64,
+    /// The list of dead and reusable indices in the buffer.
+    dead_indices: VecDeque<u64>,
+    /// The optional alignment of elements in the buffer.
+    alignment: Option<u64>,
+    /// The actual gpu buffer
+    buffer: Arc<wgpu::Buffer>,
+    _marker: PhantomData<T>,
+}
+
+impl<T: bytemuck::Pod + bytemuck::Zeroable> GpuSlotBuffer<T> {
+    /// Create a new GpuSlotBuffer with unaligned elements.
+    /// 
+    /// See [`GpuSlotBuffer::new_aligned`].
+    pub fn new(device: &wgpu::Device, label: Option<&str>, usage: wgpu::BufferUsages, capacity: u64) -> Self {
+        Self::new_impl(device, label, usage, capacity, None)
+    }
+
+    /// Create a new buffer with **aligned** elements.
+    /// 
+    /// See [`GpuSlotBuffer::new`].
+    pub fn new_aligned(device: &wgpu::Device, label: Option<&str>, usage: wgpu::BufferUsages, capacity: u64, alignment: u64) -> Self {
+        Self::new_impl(device, label, usage, capacity, Some(alignment))
+    }
+
+    fn new_impl(device: &wgpu::Device, label: Option<&str>, usage: wgpu::BufferUsages, capacity: u64, alignment: Option<u64>) -> Self {
+        let buffer = Arc::new(device.create_buffer(&wgpu::BufferDescriptor {
+            label,
+            size: capacity * mem::size_of::<T>() as u64,
+            usage,
+            mapped_at_creation: false,
+        }));
+
+        Self {
+            capacity,
+            len: 0,
+            dead_indices: VecDeque::default(),
+            buffer,
+            alignment,
+            _marker: PhantomData
+        }
+    }
+
+    /// Calculates the byte offset in the buffer of the element at `i`.
+    pub fn offset_of(&self, i: u64) -> u64 {
+        let offset = i * mem::size_of::<T>() as u64;
+
+        if let Some(align) = self.alignment {
+            round_mult::up(offset, NonZeroU64::new(align).unwrap()).unwrap()
+        } else {
+            offset
+        }
+    }
+
+    /// Set an element at `i` in the buffer to `val`.
+    pub fn set_at(&self, queue: &wgpu::Queue, i: u64, val: &T) {
+        let offset = self.offset_of(i);
+        queue.write_buffer(&self.buffer, offset, bytemuck::bytes_of(val));
+    }
+
+    /// Attempt to insert an element to the GPU buffer, returning the index it was inserted at.
+    /// 
+    /// Returns `None` when the buffer has no space to fit the element.
+    pub fn try_insert(&mut self, queue: &wgpu::Queue, val: &T) -> Option<u64> {
+        // reuse a dead index or get the next one
+        let i = match self.dead_indices.pop_front() {
+            Some(i) => i,
+            None => {
+                if self.len == self.capacity {
+                    return None;
+                }
+
+                let i = self.len;
+                self.len += 1;
+                i
+            }
+        };
+
+        self.set_at(queue, i, val);
+
+        Some(i)
+    }
+
+    /// Insert an element to the GPU buffer, returning the index it was inserted at.
+    /// 
+    /// The index is not guaranteed to be the end of the buffer since this structure reuses
+    /// indices after they're removed.
+    /// 
+    /// # Panics
+    /// Panics if the buffer does not have space to fit `val`, see [`GpuSlotBuffer::try_insert`].
+    pub fn insert(&mut self, queue: &wgpu::Queue, val: &T) -> u64 {
+        self.try_insert(queue, val)
+            .expect("GPU slot buffer ran out of slots to push elements into")
+    }
+
+    /// Remove the element at `i`, clearing the elements slot in the buffer.
+    /// 
+    /// If you do not care that the slot in the buffer is emptied, use
+    /// [`GpuSlotBuffer::remove_quick`].
+    pub fn remove(&mut self, queue: &wgpu::Queue, i: u64) {
+        let mut zeros = Vec::new();
+        zeros.resize(mem::size_of::<T>(), 0);
+        
+        let offset = self.offset_of(i);
+        queue.write_buffer(&self.buffer, offset, bytemuck::cast_slice(zeros.as_slice()));
+        self.dead_indices.push_back(i);
+    }
+
+    /// Remove the element at `i` without clearing its space in the buffer.
+    /// 
+    /// If you want to ensure that the slot in the buffer is emptied, use
+    /// [`GpuSlotBuffer::remove`].
+    pub fn remove_quick(&mut self, i: u64) {
+        self.dead_indices.push_back(i);
+    }
+
+    /// Returns the backing [`wgpu::Buffer`].
+    pub fn buffer(&self) -> &Arc<wgpu::Buffer> {
+        &self.buffer
+    }
+
+    /// Return the length of the buffer.
+    /// 
+    /// This value may not reflect the amount of elements that are actually alive in the buffer if
+    /// elements were removed and not re-added.
+    pub fn len(&self) -> u64 {
+        self.len
+    }
+
+    /// Return the amount of inuse indices in the buffer.
+    pub fn inuse_len(&self) -> u64 {
+        self.len - self.dead_indices.len() as u64
+    }
+
+    /// Returns the amount of elements the buffer can fit.
+    pub fn capacity(&self) -> u64 {
+        self.capacity
+    }
+}

lyra-game/src/render/texture_atlas.rs

@@ -1,31 +1,54 @@
-use std::sync::Arc;
+use std::{
+    collections::BTreeMap,
+    sync::Arc,
+};
 
 use glam::UVec2;
+use rectangle_pack::{pack_rects, GroupedRectsToPlace, RectToInsert, RectanglePackOk, TargetBin};
 
-#[derive(Debug, Clone, Copy)]
+#[derive(Debug, thiserror::Error)]
+pub enum AtlasPackError {
+    /// The rectangles can't be placed into the atlas. The atlas must increase in size
+    #[error("There is not enough space in the atlas for the textures")]
+    NotEnoughSpace,
+}
+
+#[repr(C)]
+#[derive(Debug, Copy, Clone, bytemuck::Pod, bytemuck::Zeroable)]
 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,
+    atlas_size: UVec2,
 
     texture_format: wgpu::TextureFormat,
     texture: Arc<wgpu::Texture>,
     view: Arc<wgpu::TextureView>,
+
+    /// The next id of the next texture that will be added to the atlas.
+    next_texture_id: u64,
+
+    rects: GroupedRectsToPlace<u64>,
+    bins: BTreeMap<u64, TargetBin>,
+    placement: Option<RectanglePackOk<u64, u64>>,
 }
 
 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;
-        
+    pub fn new(
+        device: &wgpu::Device,
+        format: wgpu::TextureFormat,
+        usages: wgpu::TextureUsages,
+        atlas_size: UVec2,
+    ) -> Self {
         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 },
+            size: wgpu::Extent3d {
+                width: atlas_size.x,
+                height: atlas_size.y,
+                depth_or_array_layers: 1,
+            },
             mip_level_count: 1,
             sample_count: 1,
             dimension: wgpu::TextureDimension::D2,
@@ -35,21 +58,98 @@ impl TextureAtlas {
         });
         let view = texture.create_view(&wgpu::TextureViewDescriptor::default());
 
+        let mut bins = BTreeMap::new();
+        // max_depth=1 for 2d
+        bins.insert(0, TargetBin::new(atlas_size.x, atlas_size.y, 1));
+
         Self {
-            texture_size,
-            texture_count,
+            atlas_size,
             texture_format: format,
             texture: Arc::new(texture),
             view: Arc::new(view),
+            next_texture_id: 0,
+            rects: GroupedRectsToPlace::new(),
+            bins,
+            placement: None,
         }
     }
 
-    /// 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;
+    /// Add a texture of `size` and pack it into the atlas, returning the id of the texture in
+    /// the atlas.
+    /// 
+    /// If you are adding multiple textures at a time and want to wait to pack the atlas, use
+    /// [`TextureAtlas::add_texture_unpacked`] and then after you're done adding them, pack them
+    /// with [`TextureAtlas::pack_atlas`].
+    pub fn pack_new_texture(&mut self, width: u32, height: u32) -> Result<u64, AtlasPackError> {
+        let id = self.next_texture_id;
+        self.next_texture_id += 1;
 
-        AtlasViewport { offset: UVec2::new(x, y), size: self.texture_size }
+        // for 2d rects, set depth to 1
+        let r = RectToInsert::new(width, height, 1);
+        self.rects.push_rect(id, None, r);
+
+        self.pack_atlas()?;
+
+        Ok(id)
+    }
+
+    /// Add a new texture and **DO NOT** pack it into the atlas.
+    /// 
+    /// <div class="warning">
+    /// 
+    /// The texture will not be packed into the atlas meaning
+    /// [`TextureAtlas::texture_viewport`] will return `None`. To pack the texture,
+    /// use [`TextureAtlas::pack_atlas`] or use [`TextureAtlas::pack_new_texture`]
+    /// when only adding a single texture.
+    /// 
+    /// </div>
+    pub fn add_texture_unpacked(&mut self, width: u32, height: u32) -> Result<u64, AtlasPackError> {
+        let id = self.next_texture_id;
+        self.next_texture_id += 1;
+
+        // for 2d rects, set depth to 1
+        let r = RectToInsert::new(width, height, 1);
+        self.rects.push_rect(id, None, r);
+
+        self.pack_atlas()?;
+
+        Ok(id)
+    }
+
+    /// Pack the textures into the atlas.
+    pub fn pack_atlas(&mut self) -> Result<(), AtlasPackError> {
+        let placement = pack_rects(
+            &self.rects,
+            &mut self.bins,
+            &rectangle_pack::volume_heuristic,
+            &rectangle_pack::contains_smallest_box,
+        )
+        .map_err(|e| match e {
+            rectangle_pack::RectanglePackError::NotEnoughBinSpace => AtlasPackError::NotEnoughSpace,
+        })?;
+        self.placement = Some(placement);
+
+        Ok(())
+    }
+
+    /// Get the viewport of a texture index in the atlas.
+    pub fn texture_viewport(&self, atlas_index: u64) -> AtlasViewport {
+        let locations = self.placement.as_ref().unwrap().packed_locations();
+        let (bin_id, loc) = locations
+            .get(&atlas_index)
+            .expect("atlas index is incorrect");
+        debug_assert_eq!(*bin_id, 0, "somehow the texture was put in some other bin");
+
+        AtlasViewport {
+            offset: UVec2 {
+                x: loc.x(),
+                y: loc.y(),
+            },
+            size: UVec2 {
+                x: loc.width(),
+                y: loc.height(),
+            },
+        }
     }
 
     pub fn view(&self) -> &Arc<wgpu::TextureView> {
@@ -64,15 +164,12 @@ impl TextureAtlas {
         &self.texture_format
     }
 
-    pub fn texture_size(&self) -> UVec2 {
-        self.texture_size
+    pub fn total_texture_count(&self) -> u64 {
+        self.next_texture_id // starts at zero, so no need to increment
     }
 
-    pub fn texture_count(&self) -> UVec2 {
-        self.texture_count
-    }
-
-    pub fn total_texture_count(&self) -> u32 {
-        self.texture_count.x * self.texture_count.y
+    /// Returns the size of the entire texture atlas.
+    pub fn atlas_size(&self) -> UVec2 {
+        self.atlas_size
     }
 }