Skip to main content

wgsl_types/
tplt.rs

1//! Built-in type-generator and function templates.
2
3#[cfg(feature = "naga-ext")]
4use crate::syntax::AccelerationStructureTag;
5use crate::{
6    Error,
7    inst::{Instance, LiteralInstance},
8    syntax::{AccessMode, AddressSpace, Enumerant, SampledType, TexelFormat},
9    ty::{TextureType, Ty, Type},
10};
11
12/// A single template parameter.
13#[derive(Clone, Debug, PartialEq)]
14pub enum TpltParam {
15    Type(Type),
16    Instance(Instance),
17    Enumerant(Enumerant),
18}
19
20type E = Error;
21
22// ------------------------
23// TYPE-GENERATOR TEMPLATES
24// ------------------------
25
26pub struct ArrayTemplate {
27    n: Option<usize>,
28    ty: Type,
29}
30
31impl ArrayTemplate {
32    pub fn new(ty: Type, n: Option<usize>) -> Self {
33        Self { n, ty }
34    }
35    pub fn parse(tplt: &[TpltParam]) -> Result<ArrayTemplate, E> {
36        let (ty, n) = match tplt {
37            [TpltParam::Type(ty)] => Ok((ty.clone(), None)),
38            [TpltParam::Type(ty), TpltParam::Instance(n)] => Ok((ty.clone(), Some(n.clone()))),
39            _ => Err(E::TemplateArgs("array")),
40        }?;
41        if let Some(n) = n {
42            let n = match n {
43                Instance::Literal(LiteralInstance::AbstractInt(n)) => (n > 0).then_some(n as usize),
44                Instance::Literal(LiteralInstance::I32(n)) => (n > 0).then_some(n as usize),
45                Instance::Literal(LiteralInstance::U32(n)) => (n > 0).then_some(n as usize),
46                #[cfg(feature = "naga-ext")]
47                Instance::Literal(LiteralInstance::I64(n)) => (n > 0).then_some(n as usize),
48                #[cfg(feature = "naga-ext")]
49                Instance::Literal(LiteralInstance::U64(n)) => (n > 0).then_some(n as usize),
50                _ => None,
51            }
52            .ok_or(E::Builtin(
53                "the array element count must evaluate to a `u32` or a `i32` greater than `0`",
54            ))?;
55            Ok(ArrayTemplate { n: Some(n), ty })
56        } else {
57            Ok(ArrayTemplate { n: None, ty })
58        }
59    }
60    pub fn ty(&self) -> Type {
61        Type::Array(Box::new(self.ty.clone()), self.n)
62    }
63    pub fn inner_ty(&self) -> Type {
64        self.ty.clone()
65    }
66    pub fn n(&self) -> Option<usize> {
67        self.n
68    }
69}
70
71#[cfg(feature = "naga-ext")]
72pub struct BindingArrayTemplate {
73    n: Option<usize>,
74    ty: Type,
75}
76
77#[cfg(feature = "naga-ext")]
78impl BindingArrayTemplate {
79    pub fn parse(tplt: &[TpltParam]) -> Result<BindingArrayTemplate, E> {
80        let (ty, n) = match tplt {
81            [TpltParam::Type(ty)] => Ok((ty.clone(), None)),
82            [TpltParam::Type(ty), TpltParam::Instance(n)] => Ok((ty.clone(), Some(n.clone()))),
83            _ => Err(E::TemplateArgs("binding_array")),
84        }?;
85        if let Some(n) = n {
86            let n = match n {
87                Instance::Literal(LiteralInstance::AbstractInt(n)) => (n > 0).then_some(n as usize),
88                Instance::Literal(LiteralInstance::I32(n)) => (n > 0).then_some(n as usize),
89                Instance::Literal(LiteralInstance::U32(n)) => (n > 0).then_some(n as usize),
90                Instance::Literal(LiteralInstance::I64(n)) => (n > 0).then_some(n as usize),
91                Instance::Literal(LiteralInstance::U64(n)) => (n > 0).then_some(n as usize),
92                _ => None,
93            }
94            .ok_or(E::Builtin(
95                "the binding_array element count must evaluate to a `u32` or a `i32` greater than `0`",
96            ))?;
97            Ok(BindingArrayTemplate { n: Some(n), ty })
98        } else {
99            Ok(BindingArrayTemplate { n: None, ty })
100        }
101    }
102    pub fn ty(&self) -> Type {
103        Type::BindingArray(Box::new(self.ty.clone()), self.n)
104    }
105    pub fn inner_ty(&self) -> Type {
106        self.ty.clone()
107    }
108    pub fn n(&self) -> Option<usize> {
109        self.n
110    }
111}
112
113pub struct VecTemplate {
114    ty: Type,
115}
116
117impl VecTemplate {
118    pub fn parse(tplt: &[TpltParam]) -> Result<VecTemplate, E> {
119        let ty = match tplt {
120            [TpltParam::Type(ty)] => Ok(ty.clone()),
121            _ => Err(E::TemplateArgs("vector")),
122        }?;
123        if ty.is_scalar() && ty.is_concrete() {
124            Ok(VecTemplate { ty })
125        } else {
126            Err(Error::Builtin("vector template type must be a scalar"))
127        }
128    }
129    pub fn ty(&self, n: u8) -> Type {
130        Type::Vec(n, self.ty.clone().into())
131    }
132    pub fn inner_ty(&self) -> &Type {
133        &self.ty
134    }
135}
136
137pub struct MatTemplate {
138    ty: Type,
139}
140
141impl MatTemplate {
142    pub fn parse(tplt: &[TpltParam]) -> Result<MatTemplate, E> {
143        let ty = match tplt {
144            [TpltParam::Type(ty)] => Ok(ty.clone()),
145            _ => Err(E::TemplateArgs("matrix")),
146        }?;
147        if ty.is_float() {
148            Ok(MatTemplate { ty })
149        } else {
150            Err(Error::Builtin("matrix template type must be f32 or f16"))
151        }
152    }
153    pub fn ty(&self, c: u8, r: u8) -> Type {
154        Type::Mat(c, r, self.ty.clone().into())
155    }
156
157    pub fn inner_ty(&self) -> &Type {
158        &self.ty
159    }
160}
161
162pub struct PtrTemplate {
163    pub space: AddressSpace,
164    pub ty: Type,
165    pub access: AccessMode,
166}
167
168impl PtrTemplate {
169    pub fn parse(tplt: &[TpltParam]) -> Result<PtrTemplate, E> {
170        let mut it = tplt.iter();
171        match (
172            it.next().cloned(),
173            it.next().cloned(),
174            it.next().cloned(),
175            it.next(),
176        ) {
177            (
178                Some(TpltParam::Enumerant(Enumerant::AddressSpace(space))),
179                Some(TpltParam::Type(ty)),
180                access,
181                None,
182            ) => {
183                if !ty.is_storable() {
184                    return Err(Error::Builtin("pointer type must be storable"));
185                }
186                let access = match access {
187                    Some(TpltParam::Enumerant(Enumerant::AccessMode(access))) => Some(access),
188                    _ => None,
189                };
190                // selecting the default access mode per address space.
191                // reference: <https://www.w3.org/TR/WGSL/#address-space>
192                let access = match (space, access) {
193                    (AddressSpace::Function, Some(access))
194                    | (AddressSpace::Private, Some(access))
195                    | (AddressSpace::Workgroup, Some(access))
196                    | (AddressSpace::Storage, Some(access)) => access,
197                    (AddressSpace::Function, None)
198                    | (AddressSpace::Private, None)
199                    | (AddressSpace::Workgroup, None) => AccessMode::ReadWrite,
200                    (AddressSpace::Uniform, Some(AccessMode::Read) | None) => AccessMode::Read,
201                    (AddressSpace::Uniform, _) => {
202                        return Err(Error::Builtin(
203                            "pointer in uniform address space must have a `read` access mode",
204                        ));
205                    }
206                    (AddressSpace::Storage, None) => AccessMode::Read,
207                    (AddressSpace::Handle, _) => {
208                        unreachable!("handle address space cannot be spelled")
209                    }
210                    (AddressSpace::Immediate, _) => {
211                        todo!("immediate")
212                    }
213                    #[cfg(feature = "naga-ext")]
214                    (AddressSpace::TaskPayload, _) => {
215                        todo!("task_payload")
216                    }
217                    #[cfg(feature = "naga-ext")]
218                    (AddressSpace::RayPayload | AddressSpace::IncomingRayPayload, access) => {
219                        access.unwrap_or(AccessMode::ReadWrite)
220                    }
221                };
222                Ok(PtrTemplate { space, ty, access })
223            }
224            _ => Err(E::TemplateArgs("pointer")),
225        }
226    }
227
228    pub fn ty(&self) -> Type {
229        Type::Ptr(self.space, self.ty.clone().into(), self.access)
230    }
231}
232
233pub struct AtomicTemplate {
234    pub ty: Type,
235}
236
237impl AtomicTemplate {
238    pub fn parse(tplt: &[TpltParam]) -> Result<AtomicTemplate, E> {
239        let ty = match tplt {
240            [TpltParam::Type(ty)] => Ok(ty.clone()),
241            _ => Err(E::TemplateArgs("atomic")),
242        }?;
243        #[cfg(feature = "naga-ext")]
244        if ty.is_f32() || ty.is_i64() || ty.is_u64() {
245            return Ok(AtomicTemplate { ty });
246        }
247        if ty.is_i32() || ty.is_u32() {
248            Ok(AtomicTemplate { ty })
249        } else {
250            Err(Error::Builtin("atomic template type must be an integer"))
251        }
252    }
253    pub fn ty(&self) -> Type {
254        Type::Atomic(self.ty.clone().into())
255    }
256    pub fn inner_ty(&self) -> Type {
257        self.ty.clone()
258    }
259}
260
261pub struct TextureTemplate {
262    ty: TextureType,
263}
264
265impl TextureTemplate {
266    pub fn parse(name: &str, tplt: &[TpltParam]) -> Result<TextureTemplate, E> {
267        let ty = match name {
268            "texture_1d" => TextureType::Sampled1D(Self::sampled_type(tplt)?),
269            "texture_2d" => TextureType::Sampled2D(Self::sampled_type(tplt)?),
270            "texture_2d_array" => TextureType::Sampled2DArray(Self::sampled_type(tplt)?),
271            "texture_3d" => TextureType::Sampled3D(Self::sampled_type(tplt)?),
272            "texture_cube" => TextureType::SampledCube(Self::sampled_type(tplt)?),
273            "texture_cube_array" => TextureType::SampledCubeArray(Self::sampled_type(tplt)?),
274            "texture_multisampled_2d" => TextureType::Multisampled2D(Self::sampled_type(tplt)?),
275            "texture_storage_1d" => {
276                let (tex, acc) = Self::texel_access(tplt)?;
277                TextureType::Storage1D(tex, acc)
278            }
279            "texture_storage_2d" => {
280                let (tex, acc) = Self::texel_access(tplt)?;
281                TextureType::Storage2D(tex, acc)
282            }
283            "texture_storage_2d_array" => {
284                let (tex, acc) = Self::texel_access(tplt)?;
285                TextureType::Storage2DArray(tex, acc)
286            }
287            "texture_storage_3d" => {
288                let (tex, acc) = Self::texel_access(tplt)?;
289                TextureType::Storage3D(tex, acc)
290            }
291            #[cfg(feature = "naga-ext")]
292            "texture_1d_array" => TextureType::Sampled1DArray(Self::sampled_type(tplt)?),
293            #[cfg(feature = "naga-ext")]
294            "texture_storage_1d_array" => {
295                let (tex, acc) = Self::texel_access(tplt)?;
296                TextureType::Storage1DArray(tex, acc)
297            }
298            #[cfg(feature = "naga-ext")]
299            "texture_multisampled_2d_array" => {
300                TextureType::Multisampled2DArray(Self::sampled_type(tplt)?)
301            }
302            _ => return Err(E::Builtin("not a templated texture type")),
303        };
304        Ok(Self { ty })
305    }
306    fn sampled_type(tplt: &[TpltParam]) -> Result<SampledType, E> {
307        match tplt {
308            [TpltParam::Type(ty)] => ty.try_into(),
309            [_] => Err(Error::Builtin(
310                "texture sampled type must be `i32`, `u32` or `f32`",
311            )),
312            _ => Err(Error::Builtin(
313                "sampled texture types take a single template parameter",
314            )),
315        }
316    }
317    fn texel_access(tplt: &[TpltParam]) -> Result<(TexelFormat, AccessMode), E> {
318        match tplt {
319            [
320                TpltParam::Enumerant(Enumerant::TexelFormat(texel)),
321                TpltParam::Enumerant(Enumerant::AccessMode(access)),
322            ] => Ok((*texel, *access)),
323            _ => Err(Error::Builtin(
324                "storage texture types take two template parameters",
325            )),
326        }
327    }
328    pub fn ty(&self) -> TextureType {
329        self.ty.clone()
330    }
331}
332
333pub struct BitcastTemplate {
334    ty: Type,
335}
336
337impl BitcastTemplate {
338    pub fn parse(tplt: &[TpltParam]) -> Result<BitcastTemplate, E> {
339        let ty = match tplt {
340            [TpltParam::Type(ty)] => Ok(ty.clone()),
341            _ => Err(E::TemplateArgs("bitcast")),
342        }?;
343        if ty.is_numeric() || ty.is_vec() && ty.inner_ty().is_numeric() {
344            Ok(BitcastTemplate { ty })
345        } else {
346            Err(Error::Builtin(
347                "bitcast template type must be a numeric scalar or numeric vector",
348            ))
349        }
350    }
351    pub fn ty(&self) -> &Type {
352        &self.ty
353    }
354    pub fn inner_ty(&self) -> Type {
355        self.ty.inner_ty()
356    }
357}
358
359#[cfg(feature = "naga-ext")]
360#[derive(Clone, Debug, PartialEq, Eq)]
361pub struct AccelerationStructureTags {
362    tags: Vec<AccelerationStructureTag>,
363}
364
365#[cfg(feature = "naga-ext")]
366impl AccelerationStructureTags {
367    pub fn parse(tplt: &[TpltParam]) -> Result<Self, E> {
368        let tags = tplt
369            .iter()
370            .map(|param| match param {
371                TpltParam::Enumerant(Enumerant::AccelerationStructureTag(tag)) => Ok(*tag),
372                TpltParam::Instance(_) => Err(Error::Builtin(
373                    "expected an acceleration structure tag, not an instance",
374                )),
375                TpltParam::Type(_) => Err(Error::Builtin(
376                    "expected an acceleration structure tag, not a type",
377                )),
378                _ => Err(Error::Builtin("unknown acceleration structure tag")),
379            })
380            .collect::<Result<Vec<_>, E>>()?;
381        Ok(Self { tags })
382    }
383
384    pub fn tags(&self) -> &[AccelerationStructureTag] {
385        &self.tags
386    }
387}