onedrop_renderer/

texture_pool.rs

//! User-loaded texture pool for the comp pass.
//!
//! MilkDrop 2 lets preset authors reference image files on disk via
//! `sampler sampler_<NAME>;` declarations in the comp shader, then sample them
//! through `tex2D(sampler_<NAME>, uv)`. The image filename (stem) becomes the
//! sampler name — `clouds.png` → `sampler_clouds`. Authors also reference
//! `sampler_rand00..15` to pull a deterministic-per-preset random texture from
//! whatever's on disk, optionally filtered by name prefix
//! (`sampler_rand02_smalltiled` → only `smalltiled*.png`).
//!
//! This module owns the disk-loaded GPU textures. The comp pipeline reads
//! through it at preset load time to populate the 8 user-texture binding slots.
//! Lookup is by canonicalised name (lowercase stem, no extension); presets in
//! the wild mix case freely (`sampler_Worms` vs `sampler_worms`).
//!
//! Pool population is best-effort:
//! - Missing directory → empty pool (engine continues with the fallback 1×1
//!   white texture). Logged once at debug level.
//! - Unreadable / undecodable file → skipped with a warn log; rest of the
//!   directory loads.
//! - Name collisions across directories → first hit wins; later directories
//!   only fill gaps. Matches MD2's "search path" semantics.

use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::sync::Arc;

/// One GPU-resident user texture.
pub struct UserTexture {
    /// Canonicalised name — lowercase stem of the source filename, no
    /// extension. `clouds.png` → `"clouds"`. Match against the `sampler_X`
    /// portion of `sampler sampler_X;` MD2 declarations (lowercased).
    pub name: String,
    pub texture: wgpu::Texture,
    pub view: wgpu::TextureView,
    pub width: u32,
    pub height: u32,
}

impl UserTexture {
    /// `vec4<f32>(w, h, 1/w, 1/h)` — the `texsize_<NAME>` constant the comp
    /// wrapper emits for this texture. Matches the MD2 convention used for
    /// the noise pack.
    pub fn texsize(&self) -> [f32; 4] {
        let w = self.width as f32;
        let h = self.height as f32;
        [w, h, 1.0 / w, 1.0 / h]
    }
}

/// Texture pool. Hold one per renderer; the comp pipeline borrows it whenever
/// a preset's user shader needs to bind disk-loaded textures.
pub struct TexturePool {
    /// All loaded textures in disk-walk order. The order is stable across
    /// runs as long as the underlying filesystem returns the same listing —
    /// used by `pick_rand` to make `sampler_rand0X` deterministic.
    textures: Vec<Arc<UserTexture>>,
    /// Lookup by canonicalised name.
    by_name: HashMap<String, usize>,
    /// 1×1 opaque white. Bound into unfilled slots and returned when a
    /// preset references a name that isn't on disk. Keeps every user
    /// shader compileable + drawable even when the texture pool is empty.
    fallback: Arc<UserTexture>,
}

impl TexturePool {
    /// Build a pool with only the 1×1 white fallback. Tests and headless CI
    /// paths use this; the real engine calls `from_dirs` after.
    pub fn new(device: &wgpu::Device, queue: &wgpu::Queue) -> Self {
        let fallback = Arc::new(make_fallback_texture(device, queue));
        Self {
            textures: Vec::new(),
            by_name: HashMap::new(),
            fallback,
        }
    }

    /// Build a pool by scanning each directory in `dirs` for image files.
    /// Earlier directories shadow later ones — pass user dirs before system
    /// dirs if you want overrides.
    ///
    /// Decode errors are logged and the offending file skipped; missing
    /// directories are silently treated as empty.
    pub fn from_dirs(device: &wgpu::Device, queue: &wgpu::Queue, dirs: &[PathBuf]) -> Self {
        let mut pool = Self::new(device, queue);
        for dir in dirs {
            pool.scan_dir(device, queue, dir);
        }
        pool
    }

