bevy_render 0.20.0-rc.1

//! Material bind group management for bindless resources.
//!
//! In bindless mode, Bevy's renderer groups materials into bind groups. This
//! allocator manages each bind group, assigning slots to materials as
//! appropriate.

use bevy_app::{App, Plugin};
use bevy_asset::{Asset, AssetId, Handle, UntypedAssetId};
use bevy_derive::{Deref, DerefMut};
use bevy_ecs::{
    resource::Resource,
    schedule::IntoScheduleConfigs as _,
    system::{Commands, Res, ResMut, SystemParamItem},
};
use bevy_platform::collections::{hash_map::Entry, HashMap, HashSet};
use bevy_reflect::{prelude::ReflectDefault, Reflect};
use bevy_utils::{default, TypeIdHashMap};
use bytemuck::{Pod, Zeroable};
use core::hash::Hash;
use core::{cmp::Ordering, iter, mem, ops::Range};
use std::any::TypeId;
use tracing::{error, trace};

use crate::{
    erased_render_asset::PrepareAssetError,
    render_asset::RenderAssets,
    render_resource::{
        AsBindGroup, AsBindGroupError, BindGroupBuilder, BindlessSlabResourceLimit,
        OwnedBindingResource, PipelineCache, UnpreparedBindingResource, UnpreparedBindingResources,
    },
    storage::ShaderBuffer,
    GpuResourceAppExt as _, Render, RenderApp, RenderStartup, RenderSystems,
};
use crate::{
    render_resource::{
        BindGroup, BindGroupEntry, BindGroupLayoutDescriptor, BindingNumber, BindingResource,
        BindlessDescriptor, BindlessIndex, BindlessIndexTableDescriptor, BindlessResourceType,
        Buffer, BufferBinding, BufferDescriptor, BufferId, BufferInitDescriptor, BufferUsages,
        CompareFunction, FilterMode, MipmapFilterMode, PreparedBindGroup, RawBufferVec, Sampler,
        SamplerDescriptor, SamplerId, TextureView, TextureViewDimension, TextureViewId,
        WgpuSampler, WgpuTextureView,
    },
    renderer::{RenderDevice, RenderQueue},
    settings::WgpuFeatures,
    storage::{GpuShaderBuffer, RenderChangedShaderBuffers},
    texture::FallbackImage,
};

/// A [`Plugin`] that provides the material bind group allocator.
///
/// The material bind group allocator is infrastructure for bindless resources.
/// It packs multiple materials into a small number of bind groups, allowing
/// Bevy to render large parts of the scene with a small number of drawcalls.
pub struct MaterialBindGroupPlugin;

/// A resource that maps each untyped material ID to its binding.
///
/// This duplicates information in `RenderAssets<M>`, but it doesn't have the
/// `M` type parameter, so it can be used in untyped contexts like
/// `collect_meshes_for_gpu_building`.
#[derive(Resource, Default, Deref, DerefMut)]
pub struct RenderMaterialBindings {
    /// The mapping from each asset ID to its location within the material bind
    /// group allocator.
    #[deref]
    bindings: HashMap<UntypedAssetId, MaterialBindingId>,

    /// The bind group we're building up.
    ///
    /// This is cleared and reused for every asset that we build in order to
    /// reuse allocations.
    bind_group_builder: BindGroupBuilder,
}

/// A resource that places materials into bind groups and tracks their
/// resources.
///
/// Internally, Bevy has separate allocators for bindless and non-bindless
/// materials. This resource provides a common interface to the specific
/// allocator in use.
#[derive(Resource, Deref, DerefMut, Default)]
pub struct MaterialBindGroupAllocators(TypeIdHashMap<MaterialBindGroupAllocator>);

/// The allocator that manages a specific material.
pub enum MaterialBindGroupAllocator {
    /// The allocator used when the material is bindless.
    Bindless(Box<MaterialBindGroupBindlessAllocator>),
    /// The allocator used when the material is non-bindless.
    NonBindless(Box<MaterialBindGroupNonBindlessAllocator>),
}

/// The allocator that places bindless materials into bind groups and tracks
/// their resources.
pub struct MaterialBindGroupBindlessAllocator {
    /// The label of the bind group allocator to use for allocated buffers.
    label: &'static str,
    /// The slabs, each of which contains a bind group.
    slabs: Vec<MaterialBindlessSlab>,
    /// The layout of the bind groups that we produce.
    bind_group_layout: BindGroupLayoutDescriptor,
    /// Information about the bindless resources in the material.
    ///
    /// We use this information to create and maintain bind groups.
    bindless_descriptor: BindlessDescriptor,

    /// Dummy buffers that we use to fill empty slots in buffer binding arrays.
    ///
    /// There's one fallback buffer for each buffer in the bind group, each
    /// appropriately sized. Each buffer contains one uninitialized element of
    /// the applicable type.
    fallback_buffers: HashMap<BindlessIndex, Buffer>,

    /// The maximum number of resources that can be stored in a slab.
    ///
    /// This corresponds to `SLAB_CAPACITY` in the `#[bindless(SLAB_CAPACITY)]`
    /// attribute, when deriving `AsBindGroup`.
    slab_capacity: u32,
}

/// A single bind group and the bookkeeping necessary to allocate into it.
pub struct MaterialBindlessSlab {
    /// The current bind group, if it's up to date.
    ///
    /// If this is `None`, then the bind group is dirty and needs to be
    /// regenerated.
    bind_group: Option<BindGroup>,

    /// The GPU-accessible buffers that hold the mapping from binding index to
    /// bindless slot.
    ///
    /// This is conventionally assigned to bind group binding 0, but it can be
    /// changed using the `#[bindless(index_table(binding(B)))]` attribute on
    /// `AsBindGroup`.
    ///
    /// Because the slab binary searches this table, the entries within must be
    /// sorted by bindless index.
    bindless_index_tables: Vec<MaterialBindlessIndexTable>,

    /// The binding arrays containing samplers.
    samplers: HashMap<BindlessResourceType, MaterialBindlessBindingArray<Sampler>>,
    /// The binding arrays containing textures.
    textures: HashMap<BindlessResourceType, MaterialBindlessBindingArray<TextureView>>,
    /// The binding arrays containing buffers.
    buffers: HashMap<BindlessIndex, MaterialBindlessBindingArray<Buffer>>,
    /// The binding arrays containing shader buffers.
    ///
    /// These differ from the binding arrays containing raw buffers because
    /// shader buffers are assets, and therefore the actual underlying buffer
    /// can change.
    shader_buffers: HashMap<BindlessIndex, MaterialBindlessBindingArray<Handle<ShaderBuffer>>>,
    /// The buffers that contain plain old data (i.e. the structure-level
    /// `#[data]` attribute of `AsBindGroup`).
    data_buffers: HashMap<BindlessIndex, MaterialDataBuffer>,

    /// A list of free slot IDs.
    free_slots: Vec<MaterialBindGroupSlot>,
    /// The total number of materials currently allocated in this slab.
    live_allocation_count: u32,
    /// The total number of resources currently allocated in the binding arrays.
    allocated_resource_count: u32,
}

/// A GPU-accessible buffer that holds the mapping from binding index to
/// bindless slot.
///
/// This is conventionally assigned to bind group binding 0, but it can be
/// changed by altering the [`Self::binding_number`], which corresponds to the
/// `#[bindless(index_table(binding(B)))]` attribute in `AsBindGroup`.
struct MaterialBindlessIndexTable {
    /// The buffer containing the mappings.
    buffer: RetainedRawBufferVec<u32>,
    /// The range of bindless indices that this bindless index table covers.
    ///
    /// If this range is M..N, then the field at index $i$ maps to bindless
    /// index $i$ + M. The size of this table is N - M.
    ///
    /// This corresponds to the `#[bindless(index_table(range(M..N)))]`
    /// attribute in `AsBindGroup`.
    index_range: Range<BindlessIndex>,
    /// The binding number that this index table is assigned to in the shader.
    binding_number: BindingNumber,
}

/// A single binding array for storing bindless resources and the bookkeeping
/// necessary to allocate into it.
struct MaterialBindlessBindingArray<R>
where
    R: GetBindingResourceId,
{
    /// The number of the binding that we attach this binding array to.
    binding_number: BindingNumber,
    /// A mapping from bindless slot index to the resource stored in that slot,
    /// if any.
    bindings: Vec<Option<MaterialBindlessBinding<R>>>,
    /// The type of resource stored in this binding array.
    resource_type: BindlessResourceType,
    /// Maps a resource ID to the slot in which it's stored.
    ///
    /// This is essentially the inverse mapping of [`Self::bindings`].
    resource_to_slot: HashMap<BindingResourceId, u32>,
    /// A list of free slots in [`Self::bindings`] that contain no binding.
    free_slots: Vec<u32>,
    /// The number of allocated objects in this binding array.
    len: u32,
}

/// A single resource (sampler, texture, or buffer) in a binding array.
///
/// Resources hold a reference count, which specifies the number of materials
/// currently allocated within the slab that refer to this resource. When the
/// reference count drops to zero, the resource is freed.
struct MaterialBindlessBinding<R>
where
    R: GetBindingResourceId,
{
    /// The sampler, texture, or buffer.
    resource: R,
    /// The number of materials currently allocated within the containing slab
    /// that use this resource.
    ref_count: u32,
}

/// The allocator that stores bind groups for non-bindless materials.
pub struct MaterialBindGroupNonBindlessAllocator {
    /// The label of the bind group allocator to use for allocated buffers.
    label: &'static str,
    /// A mapping from [`MaterialBindGroupIndex`] to the bind group allocated in
    /// each slot.
    bind_groups: Vec<Option<MaterialNonBindlessAllocatedBindGroup>>,
    /// The bind groups that are dirty and need to be prepared.
    ///
    /// To prepare the bind groups, call
    /// [`MaterialBindGroupAllocator::prepare_bind_groups`].
    to_prepare: HashSet<MaterialBindGroupIndex>,
    /// A list of free bind group indices.
    free_indices: Vec<MaterialBindGroupIndex>,
}

/// A single bind group that a [`MaterialBindGroupNonBindlessAllocator`] is
/// currently managing.
enum MaterialNonBindlessAllocatedBindGroup {
    /// An unprepared bind group.
    ///
    /// The allocator prepares all outstanding unprepared bind groups when
    /// [`MaterialBindGroupNonBindlessAllocator::prepare_bind_groups`] is
    /// called.
    Unprepared {
        /// The unprepared bind group, including extra data.
        bind_group: BindGroupBuilder,
        /// The layout of that bind group.
        layout: BindGroupLayoutDescriptor,
    },
    /// A bind group that's already been prepared.
    Prepared {
        /// The actual bind group.
        bind_group: PreparedBindGroup,
        /// The layout descriptor that reflects the layout of this bind group.
        ///
        /// Rendering doesn't use this, but we need to keep it around in case we
        /// need to demote a prepared bind group to an unprepared one. That will
        /// happen if a shader buffer changes, causing the bind group to become
        /// invalid.
        layout: BindGroupLayoutDescriptor,
        /// Uniform buffers referenced by the bind group.
        #[expect(dead_code, reason = "These buffers are only referenced by bind groups")]
        uniform_buffers: Vec<Buffer>,
    },
}

/// Dummy instances of various resources that we fill unused slots in binding
/// arrays with.
#[derive(Resource)]
pub struct FallbackBindlessResources {
    /// A dummy filtering sampler.
    filtering_sampler: Sampler,
    /// A dummy non-filtering sampler.
    non_filtering_sampler: Sampler,
    /// A dummy comparison sampler.
    comparison_sampler: Sampler,
}

/// A dummy buffer that substitutes for a shader buffer if that shader buffer
/// points to an invalid buffer.
#[derive(Resource, Deref, DerefMut)]
pub struct FallbackBuffer(Buffer);

/// The `wgpu` ID of a single bindless or non-bindless resource.
#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)]
enum BindingResourceId {
    /// A buffer.
    Buffer(BufferId),
    /// A texture view, with the given dimension.
    TextureView(TextureViewDimension, TextureViewId),
    /// A sampler.
    Sampler(SamplerId),
    /// A buffer containing plain old data.
    ///
    /// This corresponds to the `#[data]` structure-level attribute on
    /// `AsBindGroup`.
    DataBuffer,
    /// A [`ShaderBuffer`] asset.
    ShaderBuffer(AssetId<ShaderBuffer>),
}

