Index buffer specialization (#568)

index buffer specialization

crates/bevy_gltf/src/loader.rs

@@ -1,11 +1,14 @@
 use bevy_render::{
-    mesh::{Mesh, VertexAttribute},
+    mesh::{Indices, Mesh, VertexAttribute},
     pipeline::PrimitiveTopology,
 };
 
 use anyhow::Result;
 use bevy_asset::AssetLoader;
-use gltf::{buffer::Source, mesh::Mode};
+use gltf::{
+    buffer::Source,
+    mesh::{util::ReadIndices, Mode},
+};
 use std::{fs, io, path::Path};
 use thiserror::Error;
 
@@ -98,7 +101,11 @@ fn load_node(buffer_data: &[Vec<u8>], node: &gltf::Node, depth: i32) -> Result<M
             }
 
             if let Some(indices) = reader.read_indices() {
-                mesh.indices = Some(indices.into_u32().collect::<Vec<u32>>());
+                mesh.indices = match indices {
+                    ReadIndices::U8(iter) => Some(Indices::U16(iter.map(|i| i as u16).collect())),
+                    ReadIndices::U16(iter) => Some(Indices::U16(iter.collect())),
+                    ReadIndices::U32(iter) => Some(Indices::U32(iter.collect())),
+                }
             };
 
             return Ok(mesh);

crates/bevy_render/src/draw.rs

@@ -341,8 +341,7 @@ impl<'a> DrawContext<'a> {
         &self,
         draw: &mut Draw,
         render_resource_bindings: &[&RenderResourceBindings],
-    ) -> Result<Option<Range<u32>>, DrawError> {
-        let mut indices = None;
+    ) -> Result<(), DrawError> {
         let pipeline = self.current_pipeline.ok_or(DrawError::NoPipelineSet)?;
         let pipeline_descriptor = self
             .pipelines
@@ -359,13 +358,6 @@ impl<'a> DrawContext<'a> {
                 {
                     draw.set_vertex_buffer(slot as u32, vertex_buffer, 0);
                     if let Some(index_buffer) = index_buffer {
-                        if let Some(buffer_info) =
-                            self.render_resource_context.get_buffer_info(index_buffer)
-                        {
-                            indices = Some(0..(buffer_info.size / 2) as u32);
-                        } else {
-                            panic!("expected buffer type");
-                        }
                         draw.set_index_buffer(index_buffer, 0);
                     }
 
@@ -373,8 +365,7 @@ impl<'a> DrawContext<'a> {
                 }
             }
         }
-
-        Ok(indices)
+        Ok(())
     }
 }
 

crates/bevy_render/src/mesh/mesh.rs

@@ -1,8 +1,8 @@
 use super::Vertex;
 use crate::{
     pipeline::{
-        AsVertexBufferDescriptor, IndexFormat, PrimitiveTopology, RenderPipelines,
-        VertexBufferDescriptor, VertexBufferDescriptors, VertexFormat,
+        AsVertexBufferDescriptor, PrimitiveTopology, RenderPipelines, VertexBufferDescriptor,
+        VertexBufferDescriptors, VertexFormat,
     },
     renderer::{BufferInfo, BufferUsage, RenderResourceContext, RenderResourceId},
 };
@@ -106,11 +106,17 @@ pub enum MeshToVertexBufferError {
     },
 }
 