    /// Append textures from one directory. Doesn't recurse — flat layout
    /// matches MD2's `Textures/` convention.
    pub fn scan_dir(&mut self, device: &wgpu::Device, queue: &wgpu::Queue, dir: &Path) {
        let entries = match std::fs::read_dir(dir) {
            Ok(e) => e,
            Err(err) => {
                log::debug!("texture dir {} not readable: {}", dir.display(), err);
                return;
            }
        };
        for entry in entries.flatten() {
            let path = entry.path();
            if !path.is_file() {
                continue;
            }
            let Some(name) = canonical_stem(&path) else {
                continue;
            };
            if self.by_name.contains_key(&name) {
                continue;
            }
            match load_texture_file(device, queue, &path, &name) {
                Ok(tex) => {
                    let idx = self.textures.len();
                    self.textures.push(Arc::new(tex));
                    self.by_name.insert(name, idx);
                }
                Err(e) => {
                    log::warn!("failed to load texture {}: {}", path.display(), e);
                }
            }
        }
    }

    /// Lookup by canonicalised name. `name` is compared against the lowercase
    /// stem of the source filename — `get("clouds")` matches both
    /// `clouds.png` and `Clouds.PNG`.
    pub fn get(&self, name: &str) -> Option<&Arc<UserTexture>> {
        let canon = name.to_ascii_lowercase();
        self.by_name.get(&canon).map(|&i| &self.textures[i])
    }

    /// Always-non-None fallback — bound into unfilled user-texture slots so
    /// the comp pass's bind group is complete even when the pool is empty
    /// or a preset references a name we don't have.
    pub fn fallback(&self) -> &Arc<UserTexture> {
        &self.fallback
    }

    /// `true` when the pool has no user textures (only the fallback). The
    /// renderer uses this to skip the disk-load scan path entirely on
    /// `sampler_rand0X` resolution.
    pub fn is_empty(&self) -> bool {
        self.textures.is_empty()
    }

    /// Number of user textures (excludes the fallback).
    pub fn len(&self) -> usize {
        self.textures.len()
    }

    /// Deterministic-per-`seed` pick from the pool, optionally filtered by
    /// name prefix. Returns `None` when no texture matches (caller can
    /// substitute the fallback).
    ///
    /// MD2's `sampler_rand0X` slots want a stable choice for the lifetime
    /// of a preset but a different choice across presets. The caller seeds
    /// with `hash(preset_name) ^ slot_index` to achieve that.
    pub fn pick_rand(&self, seed: u64, prefix: Option<&str>) -> Option<&Arc<UserTexture>> {
        let pool_subset: Vec<usize> = match prefix {
            Some(p) => {
                let p = p.to_ascii_lowercase();
                self.textures
                    .iter()
                    .enumerate()
                    .filter(|(_, t)| t.name.starts_with(&p))
                    .map(|(i, _)| i)
                    .collect()
            }
            None => (0..self.textures.len()).collect(),
        };
        if pool_subset.is_empty() {
            return None;
        }
        let idx = pool_subset[(seed as usize) % pool_subset.len()];
        Some(&self.textures[idx])
    }
}

/// XDG-default search paths for user textures. The renderer's caller
/// (CLI / GUI) typically extends this with config-driven dirs.
///
/// Returns the candidate paths even when they don't exist on disk — the pool
/// constructor silently skips missing dirs.
pub fn default_texture_dirs() -> Vec<PathBuf> {
    let mut out = Vec::new();
    if let Some(home) = std::env::var_os("HOME") {
        let home = PathBuf::from(home);
        out.push(home.join(".config/onedrop/textures"));
        // MD2 historical layout — many users drop their texture pack under
        // `~/Music/milkdrop/textures` along with the presets.
        out.push(home.join("Music/milkdrop/textures"));
    }
    if let Some(xdg) = std::env::var_os("XDG_DATA_HOME") {
        out.push(PathBuf::from(xdg).join("onedrop/textures"));
    }
    out
}

fn canonical_stem(path: &Path) -> Option<String> {
    let stem = path.file_stem()?.to_str()?.to_ascii_lowercase();
    if stem.is_empty() {
        return None;
    }
    Some(stem)
}

fn load_texture_file(
    device: &wgpu::Device,
    queue: &wgpu::Queue,
    path: &Path,
    name: &str,
) -> Result<UserTexture, image::ImageError> {
    let img = image::open(path)?;
    let rgba = img.to_rgba8();
    let (width, height) = rgba.dimensions();
    Ok(upload_rgba8(device, queue, name, width, height, &rgba))
}

