Cleaned up the TextureAtlas API

2 files changed, 159 insertions, 181 deletions

diff --git a/alligator_resources/src/lib.rs b/alligator_resources/src/lib.rs
index 89a8f3c..9cbbba0 100644
--- a/alligator_resources/src/lib.rs
+++ b/alligator_resources/src/lib.rs
@@ -11,5 +11,3 @@ pub enum Priority {
 	Eventual(u8),
 	Urgent,
 }
-
-pub struct ResourceManager {}
diff --git a/alligator_resources/src/texture.rs b/alligator_resources/src/texture.rs
index 569f2c6..75c3889 100644
--- a/alligator_resources/src/texture.rs
+++ b/alligator_resources/src/texture.rs
@@ -1,10 +1,11 @@
+use std::cmp::Reverse;
 use std::collections::HashMap;
-use std::mem::{self, MaybeUninit};
+use std::mem::{self};
 use std::path::Path;
 use std::sync::atomic::{AtomicUsize, Ordering};
 use std::sync::Arc;
 
-use image::{GenericImage, ImageBuffer};
+use image::ImageBuffer;
 use parking_lot::Mutex;
 use texture_packer::exporter::ImageExporter;
 use texture_packer::{Frame, TexturePacker, TexturePackerConfig};
@@ -12,8 +13,10 @@ use thiserror::Error;
 
 use crate::Priority;
 
+/// The next texture ID
 static NEXT_TEXTURE_ID: AtomicUsize = AtomicUsize::new(0);
 
+/// A unique identifier for a texture
 #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
 pub struct TextureId(usize);
 
@@ -81,17 +84,20 @@ fn vec_image_to_box(vec_image: image::ImageBuffer<image::Rgba<u16>, Vec<u16>>) -
 	ImageBuffer::from_raw(width, height, buf).expect("image buffer is too small")
 }
 
+/// Get the size, in bytes, of the texture
 #[allow(clippy::missing_const_for_fn)]
 fn texture_size(image: &Rgba16Texture) -> usize {
 	image.len() * mem::size_of::<image::Rgba<u16>>()
 }
 
