diff --git a/vulkano/src/shader/mod.rs b/vulkano/src/shader/mod.rs index 1072f14a37..315c36b0f4 100644 --- a/vulkano/src/shader/mod.rs +++ b/vulkano/src/shader/mod.rs @@ -147,6 +147,7 @@ use crate::{ use ahash::{HashMap, HashSet}; use bytemuck::bytes_of; use half::f16; +use smallvec::SmallVec; use spirv::ExecutionModel; use std::{ borrow::Cow, @@ -169,8 +170,7 @@ pub struct ShaderModule { handle: ash::vk::ShaderModule, device: InstanceOwnedDebugWrapper>, id: NonZeroU64, - entry_point_map: HashMap>, - entry_point_infos: Vec, + entry_point_infos: SmallVec<[EntryPointInfo; 1]>, } impl ShaderModule { @@ -310,27 +310,11 @@ impl ShaderModule { entry_points: impl IntoIterator, ) -> Arc { let ShaderModuleCreateInfo { code: _, _ne: _ } = create_info; - - let mut entry_point_map: HashMap> = - HashMap::default(); - let entry_point_infos: Vec<_> = entry_points - .into_iter() - .enumerate() - .map(|(index, info)| { - entry_point_map - .entry(info.name.clone()) - .or_default() - .insert(ExecutionModel::from(&info.execution), index); - info - }) - .collect(); - Arc::new(ShaderModule { handle, device: InstanceOwnedDebugWrapper(device), id: Self::next_id(), - entry_point_map, - entry_point_infos, + entry_point_infos: entry_points.into_iter().collect(), }) } @@ -373,16 +357,7 @@ impl ShaderModule { /// name exist. #[inline] pub fn entry_point(self: &Arc, name: &str) -> Option { - self.entry_point_map.get(name).and_then(|infos| { - if infos.len() == 1 { - infos.iter().next().map(|(_, &info_index)| EntryPoint { - module: self.clone(), - info_index, - }) - } else { - None - } - }) + self.single_entry_point_filter(|info| info.name == name) } /// Returns information about the entry point with the provided name and execution model. @@ -393,13 +368,47 @@ impl ShaderModule { name: &str, execution: ExecutionModel, ) -> Option { - self.entry_point_map.get(name).and_then(|infos| { - infos.get(&execution).map(|&info_index| EntryPoint { - module: self.clone(), - info_index, - }) + self.single_entry_point_filter(|info| { + info.name == name && ExecutionModel::from(&info.execution) == execution + }) + } + + /// checks for *exactly* one entry point matching the `filter`, otherwise returns `None` + #[inline] + fn single_entry_point_filter

(self: &Arc, mut filter: P) -> Option + where + P: FnMut(&EntryPointInfo) -> bool, + { + let mut iter = self + .entry_point_infos + .iter() + .enumerate() + .filter(|(_, infos)| filter(infos)) + .map(|(x, _)| x); + let info_index = iter.next()?; + iter.next().is_none().then(|| EntryPoint { + module: self.clone(), + info_index, }) } + + /// Returns information about the entry point if `self` only contains a single entry point, + /// `None` otherwise. + #[inline] + pub fn single_entry_point(self: &Arc) -> Option { + self.single_entry_point_filter(|_| true) + } + + /// Returns information about the entry point if `self` only contains a single entry point + /// with the provided `ExecutionModel`. Returns `None` if no entry point was found or multiple + /// entry points have been found matching the provided `ExecutionModel`. + #[inline] + pub fn single_entry_point_of_execution( + self: &Arc, + execution: ExecutionModel, + ) -> Option { + self.single_entry_point_filter(|info| ExecutionModel::from(&info.execution) == execution) + } } impl Drop for ShaderModule {