fn upload_rgba8(
    device: &wgpu::Device,
    queue: &wgpu::Queue,
    name: &str,
    width: u32,
    height: u32,
    bytes: &[u8],
) -> UserTexture {
    let texture = device.create_texture(&wgpu::TextureDescriptor {
        label: Some(&format!("UserTexture {name}")),
        size: wgpu::Extent3d {
            width,
            height,
            depth_or_array_layers: 1,
        },
        mip_level_count: 1,
        sample_count: 1,
        dimension: wgpu::TextureDimension::D2,
        format: wgpu::TextureFormat::Rgba8Unorm,
        usage: wgpu::TextureUsages::TEXTURE_BINDING | wgpu::TextureUsages::COPY_DST,
        view_formats: &[],
    });
    queue.write_texture(
        wgpu::TexelCopyTextureInfo {
            texture: &texture,
            mip_level: 0,
            origin: wgpu::Origin3d::ZERO,
            aspect: wgpu::TextureAspect::All,
        },
        bytes,
        wgpu::TexelCopyBufferLayout {
            offset: 0,
            bytes_per_row: Some(width * 4),
            rows_per_image: Some(height),
        },
        wgpu::Extent3d {
            width,
            height,
            depth_or_array_layers: 1,
        },
    );
    let view = texture.create_view(&wgpu::TextureViewDescriptor::default());
    UserTexture {
        name: name.to_string(),
        texture,
        view,
        width,
        height,
    }
}

