feat(engine): make shader entry point names not hardcoded

1 files changed, 161 insertions, 75 deletions

diff --git a/engine/src/rendering/shader.rs b/engine/src/rendering/shader.rs
index 2a2ac7a..05fc212 100644
--- a/engine/src/rendering/shader.rs
+++ b/engine/src/rendering/shader.rs
@@ -5,8 +5,9 @@ use std::path::Path;
 use std::str::Utf8Error;
 use std::sync::Arc;
 
-use bitflags::{bitflags, bitflags_match};
+use bitflags::{bitflags, Flags as _};
 use ecs::phase::INIT as INIT_PHASE;
+use owning_ref::OwningRef;
 use shader_slang::{
     Blob as SlangBlob,
     ComponentType as SlangComponentType,
@@ -33,6 +34,7 @@ use crate::ecs::phase::Phase;
 use crate::ecs::sole::Single;
 use crate::ecs::{declare_entity, pair, Sole};
 use crate::mesh::vertex_buffer::VertexLabel;
+use crate::rendering::shader::reflection::Stage;
 
 pub mod cursor;
 pub mod reflection;
@@ -49,7 +51,7 @@ pub struct ModuleSource
     pub name: Cow<'static, str>,
     pub file_path: Cow<'static, Path>,
     pub source: Cow<'static, str>,
-    pub link_entrypoints: EntrypointFlags,
+    pub link_entry_points: EntrypointFlags,
 }
 
 impl Asset for ModuleSource
@@ -171,23 +173,23 @@ impl Program
         target_id: CompileTargetId,
     ) -> Result<VertexDescription, VertexDescCreationError>
     {
-        if !self
-            .metadata
-            .link_entrypoints
-            .contains(EntrypointFlags::VERTEX)
-        {
+        let Some(vertex_stage_entry_point_idx) =
+            self.metadata.vertex_stage_entry_point_idx
+        else {
             return Err(VertexDescCreationError::NoVertexStageEntrypointInProgram);
-        }
+        };
 
         let program_reflection = self
             .reflection(target_id)
             .map_err(VertexDescCreationError::GetProgramReflectionFailed)?;
 
-        VertexDescription::new(
-            &program_reflection
-                .get_entry_point_by_name("vertex_main")
-                .expect("Not possible"),
-        )
+        let Some(vertex_stage_entry_point) =
+            program_reflection.get_entry_point_by_index(vertex_stage_entry_point_idx)
+        else {
+            unreachable!();
+        };
+
+        VertexDescription::new(&vertex_stage_entry_point)
     }
 }
 
@@ -301,51 +303,130 @@ impl Context
         self.programs.get(asset_id)
     }
 
+    pub fn load_module(
+        &self,
+        module_source: &ModuleSource,
+    ) -> Result<Module, LoadModuleError>
+    {
+        let Some(session) = self.session.as_ref() else {
+            return Err(LoadModuleError::ContextSessionNotInitialized);
+        };
+
+        let module = session
+            .load_module_from_source_string(
+                &module_source.name,
+                &module_source.file_path.to_string_lossy(),
+                &module_source.source,
+            )
+            .map_err(Error)?;
+
+        Ok(Module { inner: module })
+    }
+
     #[tracing::instrument(skip_all)]
-    pub fn compose_into_program(
+    pub fn compose_program(
         &self,
-        module: Module,
-        link_entrypoints: EntrypointFlags,
+        modules: &[Module],
+        link_entry_points: EntrypointFlags,
+        target_id: CompileTargetId,
     ) -> Result<Program, ComposeProgramError>
     {
         let Some(session) = self.session.as_ref() else {
             return Err(ComposeProgramError::ContextSessionNotInitialized);
         };
 
-        let entry_points = link_entrypoints
+        let all_entry_points = modules
             .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))
+            .flat_map(Module::entry_points)
+            .enumerate()
+            .map(|(entry_point_index, entry_point)| {
+                Ok::<_, Error>((
+                    entry_point_index,
+                    OwningRef::new(SlangObjectRc::<SlangComponentType>::from(
+                        entry_point.inner,
+                    ))
+                    .try_map(|entry_point| {
+                        entry_point.layout(i64::from(target_id.index))
+                    })?,
+                ))
             })
             .collect::<Result<Vec<_>, _>>()?;
 
