mirage-engine 0.1.0

use std::collections::{BTreeSet, HashMap};
use std::rc::Rc;

use wgpu::util::DeviceExt;

use crate::assets::{Assets, Textures, Unresolved};
use crate::mesh::{Animation, Geometry, Meshes};

/// Frame count a mesh is kept on the GPU without being drawn.
const UNUSED_FRAMES_BEFORE_EVICTION: u64 = 240;

/// Memory the system copies hold unless the game sets another amount.
pub(crate) const DEFAULT_MEMORY: usize = 256 * 1024 * 1024;

/// Mesh data kept for the GPU, keyed by the game's mesh values.
///
/// The keys are the game's; the copies behind them are the [`Store`],
/// which nothing typed by the game reads.
pub(crate) struct MeshCache<M: Meshes> {
    assets: Rc<Assets>,
    ids: HashMap<M, MeshId>,
    store: Store,
}

impl<M: Meshes> MeshCache<M> {
    /// A cache whose system copies hold `memory` bytes at most.
    pub(crate) fn new(assets: Rc<Assets>, memory: usize) -> Self {
        Self {
            assets,
            ids: HashMap::new(),
            store: Store::new(memory),
        }
    }

    /// The copies behind the keys.
    pub(crate) fn store(&self) -> &Store {
        &self.store
    }

    /// The copies behind the keys, to upload and to end the frame over.
    pub(crate) fn store_mut(&mut self) -> &mut Store {
        &mut self.store
    }

    /// The id for `mesh`, marked used this frame; builds its mesh data on
    /// first use, and again where the cache has since dropped the copy.
    ///
    /// The errors in how a mesh was built are recorded against the assets
    /// under the mesh type's name, so the catalog run reports them.
    pub(crate) fn id_of(&mut self, mesh: &M) -> MeshId {
        let id = match self.ids.get(mesh) {
            Some(&id) => id,
            None => {
                let id = self.store.insert();
                self.ids.insert(mesh.clone(), id);
                id
            }
        };
        if self.store.geometry(id).is_some() {
            self.store.touch(id);
            return id;
        }

        let geometry = mesh.build(&self.assets).unwrap_or_else(|errors| {
            for error in errors {
                self.assets.record(Unresolved::Mesh {
                    mesh: mesh.name(),
                    error,
                });
            }
            Geometry::empty()
        });
        self.store.fill(id, geometry);
        id
    }

    /// The clips `mesh` holds, built on first use as a draw of it would be.
    pub(crate) fn clips<T>(&mut self, mesh: T) -> &[Animation]
    where
        M: From<T>,
    {
        let id = self.id_of(&M::from(mesh));

        self.store.geometry(id).map_or(&[], Geometry::clips)
    }

    /// Builds every value the set catalogs into the system copy.
    ///
    /// Everything the builds needed from assets and did not get is recorded
    /// against them, as is every error in how a mesh was built.
    pub(crate) fn build_catalog(&mut self) -> Vec<MeshId> {
        M::catalog().iter().map(|mesh| self.id_of(mesh)).collect()
    }
}

/// The copies behind a cache's keys: the system ones held within the
/// memory the game allowed, the ones used least dropped first, and the
/// GPU ones that come from them.
///
/// A dropped system copy is built again the next time the mesh is drawn.
pub(crate) struct Store {
    entries: Vec<Entry>,
    /// Every entry holding a system copy, the one used least first.
    kept: BTreeSet<(u64, MeshId)>,
    held: usize,
    memory: usize,
    frame: u64,
}

impl Store {
    fn new(memory: usize) -> Self {
        Self {
            entries: Vec::new(),
            kept: BTreeSet::new(),
            held: 0,
            memory,
            frame: 0,
        }
    }

    /// The system copy of `id`, absent where the cache has dropped it —
    /// which never happens to a mesh the frame itself resolved.
    pub(crate) fn geometry(&self, id: MeshId) -> Option<&Geometry> {
        self.entries[id.index()].geometry.as_ref()
    }

    pub(crate) fn uploaded(&self, id: MeshId) -> Option<&GpuMesh> {
        self.entries[id.index()].upload.as_ref()
    }