fn make_fallback_texture(device: &wgpu::Device, queue: &wgpu::Queue) -> UserTexture {
    // 1×1 opaque white — sampling it yields `vec4(1, 1, 1, 1)` regardless of
    // uv, so any `tex2D(sampler_missing, uv) * something` math degrades to
    // `something` (the well-behaved no-op for multiplicative texture use).
    upload_rgba8(
        device,
        queue,
        "__fallback_white",
        1,
        1,
        &[255, 255, 255, 255],
    )
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::config::RenderConfig;
    use crate::gpu_context::GpuContext;

    /// Sythesise a 4×4 RGBA8 PNG in-place. Avoids dragging a fixture file
    /// into the repo for tests that just need *some* decodable image.
    fn write_solid_png(path: &Path, color: [u8; 4]) {
        let pixels: Vec<u8> = (0..16).flat_map(|_| color).collect();
        let img =
            image::RgbaImage::from_raw(4, 4, pixels).expect("RgbaImage::from_raw must succeed");
        img.save_with_format(path, image::ImageFormat::Png).unwrap();
    }

    #[test]
    fn empty_pool_has_fallback() {
        let cfg = RenderConfig::default();
        let gpu = pollster::block_on(GpuContext::new(cfg)).unwrap();
        let pool = TexturePool::new(&gpu.device, &gpu.queue);
        assert!(pool.is_empty());
        assert_eq!(pool.len(), 0);
        assert!(pool.get("anything").is_none());
        // Fallback is always present, 1×1.
        assert_eq!(pool.fallback().width, 1);
        assert_eq!(pool.fallback().height, 1);
    }

    #[test]
    fn from_dirs_loads_pngs() {
        let cfg = RenderConfig::default();
        let gpu = pollster::block_on(GpuContext::new(cfg)).unwrap();
        let dir = tempfile::tempdir().unwrap();
        write_solid_png(&dir.path().join("Clouds.png"), [10, 20, 30, 255]);
        write_solid_png(&dir.path().join("worms.png"), [40, 50, 60, 255]);

        let pool = TexturePool::from_dirs(&gpu.device, &gpu.queue, &[dir.path().to_path_buf()]);
        assert_eq!(pool.len(), 2);
        // Case-insensitive lookup — `Clouds.png` is indexed as `clouds`.
        assert!(pool.get("clouds").is_some());
        assert!(pool.get("Clouds").is_some());
        assert!(pool.get("worms").is_some());
        assert!(pool.get("missing").is_none());
    }

    #[test]
    fn missing_dir_is_skipped() {
        let cfg = RenderConfig::default();
        let gpu = pollster::block_on(GpuContext::new(cfg)).unwrap();
        let pool = TexturePool::from_dirs(
            &gpu.device,
            &gpu.queue,
            &[PathBuf::from("/definitely/does/not/exist/onedrop-test")],
        );
        // No panic; fallback still present.
        assert!(pool.is_empty());
    }

    #[test]
    fn name_collision_first_wins() {
        let cfg = RenderConfig::default();
        let gpu = pollster::block_on(GpuContext::new(cfg)).unwrap();
        let dir_a = tempfile::tempdir().unwrap();
        let dir_b = tempfile::tempdir().unwrap();
        write_solid_png(&dir_a.path().join("clouds.png"), [10, 20, 30, 255]);
        write_solid_png(&dir_b.path().join("clouds.png"), [200, 210, 220, 255]);

        let pool = TexturePool::from_dirs(
            &gpu.device,
            &gpu.queue,
            &[dir_a.path().to_path_buf(), dir_b.path().to_path_buf()],
        );
        // Both directories had `clouds.png` but only the first one wins.
        assert_eq!(pool.len(), 1);
        assert!(pool.get("clouds").is_some());
    }

    #[test]
    fn pick_rand_is_deterministic() {
        let cfg = RenderConfig::default();
        let gpu = pollster::block_on(GpuContext::new(cfg)).unwrap();
        let dir = tempfile::tempdir().unwrap();
        for n in ["alpha", "beta", "gamma"] {
            write_solid_png(&dir.path().join(format!("{n}.png")), [0, 0, 0, 255]);
        }
        let pool = TexturePool::from_dirs(&gpu.device, &gpu.queue, &[dir.path().to_path_buf()]);
        // Same seed twice → same pick.
        let a = pool.pick_rand(42, None).map(|t| t.name.clone());
        let b = pool.pick_rand(42, None).map(|t| t.name.clone());
        assert_eq!(a, b);
        assert!(a.is_some());
    }

    #[test]
    fn pick_rand_with_prefix_filters() {
        let cfg = RenderConfig::default();
        let gpu = pollster::block_on(GpuContext::new(cfg)).unwrap();
        let dir = tempfile::tempdir().unwrap();
        write_solid_png(&dir.path().join("smalltiled_a.png"), [0; 4]);
        write_solid_png(&dir.path().join("smalltiled_b.png"), [0; 4]);
        write_solid_png(&dir.path().join("unrelated.png"), [0; 4]);
        let pool = TexturePool::from_dirs(&gpu.device, &gpu.queue, &[dir.path().to_path_buf()]);

        // Every seed must land on a `smalltiled_*` texture when the prefix
        // filter is active.
        for seed in 0..16 {
            let pick = pool.pick_rand(seed, Some("smalltiled_")).unwrap();
            assert!(
                pick.name.starts_with("smalltiled_"),
                "seed {seed} picked {}",
                pick.name
            );
        }
        // Empty prefix subset → None.
        assert!(pool.pick_rand(0, Some("nothing_matches_")).is_none());
    }

    #[test]
    fn pick_rand_on_empty_pool_is_none() {
        let cfg = RenderConfig::default();
        let gpu = pollster::block_on(GpuContext::new(cfg)).unwrap();
        let pool = TexturePool::new(&gpu.device, &gpu.queue);
        assert!(pool.pick_rand(0, None).is_none());
    }

    #[test]
    fn texsize_matches_dimensions() {
        let cfg = RenderConfig::default();
        let gpu = pollster::block_on(GpuContext::new(cfg)).unwrap();
        let dir = tempfile::tempdir().unwrap();
        write_solid_png(&dir.path().join("a.png"), [0; 4]);
        let pool = TexturePool::from_dirs(&gpu.device, &gpu.queue, &[dir.path().to_path_buf()]);
        let t = pool.get("a").unwrap();
        let v = t.texsize();
        assert_eq!(v[0], 4.0);
        assert_eq!(v[1], 4.0);
        assert!((v[2] - 0.25).abs() < 1e-6);
        assert!((v[3] - 0.25).abs() < 1e-6);
    }
}