-        let components = entry_points
-            .into_iter()
-            .map(|entry_point| {
-                SlangObjectRc::<SlangComponentType>::from(entry_point.inner)
-            })
-            .chain([SlangObjectRc::<SlangComponentType>::from(module.inner)])
-            .collect::<Vec<_>>();
+        let mut program_metadata = ProgramMetadata::default();
+
+        let mut components = Vec::with_capacity(
+            EntrypointFlags::iter_defined_names().count() + modules.len(),
+        );
+
+        components.extend(modules.iter().map(|module| {
+            SlangObjectRc::<SlangComponentType>::from(module.inner.clone())
+        }));
+
+        if link_entry_points.contains(EntrypointFlags::VERTEX) {
+            let mut entry_point_iter =
+                all_entry_points.iter().filter(|(_, entry_point)| {
+                    Stage::from_slang(
+                        entry_point.entry_point_by_index(0).unwrap().stage(),
+                    ) == Stage::Vertex
+                });
+
+            let (vertex_stage_entry_point_idx, vertex_stage_entry_point) =
+                entry_point_iter
+                    .next()
+                    .ok_or(ComposeProgramError::NoEntryPointForStage(Stage::Vertex))
+                    .cloned()?;
+
+            if entry_point_iter.next().is_some() {
+                return Err(ComposeProgramError::MultipleEntryPointsForStage(
+                    Stage::Vertex,
+                ));
+            }
+
+            // This is fine. It is very unlikely there are more than 2^32 entry points
+            program_metadata.vertex_stage_entry_point_idx =
+                Some(u32::try_from(vertex_stage_entry_point_idx).expect(
+                    "Shader's vertex stage entry point index would overflow a u32",
+                ));
+
+            components.push(vertex_stage_entry_point.as_owner().clone());
+        }
+
+        if link_entry_points.contains(EntrypointFlags::FRAGMENT) {
+            let mut entry_point_iter =
+                all_entry_points.iter().filter(|(_, entry_point)| {
+                    Stage::from_slang(
+                        entry_point.entry_point_by_index(0).unwrap().stage(),
+                    ) == Stage::Fragment
+                });
+
+            let (fragment_stage_entry_point_idx, fragment_stage_entry_point) =
+                entry_point_iter
+                    .next()
+                    .ok_or(ComposeProgramError::NoEntryPointForStage(Stage::Fragment))
+                    .cloned()?;
+
+            if entry_point_iter.next().is_some() {
+                return Err(ComposeProgramError::MultipleEntryPointsForStage(
+                    Stage::Fragment,
+                ));
+            }
+
+            // This is fine. It is very unlikely there are more than 2^32 entry points
+            program_metadata.fragment_stage_entry_point_idx =
+                Some(u32::try_from(fragment_stage_entry_point_idx).expect(
+                    "Shader's fragment stage entry point index would overflow a u32",
+                ));
+
+            components.push(fragment_stage_entry_point.as_owner().clone());
+        }
 
         let program = session
             .create_composite_component_type(&components)
-            .map_err(|err| ComposeProgramError::Other(Error(err)))?;
+            .map_err(Error)?;
 
         Ok(Program {
             inner: program,
-            metadata: ProgramMetadata { link_entrypoints },
+            metadata: program_metadata,
         })
     }
 }
@@ -524,11 +605,12 @@ impl CompileTargetFormat
     }
 }
 
-#[derive(Debug, Clone)]
+#[derive(Debug, Default, Clone)]
 #[non_exhaustive]
 pub struct ProgramMetadata
 {
-    pub link_entrypoints: EntrypointFlags,
+    pub vertex_stage_entry_point_idx: Option<u32>,
+    pub fragment_stage_entry_point_idx: Option<u32>,
 }
 
 #[derive(Debug, Clone)]
@@ -728,17 +810,14 @@ pub enum ComposeProgramError
     #[error("A shader context session has not been initialized")]
     ContextSessionNotInitialized,
 