    /// Uploads mesh `id` if the GPU does not hold it.
    pub(crate) fn upload(
        &mut self,
        device: &wgpu::Device,
        queue: &wgpu::Queue,
        textures: &Textures,
        id: MeshId,
    ) {
        let entry = &mut self.entries[id.index()];
        let Some(geometry) = entry.geometry.as_ref() else {
            return;
        };
        if geometry.indices().is_empty() {
            return;
        }

        entry
            .upload
            .get_or_insert_with(|| GpuMesh::new(device, queue, textures, geometry));
    }

    /// Ends the frame, and drops the GPU copies of meshes it did not draw
    /// for a while.
    pub(crate) fn end_frame(&mut self) {
        self.frame += 1;
        let frame = self.frame;
        for entry in &mut self.entries {
            if frame - entry.used > UNUSED_FRAMES_BEFORE_EVICTION {
                entry.upload = None;
            }
        }
    }

    /// A place for one more mesh, with no copy of it yet.
    fn insert(&mut self) -> MeshId {
        let id = MeshId(self.entries.len() as u32);
        self.entries.push(Entry {
            geometry: None,
            upload: None,
            used: self.frame,
        });
        id
    }

    /// Takes `geometry` into the cache, and drops what the memory it needs no
    /// longer leaves room for.
    fn fill(&mut self, id: MeshId, geometry: Geometry) {
        let frame = self.frame;
        let held = geometry.bytes();
        let entry = &mut self.entries[id.index()];
        entry.geometry = Some(geometry);
        entry.used = frame;
        self.held += held;
        self.kept.insert((frame, id));
        self.evict();
    }

    /// Marks `id` used this frame, which holds its copy in the cache; only
    /// ever called where the copy is there.
    fn touch(&mut self, id: MeshId) {
        let frame = self.frame;
        let entry = &mut self.entries[id.index()];
        let previous = core::mem::replace(&mut entry.used, frame);
        if previous != frame {
            self.kept.remove(&(previous, id));
            self.kept.insert((frame, id));
        }
    }

    /// Drops system copies, the one used least first, until the cache holds
    /// what its memory allows; a copy this frame used stays whatever it
    /// costs, so that a frame draws every mesh it resolved.
    fn evict(&mut self) {
        while self.held > self.memory {
            let Some(&(used, id)) = self.kept.first() else {
                return;
            };
            if used == self.frame {
                return;
            }
            self.kept.pop_first();
            self.held -= self.entries[id.index()].release();
        }
    }
}

/// A mesh's place in the cache, cheap enough to sort draws by, and to name
/// the run of the frame's palette its draws are posed into.
#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
pub(crate) struct MeshId(pub(crate) u32);

impl MeshId {
    fn index(self) -> usize {
        self.0 as usize
    }
}

/// One key's two copies: the system one while the cache holds it, the GPU
/// one while the GPU does, and the frame that last used the key.
struct Entry {
    geometry: Option<Geometry>,
    upload: Option<GpuMesh>,
    used: u64,
}

impl Entry {
    /// Drops the system copy, and returns what it held; the GPU copy draws
    /// until it is evicted.
    fn release(&mut self) -> usize {
        self.geometry.take().map_or(0, |geometry| geometry.bytes())
    }
}

/// One mesh's vertex and index buffers, what each of its parts samples,
/// and, for a skinned mesh, what each of its corners takes of its joints.
pub(crate) struct GpuMesh {
    vertices: wgpu::Buffer,
    indices: wgpu::Buffer,
    skin: Option<wgpu::Buffer>,
    slots: Vec<Option<wgpu::BindGroup>>,
}

impl GpuMesh {
    pub(crate) fn vertices(&self) -> &wgpu::Buffer {
        &self.vertices
    }

    pub(crate) fn indices(&self) -> &wgpu::Buffer {
        &self.indices
    }

    /// The joints and weights of each corner, absent for a mesh with none:
    /// the second stream the skinned stages are drawn from.
    pub(crate) fn skin(&self) -> Option<&wgpu::Buffer> {
        self.skin.as_ref()
    }

    /// The texture `part` samples, absent where it is drawn over a white
    /// default.
    pub(crate) fn part_texture(&self, part: u32) -> Option<&wgpu::BindGroup> {
        self.slots.get(part as usize)?.as_ref()
    }