/// A temporary list of references to `wgpu` bindless resources.
///
/// We need this because the `wgpu` bindless API takes a slice of references.
/// Thus we need to create intermediate vectors of bindless resources in order
/// to satisfy `wgpu`'s lifetime requirements.
enum BindingResourceArray<'a> {
    /// A list of bindings.
    Buffers(Vec<BufferBinding<'a>>),
    /// A list of texture views.
    TextureViews(Vec<&'a WgpuTextureView>),
    /// A list of samplers.
    Samplers(Vec<&'a WgpuSampler>),
}

/// The location of a material (either bindless or non-bindless) within the
/// slabs.
#[derive(Clone, Copy, Debug, Default, Pod, Zeroable, Reflect)]
#[reflect(Clone, Default)]
#[repr(C)]
pub struct MaterialBindingId {
    /// The index of the bind group (slab) where the GPU data is located.
    pub group: MaterialBindGroupIndex,
    /// The slot within that bind group.
    ///
    /// Non-bindless materials will always have a slot of 0.
    pub slot: MaterialBindGroupSlot,
}

/// The index of each material bind group.
///
/// In bindless mode, each bind group contains multiple materials. In
/// non-bindless mode, each bind group contains only one material.
#[derive(
    Clone, Copy, Debug, Default, PartialEq, Eq, Hash, Pod, Zeroable, Reflect, Deref, DerefMut,
)]
#[reflect(Default, Clone, PartialEq, Hash)]
#[repr(C)]
pub struct MaterialBindGroupIndex(pub u32);

impl From<u32> for MaterialBindGroupIndex {
    fn from(value: u32) -> Self {
        MaterialBindGroupIndex(value)
    }
}

/// The index of the slot containing material data within each material bind
/// group.
///
/// In bindless mode, this slot is needed to locate the material data in each
/// bind group, since multiple materials are packed into a single slab. In
/// non-bindless mode, this slot is always 0.
#[derive(Clone, Copy, Debug, Default, PartialEq, Pod, Zeroable, Reflect, Deref, DerefMut)]
#[reflect(Default, Clone, PartialEq)]
#[repr(C)]
pub struct MaterialBindGroupSlot(pub u32);

/// The CPU/GPU synchronization state of a buffer that we maintain.
///
/// Currently, the only buffer that we maintain is the
/// [`MaterialBindlessIndexTable`].
#[derive(Clone, Copy, PartialEq)]
enum BufferDirtyState {
    /// The buffer is currently synchronized between the CPU and GPU.
    Clean,
    /// The buffer hasn't been created yet.
    NeedsReserve,
    /// The buffer exists on both CPU and GPU, but the GPU data is out of date.
    NeedsUpload,
}

/// Information that describes a potential allocation of a [`BindGroupBuilder`]
/// into a slab.
struct BindlessAllocationCandidate {
    /// A map that, for every resource in the [`BindGroupBuilder`] that already
    /// existed in this slab, maps bindless index of that resource to its slot
    /// in the appropriate binding array.
    pre_existing_resources: HashMap<BindlessIndex, u32>,
    /// Stores the number of free slots that are needed to satisfy this
    /// allocation.
    needed_free_slots: u32,
}

/// A trait that allows fetching the [`BindingResourceId`] from a
/// [`BindlessResourceType`].
///
/// This is used when freeing bindless resources, in order to locate the IDs
/// assigned to each resource so that they can be removed from the appropriate
/// maps.
trait GetBindingResourceId {
    /// Returns the [`BindingResourceId`] for this resource.
    ///
    /// `resource_type` specifies this resource's type. This is used for
    /// textures, as a `wgpu` [`TextureView`] doesn't store enough information
    /// itself to determine its dimension.
    fn binding_resource_id(&self, resource_type: BindlessResourceType) -> BindingResourceId;
}

/// The public interface to a slab, which represents a single bind group.
pub struct MaterialSlab<'a>(MaterialSlabImpl<'a>);

/// The actual implementation of a material slab.
///
/// This has bindless and non-bindless variants.
enum MaterialSlabImpl<'a> {
    /// The implementation of the slab interface we use when the slab
    /// is bindless.
    Bindless(&'a MaterialBindlessSlab),
    /// The implementation of the slab interface we use when the slab
    /// is non-bindless.
    NonBindless(MaterialNonBindlessSlab<'a>),
}

/// A single bind group that the [`MaterialBindGroupNonBindlessAllocator`]
/// manages.
enum MaterialNonBindlessSlab<'a> {
    /// A slab that has a bind group.
    Prepared(&'a PreparedBindGroup),
    /// A slab that doesn't yet have a bind group.
    Unprepared,
}

/// Manages an array of untyped plain old data on GPU and allocates individual
/// slots within that array.
///
/// This supports the `#[data]` attribute of `AsBindGroup`.
struct MaterialDataBuffer {
    /// The number of the binding that we attach this storage buffer to.
    binding_number: BindingNumber,
    /// The actual data.
    ///
    /// Note that this is untyped (`u8`); the actual aligned size of each
    /// element is given by [`Self::aligned_element_size`];
    buffer: RetainedRawBufferVec<u8>,
    /// The size of each element in the buffer, including padding and alignment
    /// if any.
    aligned_element_size: u32,
    /// A list of free slots within the buffer.
    free_slots: Vec<u32>,
    /// The actual number of slots that have been allocated.
    len: u32,
}

/// A buffer containing plain old data, already packed into the appropriate GPU
/// format, and that can be updated incrementally.
///
/// This structure exists in order to encapsulate the lazy update
/// ([`BufferDirtyState`]) logic in a single place.
#[derive(Deref, DerefMut)]
struct RetainedRawBufferVec<T>
where
    T: Pod,
{
    /// The contents of the buffer.
    #[deref]
    buffer: RawBufferVec<T>,
    /// Whether the contents of the buffer have been uploaded to the GPU.
    dirty: BufferDirtyState,
}

/// The size of the buffer that we assign to unused buffer slots, in bytes.
///
/// This is essentially arbitrary, as it doesn't seem to matter to `wgpu` what
/// the size is.
const DEFAULT_BINDLESS_FALLBACK_BUFFER_SIZE: u64 = 16;

impl From<u32> for MaterialBindGroupSlot {
    fn from(value: u32) -> Self {
        MaterialBindGroupSlot(value)
    }
}

impl From<MaterialBindGroupSlot> for u32 {
    fn from(value: MaterialBindGroupSlot) -> Self {
        value.0
    }
}

impl<'a> From<&'a UnpreparedBindingResource> for BindingResourceId {
    fn from(value: &'a UnpreparedBindingResource) -> Self {
        match *value {
            UnpreparedBindingResource::Buffer(ref buffer) => BindingResourceId::Buffer(buffer.id()),
            UnpreparedBindingResource::ShaderBuffer(ref shader_buffer) => {
                BindingResourceId::ShaderBuffer(shader_buffer.id())
            }
            UnpreparedBindingResource::Data(_) => BindingResourceId::DataBuffer,
            UnpreparedBindingResource::TextureView(
                ref texture_view_dimension,
                ref texture_view,
            ) => BindingResourceId::TextureView(*texture_view_dimension, texture_view.id()),
            UnpreparedBindingResource::Sampler(_, ref sampler) => {
                BindingResourceId::Sampler(sampler.id())
            }
        }
    }
}

impl GetBindingResourceId for Buffer {
    fn binding_resource_id(&self, _: BindlessResourceType) -> BindingResourceId {
        BindingResourceId::Buffer(self.id())
    }
}

impl GetBindingResourceId for Sampler {
    fn binding_resource_id(&self, _: BindlessResourceType) -> BindingResourceId {
        BindingResourceId::Sampler(self.id())
    }
}

impl GetBindingResourceId for TextureView {
    fn binding_resource_id(&self, resource_type: BindlessResourceType) -> BindingResourceId {
        let texture_view_dimension = match resource_type {
            BindlessResourceType::Texture1d => TextureViewDimension::D1,
            BindlessResourceType::Texture2d => TextureViewDimension::D2,
            BindlessResourceType::Texture2dArray => TextureViewDimension::D2Array,
            BindlessResourceType::Texture3d => TextureViewDimension::D3,
            BindlessResourceType::TextureCube => TextureViewDimension::Cube,
            BindlessResourceType::TextureCubeArray => TextureViewDimension::CubeArray,
            _ => panic!("Resource type is not a texture"),
        };
        BindingResourceId::TextureView(texture_view_dimension, self.id())
    }
}

impl Plugin for MaterialBindGroupPlugin {
    fn build(&self, app: &mut App) {
        let Some(render_app) = app.get_sub_app_mut(RenderApp) else {
            return;
        };

        render_app
            .init_resource::<MaterialBindGroupAllocators>()
            .allow_ambiguous_resource::<MaterialBindGroupAllocators>()
            .init_gpu_resource::<RenderMaterialBindings>()
            .allow_ambiguous_resource::<RenderMaterialBindings>()
            .add_systems(RenderStartup, init_fallback_resources)
            .add_systems(
                Render,
                (
                    prepare_material_bind_groups,
                    write_material_bind_group_buffers,
                )
                    .chain()
                    .in_set(RenderSystems::PrepareBindGroups),
            );
    }
}

impl GetBindingResourceId for Handle<ShaderBuffer> {
    fn binding_resource_id(&self, _: BindlessResourceType) -> BindingResourceId {
        BindingResourceId::ShaderBuffer(self.id())
    }
}

impl MaterialBindGroupAllocator {
    /// Creates a new [`MaterialBindGroupAllocator`] managing the data for a
    /// single material.
    pub fn new(
        render_device: &RenderDevice,
        label: &'static str,
        bindless_descriptor: Option<BindlessDescriptor>,
        bind_group_layout: BindGroupLayoutDescriptor,
        slab_capacity: Option<BindlessSlabResourceLimit>,
    ) -> MaterialBindGroupAllocator {
        if let Some(bindless_descriptor) = bindless_descriptor {
            MaterialBindGroupAllocator::Bindless(Box::new(MaterialBindGroupBindlessAllocator::new(
                render_device,
                label,
                bindless_descriptor,
                bind_group_layout,
                slab_capacity,
            )))
        } else {
            MaterialBindGroupAllocator::NonBindless(Box::new(
                MaterialBindGroupNonBindlessAllocator::new(label),
            ))
        }
    }

    /// Returns the slab with the given index, if one exists.
    pub fn get(&self, group: MaterialBindGroupIndex) -> Option<MaterialSlab<'_>> {
        match *self {
            MaterialBindGroupAllocator::Bindless(ref bindless_allocator) => bindless_allocator
                .get(group)
                .map(|bindless_slab| MaterialSlab(MaterialSlabImpl::Bindless(bindless_slab))),
            MaterialBindGroupAllocator::NonBindless(ref non_bindless_allocator) => {
                non_bindless_allocator.get(group).map(|non_bindless_slab| {
                    MaterialSlab(MaterialSlabImpl::NonBindless(non_bindless_slab))
                })
            }
        }
    }

    /// Allocates the resources within a [`BindGroupBuilder`] and returns the
    /// resulting binding ID.
    ///
    /// This method should generally be preferred over
    /// [`Self::allocate_prepared`], because this method supports both bindless
    /// and non-bindless bind groups. Only use [`Self::allocate_prepared`] if
    /// you need to prepare the bind group yourself.
    pub fn allocate_unprepared(
        &mut self,
        unprepared_bind_group: &mut BindGroupBuilder,
        bind_group_layout: &BindGroupLayoutDescriptor,
    ) -> MaterialBindingId {
        match *self {
            MaterialBindGroupAllocator::Bindless(
                ref mut material_bind_group_bindless_allocator,
            ) => material_bind_group_bindless_allocator.allocate_unprepared(unprepared_bind_group),
            MaterialBindGroupAllocator::NonBindless(
                ref mut material_bind_group_non_bindless_allocator,
            ) => material_bind_group_non_bindless_allocator
                .allocate_unprepared(unprepared_bind_group, (*bind_group_layout).clone()),
        }
    }

    /// Places a pre-prepared bind group into a slab.
    ///
    /// For bindless materials, the allocator internally manages the bind
    /// groups, so calling this method will panic if this is a bindless
    /// allocator. Only non-bindless allocators support this method.
    ///
    /// It's generally preferred to use [`Self::allocate_unprepared`], because
    /// that method supports both bindless and non-bindless allocators. Only use
    /// this method if you need to prepare the bind group yourself.
    pub fn allocate_prepared(
        &mut self,
        prepared_bind_group: PreparedBindGroup,
        bind_group_layout: BindGroupLayoutDescriptor,
    ) -> MaterialBindingId {
        match *self {
            MaterialBindGroupAllocator::Bindless(_) => {
                panic!(
                    "Bindless resources are incompatible with implementing `as_bind_group` \
                     directly; implement `unprepared_bind_group` instead or disable bindless"
                )
            }
            MaterialBindGroupAllocator::NonBindless(ref mut non_bindless_allocator) => {
                non_bindless_allocator.allocate_prepared(prepared_bind_group, bind_group_layout)
            }
        }
    }

    /// Deallocates the material with the given binding ID.
    ///
    /// Any resources that are no longer referenced are removed from the slab.
    pub fn free(&mut self, material_binding_id: MaterialBindingId) {
        match *self {
            MaterialBindGroupAllocator::Bindless(
                ref mut material_bind_group_bindless_allocator,
            ) => material_bind_group_bindless_allocator.free(material_binding_id),
            MaterialBindGroupAllocator::NonBindless(
                ref mut material_bind_group_non_bindless_allocator,
            ) => material_bind_group_non_bindless_allocator.free(material_binding_id),
        }
    }

    /// Attempts to replace the given `existing_binding_id` with a new material,
    /// without actually reallocating any GPU resources.
    ///
    /// That is, if the given `unprepared_bind_group` represents exactly the
    /// same GPU resource bindings as the already-allocated material, this
    /// method updates the POD only and returns true. Otherwise, if the supplied
    /// bind group represents different resources, this method returns false.
    ///
    /// This is an optional fast path to handle POD-only material changes. If
    /// this method returns false, then the caller should call [`Self::free`]
    /// and [`Self::allocate_unprepared`] to reallocate the GPU resources.
    pub fn try_update_data(
        &mut self,
        existing_binding_id: MaterialBindingId,
        unprepared_bind_group: &mut BindGroupBuilder,
    ) -> bool {
        match *self {
            MaterialBindGroupAllocator::Bindless(
                ref mut material_bind_group_bindless_allocator,
            ) => material_bind_group_bindless_allocator
                .try_update_data(existing_binding_id, unprepared_bind_group),
            MaterialBindGroupAllocator::NonBindless(_) => {
                // TODO: Have a fast path for non-bindless materials.
                false
            }
        }
    }

    /// Recreates any bind groups corresponding to slabs that have been modified
    /// since last calling [`MaterialBindGroupAllocator::prepare_bind_groups`].
    pub fn prepare_bind_groups(
        &mut self,
        render_device: &RenderDevice,
        pipeline_cache: &PipelineCache,
        fallback_bindless_resources: &FallbackBindlessResources,
        fallback_image: &FallbackImage,
        fallback_buffer: &FallbackBuffer,
        shader_buffer_assets: &RenderAssets<GpuShaderBuffer>,
        changed_shader_buffers: &RenderChangedShaderBuffers,
    ) {
        match *self {
            MaterialBindGroupAllocator::Bindless(
                ref mut material_bind_group_bindless_allocator,
            ) => material_bind_group_bindless_allocator.prepare_bind_groups(
                render_device,
                pipeline_cache,
                fallback_bindless_resources,
                fallback_image,
                shader_buffer_assets,
                changed_shader_buffers,
            ),
            MaterialBindGroupAllocator::NonBindless(
                ref mut material_bind_group_non_bindless_allocator,
            ) => material_bind_group_non_bindless_allocator.prepare_bind_groups(
                render_device,
                pipeline_cache,
                fallback_buffer,
                shader_buffer_assets,
                changed_shader_buffers,
            ),
        }
    }

    /// Uploads the contents of all buffers that this
    /// [`MaterialBindGroupAllocator`] manages to the GPU.
    ///
    /// Non-bindless allocators don't currently manage any buffers, so this
    /// method only has an effect for bindless allocators.
    pub fn write_buffers(&mut self, render_device: &RenderDevice, render_queue: &RenderQueue) {
        match *self {
            MaterialBindGroupAllocator::Bindless(
                ref mut material_bind_group_bindless_allocator,
            ) => material_bind_group_bindless_allocator.write_buffers(render_device, render_queue),
            MaterialBindGroupAllocator::NonBindless(_) => {
                // Not applicable.
            }
        }
    }

    /// Get number of allocated slabs for bindless material, returns 0 if it is
    /// [`Self::NonBindless`].
    pub fn slab_count(&self) -> usize {
        match self {
            Self::Bindless(bless) => bless.slabs.len(),
            Self::NonBindless(_) => 0,
        }
    }

    /// Get total size of slabs allocated for bindless material, returns 0 if it is
    /// [`Self::NonBindless`].
    pub fn slabs_size(&self) -> usize {
        match self {
            Self::Bindless(bless) => bless
                .slabs
                .iter()
                .flat_map(|slab| {
                    slab.data_buffers
                        .iter()
                        .map(|(_, buffer)| buffer.buffer.len())
                })
                .sum(),
            Self::NonBindless(_) => 0,
        }
    }

    /// Get number of bindless material allocations in slabs, returns 0 if it is
    /// [`Self::NonBindless`].
    pub fn allocations(&self) -> u64 {
        match self {
            Self::Bindless(bless) => bless
                .slabs
                .iter()
                .map(|slab| u64::from(slab.allocated_resource_count))
                .sum(),
            Self::NonBindless(_) => 0,
        }
    }
}

