tile_chunk.wgsl

#import bevy_sprite::{
    mesh2d_functions as mesh_functions,
    mesh2d_view_bindings::view,
}

struct Vertex {
    @builtin(instance_index) instance_index: u32,
    @location(0) position: vec3<f32>,
    @location(1) uv: vec2<f32>,
    @location(5) tile_index: u32,
};

struct VertexOutput {
    @builtin(position) position: vec4<f32>,
    @location(0) uv: vec2<f32>,
    @location(1) tile_index: u32,
}

@group(2) @binding(0) var tile_atlas: texture_2d_array<f32>;
@group(2) @binding(1) var tile_atlas_sampler: sampler;
@group(2) @binding(2) var tile_indices: texture_2d<u32>;

@vertex
fn vertex(vertex: Vertex) -> VertexOutput {
    var out: VertexOutput;
    
    let world_from_local = mesh_functions::get_world_from_local(vertex.instance_index);
    let world_position = mesh_functions::mesh2d_position_local_to_world(
        world_from_local,
        vec4<f32>(vertex.position, 1.0)
    );

    out.position = mesh_functions::mesh2d_position_world_to_clip(world_position);
    out.uv = vertex.uv;
    out.tile_index = vertex.tile_index;

    return out;
}

@fragment
fn fragment(in: VertexOutput) -> @location(0) vec4<f32> {
    let chunk_size = textureDimensions(tile_indices, 0);
    let tile_xy = vec2<u32>(
        in.tile_index % chunk_size.x,
        in.tile_index / chunk_size.x
    );
    let tile_id = textureLoad(tile_indices, tile_xy, 0).r;

    return textureSample(tile_atlas, tile_atlas_sampler, in.uv, tile_id);
}

tile_entity.rs

use bevy::{
    ecs::{
        component::HookContext,
        relationship::{Relationship, RelationshipHookMode, RelationshipSourceCollection},
        world::DeferredWorld,
    },
    platform::{
        self,
        collections::{HashMap, hash_map::Entry},
    },
    prelude::*,
};

use super::{TileChunk, update_tile_chunk_indices};

pub struct TileEntityPlugin;

impl Plugin for TileEntityPlugin {
    fn build(&self, app: &mut App) {
        app.add_systems(
            PostUpdate,
            update_tile_entities.before(update_tile_chunk_indices),
        );
    }
}

#[derive(Component, Default)]
pub struct TileEntities(TileEntitiesMap);

impl RelationshipTarget for TileEntities {
    type Collection = TileEntitiesMap;
    type Relationship = Tile;

    const LINKED_SPAWN: bool = true;

    fn collection(&self) -> &Self::Collection {
        &self.0
    }

    fn collection_mut_risky(&mut self) -> &mut Self::Collection {
        &mut self.0
    }

    fn from_collection_risky(collection: Self::Collection) -> Self {
        Self(collection)
    }
}

#[derive(Component, Debug)]
#[component(
    immutable,
    on_insert = <Self as Relationship>::on_insert,
    on_replace = <Self as Relationship>::on_replace,
)]
#[require(TileIndex)]
pub struct Tile {
    pub chunk: Entity,
    pub location: UVec2,
}

impl Relationship for Tile {
    type RelationshipTarget = TileEntities;

    fn from(chunk: Entity) -> Self {
        Self {
            chunk,
            location: UVec2::ZERO,
        }
    }

    fn get(&self) -> Entity {
        self.chunk
    }

