1use super::functions::*;
5use std::collections::HashMap;
6
7pub struct VariantCache {
12 pub capacity: usize,
14 pub(super) entries: Vec<(ShaderKey, CompiledVariant, u64)>,
16 pub(super) clock: u64,
18}
19impl VariantCache {
20 pub fn new(capacity: usize) -> Self {
22 Self {
23 capacity: capacity.max(1),
24 entries: Vec::new(),
25 clock: 0,
26 }
27 }
28 pub fn insert(&mut self, variant: CompiledVariant) {
30 if let Some(pos) = self.entries.iter().position(|(k, _, _)| *k == variant.key) {
31 self.clock += 1;
32 self.entries[pos] = (variant.key.clone(), variant, self.clock);
33 return;
34 }
35 if self.entries.len() >= self.capacity {
36 let lru_pos = self
37 .entries
38 .iter()
39 .enumerate()
40 .min_by_key(|(_, (_, _, t))| *t)
41 .map(|(i, _)| i)
42 .unwrap_or(0);
43 self.entries.swap_remove(lru_pos);
44 }
45 self.clock += 1;
46 self.entries
47 .push((variant.key.clone(), variant, self.clock));
48 }
49 pub fn get(&mut self, key: &ShaderKey) -> Option<&CompiledVariant> {
51 if let Some(pos) = self.entries.iter().position(|(k, _, _)| k == key) {
52 self.clock += 1;
53 self.entries[pos].2 = self.clock;
54 return Some(&self.entries[pos].1);
55 }
56 None
57 }
58 pub fn len(&self) -> usize {
60 self.entries.len()
61 }
62 pub fn is_empty(&self) -> bool {
64 self.entries.is_empty()
65 }
66 pub fn clear(&mut self) {
68 self.entries.clear();
69 self.clock = 0;
70 }
71 pub fn total_binary_bytes(&self) -> usize {
73 self.entries
74 .iter()
75 .map(|(_, v, _)| v.binary_size_bytes)
76 .sum()
77 }
78}
79#[derive(Debug, Clone)]
81pub struct CompiledVariant {
82 pub key: ShaderKey,
84 pub wgsl: String,
86 pub workgroup_size: [u32; 3],
88 pub binary_size_bytes: usize,
90}
91impl CompiledVariant {
92 fn new(key: ShaderKey, wgsl: String, workgroup_size: [u32; 3]) -> Self {
93 let binary_size_bytes = wgsl.len();
94 Self {
95 key,
96 wgsl,
97 workgroup_size,
98 binary_size_bytes,
99 }
100 }
101}
102#[derive(Debug, Default)]
104pub struct VariantProfileRegistry {
105 pub(super) profiles: HashMap<String, VariantProfile>,
106}
107impl VariantProfileRegistry {
108 pub fn new() -> Self {
110 Self::default()
111 }
112 pub fn register(&mut self, profile: VariantProfile) {
114 self.profiles.insert(profile.name.clone(), profile);
115 }
116 pub fn resolve(&self, name: &str) -> Option<HashMap<String, String>> {
120 let profile = self.profiles.get(name)?;
121 let mut merged = if let Some(base) = &profile.base {
122 self.resolve(base)?
123 } else {
124 HashMap::new()
125 };
126 for (k, v) in &profile.defines {
127 merged.insert(k.clone(), v.clone());
128 }
129 Some(merged)
130 }
131 pub fn profile_names(&self) -> Vec<&str> {
133 let mut names: Vec<&str> = self.profiles.keys().map(|s| s.as_str()).collect();
134 names.sort_unstable();
135 names
136 }
137}
138#[derive(Debug, Default)]
143pub struct ShaderDependencyGraph {
144 pub(super) depends_on: HashMap<String, Vec<String>>,
146}
147impl ShaderDependencyGraph {
148 pub fn new() -> Self {
150 Self::default()
151 }
152 pub fn add_dependency(&mut self, dependent: impl Into<String>, dependency: impl Into<String>) {
154 self.depends_on
155 .entry(dependent.into())
156 .or_default()
157 .push(dependency.into());
158 }
159 pub fn direct_dependents(&self, dependency: &str) -> Vec<&str> {
161 self.depends_on
162 .iter()
163 .filter(|(_, deps)| deps.iter().any(|d| d == dependency))
164 .map(|(name, _)| name.as_str())
165 .collect()
166 }
167 pub fn transitive_dependents(&self, changed_shader: &str) -> Vec<String> {
170 let mut visited = std::collections::HashSet::new();
171 let mut queue = std::collections::VecDeque::new();
172 queue.push_back(changed_shader.to_string());
173 while let Some(current) = queue.pop_front() {
174 for (name, deps) in &self.depends_on {
175 if deps.contains(¤t) && !visited.contains(name.as_str()) {
176 visited.insert(name.clone());
177 queue.push_back(name.clone());
178 }
179 }
180 }
181 let mut result: Vec<String> = visited.into_iter().collect();
182 result.sort_unstable();
183 result
184 }
185 pub fn direct_dependencies(&self, shader: &str) -> &[String] {
187 self.depends_on
188 .get(shader)
189 .map(Vec::as_slice)
190 .unwrap_or(&[])
191 }
192 pub fn shaders_with_deps(&self) -> Vec<&str> {
194 let mut names: Vec<&str> = self.depends_on.keys().map(|s| s.as_str()).collect();
195 names.sort_unstable();
196 names
197 }
198}
199#[derive(Debug, Clone, PartialEq, Eq, Hash)]
203pub struct ShaderKey {
204 pub name: String,
206 pub defines: Vec<(String, String)>,
208}
209impl ShaderKey {
210 pub fn new(name: impl Into<String>, defines: &HashMap<String, String>) -> Self {
212 let mut sorted: Vec<(String, String)> = defines
213 .iter()
214 .map(|(k, v)| (k.clone(), v.clone()))
215 .collect();
216 sorted.sort_by(|a, b| a.0.cmp(&b.0));
217 Self {
218 name: name.into(),
219 defines: sorted,
220 }
221 }
222 pub fn bare(name: impl Into<String>) -> Self {
224 Self {
225 name: name.into(),
226 defines: Vec::new(),
227 }
228 }
229 pub fn fingerprint(&self) -> String {
231 if self.defines.is_empty() {
232 return self.name.clone();
233 }
234 let suffix: String = self
235 .defines
236 .iter()
237 .map(|(k, v)| format!("{k}_{v}"))
238 .collect::<Vec<_>>()
239 .join("__");
240 format!("{}__{}", self.name, suffix)
241 }
242}
243#[derive(Debug, Clone)]
248pub struct VariantProfile {
249 pub name: String,
251 pub defines: HashMap<String, String>,
253 pub base: Option<String>,
255}
256impl VariantProfile {
257 pub fn new(name: impl Into<String>) -> Self {
259 Self {
260 name: name.into(),
261 defines: HashMap::new(),
262 base: None,
263 }
264 }
265 pub fn with_base(name: impl Into<String>, base: impl Into<String>) -> Self {
267 Self {
268 name: name.into(),
269 defines: HashMap::new(),
270 base: Some(base.into()),
271 }
272 }
273 pub fn set(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
275 self.defines.insert(key.into(), value.into());
276 self
277 }
278}
279#[derive(Debug, Default)]
284pub struct HotReloadTracker {
285 pub(super) timestamps: HashMap<String, u64>,
287 pub(super) compiled_at: HashMap<String, u64>,
289}
290impl HotReloadTracker {
291 pub fn new() -> Self {
293 Self::default()
294 }
295 pub fn touch(&mut self, name: impl Into<String>, time: u64) {
297 self.timestamps.insert(name.into(), time);
298 }
299 pub fn record_compile(&mut self, name: impl Into<String>, time: u64) {
301 self.compiled_at.insert(name.into(), time);
302 }
303 pub fn needs_recompile(&self, name: &str) -> bool {
305 let modified = self.timestamps.get(name).copied().unwrap_or(0);
306 let compiled = self.compiled_at.get(name).copied().unwrap_or(0);
307 modified > compiled
308 }
309 pub fn stale_shaders(&self) -> Vec<&str> {
311 self.timestamps
312 .keys()
313 .filter(|n| self.needs_recompile(n))
314 .map(|s| s.as_str())
315 .collect()
316 }
317}
318impl HotReloadTracker {
319 pub fn touch_batch(&mut self, names: &[&str], time: u64) {
321 for &name in names {
322 self.touch(name, time);
323 }
324 }
325 pub fn flush_stale(&mut self, time: u64) {
327 let stale: Vec<String> = self
328 .timestamps
329 .keys()
330 .filter(|n| self.needs_recompile(n))
331 .cloned()
332 .collect();
333 for name in stale {
334 self.compiled_at.insert(name, time);
335 }
336 }
337 pub fn never_compiled(&self) -> Vec<&str> {
339 self.timestamps
340 .keys()
341 .filter(|n| !self.compiled_at.contains_key(n.as_str()))
342 .map(|s| s.as_str())
343 .collect()
344 }
345}
346pub struct ShaderRegistry {
352 pub(super) sources: HashMap<String, ShaderSource>,
354 pub(super) cache: VariantCache,
356 pub cache_hits: u64,
358 pub compilations: u64,
360}
361impl ShaderRegistry {
362 pub fn new(cache_capacity: usize) -> Self {
364 Self {
365 sources: HashMap::new(),
366 cache: VariantCache::new(cache_capacity),
367 cache_hits: 0,
368 compilations: 0,
369 }
370 }
371 pub fn register(&mut self, source: ShaderSource) {
374 let name = source.name.clone();
375 self.sources.insert(name.clone(), source);
376 self.cache.entries.retain(|(k, _, _)| k.name != name);
377 }
378 pub fn get_or_compile(
383 &mut self,
384 name: &str,
385 defines: &HashMap<String, String>,
386 ) -> Result<&CompiledVariant, RegistryError> {
387 let key = ShaderKey::new(name, defines);
388 if self.cache.get(&key).is_some() {
389 self.cache_hits += 1;
390 return Ok(self
391 .cache
392 .get(&key)
393 .expect("cache entry must exist after insertion"));
394 }
395 let source = self
396 .sources
397 .get(name)
398 .ok_or_else(|| RegistryError::UnknownShader(name.to_string()))?;
399 for ph in &source.placeholders {
400 if !defines.contains_key(ph.as_str()) {
401 return Err(RegistryError::MissingDefine {
402 shader: name.to_string(),
403 define: ph.clone(),
404 });
405 }
406 }
407 let wgsl = source.instantiate(defines);
408 let workgroup_size = source.workgroup_size;
409 let variant = CompiledVariant::new(key.clone(), wgsl, workgroup_size);
410 self.compilations += 1;
411 self.cache.insert(variant);
412 Ok(self
413 .cache
414 .get(&key)
415 .expect("cache entry must exist after insertion"))
416 }
417 pub fn shader_names(&self) -> Vec<&str> {
419 self.sources.keys().map(|s| s.as_str()).collect()
420 }
421 pub fn cached_count(&self) -> usize {
423 self.cache.len()
424 }
425 pub fn invalidate(&mut self, name: &str) {
427 self.cache.entries.retain(|(k, _, _)| k.name != name);
428 }
429 pub fn invalidate_all(&mut self) {
431 self.cache.clear();
432 }
433}
434impl ShaderRegistry {
435 pub fn compile_variant(
441 &mut self,
442 name: &str,
443 defines: &HashMap<String, String>,
444 opts: &ShaderCompileOptions,
445 ) -> Result<CompiledVariant, RegistryError> {
446 let source = self
447 .sources
448 .get(name)
449 .ok_or_else(|| RegistryError::UnknownShader(name.to_string()))?;
450 let mut merged = defines.clone();
451 for (k, v) in &opts.extra_defines {
452 merged.entry(k.clone()).or_insert_with(|| v.clone());
453 }
454 for ph in &source.placeholders {
455 if !merged.contains_key(ph.as_str()) {
456 return Err(RegistryError::MissingDefine {
457 shader: name.to_string(),
458 define: ph.clone(),
459 });
460 }
461 }
462 let wgsl = source.instantiate(&merged);
463 if opts.max_source_bytes > 0 && wgsl.len() > opts.max_source_bytes {
464 return Err(RegistryError::SourceTooLarge {
465 shader: name.to_string(),
466 size: wgsl.len(),
467 limit: opts.max_source_bytes,
468 });
469 }
470 let workgroup_size = opts.workgroup_size;
471 let key = ShaderKey::new(name, &merged);
472 self.compilations += 1;
473 Ok(CompiledVariant::new(key, wgsl, workgroup_size))
474 }
475}
476impl ShaderRegistry {
477 pub fn apply_hot_reload(&mut self, tracker: &HotReloadTracker) -> Vec<String> {
481 let stale: Vec<String> = tracker
482 .stale_shaders()
483 .into_iter()
484 .map(|s| s.to_string())
485 .collect();
486 for name in &stale {
487 self.invalidate(name);
488 }
489 stale
490 }
491 pub fn registered_count(&self) -> usize {
493 self.sources.len()
494 }
495 pub fn source_bytes(&self, name: &str) -> usize {
497 self.sources.get(name).map(|s| s.wgsl.len()).unwrap_or(0)
498 }
499}
500#[derive(Debug, Clone)]
502pub struct ShaderSource {
503 pub name: String,
505 pub wgsl: String,
507 pub workgroup_size: [u32; 3],
509 pub placeholders: Vec<String>,
511}
512impl ShaderSource {
513 pub fn new(name: impl Into<String>, wgsl: impl Into<String>, workgroup_size: [u32; 3]) -> Self {
515 let wgsl_str: String = wgsl.into();
516 let placeholders = collect_placeholders(&wgsl_str);
517 Self {
518 name: name.into(),
519 wgsl: wgsl_str,
520 workgroup_size,
521 placeholders,
522 }
523 }
524 pub fn instantiate(&self, defines: &HashMap<String, String>) -> String {
528 let mut out = self.wgsl.clone();
529 for (k, v) in defines {
530 let token = format!("{{{{{}}}}}", k);
531 out = out.replace(&token, v);
532 }
533 out
534 }
535 pub fn threads_per_group(&self) -> u32 {
537 self.workgroup_size[0] * self.workgroup_size[1] * self.workgroup_size[2]
538 }
539}
540#[derive(Debug, Clone, PartialEq)]
542pub enum SpecConstValue {
543 Int(i64),
545 Uint(u64),
547 Float(f64),
549 Bool(bool),
551}
552impl SpecConstValue {
553 pub fn to_wgsl(&self) -> String {
555 match self {
556 SpecConstValue::Int(v) => v.to_string(),
557 SpecConstValue::Uint(v) => format!("{v}u"),
558 SpecConstValue::Float(v) => format!("{v}"),
559 SpecConstValue::Bool(v) => v.to_string(),
560 }
561 }
562}
563#[derive(Debug, Clone, PartialEq)]
565pub struct SpecializationConstant {
566 pub name: String,
568 pub default_value: SpecConstValue,
570 pub override_value: Option<SpecConstValue>,
572}
573impl SpecializationConstant {
574 pub fn new(name: impl Into<String>, default: SpecConstValue) -> Self {
576 Self {
577 name: name.into(),
578 default_value: default,
579 override_value: None,
580 }
581 }
582 pub fn with_override(mut self, value: SpecConstValue) -> Self {
584 self.override_value = Some(value);
585 self
586 }
587 pub fn effective_value(&self) -> &SpecConstValue {
589 self.override_value.as_ref().unwrap_or(&self.default_value)
590 }
591}
592#[derive(Debug, Default)]
595pub struct PipelineCache {
596 pub(super) entries: HashMap<PipelineCacheKey, String>,
597 pub hits: u64,
599 pub misses: u64,
601}
602impl PipelineCache {
603 pub fn new() -> Self {
605 Self::default()
606 }
607 pub fn insert(&mut self, key: PipelineCacheKey, label: impl Into<String>) {
609 self.entries.insert(key, label.into());
610 }
611 pub fn get(&mut self, key: &PipelineCacheKey) -> Option<&str> {
613 if let Some(v) = self.entries.get(key) {
614 self.hits += 1;
615 Some(v.as_str())
616 } else {
617 self.misses += 1;
618 None
619 }
620 }
621 pub fn len(&self) -> usize {
623 self.entries.len()
624 }
625 pub fn is_empty(&self) -> bool {
627 self.entries.is_empty()
628 }
629 pub fn clear(&mut self) {
631 self.entries.clear();
632 }
633 pub fn hit_rate(&self) -> f64 {
635 let total = self.hits + self.misses;
636 if total == 0 {
637 return 0.0;
638 }
639 self.hits as f64 / total as f64
640 }
641}
642#[derive(Debug, Clone)]
644pub struct PipelineDescriptor {
645 pub key: ShaderKey,
647 pub bind_group_count: u32,
649 pub push_constant_bytes: u32,
651 pub label: String,
653}
654impl PipelineDescriptor {
655 pub fn new(
657 key: ShaderKey,
658 bind_group_count: u32,
659 push_constant_bytes: u32,
660 label: impl Into<String>,
661 ) -> Self {
662 Self {
663 key,
664 bind_group_count,
665 push_constant_bytes,
666 label: label.into(),
667 }
668 }
669 pub fn validate(&self) -> Result<(), RegistryError> {
671 if self.bind_group_count == 0 {
672 return Err(RegistryError::UnknownShader(
673 "pipeline has no bind groups".into(),
674 ));
675 }
676 Ok(())
677 }
678}
679#[derive(Debug, Clone, PartialEq)]
681pub struct ShaderCompileOptions {
682 pub workgroup_size: [u32; 3],
684 pub extra_defines: HashMap<String, String>,
686 pub max_source_bytes: usize,
689}
690impl ShaderCompileOptions {
691 pub fn new() -> Self {
693 Self {
694 workgroup_size: [64, 1, 1],
695 extra_defines: HashMap::new(),
696 max_source_bytes: 0,
697 }
698 }
699}
700#[derive(Debug, Clone, Default)]
702pub struct SpecConstSet {
703 pub constants: HashMap<String, SpecializationConstant>,
705}
706impl SpecConstSet {
707 pub fn new() -> Self {
709 Self::default()
710 }
711 pub fn add(&mut self, constant: SpecializationConstant) {
713 self.constants.insert(constant.name.clone(), constant);
714 }
715 pub fn to_defines(&self) -> HashMap<String, String> {
718 self.constants
719 .iter()
720 .map(|(name, c)| (name.clone(), c.effective_value().to_wgsl()))
721 .collect()
722 }
723 pub fn has(&self, name: &str) -> bool {
725 self.constants.contains_key(name)
726 }
727 pub fn get_wgsl(&self, name: &str) -> Option<String> {
729 self.constants
730 .get(name)
731 .map(|c| c.effective_value().to_wgsl())
732 }
733}
734#[derive(Debug, Clone, PartialEq, Eq)]
736pub enum RegistryError {
737 UnknownShader(String),
739 MissingDefine {
741 shader: String,
743 define: String,
745 },
746 SourceTooLarge {
748 shader: String,
750 size: usize,
752 limit: usize,
754 },
755}
756#[derive(Debug, Clone, PartialEq, Eq, Hash)]
761pub struct PipelineCacheKey {
762 pub shader_key: ShaderKey,
764 pub workgroup_size: [u32; 3],
766 pub push_constant_bytes: u32,
768 pub layout_hash: u64,
770}
771impl PipelineCacheKey {
772 pub fn new(
774 shader_key: ShaderKey,
775 workgroup_size: [u32; 3],
776 push_constant_bytes: u32,
777 layout_hash: u64,
778 ) -> Self {
779 Self {
780 shader_key,
781 workgroup_size,
782 push_constant_bytes,
783 layout_hash,
784 }
785 }
786 pub fn hash_key(&self) -> u64 {
788 let repr = format!(
789 "{}_{:?}_{}_{}",
790 self.shader_key.fingerprint(),
791 self.workgroup_size,
792 self.push_constant_bytes,
793 self.layout_hash,
794 );
795 compute_cache_key(&repr)
796 }
797}