mod.rs

use core::{iter, num::NonZero};

use bevy_camera::Camera;
use bevy_ecs::{entity::EntityHashMap, prelude::*};
use bevy_light::{
    cluster::{
        ClusterableObjectCounts, ClusterableObjects, Clusters, GlobalClusterGpuSettings,
        GlobalClusterSettings,
    },
    ClusteredDecal, EnvironmentMapLight, IrradianceVolume, PointLight, SpotLight,
};
use bevy_math::{uvec4, UVec3, UVec4, Vec4};
use bevy_render::{
    render_resource::{
        BindingResource, BufferBindingType, BufferUsages, DownlevelFlags, RawBufferVec, ShaderSize,
        ShaderType, StorageBuffer, UniformBuffer,
    },
    renderer::{RenderAdapter, RenderDevice, RenderQueue},
    sync_world::{MainEntity, RenderEntity},
    Extract,
};
use bytemuck::{Pod, Zeroable};
use tracing::{error, info, trace, warn};

use crate::{MeshPipeline, RenderViewLightProbes};

pub(crate) mod gpu;

// NOTE: this must be kept in sync with the same constants in
// `mesh_view_types.wgsl`.
pub const MAX_UNIFORM_BUFFER_CLUSTERABLE_OBJECTS: usize = 204;
// Make sure that the clusterable object buffer doesn't overflow the maximum
// size of a UBO on WebGL 2.
const _: () =
    assert!(size_of::<GpuClusteredLight>() * MAX_UNIFORM_BUFFER_CLUSTERABLE_OBJECTS <= 16384);

// NOTE: Clustered-forward rendering requires 3 storage buffer bindings so check that
// at least that many are supported using this constant and SupportedBindingType::from_device()
pub const CLUSTERED_FORWARD_STORAGE_BUFFER_COUNT: u32 = 3;

// this must match CLUSTER_COUNT_SIZE in pbr.wgsl
// and must be large enough to contain MAX_UNIFORM_BUFFER_CLUSTERABLE_OBJECTS
const CLUSTER_COUNT_SIZE: u32 = 9;

const CLUSTER_OFFSET_MASK: u32 = (1 << (32 - (CLUSTER_COUNT_SIZE * 2))) - 1;
const CLUSTER_COUNT_MASK: u32 = (1 << CLUSTER_COUNT_SIZE) - 1;

/// The initial capacity of the Z slice list.
///
/// The application can override this by setting
/// [`GlobalClusterGpuSettings::initial_z_slice_list_capacity`].
pub const GPU_CLUSTERING_INITIAL_Z_SLICE_LIST_CAPACITY: usize = 1024;

/// The initial capacity of the clustered object index list.
///
/// The application can override this by setting
/// [`GlobalClusterGpuSettings::initial_index_list_capacity`].
pub const GPU_CLUSTERING_INITIAL_INDEX_LIST_CAPACITY: usize = 65536;