    fn on_insert(
        mut world: DeferredWorld,
        HookContext {
            entity,
            caller,
            relationship_hook_mode,
            ..
        }: HookContext,
    ) {
        match relationship_hook_mode {
            RelationshipHookMode::Run => {}
            RelationshipHookMode::Skip => return,
            RelationshipHookMode::RunIfNotLinked => {
                if <Self::RelationshipTarget as RelationshipTarget>::LINKED_SPAWN {
                    return;
                }
            }
        }

        let &Tile { chunk, location } = world.get(entity).unwrap();

        if chunk == entity {
            warn!(
                "{}The {}{{ chunk: {chunk}, location: {location} }} relationship on entity {entity} points to itself. The invalid {} relationship has been removed.",
                caller
                    .map(|location| format!("{location}: "))
                    .unwrap_or_default(),
                core::any::type_name::<Self>(),
                core::any::type_name::<Self>()
            );
            world.commands().entity(entity).remove::<Self>();
            return;
        }

        let Ok(mut chunk_mut) = world.get_entity_mut(chunk) else {
            warn!(
                "{}The {}{{ chunk: {chunk}, location: {location} }} relationship on entity {entity} relates to an entity that does not exist. The invalid {} relationship has been removed.",
                caller
                    .map(|location| format!("{location}: "))
                    .unwrap_or_default(),
                core::any::type_name::<Self>(),
                core::any::type_name::<Self>()
            );
            world.commands().entity(entity).remove::<Self>();
            return;
        };

        if let Some(mut tile_entities) = chunk_mut.get_mut::<TileEntities>() {
            match tile_entities.0.0.entry(location) {
                Entry::Occupied(e) => {
                    warn!(
                        "{}The {}{{ chunk: {chunk}, location: {location} }} relationship on entity {entity} relates to a location that is already occupied by {}. The invalid {} relationship has been removed.",
                        caller
                            .map(|location| format!("{location}: "))
                            .unwrap_or_default(),
                        core::any::type_name::<Self>(),
                        e.get(),
                        core::any::type_name::<Self>(),
                    );
                    world.commands().entity(entity).remove::<Self>();
                }
                Entry::Vacant(e) => {
                    e.insert(entity);
                }
            }
        } else {
            let mut tile_entities = TileEntities::default();
            tile_entities.0.0.insert(location, entity);
            world.commands().entity(chunk).insert(tile_entities);
        }
    }

    fn on_replace(
        mut world: DeferredWorld,
        HookContext {
            entity,
            relationship_hook_mode,
            ..
        }: HookContext,
    ) {
        match relationship_hook_mode {
            RelationshipHookMode::Run => {}
            RelationshipHookMode::Skip => return,
            RelationshipHookMode::RunIfNotLinked => {
                if <Self::RelationshipTarget as RelationshipTarget>::LINKED_SPAWN {
                    return;
                }
            }
        }

        let &Tile { chunk, location } = world.get(entity).unwrap();

        if let Some(mut tile_entities) = world.get_mut::<TileEntities>(chunk) {
            if let Entry::Occupied(e) = tile_entities.0.0.entry(location) {
                if *e.get() == entity {
                    e.remove();
                }
            }
        }
    }
}

#[derive(Component, Debug, Default)]
pub struct TileIndex(pub u16);

#[derive(Default)]
pub struct TileEntitiesMap(HashMap<UVec2, Entity>);

impl RelationshipSourceCollection for TileEntitiesMap {
    type SourceIter<'a> =
        std::iter::Copied<platform::collections::hash_map::Values<'a, UVec2, Entity>>;

    fn new() -> Self {
        Self::default()
    }

    fn with_capacity(capacity: usize) -> Self {
        Self(HashMap::with_capacity(capacity))
    }

    fn reserve(&mut self, additional: usize) {
        self.0.reserve(additional);
    }

    fn add(&mut self, _: Entity) -> bool {
        false
    }

    fn remove(&mut self, _: Entity) -> bool {
        false
    }

    fn iter(&self) -> Self::SourceIter<'_> {
        self.0.values().copied()
    }

    fn len(&self) -> usize {
        self.0.len()
    }

    fn shrink_to_fit(&mut self) {
        self.0.shrink_to_fit();
    }

    fn clear(&mut self) {
        self.0.clear();
    }
}

fn update_tile_entities(
    tiles: Query<(&Tile, &TileIndex), Changed<TileIndex>>,
    mut chunks: Query<&mut TileChunk>,
) {
    for (&Tile { chunk, location }, &TileIndex(index)) in &tiles {
        let mut chunk = chunks.get_mut(chunk).unwrap();
        if chunk
            .bypass_change_detection()
            .set_index(location, index)
            .is_some_and(|old_index| old_index != index)
        {
            chunk.set_changed();
        }
    }
}

tile_render.rs

use std::mem;

use bevy::{
    asset::{AssetEvents, RenderAssetUsages},
    ecs::{component::HookContext, world::DeferredWorld},
    image::ImageSampler,
    math::FloatOrd,
    platform::collections::HashMap,
    prelude::*,
    render::{
        mesh::{
            Indices, MeshVertexAttribute, MeshVertexBufferLayoutRef, PrimitiveTopology,
            VertexFormat,
        },
        render_resource::{
            AsBindGroup, Extent3d, RenderPipelineDescriptor, ShaderRef,
            SpecializedMeshPipelineError, TextureDescriptor, TextureDimension, TextureFormat,
            TextureUsages,
        },
    },
    sprite::{AlphaMode2d, Anchor, Material2d, Material2dKey, Material2dPlugin},
};

pub mod tile_entity;
pub mod tile_set_loader;