impl MaterialBindlessIndexTable {
    /// Creates a new [`MaterialBindlessIndexTable`] for a single slab.
    fn new(
        bindless_index_table_descriptor: &BindlessIndexTableDescriptor,
    ) -> MaterialBindlessIndexTable {
        // Preallocate space for one bindings table, so that there will always be a buffer.
        let mut buffer = RetainedRawBufferVec::new(BufferUsages::STORAGE);
        for _ in *bindless_index_table_descriptor.indices.start
            ..*bindless_index_table_descriptor.indices.end
        {
            buffer.push(0);
        }

        MaterialBindlessIndexTable {
            buffer,
            index_range: bindless_index_table_descriptor.indices.clone(),
            binding_number: bindless_index_table_descriptor.binding_number,
        }
    }

    /// Returns the bindings in the binding index table.
    ///
    /// If the current [`MaterialBindlessIndexTable::index_range`] is M..N, then
    /// element *i* of the returned binding index table contains the slot of the
    /// bindless resource with bindless index *i* + M.
    fn get(&self, slot: MaterialBindGroupSlot) -> &[u32] {
        let struct_size = *self.index_range.end as usize - *self.index_range.start as usize;
        let start = struct_size * slot.0 as usize;
        &self.buffer.values()[start..(start + struct_size)]
    }

    /// Returns a single binding from the binding index table.
    fn get_binding(
        &self,
        slot: MaterialBindGroupSlot,
        bindless_index: BindlessIndex,
    ) -> Option<u32> {
        if bindless_index < self.index_range.start || bindless_index >= self.index_range.end {
            return None;
        }
        self.get(slot)
            .get((*bindless_index - *self.index_range.start) as usize)
            .copied()
    }

    fn table_length(&self) -> u32 {
        self.index_range.end.0 - self.index_range.start.0
    }

    /// Updates the binding index table for a single material.
    ///
    /// The `allocated_resource_slots` map contains a mapping from the
    /// [`BindlessIndex`] of each resource that the material references to the
    /// slot that that resource occupies in the appropriate binding array. This
    /// method serializes that map into a binding index table that the shader
    /// can read.
    fn set(
        &mut self,
        slot: MaterialBindGroupSlot,
        allocated_resource_slots: &HashMap<BindlessIndex, u32>,
    ) {
        let table_len = self.table_length() as usize;
        let range = (slot.0 as usize * table_len)..((slot.0 as usize + 1) * table_len);
        while self.buffer.len() < range.end {
            self.buffer.push(0);
        }

        for (&bindless_index, &resource_slot) in allocated_resource_slots {
            if self.index_range.contains(&bindless_index) {
                self.buffer.set(
                    *bindless_index + range.start as u32 - *self.index_range.start,
                    resource_slot,
                );
            }
        }

        // Mark the buffer as needing to be recreated, in case we grew it.
        self.buffer.dirty = BufferDirtyState::NeedsReserve;
    }

    /// Returns the [`BindGroupEntry`] for the index table itself.
    fn bind_group_entry(&self) -> BindGroupEntry<'_> {
        BindGroupEntry {
            binding: *self.binding_number,
            resource: self
                .buffer
                .buffer()
                .expect("Bindings buffer must exist")
                .as_entire_binding(),
        }
    }
}

impl<T> RetainedRawBufferVec<T>
where
    T: Pod,
{
    /// Creates a new empty [`RetainedRawBufferVec`] supporting the given
    /// [`BufferUsages`].
    fn new(buffer_usages: BufferUsages) -> RetainedRawBufferVec<T> {
        RetainedRawBufferVec {
            buffer: RawBufferVec::new(buffer_usages),
            dirty: BufferDirtyState::NeedsUpload,
        }
    }

    /// Recreates the GPU backing buffer if needed.
    fn prepare(&mut self, render_device: &RenderDevice) {
        match self.dirty {
            BufferDirtyState::Clean | BufferDirtyState::NeedsUpload => {}
            BufferDirtyState::NeedsReserve => {
                let capacity = self.buffer.len();
                self.buffer.reserve(capacity, render_device);
                self.dirty = BufferDirtyState::NeedsUpload;
            }
        }
    }

    /// Writes the current contents of the buffer to the GPU if necessary.
    fn write(&mut self, render_device: &RenderDevice, render_queue: &RenderQueue) {
        match self.dirty {
            BufferDirtyState::Clean => {}
            BufferDirtyState::NeedsReserve | BufferDirtyState::NeedsUpload => {
                self.buffer.write_buffer(render_device, render_queue);
                self.dirty = BufferDirtyState::Clean;
            }
        }
    }
}

impl MaterialBindGroupBindlessAllocator {
    /// Creates a new [`MaterialBindGroupBindlessAllocator`] managing the data
    /// for a single bindless material.
    fn new(
        render_device: &RenderDevice,
        label: &'static str,
        bindless_descriptor: BindlessDescriptor,
        bind_group_layout: BindGroupLayoutDescriptor,
        slab_capacity: Option<BindlessSlabResourceLimit>,
    ) -> MaterialBindGroupBindlessAllocator {
        let fallback_buffers = bindless_descriptor
            .buffers
            .iter()
            .map(|bindless_buffer_descriptor| {
                (
                    bindless_buffer_descriptor.bindless_index,
                    render_device.create_buffer(&BufferDescriptor {
                        label: Some("bindless fallback buffer"),
                        size: match bindless_buffer_descriptor.size {
                            Some(size) => size as u64,
                            None => DEFAULT_BINDLESS_FALLBACK_BUFFER_SIZE,
                        },
                        usage: BufferUsages::STORAGE,
                        mapped_at_creation: false,
                    }),
                )
            })
            .collect();

        MaterialBindGroupBindlessAllocator {
            label,
            slabs: vec![],
            bind_group_layout,
            bindless_descriptor,
            fallback_buffers,
            slab_capacity: slab_capacity
                .expect("Non-bindless materials should use the non-bindless allocator")
                .resolve(),
        }
    }

    /// Allocates the resources for a single material into a slab and returns
    /// the resulting ID.
    ///
    /// The returned [`MaterialBindingId`] can later be used to fetch the slab
    /// that was used.
    ///
    /// This function can't fail. If all slabs are full, then a new slab is
    /// created, and the material is allocated into it.
    fn allocate_unprepared(
        &mut self,
        unprepared_bind_group: &mut BindGroupBuilder,
    ) -> MaterialBindingId {
        for (slab_index, slab) in self.slabs.iter_mut().enumerate() {
            trace!("Trying to allocate in slab {}", slab_index);
            if let Ok(slot) = slab.try_allocate(unprepared_bind_group, self.slab_capacity) {
                return MaterialBindingId {
                    group: MaterialBindGroupIndex(slab_index as u32),
                    slot,
                };
            }
        }

        let group = MaterialBindGroupIndex(self.slabs.len() as u32);
        self.slabs
            .push(MaterialBindlessSlab::new(&self.bindless_descriptor));

        // Allocate into the newly-pushed slab.
        let Ok(slot) = self
            .slabs
            .last_mut()
            .expect("We just pushed a slab")
            .try_allocate(unprepared_bind_group, self.slab_capacity)
        else {
            panic!("An allocation into an empty slab should always succeed")
        };

        MaterialBindingId { group, slot }
    }