/// Creates the default [`GlobalClusterSettings`] resource.
pub(crate) fn make_global_cluster_settings(world: &World) -> GlobalClusterSettings {
    let device = world.resource::<RenderDevice>();
    let adapter = world.resource::<RenderAdapter>();
    let clustered_decals_are_usable =
        crate::decal::clustered::clustered_decals_are_usable(device, adapter);
    let supports_storage_buffers = matches!(
        device.get_supported_read_only_binding_type(CLUSTERED_FORWARD_STORAGE_BUFFER_COUNT),
        BufferBindingType::Storage { .. }
    );

    // We need to support compute shaders to use GPU clustering. To deal with
    // the `WGPU_SETTINGS_PRIO="webgl2"` environment setting, we check the
    // `RenderDevice` limits in addition to the `RenderAdapter`.
    //
    // Some android devices report the capabilities and limits wrong, so we can't rely on them.
    // See <https://github.com/bevyengine/bevy/issues/23208> for Android issues
    //
    // GPU clustering doesn't work properly on iOS simulator. See https://github.com/bevyengine/bevy/issues/23428
    let gpu_clustering_supported = !(cfg!(target_os = "android") || cfg!(target_abi = "sim"))
        && adapter
            .get_downlevel_capabilities()
            .flags
            .contains(DownlevelFlags::COMPUTE_SHADERS)
        && device.limits().max_storage_buffers_per_shader_stage > 0;

    let gpu_clustering = if gpu_clustering_supported {
        info!("GPU clustering is supported on this device.");
        Some(GlobalClusterGpuSettings {
            initial_z_slice_list_capacity: GPU_CLUSTERING_INITIAL_Z_SLICE_LIST_CAPACITY,
            initial_index_list_capacity: GPU_CLUSTERING_INITIAL_INDEX_LIST_CAPACITY,
        })
    } else {
        info!("GPU clustering isn't supported on this device; falling back to CPU clustering.");
        None
    };

    GlobalClusterSettings {
        supports_storage_buffers,
        clustered_decals_are_usable,
        gpu_clustering,
        max_uniform_buffer_clusterable_objects: MAX_UNIFORM_BUFFER_CLUSTERABLE_OBJECTS,
        view_cluster_bindings_max_indices: ViewClusterBindings::MAX_INDICES,
    }
}

/// The GPU-side structure that stores information about a clustered light
/// (point or spot).
///
/// This is *not* used for other clustered objects, such as light probes.
#[derive(Copy, Clone, ShaderType, Default, Pod, Zeroable, Debug)]
#[repr(C)]
pub struct GpuClusteredLight {
    // For point lights: the lower-right 2x2 values of the projection matrix [2][2] [2][3] [3][2] [3][3]
    // For spot lights: 2 components of the direction (x,z), spot_scale and spot_offset
    pub(crate) light_custom_data: Vec4,
    pub(crate) color_inverse_square_range: Vec4,
    pub(crate) position_radius: Vec4,
    pub(crate) flags: u32,
    pub(crate) shadow_depth_bias: f32,
    pub(crate) shadow_normal_bias: f32,
    pub(crate) spot_light_tan_angle: f32,
    pub(crate) soft_shadow_size: f32,
    pub(crate) shadow_map_near_z: f32,
    /// The decal applied to this light.
    ///
    /// Note that this is separate from clustered decals. Clustered decals have
    /// their own structures and don't use [`GpuClusteredLight`].
    pub(crate) decal_index: u32,
    /// The radius of the range that the light affects, used for clustering.
    pub(crate) range: f32,
}

/// Contains information about clusterable objects in the scene that's global:
/// i.e. not specific to any view.
#[derive(Resource)]
pub struct GlobalClusterableObjectMeta {
    /// GPU buffers that hold data about the clustered lights.
    ///
    /// This is only for lights. Data about other clusterable objects are stored
    /// in other buffers.
    pub gpu_clustered_lights: GpuClusteredLights,

    /// Maps a *render-world* entity to the index in the appropriate list.
    ///
    /// Only clusterable objects that have render-world entities are in this
    /// list! In particular, light probes (reflection probes and irradiance
    /// volumes) are not.
    pub entity_to_index: EntityHashMap<usize>,
}

/// GPU buffers that hold data about the clustered lights.
///
/// This is only for lights. Data about other clusterable objects are stored in
/// other buffers.
///
/// This has two variants in order to handle platforms in which storage buffers
/// aren't available.
pub struct GpuClusteredLights {
    data: RawBufferVec<GpuClusteredLight>,
    is_storage_buffer: bool,
}

#[derive(Component)]
pub struct ExtractedClusterConfig {
    /// Special near value for cluster calculations
    pub(crate) near: f32,
    pub(crate) far: f32,
    /// Number of clusters in `X` / `Y` / `Z` in the view frustum
    pub(crate) dimensions: UVec3,
}

impl<'a> From<&'a Clusters> for ExtractedClusterConfig {
    fn from(clusters: &'a Clusters) -> Self {
        Self {
            near: clusters.near,
            far: clusters.far,
            dimensions: clusters.dimensions,
        }
    }
}

