mirage-engine 0.2.0

use std::collections::{BTreeSet, HashMap};
use std::sync::Arc;

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

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

/// The game thread's meshes: the game's keys to ids, the system copies held
/// within the memory the game allowed, the ones used least dropped first,
/// and the copies the display thread has not been handed yet.
///
/// A dropped system copy is built again the next time the mesh is drawn,
/// and handed again.
pub(crate) struct MeshCatalog<M: Meshes> {
    assets: Arc<Assets>,
    ids: HashMap<M, MeshId>,
    entries: Vec<Entry>,
    /// Every entry holding a system copy, the one used least first.
    kept: BTreeSet<(u64, MeshId)>,
    held: usize,
    memory: usize,
    frame: u64,
    /// Every copy built since the last hand, which the next handed frame
    /// carries.
    unhanded: Vec<(MeshId, Arc<Geometry>)>,
    /// Every id dropped since the last hand, which the display thread drops
    /// on that word.
    dropped: Vec<MeshId>,
}

impl<M: Meshes> MeshCatalog<M> {
    /// A catalog whose system copies hold `memory` bytes at most.
    pub(crate) fn new(assets: Arc<Assets>, memory: usize) -> Self {
        Self {
            assets,
            ids: HashMap::new(),
            entries: Vec::new(),
            kept: BTreeSet::new(),
            held: 0,
            memory,
            frame: 0,
            unhanded: Vec::new(),
            dropped: Vec::new(),
        }
    }

    /// The id for `mesh`, marked used this frame; builds its mesh data on
    /// first use, and again where the catalog has since dropped the copy.
    ///
    /// What the build did not get is a debug log here: the catalog run has
    /// already reported what every cataloged value did not get.
    pub(crate) fn id_of(&mut self, mesh: &M) -> MeshId {
        let (id, unresolved) = self.built(mesh);
        unresolved.logged();
        id
    }

    /// Builds `mesh` now, instead of on its first draw, and hands it with
    /// the next frame.
    pub(crate) fn prepare(&mut self, mesh: M) {
        self.id_of(&mesh);
    }

    /// 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.geometry(id).map_or(&[], Geometry::clips)
    }

    /// The system copy of `id`, absent where the catalog 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_deref()
    }

    /// Builds every value the set catalogs into the system copy, and returns
    /// what those builds did not get, every error in how a mesh was built
    /// among it.
    pub(crate) fn build_catalog(&mut self) -> Unresolved {
        let mut unresolved = Unresolved::default();
        for mesh in M::catalog() {
            unresolved.record(self.built(&mesh).1);
        }

        unresolved
    }

    /// The id for `mesh` beside what its build did not get, building it
    /// where the catalog holds no copy.
    fn built(&mut self, mesh: &M) -> (MeshId, Unresolved) {
        let id = match self.ids.get(mesh) {
            Some(&id) => id,
            None => {
                let id = self.insert();
                self.ids.insert(mesh.clone(), id);
                id
            }
        };
        if self.entries[id.index()].geometry.is_some() {
            self.touch(id);
            return (id, Unresolved::default());
        }

        let Built {
            geometry,
            unresolved,
        } = mesh.build(&self.assets);
        self.fill(id, Arc::new(geometry));

        (id, unresolved)
    }

    /// Ends the frame, from which the copies it used may be dropped, and
    /// hands over what it built and what it dropped.
    ///
    /// The hand is the one way out of the catalog, so no copy is built
    /// without the display thread being handed it.
    pub(crate) fn end_frame(&mut self) -> HandedMeshes {
        self.frame += 1;
        HandedMeshes {
            geometry: core::mem::take(&mut self.unhanded),
            dropped: core::mem::take(&mut self.dropped),
        }
    }

    /// 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,
            used: self.frame,
        });
        id
    }

    /// Takes `geometry` into the catalog, marks it to hand, and drops what
    /// the memory it needs no longer leaves room for.
    fn fill(&mut self, id: MeshId, geometry: Arc<Geometry>) {
        let frame = self.frame;
        let held = geometry.bytes();
        let entry = &mut self.entries[id.index()];
        entry.geometry = Some(Arc::clone(&geometry));
        entry.used = frame;
        self.dropped.retain(|&dropped| dropped != id);
        self.unhanded.push((id, geometry));
        self.held += held;
        self.kept.insert((frame, id));
        self.evict();
    }

    /// Marks `id` used this frame, which holds its copy in the catalog;
    /// 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 catalog
    /// 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();
            self.unhanded.retain(|&(unhanded, _)| unhanded != id);
            self.dropped.push(id);
        }
    }
}

/// What one frame hands the display thread's meshes: every copy built since
/// the last hand, and every id dropped since it, which name no id twice.
pub(crate) struct HandedMeshes {
    geometry: Vec<(MeshId, Arc<Geometry>)>,
    dropped: Vec<MeshId>,
}

/// A mesh's place in the catalog, cheap enough to sort draws by, and to
/// name the run of the frame's palette its draws are posed into; the one
/// key the two threads share.
#[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 system copy while the catalog holds it, and the frame that
/// last used the key.
struct Entry {
    geometry: Option<Arc<Geometry>>,
    used: u64,
}

impl Entry {
    /// Drops the system copy, and returns what it held; the display thread
    /// drops its own on the same word.
    fn release(&mut self) -> usize {
        self.geometry.take().map_or(0, |geometry| geometry.bytes())
    }
}

/// The display thread's meshes, by id: the copy each id was last handed and
/// the GPU copy built from it, both held until the game thread drops the id.
pub(crate) struct GpuMeshes {
    received: Vec<Option<Received>>,
}