    fn new(
        device: &wgpu::Device,
        queue: &wgpu::Queue,
        textures: &Textures,
        geometry: &Geometry,
    ) -> Self {
        Self {
            vertices: device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
                label: Some("mirage-engine mesh vertices"),
                contents: bytemuck::cast_slice(geometry.vertices()),
                usage: wgpu::BufferUsages::VERTEX,
            }),
            indices: device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
                label: Some("mirage-engine mesh indices"),
                contents: bytemuck::cast_slice(geometry.indices()),
                usage: wgpu::BufferUsages::INDEX,
            }),
            skin: geometry.rig().skins().then(|| {
                device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
                    label: Some("mirage-engine mesh skin"),
                    contents: bytemuck::cast_slice(geometry.rig().weights()),
                    usage: wgpu::BufferUsages::VERTEX,
                })
            }),
            slots: (0..geometry.part_count())
                .map(|part| {
                    textures.bind(
                        device,
                        queue,
                        geometry.part_texture(part),
                        geometry.part_relief(part),
                        geometry.part_shading(part),
                        geometry.part_emissive(part),
                    )
                })
                .collect(),
        }
    }
}

#[cfg(test)]
mod tests {
    use core::cell::Cell;

    use super::*;
    use crate::Catalog;
    use crate::math::{Vec2, Vec3};
    use crate::mesh::{Cube, Mesh, MeshData, Part, Slot, Vertex};
    use crate::{Material, meshes};

    /// Corner count of one of these meshes.
    const CORNERS: usize = 3;

    /// And what one of them costs the cache.
    const SLAB: usize = CORNERS * size_of::<Vertex>() + CORNERS * size_of::<u32>();

    /// A vocabulary whose every value loads a mesh by name, cataloged in
    /// full.
    #[derive(Clone, Eq, Hash, PartialEq)]
    enum Ship {
        Hull,
        Thruster,
    }

    impl Catalog for Ship {
        fn catalog() -> Vec<Self> {
            vec![Self::Hull, Self::Thruster]
        }
    }

    impl Mesh for Ship {
        fn build(&self, assets: &Assets) -> MeshData {
            match self {
                Self::Hull => assets.mesh("hull"),
                Self::Thruster => assets.mesh("thruster"),
            }
        }
    }

    meshes! { enum Ships { Ship } }

    #[test]
    fn the_catalog_run_names_every_asset_it_could_not_find() {
        let assets = Rc::new(Assets::default());
        MeshCache::<Ships>::new(Rc::clone(&assets), DEFAULT_MEMORY).build_catalog();

        let error = assets
            .unresolved()
            .expect("nothing was loaded, so both pulls miss");
        assert_eq!(
            error.to_string(),
            "the game's assets did not resolve: no asset is named `hull`; \
             no asset is named `thruster`"
        );
    }

    /// The example's mesh, which names two materials, under a vocabulary
    /// that names one of them and leaves the other anonymous.
    #[derive(Clone, Eq, Hash, PartialEq)]
    struct Beacon;

    impl Catalog for Beacon {
        fn catalog() -> Vec<Self> {
            vec![Self]
        }
    }

    impl Mesh<Panels> for Beacon {
        fn build(&self, assets: &Assets) -> MeshData<Panels> {
            assets.mesh("beacon")
        }
    }

    #[derive(Clone, Debug, Eq, Hash, PartialEq)]
    struct Panels;

    impl Part for Panels {
        fn from_name(name: &str) -> Option<Self> {
            (name == "Beacon Panels").then_some(Self)
        }

        fn all() -> Vec<Self> {
            vec![Self]
        }

        fn index(&self) -> u32 {
            0
        }
    }

    /// A vocabulary that also names a part the mesh does not have — what a
    /// rename in the source leaves behind.
    #[derive(Clone, Debug, Eq, Hash, PartialEq)]
    enum Renamed {
        Panels,
        Wing,
    }

    impl Part for Renamed {
        fn from_name(name: &str) -> Option<Self> {
            match name {
                "Beacon Panels" => Some(Self::Panels),
                "Wing" => Some(Self::Wing),
                _ => None,
            }
        }

        fn all() -> Vec<Self> {
            vec![Self::Panels, Self::Wing]
        }