/// A single command in the stream that [`extract_clusters_for_cpu_clustering`]
/// produces.
enum ExtractedClusterableObjectElement {
    /// Marks the beginning of a new cluster.
    ClusterHeader(ClusterableObjectCounts),
    /// Represents a light.
    ///
    /// The given entity is the render-world entity.
    Light(Entity),
    /// Represents a reflection probe.
    ///
    /// The given entity is the main-world entity of the light probe, as light
    /// probes don't have render world entities.
    ReflectionProbe(MainEntity),
    /// Represents an irradiance volume.
    ///
    /// The given entity is the main-world entity of the light probe, as light
    /// probes don't have render world entities.
    IrradianceVolume(MainEntity),
    /// Represents a clustered decal.
    ///
    /// The given entity is the render-world entity.
    Decal(Entity),
}

#[derive(Component)]
pub struct ExtractedClusterableObjects {
    data: Vec<ExtractedClusterableObjectElement>,
}

#[derive(ShaderType)]
struct GpuClusterOffsetsAndCountsUniform {
    data: Box<[UVec4; ViewClusterBindings::MAX_UNIFORM_ITEMS]>,
}

#[derive(ShaderType, Default)]
struct GpuClusterableObjectIndexListsStorage {
    #[shader(size(runtime))]
    data: Vec<u32>,
}

#[derive(ShaderType, Default)]
struct GpuClusterOffsetsAndCountsStorage {
    /// The starting offset, followed by the number of point lights, spot
    /// lights, reflection probes, and irradiance volumes in each cluster, in
    /// that order. The remaining fields are filled with zeroes.
    #[shader(size(runtime))]
    data: Vec<GpuClusterOffsetAndCounts>,
}

/// The type we use for the offset and counts for each cluster.
type GpuClusterOffsetAndCounts = [UVec4; 2];

enum ViewClusterBuffers {
    Uniform {
        // NOTE: UVec4 is because all arrays in Std140 layout have 16-byte alignment
        clusterable_object_index_lists: UniformBuffer<GpuClusterableObjectIndexListsUniform>,
        // NOTE: UVec4 is because all arrays in Std140 layout have 16-byte alignment
        cluster_offsets_and_counts: UniformBuffer<GpuClusterOffsetsAndCountsUniform>,
    },
    Storage {
        clusterable_object_index_lists: StorageBuffer<GpuClusterableObjectIndexListsStorage>,
        cluster_offsets_and_counts: StorageBuffer<GpuClusterOffsetsAndCountsStorage>,
    },
}

#[derive(Component)]
pub struct ViewClusterBindings {
    n_indices: usize,
    n_offsets: usize,
    buffers: ViewClusterBuffers,
}

pub fn init_global_clusterable_object_meta(
    mut commands: Commands,
    render_device: Res<RenderDevice>,
) {
    commands.insert_resource(GlobalClusterableObjectMeta::new(
        render_device.get_supported_read_only_binding_type(CLUSTERED_FORWARD_STORAGE_BUFFER_COUNT),
    ));
}

impl GlobalClusterableObjectMeta {
    pub fn new(buffer_binding_type: BufferBindingType) -> Self {
        Self {
            gpu_clustered_lights: GpuClusteredLights::new(buffer_binding_type),
            entity_to_index: EntityHashMap::default(),
        }
    }
}

impl GpuClusteredLights {
    fn new(buffer_binding_type: BufferBindingType) -> Self {
        match buffer_binding_type {
            BufferBindingType::Storage { .. } => Self::storage(),
            BufferBindingType::Uniform => Self::uniform(),
        }
    }

    fn uniform() -> Self {
        GpuClusteredLights {
            data: RawBufferVec::new(BufferUsages::UNIFORM),
            is_storage_buffer: false,
        }
    }

    fn storage() -> Self {
        GpuClusteredLights {
            data: RawBufferVec::new(BufferUsages::STORAGE),
            is_storage_buffer: true,
        }
    }