    /// Attempts to replace the given `existing_binding_id` with a new material,
    /// without actually reallocating any GPU resources.
    ///
    /// See [`MaterialBindGroupAllocator::try_update_data`] for more
    /// information.
    fn try_update_data(
        &mut self,
        existing_binding_id: MaterialBindingId,
        unprepared_bind_group: &mut BindGroupBuilder,
    ) -> bool {
        self.slabs
            .get_mut(existing_binding_id.group.0 as usize)
            .is_some_and(|slab| {
                slab.try_update_data(existing_binding_id.slot, unprepared_bind_group)
            })
    }

    /// Deallocates the material with the given binding ID.
    ///
    /// Any resources that are no longer referenced are removed from the slab.
    fn free(&mut self, material_binding_id: MaterialBindingId) {
        self.slabs
            .get_mut(material_binding_id.group.0 as usize)
            .expect("Slab should exist")
            .free(material_binding_id.slot, &self.bindless_descriptor);
    }

    /// Returns the slab with the given bind group index.
    ///
    /// A [`MaterialBindGroupIndex`] can be fetched from a
    /// [`MaterialBindingId`].
    fn get(&self, group: MaterialBindGroupIndex) -> Option<&MaterialBindlessSlab> {
        self.slabs.get(group.0 as usize)
    }

    /// Recreates any bind groups corresponding to slabs that have been modified
    /// since last calling
    /// [`MaterialBindGroupBindlessAllocator::prepare_bind_groups`].
    fn prepare_bind_groups(
        &mut self,
        render_device: &RenderDevice,
        pipeline_cache: &PipelineCache,
        fallback_bindless_resources: &FallbackBindlessResources,
        fallback_image: &FallbackImage,
        shader_buffer_assets: &RenderAssets<GpuShaderBuffer>,
        changed_shader_buffers: &RenderChangedShaderBuffers,
    ) {
        for slab in &mut self.slabs {
            slab.prepare(
                render_device,
                pipeline_cache,
                self.label,
                &self.bind_group_layout,
                fallback_bindless_resources,
                &self.fallback_buffers,
                fallback_image,
                shader_buffer_assets,
                changed_shader_buffers,
                &self.bindless_descriptor,
                self.slab_capacity,
            );
        }
    }

    /// Writes any buffers that we're managing to the GPU.
    ///
    /// Currently, this only consists of the bindless index tables.
    fn write_buffers(&mut self, render_device: &RenderDevice, render_queue: &RenderQueue) {
        for slab in &mut self.slabs {
            slab.write_buffer(render_device, render_queue);
        }
    }
}

impl MaterialBindlessSlab {
    /// Attempts to allocate the given unprepared bind group in this slab.
    ///
    /// If the allocation succeeds, this method returns the slot that the
    /// allocation was placed in. If the allocation fails because the slab was
    /// full, this method returns the unprepared bind group back to the caller
    /// so that it can try to allocate again.
    fn try_allocate(
        &mut self,
        unprepared_bind_group: &mut BindGroupBuilder,
        slot_capacity: u32,
    ) -> Result<MaterialBindGroupSlot, ()> {
        // Locate pre-existing resources, and determine how many free slots we need.
        let Some(allocation_candidate) =
            self.check_allocation(&unprepared_bind_group.binding_resources)
        else {
            return Err(());
        };

        // Check to see if we have enough free space.
        //
        // As a special case, note that if *nothing* is allocated in this slab,
        // then we always allow a material to be placed in it, regardless of the
        // number of bindings the material has. This is so that, if the
        // platform's maximum bindless count is set too low to hold even a
        // single material, we can still place each material into a separate
        // slab instead of failing outright.
        if self.allocated_resource_count > 0
            && self.allocated_resource_count + allocation_candidate.needed_free_slots
                > slot_capacity
        {
            trace!("Slab is full, can't allocate");
            return Err(());
        }

        // OK, we can allocate in this slab. Assign a slot ID.
        let slot = match self.free_slots.pop() {
            Some(slot) => slot,
            None => {
                // The material bind group slot is packed into 16 bits on
                // the GPU, so spill to a new slab before we would overflow.
                if self.live_allocation_count > 0xFFFF {
                    trace!("Slab material bind group slot would overflow, can't allocate");
                    return Err(());
                }
                MaterialBindGroupSlot(self.live_allocation_count)
            }
        };

        // Bump the live allocation count.
        self.live_allocation_count += 1;

        // Insert the resources into the binding arrays.
        let allocated_resource_slots =
            self.insert_resources(unprepared_bind_group, allocation_candidate);

        // Serialize the allocated resource slots.
        for bindless_index_table in &mut self.bindless_index_tables {
            bindless_index_table.set(slot, &allocated_resource_slots);
        }

        // Invalidate the cached bind group.
        self.bind_group = None;

        Ok(slot)
    }

    /// Attempts to replace the material in the `existing_slot` with a new
    /// material, without actually reallocating any GPU resources.
    ///
    /// See [`MaterialBindGroupAllocator::try_update_data`] for more
    /// information.
    fn try_update_data(
        &mut self,
        existing_slot: MaterialBindGroupSlot,
        unprepared_bind_group: &mut BindGroupBuilder,
    ) -> bool {
        // If the GPU resources have changed, bail. The caller will need to
        // explicitly deallocate and reallocate the bind group.
        if !self
            .bindless_index_tables
            .iter()
            .all(|bindless_index_table| {
                self.gpu_resources_are_unchanged(
                    &unprepared_bind_group.binding_resources,
                    existing_slot,
                    bindless_index_table,
                )
            })
        {
            return false;
        }

        // Update all data in each bindless index table.
        // We can't directly iterate over `Self::bindless_index_tables` for
        // borrow check reasons.
        for bindless_index_table_index in 0..self.bindless_index_tables.len() {
            self.update_data(
                unprepared_bind_group,
                existing_slot,
                bindless_index_table_index,
            );
        }

        true
    }

    /// Gathers the information needed to determine whether the given unprepared
    /// bind group can be allocated in this slab.
    fn check_allocation(
        &self,
        unprepared_binding_resources: &UnpreparedBindingResources,
    ) -> Option<BindlessAllocationCandidate> {
        let mut allocation_candidate = BindlessAllocationCandidate {
            pre_existing_resources: HashMap::default(),
            needed_free_slots: 0,
        };

        for &(bindless_index, ref unprepared_binding_resource) in
            unprepared_binding_resources.iter()
        {
            let bindless_index = BindlessIndex(bindless_index);
            match *unprepared_binding_resource {
                UnpreparedBindingResource::Buffer(ref buffer) => {
                    let Some(binding_array) = self.buffers.get(&bindless_index) else {
                        error!(
                            "Binding array wasn't present for buffer at index {:?}",
                            bindless_index
                        );
                        return None;
                    };
                    match binding_array.find(BindingResourceId::Buffer(buffer.id())) {
                        Some(slot) => {
                            allocation_candidate
                                .pre_existing_resources
                                .insert(bindless_index, slot);
                        }
                        None => allocation_candidate.needed_free_slots += 1,
                    }
                }

                UnpreparedBindingResource::ShaderBuffer(ref shader_buffer) => {
                    let Some(binding_array) = self.shader_buffers.get(&bindless_index) else {
                        error!(
                            "Binding array wasn't present for shader buffer at index {:?}",
                            bindless_index
                        );
                        return None;
                    };
                    match binding_array.find(BindingResourceId::ShaderBuffer(shader_buffer.id())) {
                        Some(slot) => {
                            allocation_candidate
                                .pre_existing_resources
                                .insert(bindless_index, slot);
                        }
                        None => allocation_candidate.needed_free_slots += 1,
                    }
                }

                UnpreparedBindingResource::Data(_) => {
                    // The size of a data buffer is unlimited.
                }

                UnpreparedBindingResource::TextureView(
                    texture_view_dimension,
                    ref texture_view,
                ) => {
                    let bindless_resource_type = BindlessResourceType::from(texture_view_dimension);
                    match self
                        .textures
                        .get(&bindless_resource_type)
                        .expect("Missing binding array for texture")
                        .find(BindingResourceId::TextureView(
                            texture_view_dimension,
                            texture_view.id(),
                        )) {
                        Some(slot) => {
                            allocation_candidate
                                .pre_existing_resources
                                .insert(bindless_index, slot);
                        }
                        None => {
                            allocation_candidate.needed_free_slots += 1;
                        }
                    }
                }

                UnpreparedBindingResource::Sampler(sampler_binding_type, ref sampler) => {
                    let bindless_resource_type = BindlessResourceType::from(sampler_binding_type);
                    match self
                        .samplers
                        .get(&bindless_resource_type)
                        .expect("Missing binding array for sampler")
                        .find(BindingResourceId::Sampler(sampler.id()))
                    {
                        Some(slot) => {
                            allocation_candidate
                                .pre_existing_resources
                                .insert(bindless_index, slot);
                        }
                        None => {
                            allocation_candidate.needed_free_slots += 1;
                        }
                    }
                }
            }
        }

        Some(allocation_candidate)
    }

    /// Inserts the bind group resources described by the given
    /// [`BindGroupBuilder`] into this slab.
    ///
    /// Returns a table that maps the bindless index of each resource to its
    /// slot in its binding array.
    fn insert_resources(
        &mut self,
        unprepared_bind_group: &BindGroupBuilder,
        allocation_candidate: BindlessAllocationCandidate,
    ) -> HashMap<BindlessIndex, u32> {
        let mut allocated_resource_slots = HashMap::default();

        for (bindless_index, unprepared_binding_resource) in unprepared_bind_group.iter() {
            let bindless_index = BindlessIndex(*bindless_index);

            let pre_existing_slot = allocation_candidate
                .pre_existing_resources
                .get(&bindless_index);

            // Otherwise, we need to insert it anew.
            let binding_resource_id = BindingResourceId::from(unprepared_binding_resource);
            let increment_allocated_resource_count = match unprepared_binding_resource {
                UnpreparedBindingResource::Buffer(buffer) => {
                    let slot = self
                        .buffers
                        .get_mut(&bindless_index)
                        .expect("Buffer binding array should exist")
                        .insert(binding_resource_id, buffer.clone());
                    allocated_resource_slots.insert(bindless_index, slot);

                    debug_assert!(pre_existing_slot
                        .is_none_or(|pre_existing_slot| *pre_existing_slot == slot));
                    pre_existing_slot.is_none()
                }

                UnpreparedBindingResource::ShaderBuffer(shader_buffer) => {
                    let slot = self
                        .shader_buffers
                        .get_mut(&bindless_index)
                        .expect("Shader buffer binding array should exist")
                        .insert(binding_resource_id, shader_buffer.clone());
                    allocated_resource_slots.insert(bindless_index, slot);

                    debug_assert!(pre_existing_slot
                        .is_none_or(|pre_existing_slot| *pre_existing_slot == slot));
                    pre_existing_slot.is_none()
                }

                UnpreparedBindingResource::Data(data_range) => {
                    if pre_existing_slot.is_some() {
                        panic!("Data buffers can't be deduplicated")
                    }

                    let slot = self
                        .data_buffers
                        .get_mut(&bindless_index)
                        .expect("Data buffer binding array should exist")
                        .insert(
                            &unprepared_bind_group.data_buffer
                                [(data_range.start as usize)..(data_range.end as usize)],
                        );
                    allocated_resource_slots.insert(bindless_index, slot);
                    false
                }

                UnpreparedBindingResource::TextureView(texture_view_dimension, texture_view) => {
                    let bindless_resource_type =
                        BindlessResourceType::from(*texture_view_dimension);
                    let slot = self
                        .textures
                        .get_mut(&bindless_resource_type)
                        .expect("Texture array should exist")
                        .insert(binding_resource_id, texture_view.clone());
                    allocated_resource_slots.insert(bindless_index, slot);

                    debug_assert!(pre_existing_slot
                        .is_none_or(|pre_existing_slot| *pre_existing_slot == slot));
                    pre_existing_slot.is_none()
                }

                UnpreparedBindingResource::Sampler(sampler_binding_type, sampler) => {
                    let bindless_resource_type = BindlessResourceType::from(*sampler_binding_type);
                    let slot = self
                        .samplers
                        .get_mut(&bindless_resource_type)
                        .expect("Sampler should exist")
                        .insert(binding_resource_id, sampler.clone());
                    allocated_resource_slots.insert(bindless_index, slot);

                    debug_assert!(pre_existing_slot
                        .is_none_or(|pre_existing_slot| *pre_existing_slot == slot));
                    pre_existing_slot.is_none()
                }
            };

            // Bump the allocated resource count.
            if increment_allocated_resource_count {
                self.allocated_resource_count += 1;
            }
        }

        allocated_resource_slots
    }