        fn index(&self) -> u32 {
            match self {
                Self::Panels => 0,
                Self::Wing => 1,
            }
        }
    }

    /// The same mesh under a vocabulary that has gone stale.
    #[derive(Clone, Eq, Hash, PartialEq)]
    struct StaleBeacon;

    impl Catalog for StaleBeacon {
        fn catalog() -> Vec<Self> {
            vec![Self]
        }
    }

    impl Mesh<Renamed> for StaleBeacon {
        fn build(&self, assets: &Assets) -> MeshData<Renamed> {
            assets.mesh("beacon")
        }
    }

    /// A generated mesh whose slots stop short of its indices.
    #[derive(Clone, Eq, Hash, PartialEq)]
    struct Short;

    impl Catalog for Short {
        fn catalog() -> Vec<Self> {
            vec![Self]
        }
    }

    impl Mesh<Panels> for Short {
        fn build(&self, assets: &Assets) -> MeshData<Panels> {
            let cube = Cube.build(assets);
            MeshData::in_parts(cube.vertices().to_vec(), cube.indices().to_vec(), |_| {
                Slot::new(3, Material::default())
            })
        }
    }

    meshes! { enum Beacons { Beacon, StaleBeacon, Short } }

    fn with_model() -> Rc<Assets> {
        Rc::new(
            Assets::load([crate::assets::file("hello.glb", crate::assets::BEACON)])
                .expect("the example's model decodes"),
        )
    }

    #[test]
    fn the_catalog_run_leaves_a_material_the_vocabulary_never_names_alone() {
        let assets = with_model();
        let mut cache = MeshCache::<Beacons>::new(Rc::clone(&assets), DEFAULT_MEMORY);
        cache.id_of(&Beacons::Beacon(Beacon));

        assert!(
            assets.unresolved().is_none(),
            "the caps are simply nothing this game addresses"
        );
    }

    #[test]
    fn the_catalog_run_reports_a_part_the_mesh_turns_out_not_to_have() {
        let assets = with_model();
        let mut cache = MeshCache::<Beacons>::new(Rc::clone(&assets), DEFAULT_MEMORY);
        cache.id_of(&Beacons::StaleBeacon(StaleBeacon));

        let error = assets
            .unresolved()
            .expect("nothing the mesh holds is a wing");
        assert_eq!(
            error.to_string(),
            "the game's assets did not resolve: the asset `beacon` has no material that \
             resolves to the part Wing"
        );
    }

    #[test]
    fn the_catalog_run_reports_a_mesh_whose_slots_do_not_cover_it_under_its_own_name() {
        let assets = with_model();
        let mut cache = MeshCache::<Beacons>::new(Rc::clone(&assets), DEFAULT_MEMORY);
        let short = cache.id_of(&Beacons::Short(Short));

        let error = assets.unresolved().expect("the slots stop short");
        assert_eq!(
            error.to_string(),
            "the game's assets did not resolve: the mesh `Short` has slots covering 3 \
             indices where it holds 36"
        );
        assert_eq!(
            cache.store().geometry(short).expect("held").part_count(),
            0,
            "and nothing of it is drawn"
        );
    }

    meshes! { enum Cubes { Cube } }

    #[test]
    fn the_catalog_run_fills_the_system_tier_once() {
        let mut cache = MeshCache::<Cubes>::new(Rc::new(Assets::default()), DEFAULT_MEMORY);

        let ids = cache.build_catalog();
        assert_eq!(ids.len(), 1);
        assert_eq!(
            cache.id_of(&Cube.into()),
            ids[0],
            "a drawn cube reuses the build"
        );
        assert_eq!(cache.store().entries.len(), 1);
    }

    /// One triangle per key, every one of them the same size, over a count
    /// each build of it adds to.
    ///
    /// Two of them are the same mesh where their keys match: a test reads
    /// the count they share, never part of what the cache is keyed by.
    #[derive(Clone)]
    struct Slab {
        key: u32,
        builds: Rc<Cell<u32>>,
    }

    impl PartialEq for Slab {
        fn eq(&self, other: &Self) -> bool {
            self.key == other.key
        }
    }

    impl Eq for Slab {}

    impl core::hash::Hash for Slab {
        fn hash<H: core::hash::Hasher>(&self, state: &mut H) {
            self.key.hash(state);
        }
    }