impl GpuMeshes {
    pub(crate) fn new() -> Self {
        Self {
            received: Vec::new(),
        }
    }

    /// Takes what one frame handed: both copies of every id it dropped go,
    /// on the game thread's word, and every copy it built takes the place of
    /// what its id held, its GPU copy built again on its next draw.
    pub(crate) fn take(&mut self, handed: HandedMeshes) {
        for id in handed.dropped {
            if let Some(received) = self.received.get_mut(id.index()) {
                *received = None;
            }
        }
        for (id, geometry) in handed.geometry {
            if id.index() >= self.received.len() {
                self.received.resize_with(id.index() + 1, || None);
            }
            self.received[id.index()] = Some(Received {
                geometry,
                upload: None,
            });
        }
    }

    /// The copy `id` was handed, absent where no frame has handed it yet.
    pub(crate) fn geometry(&self, id: MeshId) -> Option<&Geometry> {
        Some(&self.received.get(id.index())?.as_ref()?.geometry)
    }

    /// The GPU copy of `id`, absent until a frame's batch of it uploads it.
    pub(crate) fn uploaded(&self, id: MeshId) -> Option<&GpuMesh> {
        self.received.get(id.index())?.as_ref()?.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 Some(Some(received)) = self.received.get_mut(id.index()) else {
            return;
        };
        if received.geometry.indices().is_empty() {
            return;
        }

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

/// One id as the display thread holds it: the copy it was handed, and the
/// GPU copy while the GPU does.
struct Received {
    geometry: Arc<Geometry>,
    upload: Option<GpuMesh>,
}

/// 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: Self::written(
                device,
                queue,
                "mirage-engine mesh vertices",
                bytemuck::cast_slice(geometry.vertices()),
                wgpu::BufferUsages::VERTEX,
            ),
            indices: Self::written(
                device,
                queue,
                "mirage-engine mesh indices",
                bytemuck::cast_slice(geometry.indices()),
                wgpu::BufferUsages::INDEX,
            ),
            skin: geometry.rig().skins().then(|| {
                Self::written(
                    device,
                    queue,
                    "mirage-engine mesh skin",
                    bytemuck::cast_slice(geometry.rig().weights()),
                    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(),
        }
    }

    /// Fills a new `usage` buffer called `label` with `contents` through the
    /// queue instead of mapping it at creation.
    fn written(
        device: &wgpu::Device,
        queue: &wgpu::Queue,
        label: &str,
        contents: &[u8],
        usage: wgpu::BufferUsages,
    ) -> wgpu::Buffer {
        let written = buffer(device, label, contents.len() as wgpu::BufferAddress, usage);
        queue.write_buffer(&written, 0, contents);

        written
    }
}

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

    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 catalog.
    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 = Arc::new(Assets::default());
        let unresolved =
            MeshCatalog::<Ships>::new(Arc::clone(&assets), DEFAULT_MEMORY).build_catalog();

        let error = unresolved
            .error()
            .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() -> Arc<Assets> {
        Arc::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 catalog = MeshCatalog::<Beacons>::new(assets, DEFAULT_MEMORY);

        assert!(
            catalog.built(&Beacons::Beacon(Beacon)).1.is_empty(),
            "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 catalog = MeshCatalog::<Beacons>::new(assets, DEFAULT_MEMORY);
        let unresolved = catalog.built(&Beacons::StaleBeacon(StaleBeacon)).1;

        let error = unresolved
            .error()
            .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 catalog = MeshCatalog::<Beacons>::new(assets, DEFAULT_MEMORY);
        let (short, unresolved) = catalog.built(&Beacons::Short(Short));

        let error = unresolved.error().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!(
            catalog.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_and_hands_it_once() {
        let mut catalog = MeshCatalog::<Cubes>::new(Arc::new(Assets::default()), DEFAULT_MEMORY);

        catalog.build_catalog();
        let handed = catalog.end_frame();
        assert_eq!(handed.geometry.len(), 1);
        assert_eq!(
            catalog.id_of(&Cube.into()),
            handed.geometry[0].0,
            "a drawn cube reuses the build"
        );
        assert_eq!(catalog.entries.len(), 1);
        assert!(
            catalog.end_frame().geometry.is_empty(),
            "and the display thread is handed nothing twice"
        );
    }

    /// 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 catalog 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 catalog of them holding `memory` bytes, and the count its builds
    /// add to.
    fn slabs(memory: usize) -> (MeshCatalog<Slabs>, Rc<Cell<u32>>) {
        let catalog = MeshCatalog::new(Arc::new(Assets::default()), memory);
        (catalog, 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_catalog_was_given() {
        let (mut catalog, builds) = slabs(3 * SLAB);

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

        assert_eq!(catalog.kept.len(), 3, "three of the eight are left");
        assert_eq!(catalog.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 catalog, builds) = slabs(2 * SLAB);
        let first = catalog.id_of(&slab(&builds, 0));
        catalog.end_frame();
        let second = catalog.id_of(&slab(&builds, 1));
        catalog.end_frame();

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

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

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

        assert!(
            catalog.geometry(first).is_none(),
            "the second took the room"
        );
        assert_eq!(
            handed.dropped,
            vec![first],
            "and the display thread is told to drop it"
        );
        assert_eq!(
            catalog.id_of(&slab(&builds, 0)),
            first,
            "and it keeps its id"
        );
        assert_eq!(builds.get(), 3, "built once more, and only once");
        assert_eq!(
            catalog.end_frame().geometry.len(),
            1,
            "and the display thread is handed the new copy"
        );

        catalog.id_of(&slab(&builds, 0));
        catalog.end_frame();
        catalog.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 catalog, builds) = slabs(DEFAULT_MEMORY);

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

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

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

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

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