    /// Removes the material allocated in the given slot, with the given
    /// descriptor, from this slab.
    fn free(&mut self, slot: MaterialBindGroupSlot, bindless_descriptor: &BindlessDescriptor) {
        // Loop through each binding.
        for (bindless_index, bindless_resource_type) in
            bindless_descriptor.resources.iter().enumerate()
        {
            let bindless_index = BindlessIndex::from(bindless_index as u32);
            let Some(bindless_index_table) = self.get_bindless_index_table(bindless_index) else {
                continue;
            };
            let Some(bindless_binding) = bindless_index_table.get_binding(slot, bindless_index)
            else {
                continue;
            };

            // Free the binding. If the resource in question was anything other
            // than a data buffer, then it has a reference count and
            // consequently we need to decrement it.
            let decrement_allocated_resource_count = match *bindless_resource_type {
                BindlessResourceType::None => false,
                BindlessResourceType::Buffer => self
                    .buffers
                    .get_mut(&bindless_index)
                    .expect("Buffer should exist with that bindless index")
                    .remove(bindless_binding),
                BindlessResourceType::ShaderBuffer => self
                    .shader_buffers
                    .get_mut(&bindless_index)
                    .expect("Shader buffer should exist with that bindless index")
                    .remove(bindless_binding),
                BindlessResourceType::DataBuffer => {
                    self.data_buffers
                        .get_mut(&bindless_index)
                        .expect("Data buffer should exist with that bindless index")
                        .remove(bindless_binding);
                    false
                }
                BindlessResourceType::SamplerFiltering
                | BindlessResourceType::SamplerNonFiltering
                | BindlessResourceType::SamplerComparison => self
                    .samplers
                    .get_mut(bindless_resource_type)
                    .expect("Sampler array should exist")
                    .remove(bindless_binding),
                BindlessResourceType::Texture1d
                | BindlessResourceType::Texture2d
                | BindlessResourceType::Texture2dArray
                | BindlessResourceType::Texture3d
                | BindlessResourceType::TextureCube
                | BindlessResourceType::TextureCubeArray => self
                    .textures
                    .get_mut(bindless_resource_type)
                    .expect("Texture array should exist")
                    .remove(bindless_binding),
            };

            // If the slot is now free, decrement the allocated resource
            // count.
            if decrement_allocated_resource_count {
                self.allocated_resource_count -= 1;
            }
        }

        // Invalidate the cached bind group.
        self.bind_group = None;

        // Release the slot ID.
        self.free_slots.push(slot);
        self.live_allocation_count -= 1;
    }

    /// Returns true if all the GPU resources in the given
    /// [`UnpreparedBindingResources`] described by the given bindless index
    /// table are identical to the resources referenced in the
    /// `bind_group_slot`.
    ///
    /// The [`Self::try_update_data`] fast path uses this method to determine
    /// whether the fast path can be used.
    fn gpu_resources_are_unchanged(
        &self,
        binding_resources: &UnpreparedBindingResources,
        bind_group_slot: MaterialBindGroupSlot,
        bindless_index_table: &MaterialBindlessIndexTable,
    ) -> bool {
        for (bindless_index, unprepared_binding_resource) in binding_resources.iter() {
            let bindless_index = BindlessIndex(*bindless_index);
            let binding_resource_id = BindingResourceId::from(unprepared_binding_resource);

            match unprepared_binding_resource {
                UnpreparedBindingResource::Buffer(_)
                | UnpreparedBindingResource::ShaderBuffer(_) => {
                    // Ignore this buffer if it's not in our bindless index
                    // table.
                    if bindless_index_table
                        .get_binding(bind_group_slot, bindless_index)
                        .is_some_and(|binding_index| {
                            // Make sure that the slot references the correct
                            // buffer.
                            !self
                                .buffers
                                .get(&bindless_index)
                                .is_some_and(|binding_array| {
                                    binding_array
                                        .slot_references(binding_index, binding_resource_id)
                                })
                        })
                    {
                        return false;
                    }
                }

                UnpreparedBindingResource::TextureView(texture_view_dimension, _) => {
                    let bindless_resource_type =
                        BindlessResourceType::from(*texture_view_dimension);
                    // Ignore this texture if it's not in our bindless index
                    // table.
                    if bindless_index_table
                        .get_binding(bind_group_slot, bindless_index)
                        .is_some_and(|binding_index| {
                            // Make sure that the slot references the correct
                            // texture.
                            !self.textures.get(&bindless_resource_type).is_some_and(
                                |binding_array| {
                                    binding_array
                                        .slot_references(binding_index, binding_resource_id)
                                },
                            )
                        })
                    {
                        return false;
                    }
                }

                UnpreparedBindingResource::Sampler(sampler_binding_type, _) => {
                    let bindless_resource_type = BindlessResourceType::from(*sampler_binding_type);
                    // Ignore this sampler if it's not in our bindless index
                    // table.
                    if bindless_index_table
                        .get_binding(bind_group_slot, bindless_index)
                        .is_some_and(|binding_index| {
                            // Make sure that the slot references the correct
                            // sampler.
                            !self.samplers.get(&bindless_resource_type).is_some_and(
                                |binding_array| {
                                    binding_array
                                        .slot_references(binding_index, binding_resource_id)
                                },
                            )
                        })
                    {
                        return false;
                    }
                }

                UnpreparedBindingResource::Data(_) => {
                    // Ignore POD.
                }
            }
        }

        true
    }

    // TODO: In the future, this could be `&self`, which would allow this method
    // to operate in parallel. We would probably want `Self::data_buffers` to be
    // `AtomicSparseBufferVec`s in that case.

    /// Replaces the POD ([`UnpreparedBindingResource::Data`]) with new data,
    /// without touching the allocation.
    ///
    /// It's the caller's responsibility to ensure that this is a reasonable
    /// thing to do (probably by calling [`Self::gpu_resources_are_unchanged`]).
    fn update_data(
        &mut self,
        unprepared_bind_group: &BindGroupBuilder,
        bind_group_slot: MaterialBindGroupSlot,
        bindless_index_table_index: usize,
    ) {
        let bindless_index_table = &self.bindless_index_tables[bindless_index_table_index];

        // Go over the index table looking for POD.
        for (bindless_index, unprepared_binding_resource) in
            unprepared_bind_group.binding_resources.iter()
        {
            let bindless_index = BindlessIndex(*bindless_index);

            match unprepared_binding_resource {
                // Ignore all resources other than data.
                UnpreparedBindingResource::Buffer(_)
                | UnpreparedBindingResource::ShaderBuffer(_)
                | UnpreparedBindingResource::TextureView(..)
                | UnpreparedBindingResource::Sampler(..) => {}

                UnpreparedBindingResource::Data(data_range) => {
                    // Ignore this resource unless it represents data we're
                    // responsible for.
                    let Some(binding_index) =
                        bindless_index_table.get_binding(bind_group_slot, bindless_index)
                    else {
                        continue;
                    };

                    // Poke the new data into the buffer.
                    self.data_buffers
                        .get_mut(&bindless_index)
                        .expect("Data buffer should exist if we're in the fast path")
                        .set(
                            binding_index,
                            &unprepared_bind_group.data_buffer
                                [(data_range.start as usize)..(data_range.end as usize)],
                        );
                }
            }
        }
    }

    /// Recreates the bind group and bindless index table buffer if necessary.
    fn prepare(
        &mut self,
        render_device: &RenderDevice,
        pipeline_cache: &PipelineCache,
        label: &'static str,
        bind_group_layout: &BindGroupLayoutDescriptor,
        fallback_bindless_resources: &FallbackBindlessResources,
        fallback_buffers: &HashMap<BindlessIndex, Buffer>,
        fallback_image: &FallbackImage,
        shader_buffer_assets: &RenderAssets<GpuShaderBuffer>,
        changed_shader_buffers: &RenderChangedShaderBuffers,
        bindless_descriptor: &BindlessDescriptor,
        slab_capacity: u32,
    ) {
        self.invalidate_bind_group_for_changed_shader_buffers_if_needed(changed_shader_buffers);

        // Create the bindless index table buffers if needed.
        for bindless_index_table in &mut self.bindless_index_tables {
            bindless_index_table.buffer.prepare(render_device);
        }

        // Create any data buffers we were managing if necessary.
        for data_buffer in self.data_buffers.values_mut() {
            data_buffer.buffer.prepare(render_device);
        }

        // Create the bind group if needed.
        self.prepare_bind_group(
            render_device,
            pipeline_cache,
            label,
            bind_group_layout,
            fallback_bindless_resources,
            fallback_buffers,
            fallback_image,
            shader_buffer_assets,
            bindless_descriptor,
            slab_capacity,
        );
    }

    /// Recreates the bind group if this slab has been changed since the last
    /// time we created it.
    fn prepare_bind_group(
        &mut self,
        render_device: &RenderDevice,
        pipeline_cache: &PipelineCache,
        label: &'static str,
        bind_group_layout: &BindGroupLayoutDescriptor,
        fallback_bindless_resources: &FallbackBindlessResources,
        fallback_buffers: &HashMap<BindlessIndex, Buffer>,
        fallback_image: &FallbackImage,
        shader_buffer_assets: &RenderAssets<GpuShaderBuffer>,
        bindless_descriptor: &BindlessDescriptor,
        slab_capacity: u32,
    ) {
        // If the bind group is clean, then do nothing.
        if self.bind_group.is_some() {
            return;
        }

        // Determine whether we need to pad out our binding arrays with dummy
        // resources.
        let required_binding_array_size = if render_device
            .features()
            .contains(WgpuFeatures::PARTIALLY_BOUND_BINDING_ARRAY)
        {
            None
        } else {
            Some(slab_capacity)
        };

        let binding_resource_arrays = self.create_binding_resource_arrays(
            fallback_bindless_resources,
            fallback_buffers,
            fallback_image,
            shader_buffer_assets,
            bindless_descriptor,
            required_binding_array_size,
        );

        let mut bind_group_entries: Vec<_> = self
            .bindless_index_tables
            .iter()
            .map(|bindless_index_table| bindless_index_table.bind_group_entry())
            .collect();

        for &(&binding, ref binding_resource_array) in binding_resource_arrays.iter() {
            bind_group_entries.push(BindGroupEntry {
                binding,
                resource: match *binding_resource_array {
                    BindingResourceArray::Buffers(ref buffer_bindings) => {
                        BindingResource::BufferArray(&buffer_bindings[..])
                    }
                    BindingResourceArray::TextureViews(ref texture_views) => {
                        BindingResource::TextureViewArray(&texture_views[..])
                    }
                    BindingResourceArray::Samplers(ref samplers) => {
                        BindingResource::SamplerArray(&samplers[..])
                    }
                },
            });
        }

        // Create bind group entries for any data buffers we're managing.
        for data_buffer in self.data_buffers.values() {
            bind_group_entries.push(BindGroupEntry {
                binding: *data_buffer.binding_number,
                resource: data_buffer
                    .buffer
                    .buffer()
                    .expect("Backing data buffer must have been uploaded by now")
                    .as_entire_binding(),
            });
        }

        self.bind_group = Some(render_device.create_bind_group(
            Some(label),
            &pipeline_cache.get_bind_group_layout(bind_group_layout),
            &bind_group_entries,
        ));
    }

    /// If this bind group contains a [`ShaderBuffer`], and the shader buffer
    /// has changed, marks the bind group as invalid so that it'll be rebuilt.
    ///
    /// Shader buffers are assets that can legally be updated to point to a
    /// different raw buffer. This will happen when, for example, the shader
    /// buffer grows. We must detect this circumstance and rebuild the bind
    /// group when it happens.
    fn invalidate_bind_group_for_changed_shader_buffers_if_needed(
        &mut self,
        changed_shader_buffers: &RenderChangedShaderBuffers,
    ) {
        if !changed_shader_buffers.is_empty()
            && self.bind_group.is_some()
            && self.shader_buffers.values().any(|buffer| {
                matches!(buffer.resource_type, BindlessResourceType::ShaderBuffer)
                    && buffer.resource_to_slot.keys().any(|binding_resource_id| {
                        matches!(
                            *binding_resource_id,
                            BindingResourceId::ShaderBuffer(ref shader_buffer) if
                                changed_shader_buffers.contains(shader_buffer)
                        )
                    })
            })
        {
            self.bind_group = None;
        }
    }

    /// Writes any buffers that we're managing to the GPU.
    ///
    /// Currently, this consists of the bindless index table plus any data
    /// buffers we're managing.
    fn write_buffer(&mut self, render_device: &RenderDevice, render_queue: &RenderQueue) {
        for bindless_index_table in &mut self.bindless_index_tables {
            bindless_index_table
                .buffer
                .write(render_device, render_queue);
        }

        for data_buffer in self.data_buffers.values_mut() {
            data_buffer.buffer.write(render_device, render_queue);
        }
    }