+#[derive(Debug)]
+pub enum Indices {
+    U16(Vec<u16>),
+    U32(Vec<u32>),
+}
+
 #[derive(Debug)]
 pub struct Mesh {
     pub primitive_topology: PrimitiveTopology,
     pub attributes: Vec<VertexAttribute>,
-    pub indices: Option<Vec<u32>>,
+    pub indices: Option<Indices>,
 }
 
 impl Mesh {
@@ -159,23 +165,17 @@ impl Mesh {
         Ok(bytes)
     }
 
-    pub fn get_index_buffer_bytes(&self, index_format: IndexFormat) -> Option<Vec<u8>> {
-        self.indices.as_ref().map(|indices| match index_format {
-            IndexFormat::Uint16 => indices
-                .iter()
-                .map(|i| *i as u16)
-                .collect::<Vec<u16>>()
-                .as_slice()
-                .as_bytes()
-                .to_vec(),
-            IndexFormat::Uint32 => indices.as_slice().as_bytes().to_vec(),
+    pub fn get_index_buffer_bytes(&self) -> Option<Vec<u8>> {
+        self.indices.as_ref().map(|indices| match &indices {
+            Indices::U16(indices) => indices.as_slice().as_bytes().to_vec(),
+            Indices::U32(indices) => indices.as_slice().as_bytes().to_vec(),
         })
     }
 }
 
 /// Generation for some primitive shape meshes.
 pub mod shape {
-    use super::{Mesh, VertexAttribute};
+    use super::{Indices, Mesh, VertexAttribute};
     use crate::pipeline::PrimitiveTopology;
     use bevy_math::*;
     use hexasphere::Hexasphere;
@@ -237,14 +237,14 @@ pub mod shape {
                 uvs.push(*uv);
             }
 
-            let indices = vec![
+            let indices = Indices::U16(vec![
                 0, 1, 2, 2, 3, 0, // top
                 4, 5, 6, 6, 7, 4, // bottom
                 8, 9, 10, 10, 11, 8, // right
                 12, 13, 14, 14, 15, 12, // left
                 16, 17, 18, 18, 19, 16, // front
                 20, 21, 22, 22, 23, 20, // back
-            ];
+            ]);
 
             Mesh {
                 primitive_topology: PrimitiveTopology::TriangleList,
@@ -333,7 +333,7 @@ pub mod shape {
                 ]
             };
 
-            let indices = vec![0, 2, 1, 0, 3, 2];
+            let indices = Indices::U16(vec![0, 2, 1, 0, 3, 2]);
 
             let mut positions = Vec::new();
             let mut normals = Vec::new();
@@ -373,7 +373,7 @@ pub mod shape {
                 ([-extent, 0.0, -extent], [0.0, 1.0, 0.0], [0.0, 1.0]),
             ];
 
-            let indices = vec![0, 2, 1, 0, 3, 2];
+            let indices = Indices::U16(vec![0, 2, 1, 0, 3, 2]);
 
             let mut positions = Vec::new();
             let mut normals = Vec::new();
@@ -455,6 +455,8 @@ pub mod shape {
                 hexasphere.get_indices(i, &mut indices);
             }
 
+            let indices = Indices::U32(indices);
+
             Mesh {
                 primitive_topology: PrimitiveTopology::TriangleList,
                 attributes: vec![
@@ -544,7 +546,7 @@ pub fn mesh_resource_provider_system(
                 &vertex_bytes,
             );
 
-            let index_bytes = mesh.get_index_buffer_bytes(IndexFormat::Uint16).unwrap();
+            let index_bytes = mesh.get_index_buffer_bytes().unwrap();
             let index_buffer = render_resource_context.create_buffer_with_data(
                 BufferInfo {
                     buffer_usage: BufferUsage::INDEX,

crates/bevy_render/src/pipeline/pipeline_compiler.rs

@@ -1,4 +1,6 @@
-use super::{state_descriptors::PrimitiveTopology, PipelineDescriptor, VertexBufferDescriptors};
+use super::{
+    state_descriptors::PrimitiveTopology, IndexFormat, PipelineDescriptor, VertexBufferDescriptors,
+};
 use crate::{
     renderer::RenderResourceContext,
     shader::{Shader, ShaderSource},
@@ -14,6 +16,7 @@ pub struct PipelineSpecialization {
     pub shader_specialization: ShaderSpecialization,
     pub primitive_topology: PrimitiveTopology,
     pub dynamic_bindings: Vec<DynamicBinding>,
+    pub index_format: IndexFormat,
     pub sample_count: u32,
 }
 
@@ -24,6 +27,7 @@ impl Default for PipelineSpecialization {
             shader_specialization: Default::default(),
             primitive_topology: Default::default(),
             dynamic_bindings: Default::default(),
+            index_format: IndexFormat::Uint16,
         }
     }
 }
@@ -161,6 +165,7 @@ impl PipelineCompiler {
 
         specialized_descriptor.sample_count = pipeline_specialization.sample_count;
         specialized_descriptor.primitive_topology = pipeline_specialization.primitive_topology;
+        specialized_descriptor.index_format = pipeline_specialization.index_format;
 
         let specialized_pipeline_handle = pipelines.add(specialized_descriptor);
         render_resource_context.create_render_pipeline(

crates/bevy_render/src/pipeline/render_pipelines.rs

@@ -1,10 +1,11 @@
-use super::{PipelineDescriptor, PipelineSpecialization};
+use super::{IndexFormat, PipelineDescriptor, PipelineSpecialization};
 use crate::{
     draw::{Draw, DrawContext},
+    mesh::{Indices, Mesh},
     prelude::Msaa,
     renderer::RenderResourceBindings,
 };
-use bevy_asset::Handle;
+use bevy_asset::{Assets, Handle};
 use bevy_ecs::{Query, Res, ResMut};
 use bevy_property::Properties;
 
@@ -75,15 +76,25 @@ pub fn draw_render_pipelines_system(
     mut draw_context: DrawContext,
     mut render_resource_bindings: ResMut<RenderResourceBindings>,
     msaa: Res<Msaa>,
-    mut query: Query<(&mut Draw, &mut RenderPipelines)>,
+    meshes: Res<Assets<Mesh>>,
+    mut query: Query<(&mut Draw, &mut RenderPipelines, &Handle<Mesh>)>,
 ) {
-    for (mut draw, mut render_pipelines) in &mut query.iter() {
+    for (mut draw, mut render_pipelines, mesh_handle) in &mut query.iter() {
         if !draw.is_visible {
             continue;
         }
+
+        let mesh = meshes.get(mesh_handle).unwrap();
+        let (index_range, index_format) = match mesh.indices.as_ref() {
+            Some(Indices::U32(indices)) => (Some(0..indices.len() as u32), IndexFormat::Uint32),
+            Some(Indices::U16(indices)) => (Some(0..indices.len() as u32), IndexFormat::Uint16),
+            None => (None, IndexFormat::Uint16),
+        };
+
         let render_pipelines = &mut *render_pipelines;
         for pipeline in render_pipelines.pipelines.iter_mut() {
             pipeline.specialization.sample_count = msaa.samples;
+            pipeline.specialization.index_format = index_format;
         }
 
         for render_pipeline in render_pipelines.pipelines.iter() {
@@ -103,10 +114,10 @@ pub fn draw_render_pipelines_system(
                     ],
                 )
                 .unwrap();
-            let indices = draw_context
+            draw_context
                 .set_vertex_buffers_from_bindings(&mut draw, &[&render_pipelines.bindings])
                 .unwrap();
-            if let Some(indices) = indices {
+            if let Some(indices) = index_range.clone() {
                 draw.draw_indexed(indices, 0, 0..1);
             }
         }

crates/bevy_render/src/pipeline/state_descriptors.rs

@@ -182,7 +182,7 @@ impl Default for BlendOperation {
     }
 }
 
-#[derive(Copy, Clone, Debug, Hash, Eq, PartialEq)]
+#[derive(Copy, Clone, Debug, Hash, Eq, PartialEq, Serialize, Deserialize, Property)]
 pub enum IndexFormat {
     Uint16 = 0,
     Uint32 = 1,