    pub(crate) fn clear(&mut self) {
        self.data.clear();
    }

    pub(crate) fn len(&self) -> usize {
        self.data.len()
    }

    pub(crate) fn add(&mut self, light: GpuClusteredLight) {
        if self.is_storage_buffer || self.data.len() < MAX_UNIFORM_BUFFER_CLUSTERABLE_OBJECTS {
            self.data.push(light);
        }
    }

    pub(crate) fn write_buffer(
        &mut self,
        render_device: &RenderDevice,
        render_queue: &RenderQueue,
    ) {
        if self.is_storage_buffer {
            if self.data.is_empty() {
                self.data.push(GpuClusteredLight::default());
            }
        } else {
            while self.data.len() < MAX_UNIFORM_BUFFER_CLUSTERABLE_OBJECTS {
                self.data.push(GpuClusteredLight::default());
            }
        }

        self.data.write_buffer(render_device, render_queue);
    }

    pub fn binding(&self) -> Option<BindingResource<'_>> {
        self.data.binding()
    }

    pub fn min_size(buffer_binding_type: BufferBindingType) -> NonZero<u64> {
        match buffer_binding_type {
            BufferBindingType::Storage { .. } => GpuClusteredLight::min_size(),
            BufferBindingType::Uniform => NonZero::try_from(
                u64::from(GpuClusteredLight::min_size())
                    * MAX_UNIFORM_BUFFER_CLUSTERABLE_OBJECTS as u64,
            )
            .unwrap(),
        }
    }

    pub fn max_clustered_lights(&self) -> Option<usize> {
        if self.is_storage_buffer {
            None
        } else {
            Some(MAX_UNIFORM_BUFFER_CLUSTERABLE_OBJECTS)
        }
    }
}

/// A shortcut for testing the type of a clusterable object.
type ClusterExtractionMapperQueryFlags = (
    Has<PointLight>,
    Has<SpotLight>,
    Has<EnvironmentMapLight>,
    Has<IrradianceVolume>,
    Has<ClusteredDecal>,
);
/// A shortcut for testing whether an entity is any type of clusterable object.
type ClusterExtractionMapperQueryFilter = Or<(
    With<PointLight>,
    With<SpotLight>,
    With<EnvironmentMapLight>,
    With<IrradianceVolume>,
    With<ClusteredDecal>,
)>;

/// A run condition that tests whether GPU clustering is enabled.
///
/// This is the version for use in extraction systems.
pub fn gpu_clustering_is_enabled_during_extraction(
    global_cluster_settings: Extract<Res<GlobalClusterSettings>>,
) -> bool {
    global_cluster_settings.gpu_clustering.is_some()
}

/// A run condition that tests whether GPU clustering is enabled.
///
/// This is the version for use in non-extraction systems.
pub fn gpu_clustering_is_enabled(global_cluster_settings: Res<GlobalClusterSettings>) -> bool {
    global_cluster_settings.gpu_clustering.is_some()
}