    /// Converts our binding arrays into binding resource arrays suitable for
    /// passing to `wgpu`.
    fn create_binding_resource_arrays<'a>(
        &'a self,
        fallback_bindless_resources: &'a FallbackBindlessResources,
        fallback_buffers: &'a HashMap<BindlessIndex, Buffer>,
        fallback_image: &'a FallbackImage,
        shader_buffer_assets: &'a RenderAssets<GpuShaderBuffer>,
        bindless_descriptor: &'a BindlessDescriptor,
        required_binding_array_size: Option<u32>,
    ) -> Vec<(&'a u32, BindingResourceArray<'a>)> {
        let mut binding_resource_arrays = vec![];

        // Build sampler bindings.
        self.create_sampler_binding_resource_arrays(
            &mut binding_resource_arrays,
            fallback_bindless_resources,
            bindless_descriptor,
            required_binding_array_size,
        );

        // Build texture bindings.
        self.create_texture_binding_resource_arrays(
            &mut binding_resource_arrays,
            fallback_image,
            bindless_descriptor,
            required_binding_array_size,
        );

        // Build buffer bindings.
        self.create_buffer_binding_resource_arrays(
            &mut binding_resource_arrays,
            fallback_buffers,
            shader_buffer_assets,
            bindless_descriptor,
            required_binding_array_size,
        );

        binding_resource_arrays
    }

    /// Accumulates sampler binding arrays into binding resource arrays suitable
    /// for passing to `wgpu`.
    fn create_sampler_binding_resource_arrays<'a, 'b>(
        &'a self,
        binding_resource_arrays: &'b mut Vec<(&'a u32, BindingResourceArray<'a>)>,
        fallback_bindless_resources: &'a FallbackBindlessResources,
        bindless_descriptor: &'a BindlessDescriptor,
        required_binding_array_size: Option<u32>,
    ) {
        // We have one binding resource array per sampler type.
        for (bindless_resource_type, fallback_sampler) in [
            (
                BindlessResourceType::SamplerFiltering,
                &fallback_bindless_resources.filtering_sampler,
            ),
            (
                BindlessResourceType::SamplerNonFiltering,
                &fallback_bindless_resources.non_filtering_sampler,
            ),
            (
                BindlessResourceType::SamplerComparison,
                &fallback_bindless_resources.comparison_sampler,
            ),
        ] {
            // Skip resource types not used by this material.
            if !bindless_descriptor
                .resources
                .contains(&bindless_resource_type)
            {
                continue;
            }

            let mut sampler_bindings = vec![];

            match self.samplers.get(&bindless_resource_type) {
                Some(sampler_bindless_binding_array) => {
                    for maybe_bindless_binding in sampler_bindless_binding_array.bindings.iter() {
                        match *maybe_bindless_binding {
                            Some(ref bindless_binding) => {
                                sampler_bindings.push(&*bindless_binding.resource);
                            }
                            None => sampler_bindings.push(&**fallback_sampler),
                        }
                    }
                }

                None => {
                    // Fill with a single fallback sampler.
                    sampler_bindings.push(&**fallback_sampler);
                }
            }

            if let Some(required_binding_array_size) = required_binding_array_size {
                sampler_bindings.extend(iter::repeat_n(
                    &**fallback_sampler,
                    required_binding_array_size as usize - sampler_bindings.len(),
                ));
            }

            let binding_number = bindless_resource_type
                .binding_number()
                .expect("Sampler bindless resource type must have a binding number");

            binding_resource_arrays.push((
                &**binding_number,
                BindingResourceArray::Samplers(sampler_bindings),
            ));
        }
    }

    /// Accumulates texture binding arrays into binding resource arrays suitable
    /// for passing to `wgpu`.
    fn create_texture_binding_resource_arrays<'a, 'b>(
        &'a self,
        binding_resource_arrays: &'b mut Vec<(&'a u32, BindingResourceArray<'a>)>,
        fallback_image: &'a FallbackImage,
        bindless_descriptor: &'a BindlessDescriptor,
        required_binding_array_size: Option<u32>,
    ) {
        for (bindless_resource_type, fallback_image) in [
            (BindlessResourceType::Texture1d, &fallback_image.d1),
            (BindlessResourceType::Texture2d, &fallback_image.d2),
            (
                BindlessResourceType::Texture2dArray,
                &fallback_image.d2_array,
            ),
            (BindlessResourceType::Texture3d, &fallback_image.d3),
            (BindlessResourceType::TextureCube, &fallback_image.cube),
            (
                BindlessResourceType::TextureCubeArray,
                &fallback_image.cube_array,
            ),
        ] {
            // Skip texture types that this material doesn't use.
            if !bindless_descriptor
                .resources
                .contains(&bindless_resource_type)
            {
                continue;
            }

            let mut texture_bindings = vec![];

            let binding_number = bindless_resource_type
                .binding_number()
                .expect("Texture bindless resource type must have a binding number");

            match self.textures.get(&bindless_resource_type) {
                Some(texture_bindless_binding_array) => {
                    for maybe_bindless_binding in texture_bindless_binding_array.bindings.iter() {
                        match *maybe_bindless_binding {
                            Some(ref bindless_binding) => {
                                texture_bindings.push(&*bindless_binding.resource);
                            }
                            None => texture_bindings.push(&*fallback_image.texture_view),
                        }
                    }
                }

                None => {
                    // Fill with a single fallback image.
                    texture_bindings.push(&*fallback_image.texture_view);
                }
            }

            if let Some(required_binding_array_size) = required_binding_array_size {
                texture_bindings.extend(iter::repeat_n(
                    &*fallback_image.texture_view,
                    required_binding_array_size as usize - texture_bindings.len(),
                ));
            }

            binding_resource_arrays.push((
                binding_number,
                BindingResourceArray::TextureViews(texture_bindings),
            ));
        }
    }

    /// Accumulates buffer binding arrays into binding resource arrays suitable
    /// for `wgpu`.
    fn create_buffer_binding_resource_arrays<'a, 'b>(
        &'a self,
        binding_resource_arrays: &'b mut Vec<(&'a u32, BindingResourceArray<'a>)>,
        fallback_buffers: &'a HashMap<BindlessIndex, Buffer>,
        shader_buffer_assets: &'a RenderAssets<GpuShaderBuffer>,
        bindless_descriptor: &'a BindlessDescriptor,
        required_binding_array_size: Option<u32>,
    ) {
        for bindless_buffer_descriptor in bindless_descriptor.buffers.iter() {
            // If this is a raw buffer, create the binding.
            if let Some(buffer_bindless_binding_array) =
                self.buffers.get(&bindless_buffer_descriptor.bindless_index)
            {
                let fallback_buffer = fallback_buffers
                    .get(&bindless_buffer_descriptor.bindless_index)
                    .expect("Fallback buffer should exist");

                let mut buffer_bindings: Vec<_> = buffer_bindless_binding_array
                    .bindings
                    .iter()
                    .map(|maybe_bindless_binding| {
                        let buffer = match *maybe_bindless_binding {
                            None => fallback_buffer,
                            Some(ref bindless_binding) => &bindless_binding.resource,
                        };
                        BufferBinding {
                            buffer,
                            offset: 0,
                            size: None,
                        }
                    })
                    .collect();

                if let Some(required_binding_array_size) = required_binding_array_size {
                    buffer_bindings.extend(iter::repeat_n(
                        BufferBinding {
                            buffer: fallback_buffer,
                            offset: 0,
                            size: None,
                        },
                        required_binding_array_size as usize - buffer_bindings.len(),
                    ));
                }

                binding_resource_arrays.push((
                    &*buffer_bindless_binding_array.binding_number,
                    BindingResourceArray::Buffers(buffer_bindings),
                ));

                continue;
            }

            // If this is a shader buffer, look up the asset ID in the
            // `shader_buffer_assets` and fetch the raw buffer from there.
            if let Some(shader_buffer_bindless_binding_array) = self
                .shader_buffers
                .get(&bindless_buffer_descriptor.bindless_index)
            {
                let fallback_buffer = fallback_buffers
                    .get(&bindless_buffer_descriptor.bindless_index)
                    .expect("Fallback buffer should exist");

                let mut buffer_bindings: Vec<_> = shader_buffer_bindless_binding_array
                    .bindings
                    .iter()
                    .map(|maybe_bindless_binding| {
                        let buffer =
                            match maybe_bindless_binding
                                .as_ref()
                                .and_then(|bindless_binding| {
                                    shader_buffer_assets.get(bindless_binding.resource.id())
                                }) {
                                None => fallback_buffer,
                                Some(gpu_shader_buffer) => &gpu_shader_buffer.buffer,
                            };
                        BufferBinding {
                            buffer,
                            offset: 0,
                            size: None,
                        }
                    })
                    .collect();

                if let Some(required_binding_array_size) = required_binding_array_size {
                    buffer_bindings.extend(iter::repeat_n(
                        BufferBinding {
                            buffer: fallback_buffer,
                            offset: 0,
                            size: None,
                        },
                        required_binding_array_size as usize - buffer_bindings.len(),
                    ));
                }

                binding_resource_arrays.push((
                    &*shader_buffer_bindless_binding_array.binding_number,
                    BindingResourceArray::Buffers(buffer_bindings),
                ));

                continue;
            };

            // If we got here, just ignore the buffer and continue. For example,
            // index buffers are present in `BindlessDescriptor::buffers` but
            // not in `BindlessDescriptor::resources`, so we need to ignore them
            // here.
        }
    }

    /// Returns the [`BindGroup`] corresponding to this slab, if it's been
    /// prepared.
    fn bind_group(&self) -> Option<&BindGroup> {
        self.bind_group.as_ref()
    }

    /// Returns the bindless index table containing the given bindless index.
    fn get_bindless_index_table(
        &self,
        bindless_index: BindlessIndex,
    ) -> Option<&MaterialBindlessIndexTable> {
        let table_index = self
            .bindless_index_tables
            .binary_search_by(|bindless_index_table| {
                if bindless_index < bindless_index_table.index_range.start {
                    Ordering::Less
                } else if bindless_index >= bindless_index_table.index_range.end {
                    Ordering::Greater
                } else {
                    Ordering::Equal
                }
            })
            .ok()?;
        self.bindless_index_tables.get(table_index)
    }
}

impl<R> MaterialBindlessBindingArray<R>
where
    R: GetBindingResourceId,
{
    /// Creates a new [`MaterialBindlessBindingArray`] with the given binding
    /// number, managing resources of the given type.
    fn new(
        binding_number: BindingNumber,
        resource_type: BindlessResourceType,
    ) -> MaterialBindlessBindingArray<R> {
        MaterialBindlessBindingArray {
            binding_number,
            bindings: vec![],
            resource_type,
            resource_to_slot: HashMap::default(),
            free_slots: vec![],
            len: 0,
        }
    }

    /// Returns the slot corresponding to the given resource, if that resource
    /// is located in this binding array.
    ///
    /// If the resource isn't in this binding array, this method returns `None`.
    fn find(&self, binding_resource_id: BindingResourceId) -> Option<u32> {
        self.resource_to_slot.get(&binding_resource_id).copied()
    }

    /// Returns true if the given slot points to the given binding resource and
    /// false otherwise.
    fn slot_references(&self, slot: u32, binding_resource_id: BindingResourceId) -> bool {
        self.bindings
            .get(slot as usize)
            .is_some_and(|maybe_binding| {
                maybe_binding.as_ref().is_some_and(|binding| {
                    binding.resource.binding_resource_id(self.resource_type) == binding_resource_id
                })
            })
    }

    /// Inserts a bindless resource into a binding array and returns the index
    /// of the slot it was inserted into.
    fn insert(&mut self, binding_resource_id: BindingResourceId, resource: R) -> u32 {
        match self.resource_to_slot.entry(binding_resource_id) {
            Entry::Occupied(o) => {
                let slot = *o.get();

                self.bindings[slot as usize]
                    .as_mut()
                    .expect("A slot in the resource_to_slot map should have a value")
                    .ref_count += 1;

                slot
            }
            Entry::Vacant(v) => {
                let slot = self.free_slots.pop().unwrap_or(self.len);
                v.insert(slot);

                if self.bindings.len() < slot as usize + 1 {
                    self.bindings.resize_with(slot as usize + 1, || None);
                }
                debug_assert!(self.bindings[slot as usize].is_none());
                self.bindings[slot as usize] = Some(MaterialBindlessBinding::new(resource));

                self.len += 1;
                slot
            }
        }
    }

    /// Removes a reference to an object from the slot.
    ///
    /// If the reference count dropped to 0 and the object was freed, this
    /// method returns true. If the object was still referenced after removing
    /// it, returns false.
    fn remove(&mut self, slot: u32) -> bool {
        let maybe_binding = &mut self.bindings[slot as usize];
        let binding = maybe_binding
            .as_mut()
            .expect("Attempted to free an already-freed binding");

        binding.ref_count -= 1;
        if binding.ref_count != 0 {
            return false;
        }

        let binding_resource_id = binding.resource.binding_resource_id(self.resource_type);
        self.resource_to_slot.remove(&binding_resource_id);

        *maybe_binding = None;
        self.free_slots.push(slot);
        self.len -= 1;
        true
    }
}

