bevy_render/mesh/
allocator.rs

1//! Manages mesh vertex and index buffers.
2
3use alloc::vec::Vec;
4use core::{
5    fmt::{self, Display, Formatter},
6    ops::Range,
7};
8use nonmax::NonMaxU32;
9
10use bevy_app::{App, Plugin};
11use bevy_asset::AssetId;
12use bevy_derive::{Deref, DerefMut};
13use bevy_ecs::{
14    resource::Resource,
15    schedule::IntoScheduleConfigs as _,
16    system::{Res, ResMut},
17    world::{FromWorld, World},
18};
19use bevy_platform::collections::{hash_map::Entry, HashMap, HashSet};
20use bevy_utils::default;
21use offset_allocator::{Allocation, Allocator};
22use tracing::error;
23use wgpu::{
24    BufferDescriptor, BufferSize, BufferUsages, CommandEncoderDescriptor, DownlevelFlags,
25    COPY_BUFFER_ALIGNMENT,
26};
27
28use crate::{
29    mesh::{Indices, Mesh, MeshVertexBufferLayouts, RenderMesh},
30    render_asset::{prepare_assets, ExtractedAssets},
31    render_resource::Buffer,
32    renderer::{RenderAdapter, RenderDevice, RenderQueue},
33    Render, RenderApp, RenderSet,
34};
35
36/// A plugin that manages GPU memory for mesh data.
37pub struct MeshAllocatorPlugin;
38
39/// Manages the assignment of mesh data to GPU buffers.
40///
41/// The Bevy renderer tries to pack vertex and index data for multiple meshes
42/// together so that multiple meshes can be drawn back-to-back without any
43/// rebinding. This resource manages these buffers.
44///
45/// Within each slab, or hardware buffer, the underlying allocation algorithm is
46/// [`offset-allocator`], a Rust port of Sebastian Aaltonen's hard-real-time C++
47/// `OffsetAllocator`. Slabs start small and then grow as their contents fill
48/// up, up to a maximum size limit. To reduce fragmentation, vertex and index
49/// buffers that are too large bypass this system and receive their own buffers.
50///
51/// The [`MeshAllocatorSettings`] allows you to tune the behavior of the
52/// allocator for better performance with your application. Most applications
53/// won't need to change the settings from their default values.
54#[derive(Resource)]
55pub struct MeshAllocator {
56    /// Holds all buffers and allocators.
57    slabs: HashMap<SlabId, Slab>,
58
59    /// Maps a layout to the slabs that hold elements of that layout.
60    ///
61    /// This is used when allocating, so that we can find the appropriate slab
62    /// to place an object in.
63    slab_layouts: HashMap<ElementLayout, Vec<SlabId>>,
64
65    /// Maps mesh asset IDs to the ID of the slabs that hold their vertex data.
66    mesh_id_to_vertex_slab: HashMap<AssetId<Mesh>, SlabId>,
67
68    /// Maps mesh asset IDs to the ID of the slabs that hold their index data.
69    mesh_id_to_index_slab: HashMap<AssetId<Mesh>, SlabId>,
70
71    /// The next slab ID to assign.
72    next_slab_id: SlabId,
73
74    /// Whether we can pack multiple vertex arrays into a single slab on this
75    /// platform.
76    ///
77    /// This corresponds to [`DownlevelFlags::BASE_VERTEX`], which is unset on
78    /// WebGL 2. On this platform, we must give each vertex array its own
79    /// buffer, because we can't adjust the first vertex when we perform a draw.
80    general_vertex_slabs_supported: bool,
81}
82
83/// Tunable parameters that customize the behavior of the allocator.
84///
85/// Generally, these parameters adjust the tradeoff between memory fragmentation
86/// and performance. You can adjust them as desired for your application. Most
87/// applications can stick with the default values.
88#[derive(Resource)]
89pub struct MeshAllocatorSettings {
90    /// The minimum size of a slab (hardware buffer), in bytes.
91    ///
92    /// The default value is 1 MiB.
93    pub min_slab_size: u64,
94
95    /// The maximum size of a slab (hardware buffer), in bytes.
96    ///
97    /// When a slab reaches this limit, a new slab is created.
98    ///
99    /// The default value is 512 MiB.
100    pub max_slab_size: u64,
101
102    /// The maximum size of vertex or index data that can be placed in a general
103    /// slab, in bytes.
104    ///
105    /// If a mesh has vertex or index data that exceeds this size limit, that
106    /// data is placed in its own slab. This reduces fragmentation, but incurs
107    /// more CPU-side binding overhead when drawing the mesh.
108    ///
109    /// The default value is 256 MiB.
110    pub large_threshold: u64,
111
112    /// The factor by which we scale a slab when growing it.
113    ///
114    /// This value must be greater than 1. Higher values result in more
115    /// fragmentation but fewer expensive copy operations when growing the
116    /// buffer.
117    ///
118    /// The default value is 1.5.
119    pub growth_factor: f64,
120}
121
122impl Default for MeshAllocatorSettings {
123    fn default() -> Self {
124        Self {
125            // 1 MiB
126            min_slab_size: 1024 * 1024,
127            // 512 MiB
128            max_slab_size: 1024 * 1024 * 512,
129            // 256 MiB
130            large_threshold: 1024 * 1024 * 256,
131            // 1.5× growth
132            growth_factor: 1.5,
133        }
134    }
135}
136
137/// The hardware buffer that mesh data lives in, as well as the range within
138/// that buffer.
139pub struct MeshBufferSlice<'a> {
140    /// The buffer that the mesh data resides in.
141    pub buffer: &'a Buffer,
142
143    /// The range of elements within this buffer that the mesh data resides in,
144    /// measured in elements.
145    ///
146    /// This is not a byte range; it's an element range. For vertex data, this
147    /// is measured in increments of a single vertex. (Thus, if a vertex is 32
148    /// bytes long, then this range is in units of 32 bytes each.) For index
149    /// data, this is measured in increments of a single index value (2 or 4
150    /// bytes). Draw commands generally take their ranges in elements, not
151    /// bytes, so this is the most convenient unit in this case.
152    pub range: Range<u32>,
153}
154
155/// The index of a single slab.
156#[derive(Clone, Copy, Default, PartialEq, Eq, PartialOrd, Ord, Hash, Debug)]
157#[repr(transparent)]
158pub struct SlabId(pub NonMaxU32);
159
160/// Data for a single slab.
161enum Slab {
162    /// A slab that can contain multiple objects.
163    General(GeneralSlab),
164    /// A slab that contains a single object.
165    LargeObject(LargeObjectSlab),
166}
167
168/// A resizable slab that can contain multiple objects.
169///
170/// This is the normal type of slab used for objects that are below the
171/// [`MeshAllocatorSettings::large_threshold`]. Slabs are divided into *slots*,
172/// which are described in detail in the [`ElementLayout`] documentation.
173struct GeneralSlab {
174    /// The [`Allocator`] that manages the objects in this slab.
175    allocator: Allocator,
176
177    /// The GPU buffer that backs this slab.
178    ///
179    /// This may be `None` if the buffer hasn't been created yet. We delay
180    /// creation of buffers until allocating all the meshes for a single frame,
181    /// so that we don't needlessly create and resize buffers when many meshes
182    /// load all at once.
183    buffer: Option<Buffer>,
184
185    /// Allocations that are on the GPU.
186    ///
187    /// The range is in slots.
188    resident_allocations: HashMap<AssetId<Mesh>, SlabAllocation>,
189
190    /// Allocations that are waiting to be uploaded to the GPU.
191    ///
192    /// The range is in slots.
193    pending_allocations: HashMap<AssetId<Mesh>, SlabAllocation>,
194
195    /// The layout of a single element (vertex or index).
196    element_layout: ElementLayout,
197
198    /// The size of this slab in slots.
199    current_slot_capacity: u32,
200}
201
202/// A slab that contains a single object.
203///
204/// Typically, this is for objects that exceed the
205/// [`MeshAllocatorSettings::large_threshold`]. This is also for objects that
206/// would ordinarily receive their own slab but can't because of platform
207/// limitations, most notably vertex arrays on WebGL 2.
208struct LargeObjectSlab {
209    /// The GPU buffer that backs this slab.
210    ///
211    /// This may be `None` if the buffer hasn't been created yet.
212    buffer: Option<Buffer>,
213
214    /// The layout of a single element (vertex or index).
215    element_layout: ElementLayout,
216}
217
218/// The type of element that a slab can store.
219#[derive(Clone, Copy, PartialEq, Eq, Hash)]
220enum ElementClass {
221    /// Data for a vertex.
222    Vertex,
223    /// A vertex index.
224    Index,
225}
226
227/// The results of [`GeneralSlab::grow_if_necessary`].
228enum SlabGrowthResult {
229    /// The mesh data already fits in the slab; the slab doesn't need to grow.
230    NoGrowthNeeded,
231    /// The slab needed to grow.
232    ///
233    /// The [`SlabToReallocate`] contains the old capacity of the slab.
234    NeededGrowth(SlabToReallocate),
235    /// The slab wanted to grow but couldn't because it hit its maximum size.
236    CantGrow,
237}
238
239/// Information about the size of individual elements (vertices or indices)
240/// within a slab.
241///
242/// Slab objects are allocated in units of *slots*. Usually, each element takes
243/// up one slot, and so elements and slots are equivalent. Occasionally,
244/// however, a slot may consist of 2 or even 4 elements. This occurs when the
245/// size of an element isn't divisible by [`COPY_BUFFER_ALIGNMENT`]. When we
246/// resize buffers, we perform GPU-to-GPU copies to shuffle the existing
247/// elements into their new positions, and such copies must be on
248/// [`COPY_BUFFER_ALIGNMENT`] boundaries. Slots solve this problem by
249/// guaranteeing that the size of an allocation quantum is divisible by both the
250/// size of an element and [`COPY_BUFFER_ALIGNMENT`], so we can relocate it
251/// freely.
252#[derive(Clone, Copy, PartialEq, Eq, Hash)]
253struct ElementLayout {
254    /// Either a vertex or an index.
255    class: ElementClass,
256
257    /// The size in bytes of a single element (vertex or index).
258    size: u64,
259
260    /// The number of elements that make up a single slot.
261    ///
262    /// Usually, this is 1, but it can be different if [`ElementLayout::size`]
263    /// isn't divisible by 4. See the comment in [`ElementLayout`] for more
264    /// details.
265    elements_per_slot: u32,
266}
267
268/// The location of an allocation and the slab it's contained in.
269struct MeshAllocation {
270    /// The ID of the slab.
271    slab_id: SlabId,
272    /// Holds the actual allocation.
273    slab_allocation: SlabAllocation,
274}
275
276/// An allocation within a slab.
277#[derive(Clone)]
278struct SlabAllocation {
279    /// The actual [`Allocator`] handle, needed to free the allocation.
280    allocation: Allocation,
281    /// The number of slots that this allocation takes up.
282    slot_count: u32,
283}
284
285/// Holds information about all slabs scheduled to be allocated or reallocated.
286#[derive(Default, Deref, DerefMut)]
287struct SlabsToReallocate(HashMap<SlabId, SlabToReallocate>);
288
289/// Holds information about a slab that's scheduled to be allocated or
290/// reallocated.
291#[derive(Default)]
292struct SlabToReallocate {
293    /// The capacity of the slab before we decided to grow it.
294    old_slot_capacity: u32,
295}
296
297impl Display for SlabId {
298    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
299        self.0.fmt(f)
300    }
301}
302
303impl Plugin for MeshAllocatorPlugin {
304    fn build(&self, app: &mut App) {
305        let Some(render_app) = app.get_sub_app_mut(RenderApp) else {
306            return;
307        };
308
309        render_app
310            .init_resource::<MeshAllocatorSettings>()
311            .add_systems(
312                Render,
313                allocate_and_free_meshes
314                    .in_set(RenderSet::PrepareAssets)
315                    .before(prepare_assets::<RenderMesh>),
316            );
317    }
318
319    fn finish(&self, app: &mut App) {
320        let Some(render_app) = app.get_sub_app_mut(RenderApp) else {
321            return;
322        };
323
324        // The `RenderAdapter` isn't available until now, so we can't do this in
325        // [`Plugin::build`].
326        render_app.init_resource::<MeshAllocator>();
327    }
328}
329
330impl FromWorld for MeshAllocator {
331    fn from_world(world: &mut World) -> Self {
332        // Note whether we're on WebGL 2. In this case, we must give every
333        // vertex array its own slab.
334        let render_adapter = world.resource::<RenderAdapter>();
335        let general_vertex_slabs_supported = render_adapter
336            .get_downlevel_capabilities()
337            .flags
338            .contains(DownlevelFlags::BASE_VERTEX);
339
340        Self {
341            slabs: HashMap::default(),
342            slab_layouts: HashMap::default(),
343            mesh_id_to_vertex_slab: HashMap::default(),
344            mesh_id_to_index_slab: HashMap::default(),
345            next_slab_id: default(),
346            general_vertex_slabs_supported,
347        }
348    }
349}
350
351/// A system that processes newly-extracted or newly-removed meshes and writes
352/// their data into buffers or frees their data as appropriate.
353pub fn allocate_and_free_meshes(
354    mut mesh_allocator: ResMut<MeshAllocator>,
355    mesh_allocator_settings: Res<MeshAllocatorSettings>,
356    extracted_meshes: Res<ExtractedAssets<RenderMesh>>,
357    mut mesh_vertex_buffer_layouts: ResMut<MeshVertexBufferLayouts>,
358    render_device: Res<RenderDevice>,
359    render_queue: Res<RenderQueue>,
360) {
361    // Process removed or modified meshes.
362    mesh_allocator.free_meshes(&extracted_meshes);
363
364    // Process newly-added or modified meshes.
365    mesh_allocator.allocate_meshes(
366        &mesh_allocator_settings,
367        &extracted_meshes,
368        &mut mesh_vertex_buffer_layouts,
369        &render_device,
370        &render_queue,
371    );
372}
373
374impl MeshAllocator {
375    /// Returns the buffer and range within that buffer of the vertex data for
376    /// the mesh with the given ID.
377    ///
378    /// If the mesh wasn't allocated, returns None.
379    pub fn mesh_vertex_slice(&self, mesh_id: &AssetId<Mesh>) -> Option<MeshBufferSlice> {
380        self.mesh_slice_in_slab(mesh_id, *self.mesh_id_to_vertex_slab.get(mesh_id)?)
381    }
382
383    /// Returns the buffer and range within that buffer of the index data for
384    /// the mesh with the given ID.
385    ///
386    /// If the mesh has no index data or wasn't allocated, returns None.
387    pub fn mesh_index_slice(&self, mesh_id: &AssetId<Mesh>) -> Option<MeshBufferSlice> {
388        self.mesh_slice_in_slab(mesh_id, *self.mesh_id_to_index_slab.get(mesh_id)?)
389    }
390
391    /// Returns the IDs of the vertex buffer and index buffer respectively for
392    /// the mesh with the given ID.
393    ///
394    /// If the mesh wasn't allocated, or has no index data in the case of the
395    /// index buffer, the corresponding element in the returned tuple will be
396    /// None.
397    pub fn mesh_slabs(&self, mesh_id: &AssetId<Mesh>) -> (Option<SlabId>, Option<SlabId>) {
398        (
399            self.mesh_id_to_vertex_slab.get(mesh_id).cloned(),
400            self.mesh_id_to_index_slab.get(mesh_id).cloned(),
401        )
402    }
403
404    /// Given a slab and a mesh with data located with it, returns the buffer
405    /// and range of that mesh data within the slab.
406    fn mesh_slice_in_slab(
407        &self,
408        mesh_id: &AssetId<Mesh>,
409        slab_id: SlabId,
410    ) -> Option<MeshBufferSlice> {
411        match self.slabs.get(&slab_id)? {
412            Slab::General(general_slab) => {
413                let slab_allocation = general_slab.resident_allocations.get(mesh_id)?;
414                Some(MeshBufferSlice {
415                    buffer: general_slab.buffer.as_ref()?,
416                    range: (slab_allocation.allocation.offset
417                        * general_slab.element_layout.elements_per_slot)
418                        ..((slab_allocation.allocation.offset + slab_allocation.slot_count)
419                            * general_slab.element_layout.elements_per_slot),
420                })
421            }
422
423            Slab::LargeObject(large_object_slab) => {
424                let buffer = large_object_slab.buffer.as_ref()?;
425                Some(MeshBufferSlice {
426                    buffer,
427                    range: 0..((buffer.size() / large_object_slab.element_layout.size) as u32),
428                })
429            }
430        }
431    }
432
433    /// Processes newly-loaded meshes, allocating room in the slabs for their
434    /// mesh data and performing upload operations as appropriate.
435    fn allocate_meshes(
436        &mut self,
437        mesh_allocator_settings: &MeshAllocatorSettings,
438        extracted_meshes: &ExtractedAssets<RenderMesh>,
439        mesh_vertex_buffer_layouts: &mut MeshVertexBufferLayouts,
440        render_device: &RenderDevice,
441        render_queue: &RenderQueue,
442    ) {
443        let mut slabs_to_grow = SlabsToReallocate::default();
444
445        // Allocate.
446        for (mesh_id, mesh) in &extracted_meshes.extracted {
447            // Allocate vertex data. Note that we can only pack mesh vertex data
448            // together if the platform supports it.
449            let vertex_element_layout = ElementLayout::vertex(mesh_vertex_buffer_layouts, mesh);
450            if self.general_vertex_slabs_supported {
451                self.allocate(
452                    mesh_id,
453                    mesh.get_vertex_buffer_size() as u64,
454                    vertex_element_layout,
455                    &mut slabs_to_grow,
456                    mesh_allocator_settings,
457                );
458            } else {
459                self.allocate_large(mesh_id, vertex_element_layout);
460            }
461
462            // Allocate index data.
463            if let (Some(index_buffer_data), Some(index_element_layout)) =
464                (mesh.get_index_buffer_bytes(), ElementLayout::index(mesh))
465            {
466                self.allocate(
467                    mesh_id,
468                    index_buffer_data.len() as u64,
469                    index_element_layout,
470                    &mut slabs_to_grow,
471                    mesh_allocator_settings,
472                );
473            }
474        }
475
476        // Perform growth.
477        for (slab_id, slab_to_grow) in slabs_to_grow.0 {
478            self.reallocate_slab(render_device, render_queue, slab_id, slab_to_grow);
479        }
480
481        // Copy new mesh data in.
482        for (mesh_id, mesh) in &extracted_meshes.extracted {
483            self.copy_mesh_vertex_data(mesh_id, mesh, render_device, render_queue);
484            self.copy_mesh_index_data(mesh_id, mesh, render_device, render_queue);
485        }
486    }
487
488    /// Copies vertex array data from a mesh into the appropriate spot in the
489    /// slab.
490    fn copy_mesh_vertex_data(
491        &mut self,
492        mesh_id: &AssetId<Mesh>,
493        mesh: &Mesh,
494        render_device: &RenderDevice,
495        render_queue: &RenderQueue,
496    ) {
497        let Some(&slab_id) = self.mesh_id_to_vertex_slab.get(mesh_id) else {
498            return;
499        };
500
501        // Call the generic function.
502        self.copy_element_data(
503            mesh_id,
504            mesh.get_vertex_buffer_size(),
505            |slice| mesh.write_packed_vertex_buffer_data(slice),
506            BufferUsages::VERTEX,
507            slab_id,
508            render_device,
509            render_queue,
510        );
511    }
512
513    /// Copies index array data from a mesh into the appropriate spot in the
514    /// slab.
515    fn copy_mesh_index_data(
516        &mut self,
517        mesh_id: &AssetId<Mesh>,
518        mesh: &Mesh,
519        render_device: &RenderDevice,
520        render_queue: &RenderQueue,
521    ) {
522        let Some(&slab_id) = self.mesh_id_to_index_slab.get(mesh_id) else {
523            return;
524        };
525        let Some(index_data) = mesh.get_index_buffer_bytes() else {
526            return;
527        };
528
529        // Call the generic function.
530        self.copy_element_data(
531            mesh_id,
532            index_data.len(),
533            |slice| slice.copy_from_slice(index_data),
534            BufferUsages::INDEX,
535            slab_id,
536            render_device,
537            render_queue,
538        );
539    }
540
541    /// A generic function that copies either vertex or index data into a slab.
542    fn copy_element_data(
543        &mut self,
544        mesh_id: &AssetId<Mesh>,
545        len: usize,
546        fill_data: impl Fn(&mut [u8]),
547        buffer_usages: BufferUsages,
548        slab_id: SlabId,
549        render_device: &RenderDevice,
550        render_queue: &RenderQueue,
551    ) {
552        let Some(slab) = self.slabs.get_mut(&slab_id) else {
553            return;
554        };
555
556        match *slab {
557            Slab::General(ref mut general_slab) => {
558                let (Some(buffer), Some(allocated_range)) = (
559                    &general_slab.buffer,
560                    general_slab.pending_allocations.remove(mesh_id),
561                ) else {
562                    return;
563                };
564
565                let slot_size = general_slab.element_layout.slot_size();
566
567                // round up size to a multiple of the slot size to satisfy wgpu alignment requirements
568                if let Some(size) = BufferSize::new((len as u64).next_multiple_of(slot_size)) {
569                    // Write the data in.
570                    if let Some(mut buffer) = render_queue.write_buffer_with(
571                        buffer,
572                        allocated_range.allocation.offset as u64 * slot_size,
573                        size,
574                    ) {
575                        let slice = &mut buffer.as_mut()[..len];
576                        fill_data(slice);
577                    }
578                }
579
580                // Mark the allocation as resident.
581                general_slab
582                    .resident_allocations
583                    .insert(*mesh_id, allocated_range);
584            }
585
586            Slab::LargeObject(ref mut large_object_slab) => {
587                debug_assert!(large_object_slab.buffer.is_none());
588
589                // Create the buffer and its data in one go.
590                let buffer = render_device.create_buffer(&BufferDescriptor {
591                    label: Some(&format!(
592                        "large mesh slab {} ({}buffer)",
593                        slab_id,
594                        buffer_usages_to_str(buffer_usages)
595                    )),
596                    size: len as u64,
597                    usage: buffer_usages | BufferUsages::COPY_DST,
598                    mapped_at_creation: true,
599                });
600                {
601                    let slice = &mut buffer.slice(..).get_mapped_range_mut()[..len];
602                    fill_data(slice);
603                }
604                buffer.unmap();
605                large_object_slab.buffer = Some(buffer);
606            }
607        }
608    }
609
610    /// Frees allocations for meshes that were removed or modified this frame.
611    fn free_meshes(&mut self, extracted_meshes: &ExtractedAssets<RenderMesh>) {
612        let mut empty_slabs = <HashSet<_>>::default();
613
614        // TODO: Consider explicitly reusing allocations for changed meshes of the same size
615        let meshes_to_free = extracted_meshes
616            .removed
617            .iter()
618            .chain(extracted_meshes.modified.iter());
619
620        for mesh_id in meshes_to_free {
621            if let Some(slab_id) = self.mesh_id_to_vertex_slab.remove(mesh_id) {
622                self.free_allocation_in_slab(mesh_id, slab_id, &mut empty_slabs);
623            }
624            if let Some(slab_id) = self.mesh_id_to_index_slab.remove(mesh_id) {
625                self.free_allocation_in_slab(mesh_id, slab_id, &mut empty_slabs);
626            }
627        }
628
629        for empty_slab in empty_slabs {
630            self.slab_layouts.values_mut().for_each(|slab_ids| {
631                let idx = slab_ids.iter().position(|&slab_id| slab_id == empty_slab);
632                if let Some(idx) = idx {
633                    slab_ids.remove(idx);
634                }
635            });
636            self.slabs.remove(&empty_slab);
637        }
638    }
639
640    /// Given a slab and the ID of a mesh containing data in it, marks the
641    /// allocation as free.
642    ///
643    /// If this results in the slab becoming empty, this function adds the slab
644    /// to the `empty_slabs` set.
645    fn free_allocation_in_slab(
646        &mut self,
647        mesh_id: &AssetId<Mesh>,
648        slab_id: SlabId,
649        empty_slabs: &mut HashSet<SlabId>,
650    ) {
651        let Some(slab) = self.slabs.get_mut(&slab_id) else {
652            return;
653        };
654
655        match *slab {
656            Slab::General(ref mut general_slab) => {
657                let Some(slab_allocation) = general_slab
658                    .resident_allocations
659                    .remove(mesh_id)
660                    .or_else(|| general_slab.pending_allocations.remove(mesh_id))
661                else {
662                    return;
663                };
664
665                general_slab.allocator.free(slab_allocation.allocation);
666
667                if general_slab.is_empty() {
668                    empty_slabs.insert(slab_id);
669                }
670            }
671            Slab::LargeObject(_) => {
672                empty_slabs.insert(slab_id);
673            }
674        }
675    }
676
677    /// Allocates space for mesh data with the given byte size and layout in the
678    /// appropriate slab, creating that slab if necessary.
679    fn allocate(
680        &mut self,
681        mesh_id: &AssetId<Mesh>,
682        data_byte_len: u64,
683        layout: ElementLayout,
684        slabs_to_grow: &mut SlabsToReallocate,
685        settings: &MeshAllocatorSettings,
686    ) {
687        let data_element_count = data_byte_len.div_ceil(layout.size) as u32;
688        let data_slot_count = data_element_count.div_ceil(layout.elements_per_slot);
689
690        // If the mesh data is too large for a slab, give it a slab of its own.
691        if data_slot_count as u64 * layout.slot_size()
692            >= settings.large_threshold.min(settings.max_slab_size)
693        {
694            self.allocate_large(mesh_id, layout);
695        } else {
696            self.allocate_general(mesh_id, data_slot_count, layout, slabs_to_grow, settings);
697        }
698    }
699
700    /// Allocates space for mesh data with the given slot size and layout in the
701    /// appropriate general slab.
702    fn allocate_general(
703        &mut self,
704        mesh_id: &AssetId<Mesh>,
705        data_slot_count: u32,
706        layout: ElementLayout,
707        slabs_to_grow: &mut SlabsToReallocate,
708        settings: &MeshAllocatorSettings,
709    ) {
710        let candidate_slabs = self.slab_layouts.entry(layout).or_default();
711
712        // Loop through the slabs that accept elements of the appropriate type
713        // and try to allocate the mesh inside them. We go with the first one
714        // that succeeds.
715        let mut mesh_allocation = None;
716        for &slab_id in &*candidate_slabs {
717            let Some(Slab::General(slab)) = self.slabs.get_mut(&slab_id) else {
718                unreachable!("Slab not found")
719            };
720
721            let Some(allocation) = slab.allocator.allocate(data_slot_count) else {
722                continue;
723            };
724
725            // Try to fit the object in the slab, growing if necessary.
726            match slab.grow_if_necessary(allocation.offset + data_slot_count, settings) {
727                SlabGrowthResult::NoGrowthNeeded => {}
728                SlabGrowthResult::NeededGrowth(slab_to_reallocate) => {
729                    // If we already grew the slab this frame, don't replace the
730                    // `SlabToReallocate` entry. We want to keep the entry
731                    // corresponding to the size that the slab had at the start
732                    // of the frame, so that we can copy only the used portion
733                    // of the initial buffer to the new one.
734                    if let Entry::Vacant(vacant_entry) = slabs_to_grow.entry(slab_id) {
735                        vacant_entry.insert(slab_to_reallocate);
736                    }
737                }
738                SlabGrowthResult::CantGrow => continue,
739            }
740
741            mesh_allocation = Some(MeshAllocation {
742                slab_id,
743                slab_allocation: SlabAllocation {
744                    allocation,
745                    slot_count: data_slot_count,
746                },
747            });
748            break;
749        }
750
751        // If we still have no allocation, make a new slab.
752        if mesh_allocation.is_none() {
753            let new_slab_id = self.next_slab_id;
754            self.next_slab_id.0 = NonMaxU32::new(self.next_slab_id.0.get() + 1).unwrap_or_default();
755
756            let new_slab = GeneralSlab::new(
757                new_slab_id,
758                &mut mesh_allocation,
759                settings,
760                layout,
761                data_slot_count,
762            );
763
764            self.slabs.insert(new_slab_id, Slab::General(new_slab));
765            candidate_slabs.push(new_slab_id);
766            slabs_to_grow.insert(new_slab_id, SlabToReallocate::default());
767        }
768
769        let mesh_allocation = mesh_allocation.expect("Should have been able to allocate");
770
771        // Mark the allocation as pending. Don't copy it in just yet; further
772        // meshes loaded this frame may result in its final allocation location
773        // changing.
774        if let Some(Slab::General(general_slab)) = self.slabs.get_mut(&mesh_allocation.slab_id) {
775            general_slab
776                .pending_allocations
777                .insert(*mesh_id, mesh_allocation.slab_allocation);
778        };
779
780        self.record_allocation(mesh_id, mesh_allocation.slab_id, layout.class);
781    }
782
783    /// Allocates an object into its own dedicated slab.
784    fn allocate_large(&mut self, mesh_id: &AssetId<Mesh>, layout: ElementLayout) {
785        let new_slab_id = self.next_slab_id;
786        self.next_slab_id.0 = NonMaxU32::new(self.next_slab_id.0.get() + 1).unwrap_or_default();
787
788        self.record_allocation(mesh_id, new_slab_id, layout.class);
789
790        self.slabs.insert(
791            new_slab_id,
792            Slab::LargeObject(LargeObjectSlab {
793                buffer: None,
794                element_layout: layout,
795            }),
796        );
797    }
798
799    /// Reallocates a slab that needs to be resized, or allocates a new slab.
800    ///
801    /// This performs the actual growth operation that
802    /// [`GeneralSlab::grow_if_necessary`] scheduled. We do the growth in two
803    /// phases so that, if a slab grows multiple times in the same frame, only
804    /// one new buffer is reallocated, rather than reallocating the buffer
805    /// multiple times.
806    fn reallocate_slab(
807        &mut self,
808        render_device: &RenderDevice,
809        render_queue: &RenderQueue,
810        slab_id: SlabId,
811        slab_to_grow: SlabToReallocate,
812    ) {
813        let Some(Slab::General(slab)) = self.slabs.get_mut(&slab_id) else {
814            error!("Couldn't find slab {} to grow", slab_id);
815            return;
816        };
817
818        let old_buffer = slab.buffer.take();
819
820        let mut buffer_usages = BufferUsages::COPY_SRC | BufferUsages::COPY_DST;
821        match slab.element_layout.class {
822            ElementClass::Vertex => buffer_usages |= BufferUsages::VERTEX,
823            ElementClass::Index => buffer_usages |= BufferUsages::INDEX,
824        };
825
826        // Create the buffer.
827        let new_buffer = render_device.create_buffer(&BufferDescriptor {
828            label: Some(&format!(
829                "general mesh slab {} ({}buffer)",
830                slab_id,
831                buffer_usages_to_str(buffer_usages)
832            )),
833            size: slab.current_slot_capacity as u64 * slab.element_layout.slot_size(),
834            usage: buffer_usages,
835            mapped_at_creation: false,
836        });
837
838        slab.buffer = Some(new_buffer.clone());
839
840        let Some(old_buffer) = old_buffer else { return };
841
842        // In order to do buffer copies, we need a command encoder.
843        let mut encoder = render_device.create_command_encoder(&CommandEncoderDescriptor {
844            label: Some("slab resize encoder"),
845        });
846
847        // Copy the data from the old buffer into the new one.
848        encoder.copy_buffer_to_buffer(
849            &old_buffer,
850            0,
851            &new_buffer,
852            0,
853            slab_to_grow.old_slot_capacity as u64 * slab.element_layout.slot_size(),
854        );
855
856        let command_buffer = encoder.finish();
857        render_queue.submit([command_buffer]);
858    }
859
860    /// Records the location of the given newly-allocated mesh data in the
861    /// [`Self::mesh_id_to_vertex_slab`] or [`Self::mesh_id_to_index_slab`]
862    /// tables as appropriate.
863    fn record_allocation(
864        &mut self,
865        mesh_id: &AssetId<Mesh>,
866        slab_id: SlabId,
867        element_class: ElementClass,
868    ) {
869        match element_class {
870            ElementClass::Vertex => {
871                self.mesh_id_to_vertex_slab.insert(*mesh_id, slab_id);
872            }
873            ElementClass::Index => {
874                self.mesh_id_to_index_slab.insert(*mesh_id, slab_id);
875            }
876        }
877    }
878}
879
880impl GeneralSlab {
881    /// Creates a new growable slab big enough to hold a single element of
882    /// `data_slot_count` size with the given `layout`.
883    fn new(
884        new_slab_id: SlabId,
885        mesh_allocation: &mut Option<MeshAllocation>,
886        settings: &MeshAllocatorSettings,
887        layout: ElementLayout,
888        data_slot_count: u32,
889    ) -> GeneralSlab {
890        let initial_slab_slot_capacity = (settings.min_slab_size.div_ceil(layout.slot_size())
891            as u32)
892            .max(offset_allocator::ext::min_allocator_size(data_slot_count));
893        let max_slab_slot_capacity = (settings.max_slab_size.div_ceil(layout.slot_size()) as u32)
894            .max(offset_allocator::ext::min_allocator_size(data_slot_count));
895
896        let mut new_slab = GeneralSlab {
897            allocator: Allocator::new(max_slab_slot_capacity),
898            buffer: None,
899            resident_allocations: HashMap::default(),
900            pending_allocations: HashMap::default(),
901            element_layout: layout,
902            current_slot_capacity: initial_slab_slot_capacity,
903        };
904
905        // This should never fail.
906        if let Some(allocation) = new_slab.allocator.allocate(data_slot_count) {
907            *mesh_allocation = Some(MeshAllocation {
908                slab_id: new_slab_id,
909                slab_allocation: SlabAllocation {
910                    slot_count: data_slot_count,
911                    allocation,
912                },
913            });
914        }
915
916        new_slab
917    }
918
919    /// Checks to see if the size of this slab is at least `new_size_in_slots`
920    /// and grows the slab if it isn't.
921    ///
922    /// The returned [`SlabGrowthResult`] describes whether the slab needed to
923    /// grow and whether, if so, it was successful in doing so.
924    fn grow_if_necessary(
925        &mut self,
926        new_size_in_slots: u32,
927        settings: &MeshAllocatorSettings,
928    ) -> SlabGrowthResult {
929        // Is the slab big enough already?
930        let initial_slot_capacity = self.current_slot_capacity;
931        if self.current_slot_capacity >= new_size_in_slots {
932            return SlabGrowthResult::NoGrowthNeeded;
933        }
934
935        // Try to grow in increments of `MeshAllocatorSettings::growth_factor`
936        // until we're big enough.
937        while self.current_slot_capacity < new_size_in_slots {
938            let new_slab_slot_capacity =
939                ((self.current_slot_capacity as f64 * settings.growth_factor).ceil() as u32)
940                    .min((settings.max_slab_size / self.element_layout.slot_size()) as u32);
941            if new_slab_slot_capacity == self.current_slot_capacity {
942                // The slab is full.
943                return SlabGrowthResult::CantGrow;
944            }
945
946            self.current_slot_capacity = new_slab_slot_capacity;
947        }
948
949        // Tell our caller what we did.
950        SlabGrowthResult::NeededGrowth(SlabToReallocate {
951            old_slot_capacity: initial_slot_capacity,
952        })
953    }
954}
955
956impl ElementLayout {
957    /// Creates an [`ElementLayout`] for mesh data of the given class (vertex or
958    /// index) with the given byte size.
959    fn new(class: ElementClass, size: u64) -> ElementLayout {
960        const {
961            assert!(4 == COPY_BUFFER_ALIGNMENT);
962        }
963        // this is equivalent to `4 / gcd(4,size)` but lets us not implement gcd.
964        // ping @atlv if above assert ever fails (likely never)
965        let elements_per_slot = [1, 4, 2, 4][size as usize & 3];
966        ElementLayout {
967            class,
968            size,
969            // Make sure that slot boundaries begin and end on
970            // `COPY_BUFFER_ALIGNMENT`-byte (4-byte) boundaries.
971            elements_per_slot,
972        }
973    }
974
975    fn slot_size(&self) -> u64 {
976        self.size * self.elements_per_slot as u64
977    }
978
979    /// Creates the appropriate [`ElementLayout`] for the given mesh's vertex
980    /// data.
981    fn vertex(
982        mesh_vertex_buffer_layouts: &mut MeshVertexBufferLayouts,
983        mesh: &Mesh,
984    ) -> ElementLayout {
985        let mesh_vertex_buffer_layout =
986            mesh.get_mesh_vertex_buffer_layout(mesh_vertex_buffer_layouts);
987        ElementLayout::new(
988            ElementClass::Vertex,
989            mesh_vertex_buffer_layout.0.layout().array_stride,
990        )
991    }
992
993    /// Creates the appropriate [`ElementLayout`] for the given mesh's index
994    /// data.
995    fn index(mesh: &Mesh) -> Option<ElementLayout> {
996        let size = match mesh.indices()? {
997            Indices::U16(_) => 2,
998            Indices::U32(_) => 4,
999        };
1000        Some(ElementLayout::new(ElementClass::Index, size))
1001    }
1002}
1003
1004impl GeneralSlab {
1005    /// Returns true if this slab is empty.
1006    fn is_empty(&self) -> bool {
1007        self.resident_allocations.is_empty() && self.pending_allocations.is_empty()
1008    }
1009}
1010
1011/// Returns a string describing the given buffer usages.
1012fn buffer_usages_to_str(buffer_usages: BufferUsages) -> &'static str {
1013    if buffer_usages.contains(BufferUsages::VERTEX) {
1014        "vertex "
1015    } else if buffer_usages.contains(BufferUsages::INDEX) {
1016        "index "
1017    } else {
1018        ""
1019    }
1020}