use tile_entity::TileEntityPlugin;
use tile_set_loader::TileSetLoader;

pub const ATTRIBUTE_TILE_INDEX: MeshVertexAttribute =
    MeshVertexAttribute::new("Vertex_TileIndex", 264043692, VertexFormat::Uint32);

pub struct TileRenderPlugin;

impl Plugin for TileRenderPlugin {
    fn build(&self, app: &mut App) {
        app.add_plugins((
            Material2dPlugin::<TileMaterial>::default(),
            TileEntityPlugin,
        ))
        .init_asset_loader::<TileSetLoader>()
        .init_resource::<ChunkMeshCache>()
        .add_systems(PostUpdate, update_tile_chunk_indices.before(AssetEvents));
    }
}

#[derive(Component)]
#[component(
    on_insert = on_tile_chunk_inserted,
    on_remove = on_tile_chunk_removed,
)]
#[require(Mesh2d, MeshMaterial2d<TileMaterial>)]
pub struct TileChunk {
    pub size: UVec2,
    pub custom_size: Option<Vec2>,
    pub anchor: Anchor,
    pub alpha_mode: AlphaMode2d,
    pub tile_atlas: Handle<Image>,
    pub indices: Vec<u16>,
}

impl TileChunk {
    pub fn size(&self) -> UVec2 {
        self.size
    }

    pub fn set_index(&mut self, location: UVec2, index: u16) -> Option<u16> {
        if location.cmpge(self.size).any() {
            return None;
        }
        let i = location.x + location.y * self.size.x;
        Some(mem::replace(&mut self.indices[i as usize], index))
    }
}

#[derive(Asset, TypePath, AsBindGroup, Debug, Clone)]
pub struct TileMaterial {
    pub alpha_mode: AlphaMode2d,
    #[texture(0, dimension = "2d_array")]
    #[sampler(1)]
    pub tile_atlas: Handle<Image>,
    #[texture(2, sample_type = "u_int")]
    pub tile_indices: Handle<Image>,
}

impl Material2d for TileMaterial {
    fn fragment_shader() -> ShaderRef {
        "shaders/tile_chunk.wgsl".into()
    }

    fn vertex_shader() -> ShaderRef {
        "shaders/tile_chunk.wgsl".into()
    }

    fn alpha_mode(&self) -> AlphaMode2d {
        self.alpha_mode
    }

    fn specialize(
        descriptor: &mut RenderPipelineDescriptor,
        layout: &MeshVertexBufferLayoutRef,
        _key: Material2dKey<Self>,
    ) -> Result<(), SpecializedMeshPipelineError> {
        let vertex_layout = layout.0.get_layout(&[
            Mesh::ATTRIBUTE_POSITION.at_shader_location(0),
            Mesh::ATTRIBUTE_UV_0.at_shader_location(1),
            ATTRIBUTE_TILE_INDEX.at_shader_location(5),
        ])?;
        descriptor.vertex.buffers = vec![vertex_layout];
        Ok(())
    }
}

type MeshKey = (UVec2, FloatOrd, FloatOrd, FloatOrd, FloatOrd);

#[derive(Resource, Default)]
struct ChunkMeshCache {
    chunk_meshes: HashMap<MeshKey, Handle<Mesh>>,
}

fn on_tile_chunk_inserted(mut world: DeferredWorld, HookContext { entity, .. }: HookContext) {
    let tile_chunk = world.get::<TileChunk>(entity).unwrap();

    let chunk_size = tile_chunk.size();
    let display_size = tile_chunk.custom_size.unwrap_or(chunk_size.as_vec2());
    let alpha_mode = tile_chunk.alpha_mode;
    let anchor = tile_chunk.anchor;
    let tile_atlas = tile_chunk.tile_atlas.clone();

    let mesh_key = (
        chunk_size,
        FloatOrd(display_size.x),
        FloatOrd(display_size.y),
        FloatOrd(anchor.as_vec().x),
        FloatOrd(anchor.as_vec().y),
    );

    let tile_indices_image = make_chunk_image(tile_chunk.size(), &tile_chunk.indices);

    let tile_indices = world
        .resource_mut::<Assets<Image>>()
        .add(tile_indices_image);

    let material_handle = world
        .resource_mut::<Assets<TileMaterial>>()
        .add(TileMaterial {
            alpha_mode,
            tile_atlas,
            tile_indices,
        });

    let mesh_handle = match world
        .resource::<ChunkMeshCache>()
        .chunk_meshes
        .get(&mesh_key)
    {
        Some(mesh) => mesh.clone(),
        None => {
            let mesh = make_chunk_mesh(chunk_size, display_size, anchor);
            let mesh_handle = world.resource_mut::<Assets<Mesh>>().add(mesh);
            world
                .resource_mut::<ChunkMeshCache>()
                .chunk_meshes
                .insert(mesh_key, mesh_handle.clone());
            mesh_handle
        }
    };

    let mut chunk_mut = world.entity_mut(entity);
    chunk_mut.get_mut::<Mesh2d>().unwrap().0 = mesh_handle;
    chunk_mut
        .get_mut::<MeshMaterial2d<TileMaterial>>()
        .unwrap()
        .0 = material_handle;
}