impl<R> MaterialBindlessBinding<R>
where
    R: GetBindingResourceId,
{
    /// Creates a new [`MaterialBindlessBinding`] for a freshly-added resource.
    ///
    /// The reference count is initialized to 1.
    fn new(resource: R) -> MaterialBindlessBinding<R> {
        MaterialBindlessBinding {
            resource,
            ref_count: 1,
        }
    }
}

/// Returns true if the material will *actually* use bindless resources or false
/// if it won't.
///
/// This takes the platform support (or lack thereof) for bindless resources
/// into account.
pub fn material_uses_bindless_resources<M>(render_device: &RenderDevice) -> bool
where
    M: AsBindGroup,
{
    M::bindless_slot_count().is_some_and(|bindless_slot_count| {
        M::bindless_supported(render_device) && bindless_slot_count.resolve() > 1
    })
}

impl MaterialBindlessSlab {
    /// Creates a new [`MaterialBindlessSlab`] for a material with the given
    /// bindless descriptor.
    ///
    /// We use this when no existing slab could hold a material to be allocated.
    fn new(bindless_descriptor: &BindlessDescriptor) -> MaterialBindlessSlab {
        let mut buffers = HashMap::default();
        let mut samplers = HashMap::default();
        let mut textures = HashMap::default();
        let mut data_buffers = HashMap::default();
        let mut shader_buffers = HashMap::default();

        for (bindless_index, bindless_resource_type) in
            bindless_descriptor.resources.iter().enumerate()
        {
            let bindless_index = BindlessIndex(bindless_index as u32);
            match *bindless_resource_type {
                BindlessResourceType::None => {}
                BindlessResourceType::Buffer => {
                    let binding_number = bindless_descriptor
                        .buffers
                        .iter()
                        .find(|bindless_buffer_descriptor| {
                            bindless_buffer_descriptor.bindless_index == bindless_index
                        })
                        .expect(
                            "Bindless buffer descriptor matching that bindless index should be \
                             present",
                        )
                        .binding_number;
                    buffers.insert(
                        bindless_index,
                        MaterialBindlessBindingArray::new(binding_number, *bindless_resource_type),
                    );
                }
                BindlessResourceType::ShaderBuffer => {
                    let binding_number = bindless_descriptor
                        .buffers
                        .iter()
                        .find(|bindless_buffer_descriptor| {
                            bindless_buffer_descriptor.bindless_index == bindless_index
                        })
                        .expect(
                            "Bindless buffer descriptor matching that bindless index should be \
                             present",
                        )
                        .binding_number;
                    shader_buffers.insert(
                        bindless_index,
                        MaterialBindlessBindingArray::new(binding_number, *bindless_resource_type),
                    );
                }
                BindlessResourceType::DataBuffer => {
                    // Copy the data in.
                    let buffer_descriptor = bindless_descriptor
                        .buffers
                        .iter()
                        .find(|bindless_buffer_descriptor| {
                            bindless_buffer_descriptor.bindless_index == bindless_index
                        })
                        .expect(
                            "Bindless buffer descriptor matching that bindless index should be \
                             present",
                        );
                    data_buffers.insert(
                        bindless_index,
                        MaterialDataBuffer::new(
                            buffer_descriptor.binding_number,
                            buffer_descriptor
                                .size
                                .expect("Data buffers should have a size")
                                as u32,
                        ),
                    );
                }
                BindlessResourceType::SamplerFiltering
                | BindlessResourceType::SamplerNonFiltering
                | BindlessResourceType::SamplerComparison => {
                    samplers.insert(
                        *bindless_resource_type,
                        MaterialBindlessBindingArray::new(
                            *bindless_resource_type.binding_number().unwrap(),
                            *bindless_resource_type,
                        ),
                    );
                }
                BindlessResourceType::Texture1d
                | BindlessResourceType::Texture2d
                | BindlessResourceType::Texture2dArray
                | BindlessResourceType::Texture3d
                | BindlessResourceType::TextureCube
                | BindlessResourceType::TextureCubeArray => {
                    textures.insert(
                        *bindless_resource_type,
                        MaterialBindlessBindingArray::new(
                            *bindless_resource_type.binding_number().unwrap(),
                            *bindless_resource_type,
                        ),
                    );
                }
            }
        }

        let bindless_index_tables = bindless_descriptor
            .index_tables
            .iter()
            .map(MaterialBindlessIndexTable::new)
            .collect();

        MaterialBindlessSlab {
            bind_group: None,
            bindless_index_tables,
            samplers,
            textures,
            buffers,
            shader_buffers,
            data_buffers,
            free_slots: vec![],
            live_allocation_count: 0,
            allocated_resource_count: 0,
        }
    }
}

/// Creates and inserts the [`FallbackBindlessResources`] and [`FallbackBuffer`].
fn init_fallback_resources(mut commands: Commands, render_device: Res<RenderDevice>) {
    // Create the `FallbackBindlessResources`.
    commands.insert_resource(FallbackBindlessResources {
        filtering_sampler: render_device.create_sampler(&SamplerDescriptor {
            label: Some("fallback filtering sampler"),
            ..default()
        }),
        non_filtering_sampler: render_device.create_sampler(&SamplerDescriptor {
            label: Some("fallback non-filtering sampler"),
            mag_filter: FilterMode::Nearest,
            min_filter: FilterMode::Nearest,
            mipmap_filter: MipmapFilterMode::Nearest,
            ..default()
        }),
        comparison_sampler: render_device.create_sampler(&SamplerDescriptor {
            label: Some("fallback comparison sampler"),
            compare: Some(CompareFunction::Always),
            ..default()
        }),
    });

    // Creates the `FallbackBuffer`.
    commands.insert_resource(FallbackBuffer(render_device.create_buffer(
        &BufferDescriptor {
            label: Some("fallback buffer"),
            size: 1,
            usage: BufferUsages::COPY_SRC
                | BufferUsages::COPY_DST
                | BufferUsages::STORAGE
                | BufferUsages::UNIFORM,
            mapped_at_creation: false,
        },
    )));
}

impl MaterialBindGroupNonBindlessAllocator {
    /// Creates a new [`MaterialBindGroupNonBindlessAllocator`] managing the
    /// bind groups for a single non-bindless material.
    fn new(label: &'static str) -> MaterialBindGroupNonBindlessAllocator {
        MaterialBindGroupNonBindlessAllocator {
            label,
            bind_groups: vec![],
            to_prepare: HashSet::default(),
            free_indices: vec![],
        }
    }

    /// Inserts a bind group, either unprepared or prepared, into this allocator
    /// and returns a [`MaterialBindingId`].
    ///
    /// The returned [`MaterialBindingId`] can later be used to fetch the bind
    /// group.
    fn allocate(&mut self, bind_group: MaterialNonBindlessAllocatedBindGroup) -> MaterialBindingId {
        let group_id = self
            .free_indices
            .pop()
            .unwrap_or(MaterialBindGroupIndex(self.bind_groups.len() as u32));
        if self.bind_groups.len() < *group_id as usize + 1 {
            self.bind_groups
                .resize_with(*group_id as usize + 1, || None);
        }

        if matches!(
            bind_group,
            MaterialNonBindlessAllocatedBindGroup::Unprepared { .. }
        ) {
            self.to_prepare.insert(group_id);
        }

        self.bind_groups[*group_id as usize] = Some(bind_group);

        MaterialBindingId {
            group: group_id,
            slot: default(),
        }
    }

    /// Inserts an unprepared bind group into this allocator and returns a
    /// [`MaterialBindingId`].
    fn allocate_unprepared(
        &mut self,
        unprepared_bind_group: &mut BindGroupBuilder,
        bind_group_layout: BindGroupLayoutDescriptor,
    ) -> MaterialBindingId {
        self.allocate(MaterialNonBindlessAllocatedBindGroup::Unprepared {
            // FIXME: This is slow as it takes all the allocations, meaning that
            // the next material to be prepared is going to have to reallocate
            // the unprepared bind group. But there's a reason for it: we have
            // to hold onto the data allocations so that we can upload them
            // later, in `Self::prepare_bind_groups`. Addressing this will
            // require either moving those allocations to happen earlier or
            // setting up some sort of buffer recycling scheme.
            bind_group: mem::take(unprepared_bind_group),
            layout: bind_group_layout,
        })
    }

    /// Inserts an prepared bind group into this allocator and returns a
    /// [`MaterialBindingId`].
    fn allocate_prepared(
        &mut self,
        prepared_bind_group: PreparedBindGroup,
        bind_group_layout: BindGroupLayoutDescriptor,
    ) -> MaterialBindingId {
        self.allocate(MaterialNonBindlessAllocatedBindGroup::Prepared {
            bind_group: prepared_bind_group,
            layout: bind_group_layout,
            uniform_buffers: vec![],
        })
    }

    /// Deallocates the bind group with the given binding ID.
    fn free(&mut self, binding_id: MaterialBindingId) {
        debug_assert_eq!(binding_id.slot, MaterialBindGroupSlot(0));
        debug_assert!(self.bind_groups[*binding_id.group as usize].is_some());
        self.bind_groups[*binding_id.group as usize] = None;
        self.to_prepare.remove(&binding_id.group);
        self.free_indices.push(binding_id.group);
    }

