feat(engine): set GPU vertex attrs on DrawMesh command if necessaryHEADmaster

1 files changed, 95 insertions, 96 deletions

diff --git a/engine/src/rendering/shader.rs b/engine/src/rendering/shader.rs
index b508670..3bcf73e 100644
--- a/engine/src/rendering/shader.rs
+++ b/engine/src/rendering/shader.rs
@@ -4,6 +4,7 @@ use std::collections::HashMap;
 use std::fmt::Debug;
 use std::path::Path;
 use std::str::Utf8Error;
+use std::sync::Arc;
 
 use bitflags::{bitflags, bitflags_match};
 use ecs::phase::INIT as INIT_PHASE;
@@ -156,15 +157,24 @@ impl<'a> EntryPointReflection<'a>
 pub struct Program
 {
     inner: SlangComponentType,
+    metadata: ProgramMetadata,
 }
 
 impl Program
 {
-    pub fn link(&self) -> Result<Program, Error>
+    pub fn into_linked(self) -> Result<Program, Error>
     {
         let linked_program = self.inner.link()?;
 
-        Ok(Program { inner: linked_program })
+        Ok(Program {
+            inner: linked_program,
+            metadata: self.metadata,
+        })
+    }
+
+    pub fn metadata(&self) -> &ProgramMetadata
+    {
+        &self.metadata
     }
 
     pub fn get_entry_point_code(&self, entry_point_index: u32) -> Result<Blob, Error>
@@ -791,7 +801,7 @@ pub struct Context
     _global_session: SlangGlobalSession,
     session: SlangSession,
     modules: HashMap<AssetId, Module>,
-    programs: HashMap<AssetId, (Program, ProgramMetadata)>,
+    programs: HashMap<AssetId, Program>,
 }
 
 impl Context
@@ -803,38 +813,70 @@ impl Context
 
     pub fn get_program(&self, asset_id: &AssetId) -> Option<&Program>
     {
-        self.programs.get(asset_id).map(|(program, _)| program)
-    }
-
-    pub fn get_program_metadata(&self, asset_id: &AssetId) -> Option<&ProgramMetadata>
-    {
-        self.programs
-            .get(asset_id)
-            .map(|(_, program_metadata)| program_metadata)
+        self.programs.get(asset_id)
     }
 
     pub fn compose_into_program(
         &self,
-        modules: impl IntoIterator<Item = Module>,
-        entry_points: impl IntoIterator<Item = EntryPoint>,
-    ) -> Result<Program, Error>
+        module: Module,
+        link_entrypoints: EntrypointFlags,
+    ) -> Result<Program, ComposeProgramError>
     {
-        let components =
-            modules
-                .into_iter()
-                .map(|module| SlangComponentType::from(module.inner.clone()))
-                .chain(entry_points.into_iter().map(|entry_point| {
-                    SlangComponentType::from(entry_point.inner.clone())
-                }))
-                .collect::<Vec<_>>();
-
-        let program = self.session.create_composite_component_type(&components)?;
-
-        Ok(Program { inner: program })
+        let entry_points = link_entrypoints
+            .iter()
+            .filter_map(|entrypoint_flag| {
+                let entrypoint_name = bitflags_match!(entrypoint_flag, {
+                    EntrypointFlags::VERTEX => Some("vertex_main"),
+                    EntrypointFlags::FRAGMENT => Some("fragment_main"),
+                    _ => None
+                })?;
+
+                let Some(entry_point) = module.get_entry_point(entrypoint_name) else {
+                    return Some(Err(ComposeProgramError::EntrypointNotFoundError {
+                        entrypoint_name,
+                    }));
+                };
+
+                Some(Ok(entry_point))
+            })
+            .collect::<Result<Vec<_>, _>>()?;
+
+        let components = entry_points
+            .into_iter()
+            .map(|entry_point| SlangComponentType::from(entry_point.inner))
+            .chain([SlangComponentType::from(module.inner)])
+            .collect::<Vec<_>>();
+
+        let program = self
+            .session
+            .create_composite_component_type(&components)
+            .map_err(|err| ComposeProgramError::Other(Error(err)))?;
+
+        let vertex_desc = if link_entrypoints.contains(EntrypointFlags::VERTEX) {
+            // TODO: Do not have hard coded target here
+            let Ok(program_reflection) = program.layout(0) else {
+                unreachable!();
+            };
+
+            let program_reflection = ProgramReflection { inner: program_reflection };
+
+            Some(VertexDescription::new(
+                &program_reflection
+                    .get_entry_point_by_name("vertex_main")
+                    .expect("Not possible"),
+            )?)
+        } else {
+            None
+        };
+
+        Ok(Program {
+            inner: program,
+            metadata: ProgramMetadata { vertex_desc },
+        })
     }
 }
 
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 #[non_exhaustive]
 pub struct ProgramMetadata
 {
@@ -843,11 +885,11 @@ pub struct ProgramMetadata
     pub vertex_desc: Option<VertexDescription>,
 }
 
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 #[non_exhaustive]
 pub struct VertexDescription
 {
-    pub fields: Box<[VertexFieldDescription]>,
+    pub fields: Arc<[VertexFieldDescription]>,
 }
 
 impl VertexDescription