    impl Catalog for Slab {
        fn catalog() -> Vec<Self> {
            Vec::new()
        }
    }

    impl Mesh for Slab {
        fn build(&self, _assets: &Assets) -> MeshData {
            self.builds.set(self.builds.get() + 1);
            MeshData::new(
                vec![Vertex::new(Vec3::ZERO, Vec3::Y, Vec2::ZERO); CORNERS],
                (0..CORNERS as u32).collect(),
            )
        }
    }

    meshes! { enum Slabs { Slab } }

    /// A cache of them holding `memory` bytes, and the count its builds add to.
    fn slabs(memory: usize) -> (MeshCache<Slabs>, Rc<Cell<u32>>) {
        let cache = MeshCache::new(Rc::new(Assets::default()), memory);
        (cache, Rc::new(Cell::new(0)))
    }

    fn slab(builds: &Rc<Cell<u32>>, key: u32) -> Slabs {
        Slab {
            key,
            builds: Rc::clone(builds),
        }
        .into()
    }

    #[test]
    fn the_system_copies_are_held_inside_the_memory_the_cache_was_given() {
        let (mut cache, builds) = slabs(3 * SLAB);

        for key in 0..8 {
            cache.id_of(&slab(&builds, key));
            cache.store_mut().end_frame();
        }

        assert_eq!(cache.store().kept.len(), 3, "three of the eight are left");
        assert_eq!(cache.store().held, 3 * SLAB);
        assert_eq!(builds.get(), 8, "and every key was built once");
    }

    #[test]
    fn the_copy_the_frame_left_alone_is_the_one_that_drops() {
        let (mut cache, builds) = slabs(2 * SLAB);
        let first = cache.id_of(&slab(&builds, 0));
        cache.store_mut().end_frame();
        let second = cache.id_of(&slab(&builds, 1));
        cache.store_mut().end_frame();

        cache.id_of(&slab(&builds, 0));
        cache.id_of(&slab(&builds, 2));

        assert!(
            cache.store().geometry(first).is_some(),
            "the frame used this one"
        );
        assert!(cache.store().geometry(second).is_none(), "and not this one");
        assert_eq!(builds.get(), 3);
    }

    #[test]
    fn a_dropped_mesh_is_built_again_once_and_holds_while_it_is_used() {
        let (mut cache, builds) = slabs(SLAB);
        let first = cache.id_of(&slab(&builds, 0));
        cache.store_mut().end_frame();
        cache.id_of(&slab(&builds, 1));
        cache.store_mut().end_frame();

        assert!(
            cache.store().geometry(first).is_none(),
            "the second took the room"
        );
        assert_eq!(cache.id_of(&slab(&builds, 0)), first, "and it keeps its id");
        assert_eq!(builds.get(), 3, "built once more, and only once");

        cache.id_of(&slab(&builds, 0));
        cache.store_mut().end_frame();
        cache.id_of(&slab(&builds, 0));

        assert_eq!(builds.get(), 3, "the copy stands while the game draws it");
    }

    #[test]
    fn a_run_inside_the_memory_never_builds_a_mesh_twice() {
        let (mut cache, builds) = slabs(DEFAULT_MEMORY);

        for _ in 0..64 {
            for key in 0..4 {
                cache.id_of(&slab(&builds, key));
            }
            cache.store_mut().end_frame();
        }

        assert_eq!(builds.get(), 4);
        assert_eq!(cache.store().kept.len(), 4);
    }

    #[test]
    fn an_insert_never_drops_what_the_frame_it_lands_in_took() {
        let (mut cache, builds) = slabs(SLAB / 2);
        let alone = cache.id_of(&slab(&builds, 0));
        assert!(
            cache.store().geometry(alone).is_some(),
            "a mesh past the memory on its own draws the frame that asked for it"
        );

        let second = cache.id_of(&slab(&builds, 1));
        assert!(
            cache.store().geometry(alone).is_some() && cache.store().geometry(second).is_some(),
            "and so does everything else that frame resolved"
        );

        cache.store_mut().end_frame();
        cache.id_of(&slab(&builds, 2));
        assert_eq!(
            cache.store().kept.len(),
            1,
            "the frame after takes the room back"
        );
    }
}