/// Extracts the clusters that the CPU produced into the render world.
pub fn extract_clusters_for_cpu_clustering(
    mut commands: Commands,
    views: Extract<Query<(RenderEntity, &Clusters, &Camera)>>,
    mapper: Extract<
        Query<
            (Option<&RenderEntity>, ClusterExtractionMapperQueryFlags),
            ClusterExtractionMapperQueryFilter,
        >,
    >,
    global_cluster_settings: Extract<Res<GlobalClusterSettings>>,
) {
    for (entity, clusters, camera) in &views {
        let mut entity_commands = commands
            .get_entity(entity)
            .expect("Clusters entity wasn't synced.");
        if !camera.is_active {
            entity_commands.remove::<(ExtractedClusterableObjects, ExtractedClusterConfig)>();
            continue;
        }

        let clusterable_objects = match clusters.clusterable_objects {
            ClusterableObjects::Cpu(ref cpu_clusterable_objects) => cpu_clusterable_objects,
            ClusterableObjects::Gpu => {
                error!("Clusterable objects must have been in CPU mode if doing CPU clustering");
                continue;
            }
        };

        let mut data = vec![];
        for cluster_objects in clusterable_objects {
            data.push(ExtractedClusterableObjectElement::ClusterHeader(
                cluster_objects.counts,
            ));
            for clusterable_entity in cluster_objects.iter() {
                let Ok((
                    maybe_render_entity,
                    (
                        is_point_light,
                        is_spot_light,
                        is_reflection_probe,
                        is_irradiance_volume,
                        is_clustered_decal,
                    ),
                )) = mapper.get(*clusterable_entity)
                else {
                    error!(
                        "Couldn't find clustered object {:?} in the main world",
                        clusterable_entity
                    );
                    continue;
                };

                if let Some(render_entity) = maybe_render_entity {
                    if is_clustered_decal {
                        data.push(ExtractedClusterableObjectElement::Decal(**render_entity));
                    } else if is_point_light || is_spot_light {
                        data.push(ExtractedClusterableObjectElement::Light(**render_entity));
                    }
                }
                if is_reflection_probe {
                    data.push(ExtractedClusterableObjectElement::ReflectionProbe(
                        MainEntity::from(*clusterable_entity),
                    ));
                }
                if is_irradiance_volume {
                    data.push(ExtractedClusterableObjectElement::IrradianceVolume(
                        MainEntity::from(*clusterable_entity),
                    ));
                }
            }
        }

        entity_commands.insert((
            ExtractedClusterableObjects { data },
            ExtractedClusterConfig::from(clusters),
        ));
    }

    commands.insert_resource(global_cluster_settings.clone());
}

/// Creates and populates the GPU buffers that store clusters when CPU
/// clustering is being used.
pub fn prepare_clusters_for_cpu_clustering(
    mut commands: Commands,
    render_device: Res<RenderDevice>,
    render_queue: Res<RenderQueue>,
    mesh_pipeline: Res<MeshPipeline>,
    global_clusterable_object_meta: Res<GlobalClusterableObjectMeta>,
    views: Query<(
        Entity,
        &ExtractedClusterableObjects,
        Option<&RenderViewLightProbes<EnvironmentMapLight>>,
        Option<&RenderViewLightProbes<IrradianceVolume>>,
    )>,
) {
    let render_device = render_device.into_inner();
    let supports_storage_buffers = matches!(
        mesh_pipeline.clustered_forward_buffer_binding_type,
        BufferBindingType::Storage { .. }
    );
    for (entity, extracted_clusters, maybe_environment_maps, maybe_irradiance_volumes) in &views {
        let mut view_clusters_bindings =
            ViewClusterBindings::new(mesh_pipeline.clustered_forward_buffer_binding_type);
        view_clusters_bindings.clear();

        for record in &extracted_clusters.data {
            match record {
                ExtractedClusterableObjectElement::ClusterHeader(counts) => {
                    let offset = view_clusters_bindings.n_indices();
                    view_clusters_bindings.push_offset_and_counts(offset, counts);
                }

                ExtractedClusterableObjectElement::Light(entity)
                | ExtractedClusterableObjectElement::Decal(entity) => {
                    if let Some(clusterable_object_index) =
                        global_clusterable_object_meta.entity_to_index.get(entity)
                    {
                        if view_clusters_bindings.n_indices() >= ViewClusterBindings::MAX_INDICES
                            && !supports_storage_buffers
                        {
                            warn!(
                                "Clusterable object index lists are full! The clusterable \
                                 objects in the view are present in too many clusters."
                            );
                            break;
                        }
                        view_clusters_bindings.push_index(*clusterable_object_index);
                    } else {
                        // This should never happen. The appropriate systems
                        // should have populated
                        // `global_clusterable_object_meta` by now.
                        error!(
                            "Clustered light or decal {:?} had no assigned index!",
                            entity
                        );
                        // Things that should never happen won't happen in debug mode.
                        debug_assert!(false);
                        view_clusters_bindings.push_dummy_index();
                    }
                }

                ExtractedClusterableObjectElement::ReflectionProbe(main_entity) => {
                    match maybe_environment_maps.and_then(|environment_maps| {
                        environment_maps
                            .main_entity_to_render_light_probe_index
                            .get(main_entity)
                    }) {
                        Some(render_light_probe_index) => {
                            view_clusters_bindings.push_index(*render_light_probe_index as usize);
                        }
                        None => {
                            // This can happen while the reflection probe is loading.
                            trace!(
                                "Clustered reflection probe {:?} had no assigned index",
                                main_entity,
                            );
                            view_clusters_bindings.push_dummy_index();
                        }
                    }
                }

                ExtractedClusterableObjectElement::IrradianceVolume(main_entity) => {
                    match maybe_irradiance_volumes.and_then(|irradiance_volumes| {
                        irradiance_volumes
                            .main_entity_to_render_light_probe_index
                            .get(main_entity)
                    }) {
                        Some(render_light_probe_index) => {
                            view_clusters_bindings.push_index(*render_light_probe_index as usize);
                        }
                        None => {
                            trace!(
                                "Clustered irradiance volume {:?} had no assigned index",
                                main_entity
                            );
                            view_clusters_bindings.push_dummy_index();
                        }
                    }
                }
            }
        }

        view_clusters_bindings.write_buffers(render_device, &render_queue);

        commands.entity(entity).insert(view_clusters_bindings);
    }
}