+/// A texture from disk
 struct TextureFile {
 	path: Box<Path>,
 	texture: Option<Arc<Rgba16Texture>>,
 }
 
 impl TextureFile {
+	/// This doesn't load the texture
 	#[allow(clippy::missing_const_for_fn)]
 	fn new(path: impl AsRef<Path>) -> Self {
 		Self {
@@ -100,13 +106,6 @@ impl TextureFile {
 		}
 	}
 
-	fn open(path: impl AsRef<Path>) -> Result<Self, LoadError> {
-		let mut this = Self::new(path);
-		this.load()?;
-
-		Ok(this)
-	}
-
 	const fn is_loaded(&self) -> bool {
 		self.texture.is_some()
 	}
@@ -122,17 +121,24 @@ impl TextureFile {
 		Ok(self.texture.as_ref().expect("the texture wasn't loaded"))
 	}
 
+	fn loaded_texture(&self) -> Option<&Rgba16Texture> {
+		self.texture.as_deref()
+	}
+
 	fn is_used(&self) -> bool {
 		let Some(arc) = &self.texture else { return false };
 		Arc::strong_count(arc) > 1
 	}
 
+	/// Unloads the texture from memory if it isn't being used
 	fn unload(&mut self) {
 		if !self.is_used() {
 			self.texture = None;
 		}
 	}
 
+	/// The amount of heap memory used, in bytes. This returns 0 if the texture
+	/// hasn't been loaded yet.
 	fn allocated_size(&self) -> usize {
 		self.texture.as_ref().map_or(0, |t| texture_size(t))
 	}
@@ -200,12 +206,21 @@ impl Texture {
 		}
 	}
 
+	/// If the texture is loaded, return it.
+	fn loaded_texture(&self) -> Option<&Rgba16Texture> {
+		match &self.buffer {
+			TextureBuffer::Memory(ref texture) => Some(texture),
+			TextureBuffer::Disk(file) => file.loaded_texture(),
+		}
+	}
+
 	fn unload(&mut self) {
 		if let TextureBuffer::Disk(file) = &mut self.buffer {
 			file.unload();
 		}
 	}
 
+	/// The amount of heap memory used for the texture, if any
 	fn allocated_size(&self) -> usize {
 		match &self.buffer {
 			TextureBuffer::Memory(texture) => texture_size(texture),
@@ -221,56 +236,108 @@ impl Texture {
 	}
 }
 
-pub struct TextureRef {
-	id: TextureId,
-	// TODO we don't need this. Just the width and height will do
-	texture: Arc<Rgba16Texture>,
-	queued_priority: Arc<Mutex<Option<Priority>>>,
+pub struct TextureAtlas<'a> {
+	width: u32,
+	height: u32,
+	packer: TexturePacker<'a, Rgba16Texture, TextureId>,
 }
 
-impl TextureRef {
+impl<'a> TextureAtlas<'a> {
+	fn new(width: u32, height: u32, textures: &HashMap<TextureId, Texture>) -> Self {
+		profiling::scope!("new atlas");
+
+		let mut packer = TexturePacker::new_skyline(TexturePackerConfig {
+			max_width: width,
+			max_height: height,
+			allow_rotation: false,
+			trim: false,
+			texture_padding: 0,
+			..Default::default()
+		});
+
+		for (id, texture) in textures {
+			if texture.is_loaded() {
+				let texture = texture
+					.loaded_texture()
+					.expect("texture couldn't be loaded");
+
+				// if the textures don't fit, make a bigger packer
+				if packer.pack_own(*id, texture.clone()).is_err() {
+					return Self::new(width * 2, height * 2, textures);
+				}
+			}
+		}
+
+		Self {
+			width,
+			height,
+			packer,
+		}
+	}
+
+	fn subtexture(&self, id: TextureId) -> Option<&Frame<TextureId>> {
+		self.packer.get_frame(&id)
+	}
+
 	#[must_use]
-	pub const fn id(&self) -> TextureId {
-		self.id
+	pub const fn atlas_width(&self) -> u32 {
+		self.width
 	}
 
-	// TODO: it's safer to replace this with a position thingy
 	#[must_use]
-	pub fn atlas_x(&self, manager: &TextureManager) -> f32 {
-		manager.subtexture_x(self.id).expect("not in texture atlas")
+	pub const fn atlas_height(&self) -> u32 {
+		self.height
 	}
 
+	/// Get the x-position of a texture, if it is in the texture atlas
 	#[must_use]
-	pub fn atlas_y(&self, manager: &TextureManager) -> f32 {
-		manager.subtexture_y(self.id).expect("not in texture atlas")
+	#[allow(clippy::cast_precision_loss)] // TODO remove this
+	pub fn subtexture_x(&self, id: TextureId) -> Option<f32> {
+		let x = self.subtexture(id)?.frame.x;
+		Some(x as f32 / self.width as f32)
 	}
 
-	// TODO: it's safer to replace this with a position thingy
+	/// Get the y-position of a texture, if it is in the texture atlas
 	#[must_use]
 	#[allow(clippy::cast_precision_loss)] // TODO remove this
-	pub fn width(&self, manager: &TextureManager) -> f32 {
-		self.texture.width() as f32 / manager.atlas_width() as f32
+	pub fn subtexture_y(&self, id: TextureId) -> Option<f32> {
+		let y = self.subtexture(id)?.frame.y;
+		Some(y as f32 / self.height as f32)
 	}
 
+	/// Get the width of a texture, if it is in the texture atlas
 	#[must_use]
 	#[allow(clippy::cast_precision_loss)] // TODO remove this
-	pub fn height(&self, manager: &TextureManager) -> f32 {
-		self.texture.height() as f32 / manager.atlas_height() as f32
+	pub fn subtexture_width(&self, id: TextureId) -> Option<f32> {
+		let width = self.subtexture(id)?.frame.w;
+		Some(width as f32 / self.width as f32)
 	}
 
-	pub fn queue_priority(&self, priority: Priority) {
-		let mut queued_priority = self.queued_priority.lock();
-		*queued_priority = Some(priority);
+	/// Get the height of a texture, if it is in the texture atlas
+	#[must_use]
+	#[allow(clippy::cast_precision_loss)] // TODO remove this
+	pub fn subtexture_height(&self, id: TextureId) -> Option<f32> {
+		let height = self.subtexture(id)?.frame.h;
+		Some(height as f32 / self.height as f32)
+	}
+
+	#[must_use]
+	pub fn to_texture(&self) -> Rgba16Texture {
+		profiling::scope!("export atlas");
+		vec_image_to_box(
+			ImageExporter::export(&self.packer)
+				.expect("ImageExporter error?")
+				.into_rgba16(),
+		)
 	}
 }
 
 pub struct TextureManager {
 	textures: HashMap<TextureId, Texture>,
-	packer: TexturePacker<'static, Rgba16Texture, TextureId>,
-	atlas: Rgba16Texture, // cached texture atlas
 	max_size: usize,
-	width: u32,
-	height: u32,
+	atlas_width: u32,
+	atlas_height: u32,
+	needs_atlas_update: bool,
 }
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
@@ -281,63 +348,55 @@ pub struct TextureConfig {
 	pub atlas_height: u32,
 }
 
-fn packer(width: u32, height: u32) -> TexturePacker<'static, Rgba16Texture, TextureId> {
-	TexturePacker::new_skyline(TexturePackerConfig {
-		max_width: width,
-		max_height: height,
-		allow_rotation: false,
-		trim: false,
-		texture_padding: 0,
-		..Default::default()
-	})
+impl Default for TextureConfig {
+	fn default() -> Self {
+		Self {
+			initial_capacity: 500,
+			max_size: 10 * 1024 * 1024, // 10 MiB
+			atlas_width: 3980,
+			atlas_height: 2160, // 4K resolution
+		}
+	}
 }
 
 impl TextureManager {
 	/// Create a new `TextureManager` with the given config options.
 	#[must_use]
 	pub fn new(config: TextureConfig) -> Self {
-		let width = config.atlas_width;
-		let height = config.atlas_height;
-		let max_size = config.max_size;
 		let textures = HashMap::with_capacity(config.initial_capacity);
-		let packer = packer(width, height);
-
-		let atlas: Box<[MaybeUninit<u16>]> = Box::new_zeroed_slice((4 * width * height) as _);
-		let atlas = unsafe { atlas.assume_init() };
-		let atlas = Rgba16Texture::from_raw(width, height, atlas);
-		let atlas = atlas.expect("atlas cache is too small");
 
 		Self {
 			textures,
-			packer,
-			atlas,
-			max_size,
-			width,
-			height,
+			max_size: config.max_size,
+			atlas_width: config.atlas_width,
+			atlas_height: config.atlas_height,
+			needs_atlas_update: false,
 		}
 	}
 
-	fn can_load(&mut self, size: usize, priority: Priority) -> bool {
-		let mut textures: Vec<&mut Texture> = self.textures.values_mut().collect();
-		textures.sort_by_key(|a| a.priority());
-		textures.reverse();
+	/// Load textures into memory that will be needed soon. Unload unnecessary textures
+	pub fn cache_files(&mut self) {
+		let mut textures: Vec<&mut Texture> = self
+			.textures
+			.values_mut()
+			.map(|t| {
+				t.unqueue_priority();
+				t
+			})
+			.collect();
+		textures.sort_by_key(|t2| Reverse(t2.priority()));
 
 		let max_size = self.max_size;
-		let priority = priority;
-		let mut total_size = size;
+		let mut total_size = 0;
 
 		for texture in textures {
-			texture.unqueue_priority();
-			if total_size + texture.allocated_size() < max_size {
-				total_size += texture.allocated_size();
-			} else if texture.priority() < priority {
+			drop(texture.load_texture());
+			total_size += texture.allocated_size();
+			if total_size > max_size && texture.priority() != Priority::Urgent {
 				texture.unload();
-			} else {
-				return false;
+				return;
 			}
 		}
-
-		true
 	}
 
 	/// Loads a texture from memory in the given format.
@@ -359,60 +418,17 @@ impl TextureManager {
 		let texture = Texture::from_buffer(texture);
 
 		self.textures.insert(id, texture);
+		self.needs_atlas_update = true;
 
 		Ok(id)
 	}
 
-	fn resize_atlas(&mut self, width: u32, height: u32) {
-		self.packer = packer(width, height);
-
-		for (id, texture) in &mut self.textures {
-			if texture.priority == Priority::Urgent {
-				let texture = texture
-					.load_texture()
-					.expect("unable to load texture when putting it in the atlas");
-
-				self.packer
-					.pack_own(*id, texture.clone())
-					.expect("resized atlas is too small");
-			}
-		}
-	}
-
-	fn reset_atlas(&mut self, fallback_width: u32, fallback_height: u32) {
-		self.packer = packer(self.width, self.height);
-
-		let mut failed = false;
-		for (id, texture) in &mut self.textures {
-			if texture.priority == Priority::Urgent {
-				let texture = texture
-					.load_texture()
-					.expect("unable to load texture when putting it in the atlas");
-
-				if self.packer.pack_own(*id, texture.clone()).is_err() {
-					failed = true;
-				}
-			}
-		}
-
-		if failed {
-			self.resize_atlas(fallback_width, fallback_height);
-		}
-	}
-
-	fn load_to_atlas(&mut self, id: TextureId) {
-		let texture = self.textures.get_mut(&id).expect("invalid texture id");
-		let texture = texture
-			.load_texture()
-			.expect("unable to load texture when putting it in the atlas");
-
-		if self.packer.pack_own(id, texture.clone()).is_err() {
-			let fallback_width = self.width * 2 + texture.width();
-			let fallback_height = self.height * 2 + texture.height();
-			self.reset_atlas(fallback_width, fallback_height);
-		}
-	}
-
+	/// Loads a texture from disk.
+	///
+	/// # Errors
+	///
+	/// This returns an error if `priority` is set to [`Priority::Urgent`] but
+	/// there was an error in loading the file to a texture.
 	pub fn load_from_file(
 		&mut self,
 		path: impl AsRef<Path>,
@@ -420,14 +436,12 @@ impl TextureManager {
 	) -> Result<TextureId, LoadError> {
 		let id = TextureId::new();
 		let mut texture = Texture::from_path(path, priority);
-		let size = texture.allocated_size();
 
-		if priority != Priority::Unnecessary
-			&& (priority == Priority::Urgent || self.can_load(size, texture.priority()))
-		{
+		if priority == Priority::Urgent {
 			match texture.load_texture() {
 				Ok(_) => {
 					self.textures.insert(id, texture);
+					self.needs_atlas_update = true;
 				}
 				Err(e) => {
 					self.textures.insert(id, texture);
@@ -438,75 +452,41 @@ impl TextureManager {
 			self.textures.insert(id, texture);
 		}
 
-		if priority == Priority::Urgent {
-			self.load_to_atlas(id);
-		}
-
 		Ok(id)
 	}
 
+	/// Loads a texture from disk.
+	///
+	/// # Errors
+	///
+	/// This returns an error if `priority` is set to [`Priority::Urgent`] but
+	/// there was an error in loading the file to a texture.
 	pub fn set_priority(&mut self, id: TextureId, priority: Priority) -> Result<(), LoadError> {
 		let texture = self.textures.get_mut(&id).expect("invalid texture id");
-		let old_priority = texture.priority();
 		texture.set_priority(priority);
-		let size = texture.allocated_size();
 
-		if priority > old_priority
-			&& priority != Priority::Unnecessary
-			&& !texture.is_loaded()
-			&& (priority == Priority::Urgent || self.can_load(size, priority))
-		{
+		if !texture.is_loaded() && priority == Priority::Urgent {
 			let texture = self.textures.get_mut(&id).expect("invalid texture id");
 			texture.load_texture()?;
-		}
-
-		if priority == Priority::Urgent {
-			self.load_to_atlas(id);
+			self.needs_atlas_update = true;
 		}
 
 		Ok(())
 	}
 
-	pub fn atlas(&mut self) -> &Rgba16Texture {
-		let atlas = {
-			profiling::scope!("export atlas");
-			ImageExporter::export(&self.packer).expect("ImageExporter error?")
-		};
-
-		profiling::scope!("copy image");
-		self.atlas
-			.copy_from(&atlas.into_rgba16(), 0, 0)
-			.expect("image cache was too small");
-
-		&self.atlas
-	}
-
-	/// Get the subtexture in the texture atlas
-	fn subtexture(&self, id: TextureId) -> Option<&Frame<TextureId>> {
-		self.packer.get_frame(&id)
-	}
-
-	const fn atlas_width(&self) -> u32 {
-		self.width
-	}
-
-	const fn atlas_height(&self) -> u32 {
-		self.height
+	/// This returns `true` if a texture has been set to have an urgent
+	/// priority since the last time this function was called.
+	pub fn needs_atlas_update(&mut self) -> bool {
+		let needs_update = self.needs_atlas_update;
+		self.needs_atlas_update = false;
+		needs_update
 	}
 
-	/// Get the x-position of a texture, if it is in the texture atlas
-	#[must_use]
-	#[allow(clippy::cast_precision_loss)] // TODO remove this
-	pub fn subtexture_x(&self, id: TextureId) -> Option<f32> {
-		let x = self.subtexture(id)?.frame.x;
-		Some(x as f32 / self.width as f32)
-	}
-
-	/// Get the y-position of a texture, if it is in the texture atlas
-	#[must_use]
-	#[allow(clippy::cast_precision_loss)] // TODO remove this
-	pub fn subtexture_y(&self, id: TextureId) -> Option<f32> {
-		let y = self.subtexture(id)?.frame.y;
-		Some(y as f32 / self.height as f32)
+	/// Create a texture atlas
+	pub fn atlas(&mut self) -> TextureAtlas<'_> {
+		let atlas = TextureAtlas::new(self.atlas_width, self.atlas_height, &self.textures);
+		self.atlas_width = atlas.atlas_width();
+		self.atlas_height = atlas.atlas_height();
+		atlas
 	}
 }