    /// Returns a wrapper around the bind group with the given index.
    fn get(&self, group: MaterialBindGroupIndex) -> Option<MaterialNonBindlessSlab<'_>> {
        self.bind_groups[group.0 as usize]
            .as_ref()
            .map(|bind_group| match bind_group {
                MaterialNonBindlessAllocatedBindGroup::Prepared { bind_group, .. } => {
                    MaterialNonBindlessSlab::Prepared(bind_group)
                }
                MaterialNonBindlessAllocatedBindGroup::Unprepared { .. } => {
                    MaterialNonBindlessSlab::Unprepared
                }
            })
    }

    /// Prepares any as-yet unprepared bind groups that this allocator is
    /// managing.
    ///
    /// Unprepared bind groups can be added to this allocator with
    /// [`Self::allocate_unprepared`]. Such bind groups will defer being
    /// prepared until the next time this method is called.
    fn prepare_bind_groups(
        &mut self,
        render_device: &RenderDevice,
        pipeline_cache: &PipelineCache,
        fallback_buffer: &FallbackBuffer,
        shader_buffer_assets: &RenderAssets<GpuShaderBuffer>,
        changed_shader_buffers: &RenderChangedShaderBuffers,
    ) {
        self.invalidate_bind_group_for_changed_shader_buffers_if_needed(changed_shader_buffers);

        for bind_group_index in mem::take(&mut self.to_prepare) {
            let Some(MaterialNonBindlessAllocatedBindGroup::Unprepared {
                bind_group: unprepared_bind_group,
                layout: bind_group_layout,
            }) = mem::take(&mut self.bind_groups[*bind_group_index as usize])
            else {
                panic!("Allocation didn't exist or was already prepared");
            };

            // Pack any `Data` into uniform buffers.
            let mut uniform_buffers = vec![];
            for (index, binding) in unprepared_bind_group.binding_resources.iter() {
                let UnpreparedBindingResource::Data(ref data_range) = *binding else {
                    continue;
                };
                let label = format!("material uniform data {}", *index);
                let uniform_buffer = render_device.create_buffer_with_data(&BufferInitDescriptor {
                    label: Some(&label),
                    contents: &unprepared_bind_group.data_buffer
                        [(data_range.start as usize)..(data_range.end as usize)],
                    usage: BufferUsages::COPY_DST | BufferUsages::UNIFORM,
                });
                uniform_buffers.push(uniform_buffer);
            }

            // Create bind group entries.
            let mut bind_group_entries = vec![];
            let mut uniform_buffers_iter = uniform_buffers.iter();
            for (index, binding) in unprepared_bind_group.binding_resources.iter() {
                match *binding {
                    UnpreparedBindingResource::Data(_) => {
                        bind_group_entries.push(BindGroupEntry {
                            binding: *index,
                            resource: uniform_buffers_iter
                                .next()
                                .expect("We should have created uniform buffers for each `Data`")
                                .as_entire_binding(),
                        });
                    }
                    _ => bind_group_entries.push(BindGroupEntry {
                        binding: *index,
                        resource: binding.get_binding(fallback_buffer, shader_buffer_assets),
                    }),
                }
            }

            // Create the bind group.
            let bind_group = render_device.create_bind_group(
                self.label,
                &pipeline_cache.get_bind_group_layout(&bind_group_layout),
                &bind_group_entries,
            );

            self.bind_groups[*bind_group_index as usize] =
                Some(MaterialNonBindlessAllocatedBindGroup::Prepared {
                    bind_group: PreparedBindGroup {
                        bindings: unprepared_bind_group.into_binding_resources(),
                        bind_group,
                    },
                    layout: bind_group_layout,
                    uniform_buffers,
                });
        }
    }

    /// Invalidates any bind group that contains a shader buffer if the shader
    /// buffer has changed identity.
    fn invalidate_bind_group_for_changed_shader_buffers_if_needed(
        &mut self,
        changed_shader_buffers: &RenderChangedShaderBuffers,
    ) {
        // Early out if there are no changed shader buffers. This is important,
        // as scanning for invalidated bind groups can be expensive.
        if changed_shader_buffers.is_empty() {
            return;
        }

        // Look for bind groups that are now invalid.
        for (allocated_bind_group_index, maybe_allocated_bind_group) in
            self.bind_groups.iter_mut().enumerate()
        {
            // We only care about prepared bind groups. Unprepared bind groups
            // store only the asset ID of the shader buffer, which is stable,
            // rather than the buffer itself, which is not.
            let Some(MaterialNonBindlessAllocatedBindGroup::Prepared {
                bind_group: ref mut prepared_bind_group,
                ref mut layout,
                ..
            }) = *maybe_allocated_bind_group
            else {
                continue;
            };

            // Check to see if any shader buffer within the bind group was
            // invalidated.
            if prepared_bind_group
                .bindings
                .iter()
                .all(|(_, binding_resource)| match *binding_resource {
                    OwnedBindingResource::ShaderBuffer(ref shader_buffer) => {
                        !changed_shader_buffers.contains(&shader_buffer.id())
                    }
                    _ => true,
                })
            {
                continue;
            }

            // If we got here, we know the bind group is invalid. Mark it as
            // such by "un-preparing" it. The logic in `prepare_bind_groups`
            // will detect that the bind group is now unprepared and will
            // rebuild it.
            let layout = mem::take(layout);
            *maybe_allocated_bind_group = Some(MaterialNonBindlessAllocatedBindGroup::Unprepared {
                bind_group: prepared_bind_group.unprepare(),
                layout,
            });
            self.to_prepare
                .insert(MaterialBindGroupIndex(allocated_bind_group_index as u32));
        }
    }
}

impl<'a> MaterialSlab<'a> {
    /// Returns the [`BindGroup`] corresponding to this slab, if it's been
    /// prepared.
    ///
    /// You can prepare bind groups by calling
    /// [`MaterialBindGroupAllocator::prepare_bind_groups`]. If the bind group
    /// isn't ready, this method returns `None`.
    pub fn bind_group(&self) -> Option<&'a BindGroup> {
        match self.0 {
            MaterialSlabImpl::Bindless(material_bindless_slab) => {
                material_bindless_slab.bind_group()
            }
            MaterialSlabImpl::NonBindless(MaterialNonBindlessSlab::Prepared(
                prepared_bind_group,
            )) => Some(&prepared_bind_group.bind_group),
            MaterialSlabImpl::NonBindless(MaterialNonBindlessSlab::Unprepared) => None,
        }
    }
}

impl MaterialDataBuffer {
    /// Creates a new [`MaterialDataBuffer`] managing a buffer of elements of
    /// size `aligned_element_size` that will be bound to the given binding
    /// number.
    fn new(binding_number: BindingNumber, aligned_element_size: u32) -> MaterialDataBuffer {
        MaterialDataBuffer {
            binding_number,
            buffer: RetainedRawBufferVec::new(BufferUsages::STORAGE),
            aligned_element_size,
            free_slots: vec![],
            len: 0,
        }
    }

    /// Allocates a slot for a new piece of data, copies the data into that
    /// slot, and returns the slot ID.
    ///
    /// The size of the piece of data supplied to this method must equal the
    /// [`Self::aligned_element_size`] provided to [`MaterialDataBuffer::new`].
    fn insert(&mut self, data: &[u8]) -> u32 {
        // Make sure the data is of the right length.
        debug_assert_eq!(data.len(), self.aligned_element_size as usize);

        // Grab a slot.
        let slot = self.free_slots.pop().unwrap_or(self.len);

        // Calculate the range we're going to copy to.
        let start = slot as usize * self.aligned_element_size as usize;
        let end = (slot as usize + 1) * self.aligned_element_size as usize;

        // Resize the buffer if necessary.
        if self.buffer.len() < end {
            self.buffer.reserve_internal(end);
        }
        while self.buffer.values().len() < end {
            self.buffer.push(0);
        }

        // Copy in the data.
        self.buffer.values_mut()[start..end].copy_from_slice(data);

        // Mark the buffer dirty, and finish up.
        self.len += 1;
        self.buffer.dirty = BufferDirtyState::NeedsReserve;
        slot
    }

    /// Updates the given slot with new data.
    fn set(&mut self, slot: u32, data: &[u8]) {
        // Make sure the data is of the right length.
        debug_assert_eq!(data.len(), self.aligned_element_size as usize);

        // Calculate the range we're going to copy to.
        let start = slot as usize * self.aligned_element_size as usize;
        let end = (slot as usize + 1) * self.aligned_element_size as usize;

        self.buffer.values_mut()[start..end].copy_from_slice(data);

        if self.buffer.dirty == BufferDirtyState::Clean {
            self.buffer.dirty = BufferDirtyState::NeedsUpload;
        }
    }

    /// Marks the given slot as free.
    fn remove(&mut self, slot: u32) {
        self.free_slots.push(slot);
        self.len -= 1;
    }
}

/// Creates and/or recreates any bind groups that contain materials that were
/// modified this frame.
pub fn prepare_material_bind_groups(
    mut allocators: ResMut<MaterialBindGroupAllocators>,
    render_device: Res<RenderDevice>,
    pipeline_cache: Res<PipelineCache>,
    fallback_image: Res<FallbackImage>,
    fallback_buffer: Res<FallbackBuffer>,
    fallback_resources: Res<FallbackBindlessResources>,
    shader_buffer_assets: Res<RenderAssets<GpuShaderBuffer>>,
    changed_shader_buffers: Res<RenderChangedShaderBuffers>,
) {
    for (_, allocator) in allocators.iter_mut() {
        allocator.prepare_bind_groups(
            &render_device,
            &pipeline_cache,
            &fallback_resources,
            &fallback_image,
            &fallback_buffer,
            &shader_buffer_assets,
            &changed_shader_buffers,
        );
    }
}

/// Uploads the contents of all buffers that the [`MaterialBindGroupAllocator`]
/// manages to the GPU.
///
/// Non-bindless allocators don't currently manage any buffers, so this method
/// only has an effect for bindless allocators.
pub fn write_material_bind_group_buffers(
    mut allocators: ResMut<MaterialBindGroupAllocators>,
    render_device: Res<RenderDevice>,
    render_queue: Res<RenderQueue>,
) {
    for (_, allocator) in allocators.iter_mut() {
        allocator.write_buffers(&render_device, &render_queue);
    }
}

impl MaterialBindGroupAllocators {
    /// Adds a new [`MaterialBindGroupAllocator`] to this resource to manage
    /// materials of the given type.
    pub fn add<M>(&mut self, render_device: &RenderDevice)
    where
        M: AsBindGroup + 'static,
    {
        self.insert(
            TypeId::of::<M>(),
            MaterialBindGroupAllocator::new(
                render_device,
                M::label(),
                material_uses_bindless_resources::<M>(render_device)
                    .then(|| M::bindless_descriptor())
                    .flatten(),
                M::bind_group_layout_descriptor(render_device),
                M::bindless_slot_count(),
            ),
        );
    }
}

impl RenderMaterialBindings {
    /// Prepares a material asset for the render world.
    ///
    /// This method allocates the material in the bind group allocators and
    /// returns its [`MaterialBindingId`].
    pub fn prepare_material<M>(
        &mut self,
        material: &M,
        material_id: AssetId<M>,
        material_param: &mut SystemParamItem<'_, '_, M::Param>,
        material_layout: &BindGroupLayoutDescriptor,
        bind_group_allocators: &mut MaterialBindGroupAllocators,
        render_device: &RenderDevice,
        pipeline_cache: &PipelineCache,
        fallback_buffer: &FallbackBuffer,
        shader_buffer_assets: &RenderAssets<GpuShaderBuffer>,
    ) -> Result<MaterialBindingId, PrepareAssetError<M>>
    where
        M: AsBindGroup + Asset + Clone,
    {
        let actual_material_layout = pipeline_cache.get_bind_group_layout(material_layout);

        // Ask the material to build the bind group.
        let result = match material.build_bind_group(
            &actual_material_layout,
            render_device,
            material_param,
            false,
            &mut self.bind_group_builder,
        ) {
            Ok(()) => {
                let bind_group_allocator =
                    bind_group_allocators.get_mut(&TypeId::of::<M>()).unwrap();
                // Allocate or update the material.
                match self.bindings.entry(material_id.into()) {
                    Entry::Occupied(mut occupied_entry) => {
                        // First, try a fast path. If none of the GPU resource
                        // bindings have changed, then we can just update the
                        // POD and don't have to update any of the bindings in
                        // the bind group allocator.
                        let old_binding = *occupied_entry.get();
                        if bind_group_allocator
                            .try_update_data(old_binding, &mut self.bind_group_builder)
                        {
                            Ok(old_binding)
                        } else {
                            // Otherwise, fall back to the slow path. Deallocate
                            // the GPU resource bindings, and reallocate them.
                            bind_group_allocator.free(old_binding);
                            let new_binding = bind_group_allocator
                                .allocate_unprepared(&mut self.bind_group_builder, material_layout);
                            *occupied_entry.get_mut() = new_binding;
                            Ok(new_binding)
                        }
                    }
                    Entry::Vacant(vacant_entry) => Ok(*vacant_entry.insert(
                        bind_group_allocator
                            .allocate_unprepared(&mut self.bind_group_builder, material_layout),
                    )),
                }
            }
            Err(AsBindGroupError::RetryNextUpdate) => {
                Err(PrepareAssetError::RetryNextUpdate((*material).clone()))
            }
            Err(AsBindGroupError::CreateBindGroupDirectly) => {
                match material.as_bind_group(
                    material_layout,
                    render_device,
                    pipeline_cache,
                    fallback_buffer,
                    shader_buffer_assets,
                    material_param,
                ) {
                    Ok(prepared_bind_group) => {
                        let bind_group_allocator =
                            bind_group_allocators.get_mut(&TypeId::of::<M>()).unwrap();
                        // Store the resulting bind group directly in the slot.
                        let material_binding_id = bind_group_allocator
                            .allocate_prepared(prepared_bind_group, material_layout.clone());
                        self.insert(material_id.into(), material_binding_id);
                        Ok(material_binding_id)
                    }
                    Err(AsBindGroupError::RetryNextUpdate) => {
                        Err(PrepareAssetError::RetryNextUpdate((*material).clone()))
                    }
                    Err(other) => Err(PrepareAssetError::AsBindGroupError(other)),
                }
            }
            Err(other) => Err(PrepareAssetError::AsBindGroupError(other)),
        };

        // Clear out the bind group builder (but retain its heap allocations; we
        // don't want to allocate here).
        self.bind_group_builder.clear();

        result
    }

    /// Removes a material asset from the render world.
    pub fn unload_material<M>(
        &mut self,
        source_asset: AssetId<M>,
        bind_group_allocators: &mut MaterialBindGroupAllocators,
    ) where
        M: Asset + AsBindGroup + Clone,
    {
        if let Some(material_binding_id) = self.remove(&source_asset.untyped()) {
            let bind_group_allocator = bind_group_allocators.get_mut(&TypeId::of::<M>()).unwrap();
            bind_group_allocator.free(material_binding_id);
        }
    }
}