impl ViewClusterBindings {
    pub const MAX_OFFSETS: usize = 16384 / 4;
    const MAX_UNIFORM_ITEMS: usize = Self::MAX_OFFSETS / 4;
    pub const MAX_INDICES: usize = 16384;

    pub fn new(buffer_binding_type: BufferBindingType) -> Self {
        Self {
            n_indices: 0,
            n_offsets: 0,
            buffers: ViewClusterBuffers::new(buffer_binding_type),
        }
    }

    pub fn clear(&mut self) {
        match &mut self.buffers {
            ViewClusterBuffers::Uniform {
                clusterable_object_index_lists,
                cluster_offsets_and_counts,
            } => {
                *clusterable_object_index_lists.get_mut().data =
                    [UVec4::ZERO; Self::MAX_UNIFORM_ITEMS];
                *cluster_offsets_and_counts.get_mut().data = [UVec4::ZERO; Self::MAX_UNIFORM_ITEMS];
            }
            ViewClusterBuffers::Storage {
                clusterable_object_index_lists,
                cluster_offsets_and_counts,
                ..
            } => {
                clusterable_object_index_lists.get_mut().data.clear();
                cluster_offsets_and_counts.get_mut().data.clear();
            }
        }
    }

    fn push_offset_and_counts(&mut self, offset: usize, counts: &ClusterableObjectCounts) {
        match &mut self.buffers {
            ViewClusterBuffers::Uniform {
                cluster_offsets_and_counts,
                ..
            } => {
                let array_index = self.n_offsets >> 2; // >> 2 is equivalent to / 4
                if array_index >= Self::MAX_UNIFORM_ITEMS {
                    warn!("cluster offset and count out of bounds!");
                    return;
                }
                let component = self.n_offsets & ((1 << 2) - 1);
                let packed =
                    pack_offset_and_counts(offset, counts.point_lights, counts.spot_lights);

                cluster_offsets_and_counts.get_mut().data[array_index][component] = packed;
            }
            ViewClusterBuffers::Storage {
                cluster_offsets_and_counts,
                ..
            } => {
                cluster_offsets_and_counts.get_mut().data.push([
                    uvec4(
                        offset as u32,
                        counts.point_lights,
                        counts.spot_lights,
                        counts.reflection_probes,
                    ),
                    uvec4(counts.irradiance_volumes, counts.decals, 0, 0),
                ]);
            }
        }

        self.n_offsets += 1;
    }

    pub fn n_indices(&self) -> usize {
        self.n_indices
    }

