primase_core/shader/
specialise.rs

1//! Explicit, generic material feature specialisations, retained alongside the original lift.
2use crate::shader::{LiftedShader, ShaderBinding, StageBinding};
3use crate::{Component, Entity, Geometry, World};
4use iridia_interface::ResourceId;
5use std::collections::{BTreeMap, BTreeSet};
6use std::sync::Arc;
7/// A material's explicit request to specialise its source programs on their stated feature words.
8/// A producer may write this only from the exact lift and the material's binding. Data constants
9/// remain uniforms; consumers keep the original shader and lift as provenance.
10#[derive(Clone, Debug, Default)]
11pub struct ShaderSpecialisations(pub Vec<StageSpecialisation>);
12/// One stage's specialised copy and the complete audit of its feature decisions.
13#[derive(Clone, Debug)]
14pub struct StageSpecialisation {
15    pub shader: Entity,
16    pub program: Option<Arc<iridia::Program>>,
17    pub words: BTreeMap<(ResourceId, u32), u32>,
18    pub removed: Vec<ResourceId>,
19    pub folded: usize,
20    pub refusal: Option<String>,
21}
22impl ShaderSpecialisations {
23    /// Specialise only on the known words deciding a fully stated branch/select. A missing exact
24    /// lift or a refused transformation is a recorded refusal, never a guessed program.
25    pub fn of(world: &World, binding: &ShaderBinding) -> Self {
26        Self::shared(world, binding, &mut SpecialisationCache::default())
27    }
28    /// The specialised program of the stage that draws with `shader`, where its specialisation was not refused.
29    pub fn program(&self, shader: Entity) -> Option<&Arc<iridia::Program>> {
30        self.0
31            .iter()
32            .find(|s| s.shader == shader)
33            .and_then(|s| s.program.as_ref())
34    }
35    /// [`ShaderSpecialisations::of`], sharing `cache` with the other materials of the same world: a stage whose
36    /// program and selected feature words (each with its value) are another material's gets that material's program.
37    pub fn shared(world: &World, binding: &ShaderBinding, cache: &mut SpecialisationCache) -> Self {
38        Self(
39            binding
40                .stages
41                .iter()
42                .map(|stage| Self::stage(world, stage, cache))
43                .collect(),
44        )
45    }
46    fn stage(world: &World, stage: &StageBinding, cache: &mut SpecialisationCache) -> StageSpecialisation {
47        let refused = |why: &str| StageSpecialisation {
48            shader: stage.shader,
49            program: None,
50            words: Default::default(),
51            removed: vec![],
52            folded: 0,
53            refusal: Some(why.to_owned()),
54        };
55        let Some(source) = world.get::<LiftedShader>(stage.shader) else {
56            return refused("source has not supplied an exact lift");
57        };
58        let known = stage.known_words();
59        // which words decide a branch or select depends on which are known, not on their values
60        let keys: BTreeSet<_> = known.keys().copied().collect();
61        let selected = match cache.selected.entry((stage.shader, keys)).or_insert_with(|| {
62            iridia::specialise::feature_words(&source.0, &known).map(|w| w.into_keys().collect())
63        }) {
64            Ok(selected) => selected,
65            Err(why) => return refused(why),
66        };
67        let words: Words = known
68            .into_iter()
69            .filter(|(key, _)| selected.contains(key))
70            .collect();
71        match cache
72            .programs
73            .entry((stage.shader, words))
74            .or_insert_with_key(|(_, words)| {
75                iridia::specialise::specialise(&source.0, words).map(|result| Specialised {
76                    program: Arc::new(result.program),
77                    words: result.words,
78                    removed: result.unread,
79                    folded: result.folded,
80                })
81            }) {
82            Ok(shared) => StageSpecialisation {
83                shader: stage.shader,
84                program: Some(shared.program.clone()),
85                words: shared.words.clone(),
86                removed: shared.removed.clone(),
87                folded: shared.folded,
88                refusal: None,
89            },
90            Err(why) => refused(why),
91        }
92    }
93}
94
95/// A uniform word, by its buffer and byte offset.
96type Word = (ResourceId, u32);
97type Words = BTreeMap<Word, u32>;
98
99/// One specialised program and its audit, shared by every stage specialised on the same words.
100#[derive(Debug)]
101struct Specialised {
102    program: Arc<iridia::Program>,
103    words: Words,
104    removed: Vec<ResourceId>,
105    folded: usize,
106}
107
108/// The specialisations made for the materials of one world, so that materials of one program and the same selected
109/// feature words (each with its value) share one specialised program rather than each holding a copy.
110#[derive(Debug, Default)]
111pub struct SpecialisationCache {
112    /// Per program and set of known words, the words that decide a branch or select (or why none can be chosen).
113    selected: BTreeMap<(Entity, BTreeSet<Word>), Result<BTreeSet<Word>, String>>,
114    /// Per program and selected words with their values, its specialisation (or why it was refused).
115    programs: BTreeMap<(Entity, Words), Result<Specialised, String>>,
116}
117impl SpecialisationCache {
118    /// The specialised programs made so far.
119    pub fn programs(&self) -> usize {
120        self.programs.values().filter(|p| p.is_ok()).count()
121    }
122}
123impl Component for ShaderSpecialisations {
124    const NAME: &'static str = "core.shader_specialisations";
125    const GEOMETRY: Geometry = Geometry::None;
126    fn references(&self, out: &mut Vec<Entity>) {
127        out.extend(self.0.iter().map(|s| s.shader));
128    }
129    fn to_json(&self) -> Result<Option<serde_json::Value>, crate::JsonError> {
130        crate::json::to_value(&self.0.iter().map(|s|serde_json::json!({
131            "shader":s.shader,"words":s.words.iter().map(|((resource,offset),word)|(*resource,*offset,*word)).collect::<Vec<_>>(),
132            "resources_removed":s.removed,"expressions_folded":s.folded,"refusal":s.refusal,
133        })).collect::<Vec<_>>()).map(Some)
134    }
135    fn verify(&self, entity: Entity, world: &World) -> Result<(), String> {
136        let binding = world
137            .get::<ShaderBinding>(entity)
138            .ok_or("specialisations have no source binding")?;
139        for (index, stage) in self.0.iter().enumerate() {
140            let source = binding
141                .stages
142                .iter()
143                .find(|s| s.shader == stage.shader)
144                .ok_or("specialisation names an unbound shader")?;
145            if self.0.iter().take(index).any(|s| s.shader == stage.shader) {
146                return Err("duplicate stage specialisation".into());
147            }
148            let known = source.known_words();
149            if stage.words.iter().any(|(k, v)| known.get(k) != Some(v)) {
150                return Err("specialisation freezes a word the binding does not state".into());
151            }
152            if let Some(program) = &stage.program {
153                program.validate().map_err(|e| format!("{e:?}"))?;
154            }
155        }
156        Ok(())
157    }
158    fn apply_linear(&self, _: &crate::Linear) -> crate::world::Linearity<Self> {
159        crate::world::Linearity::Invariant
160    }
161}
162
163#[cfg(test)]
164mod tests {
165    use super::*;
166    use crate::shader::{ConstantBinding, ConstantRole};
167    use iridia::{ResourceKind, Stage};
168
169    /// A pixel program whose first constant word switches between a texture sample scaled by the second word and a
170    /// coordinate.
171    fn program() -> iridia::Program {
172        iridia_mojo::lift(
173            "uniform vec4 ps_uniforms_vec4[1];\nuniform sampler2D ps_s0;\nvarying vec4 vTexCoord0;\nvec4 ps_r0;\nvoid main() {\nps_r0 = texture2D(ps_s0,vTexCoord0.xy);\nif (ps_uniforms_vec4[0].x != 0.0) {\nps_r0 = ps_r0 * ps_uniforms_vec4[0].yyyy;\n} else {\nps_r0 = vTexCoord0;\n}\ngl_FragColor = ps_r0;\n}",
174            Stage::Pixel,
175            &iridia_mojo::Hints {
176                target: "fixture".into(),
177                constants: vec![iridia_mojo::Constant {
178                    name: "SwitchAndScale".into(),
179                    offset: 0,
180                    bytes: 16,
181                    width: 4,
182                }],
183                samplers: vec![iridia_mojo::Sampler {
184                    glsl_name: "ps_s0".into(),
185                    name: "image".into(),
186                    slot: 0,
187                }],
188                attributes: vec![],
189            },
190        )
191        .unwrap()
192    }
193
194    /// A material binding `shader` with the switch and scale words stated.
195    fn binding(shader: Entity, buffer: ResourceId, switch: f32, scale: f32) -> ShaderBinding {
196        ShaderBinding {
197            stages: vec![StageBinding {
198                constants: vec![ConstantBinding {
199                    resource: buffer,
200                    offset: 0,
201                    size: 8,
202                    name: None,
203                    role: ConstantRole::Material(vec![switch.to_bits(), scale.to_bits()]),
204                    lines: None,
205                }],
206                ..StageBinding::new(Stage::Pixel, shader)
207            }],
208            state: Default::default(),
209        }
210    }
211
212    /// Materials of one program share its specialisation exactly when every selected word has the same value: a
213    /// word that decides nothing (the scale) does not split them, and the switch with another value does.
214    #[test]
215    fn materials_share_a_specialisation_only_when_the_selected_words_are_equal() {
216        let source = program();
217        let buffer = source
218            .interface
219            .resources
220            .iter()
221            .find(|r| matches!(r.kind, ResourceKind::ConstantBuffer { .. }))
222            .unwrap()
223            .id;
224        let mut world = World::new();
225        let shader = world.spawn();
226        world.insert(shader, LiftedShader(Arc::new(source)));
227        let mut cache = SpecialisationCache::default();
228        let mut of = |switch: f32, scale: f32| {
229            let set =
230                ShaderSpecialisations::shared(&world, &binding(shader, buffer, switch, scale), &mut cache);
231            let [stage] = &set.0[..] else { panic!("one stage") };
232            assert_eq!(stage.refusal, None);
233            (stage.program.clone().unwrap(), stage.words.clone())
234        };
235        let (on, on_words) = of(1.0, 2.0);
236        let (on_again, _) = of(1.0, 3.0);
237        let (off, off_words) = of(0.0, 2.0);
238        // the switch alone is selected, with its value
239        assert_eq!(on_words, Words::from([((buffer, 0), 1f32.to_bits())]));
240        assert_eq!(off_words, Words::from([((buffer, 0), 0f32.to_bits())]));
241        assert!(Arc::ptr_eq(&on, &on_again));
242        assert!(!Arc::ptr_eq(&on, &off));
243        // the switch off leaves the sample unread: the two programs differ
244        let textures = |p: &iridia::Program| {
245            p.interface
246                .resources
247                .iter()
248                .filter(|r| matches!(r.kind, ResourceKind::Texture { .. }))
249                .count()
250        };
251        assert_eq!((textures(&on), textures(&off)), (1, 0));
252        assert_eq!(cache.programs(), 2);
253    }
254
255    /// A stage whose program has no exact lift is refused, by reason.
256    #[test]
257    fn a_stage_without_an_exact_lift_is_refused() {
258        let mut world = World::new();
259        let shader = world.spawn();
260        let set = ShaderSpecialisations::of(&world, &binding(shader, ResourceId(0), 1.0, 1.0));
261        assert_eq!(
262            set.0[0].refusal.as_deref(),
263            Some("source has not supplied an exact lift")
264        );
265        assert!(set.0[0].program.is_none());
266    }
267}
primase_core/shader/specialise.rs

primase_core/shader/
specialise.rs