-    #[error("Shader module does not have a '{entrypoint_name}' entry point")]
-    EntrypointNotFoundError
-    {
-        entrypoint_name: &'static str
-    },
+    #[error("Multiple entry points found for shader stage {0:?}")]
+    MultipleEntryPointsForStage(Stage),
 
-    #[error("Failed to create vertex description")]
-    VertexDescriptionCreationFailed(#[from] VertexDescCreationError),
+    #[error("No entry point was found for shader stage {0:?}")]
+    NoEntryPointForStage(Stage),
 
     #[error(transparent)]
-    Other(Error),
+    Other(#[from] Error),
 }
 
 #[derive(Debug, thiserror::Error)]
@@ -755,6 +834,16 @@ pub enum ContextSessionInitError
 }
 
 #[derive(Debug, thiserror::Error)]
+pub enum LoadModuleError
+{
+    #[error("A shader context session has not been initialized")]
+    ContextSessionNotInitialized,
+
+    #[error(transparent)]
+    Other(#[from] Error),
+}
+
+#[derive(Debug, thiserror::Error)]
 #[error(transparent)]
 pub struct Error(#[from] shader_slang::Error);
 
@@ -850,7 +939,7 @@ fn import(
         source: std::fs::read_to_string(file_path)
             .map_err(ImportError::ReadFileFailed)?
             .into(),
-        link_entrypoints: settings.map_or_default(|settings| settings.link_entrypoints),
+        link_entry_points: settings.map_or_default(|settings| settings.link_entrypoints),
     });
 
     Ok(())
@@ -914,6 +1003,7 @@ fn load_modules(
     mut context: Single<Context>,
     assets: Single<Assets>,
     asset_events: Single<AssetEvents>,
+    rendering_backend_properties: Single<crate::rendering::BackendProperties>,
 ) -> Result<(), crate::Error>
 {
     let Ok(context) = context.get_mut() else {
@@ -923,7 +1013,7 @@ fn load_modules(
     let assets = assets.get()?;
     let asset_events = asset_events.get()?;
 
-    let Some(session) = context.session.as_ref() else {
+    let Ok(rendering_backend_properties) = rendering_backend_properties.get() else {
         return Ok(());
     };
 
@@ -938,25 +1028,35 @@ fn load_modules(
             unreachable!();
         };
 
-        tracing::debug!(asset_label=?asset_label, "Loading shader module");
+        tracing::debug!(%asset_label, "Loading shader module");
 
-        let module = match load_module(session, module_source) {
+        let module = match context.load_module(module_source) {
             Ok(module) => module,
             Err(err) => {
-                tracing::error!("Failed to load shader module: {err}");
+                tracing::error!(%asset_label, "Failed to load shader module: {err}");
                 continue;
             }
         };
 
-        if !module_source.link_entrypoints.is_empty() {
-            assert!(!context.programs.contains_key(asset_id));
+        if !module_source.link_entry_points.is_empty() {
+            if context.programs.contains_key(asset_id) {
+                tracing::debug!(
+                    %asset_label,
+                    "Removing old shader program associated with asset"
+                );
 
-            let shader_program = match context
-                .compose_into_program(module, module_source.link_entrypoints)
-            {
+                context.programs.remove(asset_id);
+            }
+
+            let shader_program = match context.compose_program(
+                &[module],
+                module_source.link_entry_points,
+                rendering_backend_properties.shader_target_id,
+            ) {
                 Ok(shader_program) => shader_program,
                 Err(err) => {
                     tracing::error!(
+                        %asset_label,
                         "Failed to compose shader into program: {:#}",
                         crate::Error::new(err)
                     );
@@ -967,7 +1067,7 @@ fn load_modules(
             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}");
+                    tracing::error!(%asset_label, "Failed to link shader: {err}");
                     continue;
                 }
             };
@@ -978,17 +1078,3 @@ fn load_modules(
 
     Ok(())
 }
-
-fn load_module(
-    session: &SlangSession,
-    module_source: &ModuleSource,
-) -> Result<Module, Error>
-{
-    let module = session.load_module_from_source_string(
-        &module_source.name,
-        &module_source.file_path.to_string_lossy(),
-        &module_source.source,
-    )?;
-
-    Ok(Module { inner: module })
-}