    // An internal helper method that pushes a raw clustered object index to the
    // GPU buffer.
    fn push_raw_index(&mut self, index: u32) {
        match &mut self.buffers {
            ViewClusterBuffers::Uniform {
                clusterable_object_index_lists,
                ..
            } => {
                let array_index = self.n_indices >> 4; // >> 4 is equivalent to / 16
                let component = (self.n_indices >> 2) & ((1 << 2) - 1);
                let sub_index = self.n_indices & ((1 << 2) - 1);

                clusterable_object_index_lists.get_mut().data[array_index][component] |=
                    index << (8 * sub_index);
            }
            ViewClusterBuffers::Storage {
                clusterable_object_index_lists,
                ..
            } => {
                clusterable_object_index_lists.get_mut().data.push(index);
            }
        }

        self.n_indices += 1;
    }

    /// Pushes the index of a clustered object to the GPU buffer.
    pub fn push_index(&mut self, index: usize) {
        self.push_raw_index(index as u32);
    }

    /// Pushes a placeholder -1 index to the GPU buffer.
    ///
    /// This is used when processing reflection probes and irradiance volumes
    /// that haven't loaded yet.
    pub fn push_dummy_index(&mut self) {
        self.push_raw_index(!0);
    }

    /// Reserves space in the cluster offsets-and-counts list for `clusters`
    /// clusters.
    pub fn reserve_clusters(&mut self, clusters: usize) {
        match &mut self.buffers {
            ViewClusterBuffers::Uniform { .. } => {
                error!("`reserve_clusters` should only be called in GPU clustering, which requires a storage buffer");
            }
            ViewClusterBuffers::Storage {
                cluster_offsets_and_counts,
                ..
            } => {
                cluster_offsets_and_counts
                    .get_mut()
                    .data
                    .extend(iter::repeat_n(
                        GpuClusterOffsetAndCounts::default(),
                        clusters,
                    ));
                self.n_offsets += clusters;
            }
        }
    }

    /// Reserves space in the index lists for `elements` indices.
    pub fn reserve_indices(&mut self, elements: usize) {
        match &mut self.buffers {
            ViewClusterBuffers::Uniform { .. } => {
                error!("`reserve_indices` should only be called in GPU clustering, which requires a storage buffer");
            }
            ViewClusterBuffers::Storage {
                clusterable_object_index_lists,
                ..
            } => {
                clusterable_object_index_lists
                    .get_mut()
                    .data
                    .extend(iter::repeat_n(0, elements));
                self.n_indices += elements;
            }
        }
    }

    pub fn write_buffers(&mut self, render_device: &RenderDevice, render_queue: &RenderQueue) {
        match &mut self.buffers {
            ViewClusterBuffers::Uniform {
                clusterable_object_index_lists,
                cluster_offsets_and_counts,
            } => {
                clusterable_object_index_lists.write_buffer(render_device, render_queue);
                cluster_offsets_and_counts.write_buffer(render_device, render_queue);
            }
            ViewClusterBuffers::Storage {
                clusterable_object_index_lists,
                cluster_offsets_and_counts,
            } => {
                clusterable_object_index_lists.write_buffer(render_device, render_queue);
                cluster_offsets_and_counts.write_buffer(render_device, render_queue);
            }
        }
    }

