bevy-aabb-instancing 0.1.0

use bevy::{
    prelude::*,
    render::{
        primitives::Aabb, render_component::ExtractComponent, render_resource::*,
        renderer::RenderDevice,
    },
};
use bitvec::boxed::BitBox;
use bytemuck::{Pod, Zeroable};

/// Instance-rate vertex attributes for rendered cuboid.
#[derive(Clone, Copy, Pod, Zeroable)]
#[repr(C)]
pub struct CuboidInstance {
    pub minimum: Vec3,
    pub maximum: Vec3,
    pub color_rgba: [f32; 4],
}

impl CuboidInstance {
    pub fn new(minimum: Vec3, maximum: Vec3, color_rgba: [f32; 4]) -> Self {
        Self {
            minimum,
            maximum,
            color_rgba,
        }
    }

    pub(crate) const VERTEX_ATTRIBUTES: [VertexAttribute; 3] = [
        // Min
        VertexAttribute {
            format: VertexFormat::Float32x3,
            offset: 0,
            shader_location: 3, // shader locations 0-2 are taken up by Position, Normal and UV attributes
        },
        // Shape
        VertexAttribute {
            format: VertexFormat::Float32x3,
            offset: VertexFormat::Float32x3.size(),
            shader_location: 4,
        },
        // Color
        VertexAttribute {
            format: VertexFormat::Float32x4,
            offset: 2 * VertexFormat::Float32x3.size(),
            shader_location: 5,
        },
    ];
}

/// The set of cuboids to be extracted for rendering.
#[derive(Clone, Component)]
pub struct CuboidInstances {
    /// Instances to be rendered. These can be masked on/off by creating an optional [`CuboidInstancesMask`] component.
    pub instances: Box<[CuboidInstance]>,
    /// The bounding volume for all `instances`.
    pub aabb: Aabb,
}

impl CuboidInstances {
    pub fn new(instances: Box<[CuboidInstance]>, aabb: Aabb) -> Self {
        Self { instances, aabb }
    }

    /// Automatically creates an [`Aabb`] that bounds all `instances`.
    pub fn new_with_computed_aabb(instances: Box<[CuboidInstance]>) -> Self {
        let mut min = Vec3::splat(f32::MAX);
        let mut max = Vec3::splat(f32::MIN);
        for i in instances.iter() {
            min = min.min(i.minimum);
            max = max.max(i.maximum);
        }
        Self::new(instances, Aabb::from_min_max(min, max))
    }
}

/// Won't be rendered.
pub const VISIBLE: bool = false;

/// Will be rendered.
pub const INVISIBLE: bool = true;

/// An optional component to accompany [`CuboidInstances`] and set the visibility of individual instances.
#[derive(Component)]
pub struct CuboidInstancesMask {
    /// Parallel to the `instances` on [`CuboidInstances`]. Only [`VISIBLE`] bits are rendered. [`INVISIBLE`] bits are not.
    pub bitmask: BitBox,
}

impl CuboidInstancesMask {
    pub fn new(bitmask: BitBox) -> Self {
        Self { bitmask }
    }
}

#[derive(Clone, Component, Debug)]
pub(crate) struct CuboidInstancesBuffer {
    pub buffer: Buffer,
    pub length: usize,
}

impl ExtractComponent for CuboidInstancesBuffer {
    type Query = &'static CuboidInstancesBuffer;
    type Filter = ();

    fn extract_component(buffer: bevy::ecs::query::QueryItem<Self::Query>) -> Self {
        buffer.clone()
    }
}

pub(crate) fn create_new_instance_buffers(
    mut commands: Commands,
    device: Res<RenderDevice>,
    query: Query<
        (Entity, &CuboidInstances, Option<&CuboidInstancesMask>),
        Or<(
            Added<CuboidInstances>,
            Changed<CuboidInstances>,
            Added<CuboidInstancesMask>,
            Changed<CuboidInstancesMask>,
        )>,
    >,
) {
    for (entity, instances, maybe_mask) in query.iter() {
        let mut visible_instances = Vec::new();
        let render_instances: &[CuboidInstance] = if let Some(mask) = maybe_mask {
            assert_eq!(mask.bitmask.len(), instances.instances.len());
            visible_instances.reserve(mask.bitmask.count_zeros());
            for index in mask.bitmask.iter_zeros() {
                visible_instances.push(instances.instances[index]);
            }
            &visible_instances
        } else {
            &instances.instances
        };
        let buffer = device.create_buffer_with_data(&BufferInitDescriptor {
            label: Some("instance data buffer"),
            contents: bytemuck::cast_slice(render_instances),
            usage: BufferUsages::VERTEX | BufferUsages::COPY_DST,
        });
        // HACK: get_or_spawn might help with the error: "Could not add a component ... to entity X because it doesn't exist in
        // this World." Not really sure about the root cause honestly.
        commands.get_or_spawn(entity).insert(CuboidInstancesBuffer {
            buffer,
            length: render_instances.len(),
        });
    }
}