fn on_tile_chunk_removed(mut world: DeferredWorld, HookContext { entity, .. }: HookContext) {
    world
        .commands()
        .entity(entity)
        .try_remove::<(Mesh2d, MeshMaterial2d<TileMaterial>)>();
}

fn update_tile_chunk_indices(
    mut materials: ResMut<Assets<TileMaterial>>,
    mut images: ResMut<Assets<Image>>,
    chunks: Query<(&TileChunk, &MeshMaterial2d<TileMaterial>), Changed<TileChunk>>,
) {
    for (chunk, material) in &chunks {
        // We get the material mutably here so that change detection is triggered. We don't mutate
        // the material itself.
        if let Some(material) = materials.get_mut(material.id()) {
            if let Some(image) = images.get_mut(&material.tile_indices) {
                let data = image.data.as_mut().unwrap();
                data.clear();
                data.extend(chunk.indices.iter().copied().flat_map(u16::to_ne_bytes));
            }
        }
    }
}

fn make_chunk_image(size: UVec2, indices: &[u16]) -> Image {
    Image {
        data: Some(indices.iter().copied().flat_map(u16::to_ne_bytes).collect()),
        texture_descriptor: TextureDescriptor {
            size: Extent3d {
                width: size.x,
                height: size.y,
                depth_or_array_layers: 1,
            },
            dimension: TextureDimension::D2,
            format: TextureFormat::R16Uint,
            label: None,
            mip_level_count: 1,
            sample_count: 1,
            usage: TextureUsages::TEXTURE_BINDING | TextureUsages::COPY_DST,
            view_formats: &[],
        },
        sampler: ImageSampler::nearest(),
        texture_view_descriptor: None,
        asset_usage: RenderAssetUsages::RENDER_WORLD | RenderAssetUsages::MAIN_WORLD,
    }
}

/// Creates a mesh of `size` quads, with vertices placed according to `display_size` and `anchor`.
fn make_chunk_mesh(size: UVec2, display_size: Vec2, anchor: Anchor) -> Mesh {
    let mut mesh = Mesh::new(
        PrimitiveTopology::TriangleList,
        RenderAssetUsages::RENDER_WORLD | RenderAssetUsages::MAIN_WORLD,
    );

    let offset = display_size * (Vec2::splat(-0.5) - anchor.as_vec());

    let num_quads = size.element_product() as usize;
    let quad_size = display_size / size.as_vec2();

    let mut positions = Vec::with_capacity(4 * num_quads);
    let mut uvs = Vec::with_capacity(4 * num_quads);
    let mut indices = Vec::with_capacity(6 * num_quads);

    for y in 0..size.y {
        for x in 0..size.x {
            let i = positions.len() as u32;

            let p0 = offset + quad_size * UVec2::new(x, y).as_vec2();
            let p1 = p0 + quad_size;

            positions.extend([
                Vec3::new(p0.x, p0.y, 0.0),
                Vec3::new(p1.x, p0.y, 0.0),
                Vec3::new(p0.x, p1.y, 0.0),
                Vec3::new(p1.x, p1.y, 0.0),
            ]);

            uvs.extend([
                Vec2::new(0.0, 1.0),
                Vec2::new(1.0, 1.0),
                Vec2::new(0.0, 0.0),
                Vec2::new(1.0, 0.0),
            ]);

            indices.extend([i, i + 2, i + 1]);
            indices.extend([i + 3, i + 1, i + 2]);
        }
    }

    mesh.insert_attribute(Mesh::ATTRIBUTE_POSITION, positions);
    mesh.insert_attribute(Mesh::ATTRIBUTE_UV_0, uvs);
    mesh.insert_attribute(
        ATTRIBUTE_TILE_INDEX,
        (0..size.element_product())
            .flat_map(|i| [i; 4])
            .collect::<Vec<u32>>(),
    );
    mesh.insert_indices(Indices::U32(indices));

    mesh
}