    pub fn clusterable_object_index_lists_binding(&self) -> Option<BindingResource<'_>> {
        match &self.buffers {
            ViewClusterBuffers::Uniform {
                clusterable_object_index_lists,
                ..
            } => clusterable_object_index_lists.binding(),
            ViewClusterBuffers::Storage {
                clusterable_object_index_lists,
                ..
            } => clusterable_object_index_lists.binding(),
        }
    }

    pub fn offsets_and_counts_binding(&self) -> Option<BindingResource<'_>> {
        match &self.buffers {
            ViewClusterBuffers::Uniform {
                cluster_offsets_and_counts,
                ..
            } => cluster_offsets_and_counts.binding(),
            ViewClusterBuffers::Storage {
                cluster_offsets_and_counts,
                ..
            } => cluster_offsets_and_counts.binding(),
        }
    }

    pub fn min_size_clusterable_object_index_lists(
        buffer_binding_type: BufferBindingType,
    ) -> NonZero<u64> {
        match buffer_binding_type {
            BufferBindingType::Storage { .. } => GpuClusterableObjectIndexListsStorage::min_size(),
            BufferBindingType::Uniform => GpuClusterableObjectIndexListsUniform::min_size(),
        }
    }

    pub fn min_size_cluster_offsets_and_counts(
        buffer_binding_type: BufferBindingType,
    ) -> NonZero<u64> {
        match buffer_binding_type {
            BufferBindingType::Storage { .. } => GpuClusterOffsetsAndCountsStorage::min_size(),
            BufferBindingType::Uniform => GpuClusterOffsetsAndCountsUniform::min_size(),
        }
    }
}

impl ViewClusterBuffers {
    fn new(buffer_binding_type: BufferBindingType) -> Self {
        match buffer_binding_type {
            BufferBindingType::Storage { .. } => Self::storage(),
            BufferBindingType::Uniform => Self::uniform(),
        }
    }

    fn uniform() -> Self {
        ViewClusterBuffers::Uniform {
            clusterable_object_index_lists: UniformBuffer::default(),
            cluster_offsets_and_counts: UniformBuffer::default(),
        }
    }

    fn storage() -> Self {
        ViewClusterBuffers::Storage {
            clusterable_object_index_lists: StorageBuffer::default(),
            cluster_offsets_and_counts: StorageBuffer::default(),
        }
    }
}

// Compresses the offset and counts of point and spot lights so that they fit in
// a UBO.
//
// This function is only used if storage buffers are unavailable on this
// platform: typically, on WebGL 2.
//
// NOTE: With uniform buffer max binding size as 16384 bytes
// that means we can fit 204 clusterable objects in one uniform
// buffer, which means the count can be at most 204 so it
// needs 9 bits.
// The array of indices can also use u8 and that means the
// offset in to the array of indices needs to be able to address
// 16384 values. log2(16384) = 14 bits.
// We use 32 bits to store the offset and counts so
// we pack the offset into the upper 14 bits of a u32,
// the point light count into bits 9-17, and the spot light count into bits 0-8.
//  [ 31     ..     18 | 17      ..      9 | 8       ..     0 ]
//  [      offset      | point light count | spot light count ]
//
// NOTE: This assumes CPU and GPU endianness are the same which is true
// for all common and tested x86/ARM CPUs and AMD/NVIDIA/Intel/Apple/etc GPUs
//
// NOTE: On platforms that use this function, we don't cluster light probes, so
// the number of light probes is irrelevant.
fn pack_offset_and_counts(offset: usize, point_count: u32, spot_count: u32) -> u32 {
    ((offset as u32 & CLUSTER_OFFSET_MASK) << (CLUSTER_COUNT_SIZE * 2))
        | ((point_count & CLUSTER_COUNT_MASK) << CLUSTER_COUNT_SIZE)
        | (spot_count & CLUSTER_COUNT_MASK)
}

#[derive(ShaderType)]
struct GpuClusterableObjectIndexListsUniform {
    data: Box<[UVec4; ViewClusterBindings::MAX_UNIFORM_ITEMS]>,
}

// NOTE: Assert at compile time that GpuClusterableObjectIndexListsUniform
// fits within the maximum uniform buffer binding size
const _: () = assert!(GpuClusterableObjectIndexListsUniform::SHADER_SIZE.get() <= 16384);

impl Default for GpuClusterableObjectIndexListsUniform {
    fn default() -> Self {
        Self {
            data: Box::new([UVec4::ZERO; ViewClusterBindings::MAX_UNIFORM_ITEMS]),
        }
    }
}

impl Default for GpuClusterOffsetsAndCountsUniform {
    fn default() -> Self {
        Self {
            data: Box::new([UVec4::ZERO; ViewClusterBindings::MAX_UNIFORM_ITEMS]),
        }
    }
}