@@ -916,7 +958,7 @@ impl VertexDescription
             })
             .collect::<Result<Vec<_>, _>>()?;
 
-        Ok(Self { fields: fields.into_boxed_slice() })
+        Ok(Self { fields: fields.into() })
     }
 }
 
@@ -955,6 +997,22 @@ pub enum VertexDescriptionError
 }
 
 #[derive(Debug, thiserror::Error)]
+pub enum ComposeProgramError
+{
+    #[error("Shader module does not have a '{entrypoint_name}' entry point")]
+    EntrypointNotFoundError
+    {
+        entrypoint_name: &'static str
+    },
+
+    #[error("Failed to create vertex description")]
+    VertexDescriptionCreationFailed(#[from] VertexDescriptionError),
+
+    #[error(transparent)]
+    Other(Error),
+}
+
+#[derive(Debug, thiserror::Error)]
 #[error(transparent)]
 pub struct Error(#[from] shader_slang::Error);
 
@@ -1108,44 +1166,17 @@ fn load_modules(
         if !module_source.link_entrypoints.is_empty() {
             assert!(context.programs.get(asset_id).is_none());
 
-            let entry_points = match module_source
-                .link_entrypoints
-                .iter()
-                .filter_map(|entrypoint_flag| {
-                    let entrypoint_name = bitflags_match!(entrypoint_flag, {
-                        EntrypointFlags::VERTEX => Some("vertex_main"),
-                        EntrypointFlags::FRAGMENT => Some("fragment_main"),
-                        _ => None
-                    })?;
-
-                    let Some(entry_point) = module.get_entry_point(entrypoint_name)
-                    else {
-                        return Some(Err(EntrypointNotFoundError { entrypoint_name }));
-                    };
-
-                    Some(Ok(entry_point))
-                })
-                .collect::<Result<Vec<_>, EntrypointNotFoundError>>()
+            let shader_program = match context
+                .compose_into_program(module, module_source.link_entrypoints)
             {
-                Ok(entry_points) => entry_points,
-                Err(EntrypointNotFoundError { entrypoint_name }) => {
-                    tracing::error!(
-                        "Shader module does not have a '{entrypoint_name}' entry point"
-                    );
+                Ok(shader_program) => shader_program,
+                Err(err) => {
+                    tracing::error!("Failed to compose shader into program: {err}");
                     continue;
                 }
             };
 
-            let shader_program =
-                match context.compose_into_program([module], entry_points) {
-                    Ok(shader_program) => shader_program,
-                    Err(err) => {
-                        tracing::error!("Failed to compose shader into program: {err}");
-                        continue;
-                    }
-                };
-
-            let linked_shader_program = match shader_program.link() {
+            let linked_shader_program = match shader_program.into_linked() {
                 Ok(linked_shader_program) => linked_shader_program,
                 Err(err) => {
                     tracing::error!("Failed to link shader: {err}");
@@ -1153,33 +1184,7 @@ fn load_modules(
                 }
             };
 
-            let vertex_desc = if module_source
-                .link_entrypoints
-                .contains(EntrypointFlags::VERTEX)
-            {
-                VertexDescription::new(
-                    &shader_program
-                        .reflection(0)
-                        .expect("Not possible")
-                        .get_entry_point_by_name("vertex_main")
-                        .expect("Not possible"),
-                )
-                .inspect_err(|err| {
-                    tracing::error!(
-                        "Failed to create a vertex description for shader {}: {}",
-                        asset_label,
-                        err
-                    );
-                })
-                .ok()
-            } else {
-                None
-            };
-
-            context.programs.insert(
-                *asset_id,
-                (linked_shader_program, ProgramMetadata { vertex_desc }),
-            );
+            context.programs.insert(*asset_id, linked_shader_program);
         }
     }
 
@@ -1199,9 +1204,3 @@ fn load_module(
 
     Ok(Module { inner: module })
 }
-
-#[derive(Debug)]
-struct EntrypointNotFoundError
-{
-    entrypoint_name: &'static str,
-}