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 new_unchecked(ty: Type) -> Self {
119        Self { ty }
120    }
121    pub fn parse(tplt: &[TpltParam]) -> Result<VecTemplate, E> {
122        let ty = match tplt {
123            [TpltParam::Type(ty)] => Ok(ty.clone()),
124            _ => Err(E::TemplateArgs("vector")),
125        }?;
126        if ty.is_scalar() && ty.is_concrete() {
127            Ok(VecTemplate { ty })
128        } else {
129            Err(Error::Builtin("vector template type must be a scalar"))
130        }
131    }
132    pub fn ty(&self, n: u8) -> Type {
133        Type::Vec(n, self.ty.clone().into())
134    }
135    pub fn inner_ty(&self) -> &Type {
136        &self.ty
137    }
138}
139
140pub struct MatTemplate {
141    ty: Type,
142}
143
144impl MatTemplate {
145    pub fn parse(tplt: &[TpltParam]) -> Result<MatTemplate, E> {
146        let ty = match tplt {
147            [TpltParam::Type(ty)] => Ok(ty.clone()),
148            _ => Err(E::TemplateArgs("matrix")),
149        }?;
150        if ty.is_float() {
151            Ok(MatTemplate { ty })
152        } else {
153            Err(Error::Builtin("matrix template type must be f32 or f16"))
154        }
155    }
156    pub fn ty(&self, c: u8, r: u8) -> Type {
157        Type::Mat(c, r, self.ty.clone().into())
158    }
159
160    pub fn inner_ty(&self) -> &Type {
161        &self.ty
162    }
163}
164
165pub struct PtrTemplate {
166    pub space: AddressSpace,
167    pub ty: Type,
168    pub access: AccessMode,
169}
170
171impl PtrTemplate {
172    pub fn parse(tplt: &[TpltParam]) -> Result<PtrTemplate, E> {
173        let mut it = tplt.iter();
174        match (
175            it.next().cloned(),
176            it.next().cloned(),
177            it.next().cloned(),
178            it.next(),
179        ) {
180            (
181                Some(TpltParam::Enumerant(Enumerant::AddressSpace(space))),
182                Some(TpltParam::Type(ty)),
183                access,
184                None,
185            ) => {
186                if !ty.is_storable() {
187                    return Err(Error::Builtin("pointer type must be storable"));
188                }
189                let access = match access {
190                    Some(TpltParam::Enumerant(Enumerant::AccessMode(access))) => Some(access),
191                    _ => None,
192                };
193                // selecting the default access mode per address space.
194                // reference: <https://www.w3.org/TR/WGSL/#address-space>
195                let access = match (space, access) {
196                    (AddressSpace::Function, Some(access))
197                    | (AddressSpace::Private, Some(access))
198                    | (AddressSpace::Workgroup, Some(access))
199                    | (AddressSpace::Storage, Some(access)) => access,
200                    (AddressSpace::Function, None)
201                    | (AddressSpace::Private, None)
202                    | (AddressSpace::Workgroup, None) => AccessMode::ReadWrite,
203                    (AddressSpace::Uniform, Some(AccessMode::Read) | None) => AccessMode::Read,
204                    (AddressSpace::Uniform, _) => {
205                        return Err(Error::Builtin(
206                            "pointer in uniform address space must have a `read` access mode",
207                        ));
208                    }
209                    (AddressSpace::Storage, None) => AccessMode::Read,
210                    (AddressSpace::Handle, _) => {
211                        unreachable!("handle address space cannot be spelled")
212                    }
213                    (AddressSpace::Immediate, _) => {
214                        todo!("immediate")
215                    }
216                    #[cfg(feature = "naga-ext")]
217                    (AddressSpace::TaskPayload, _) => {
218                        todo!("task_payload")
219                    }
220                    #[cfg(feature = "naga-ext")]
221                    (AddressSpace::RayPayload | AddressSpace::IncomingRayPayload, access) => {
222                        access.unwrap_or(AccessMode::ReadWrite)
223                    }
224                };
225                Ok(PtrTemplate { space, ty, access })
226            }
227            _ => Err(E::TemplateArgs("pointer")),
228        }
229    }
230
231    pub fn ty(&self) -> Type {
232        Type::Ptr(self.space, self.ty.clone().into(), self.access)
233    }
234}
235
236pub struct AtomicTemplate {
237    pub ty: Type,
238}
239
240impl AtomicTemplate {
241    pub fn parse(tplt: &[TpltParam]) -> Result<AtomicTemplate, E> {
242        let ty = match tplt {
243            [TpltParam::Type(ty)] => Ok(ty.clone()),
244            _ => Err(E::TemplateArgs("atomic")),
245        }?;
246        #[cfg(feature = "naga-ext")]
247        if ty.is_f32() || ty.is_i64() || ty.is_u64() {
248            return Ok(AtomicTemplate { ty });
249        }
250        if ty.is_i32() || ty.is_u32() {
251            Ok(AtomicTemplate { ty })
252        } else {
253            Err(Error::Builtin("atomic template type must be an integer"))
254        }
255    }
256    pub fn ty(&self) -> Type {
257        Type::Atomic(self.ty.clone().into())
258    }
259    pub fn inner_ty(&self) -> Type {
260        self.ty.clone()
261    }
262}
263
264pub struct TextureTemplate {
265    ty: TextureType,
266}
267
268impl TextureTemplate {
269    pub fn parse(name: &str, tplt: &[TpltParam]) -> Result<TextureTemplate, E> {
270        let ty = match name {
271            "texture_1d" => TextureType::Sampled1D(Self::sampled_type(tplt)?),
272            "texture_2d" => TextureType::Sampled2D(Self::sampled_type(tplt)?),
273            "texture_2d_array" => TextureType::Sampled2DArray(Self::sampled_type(tplt)?),
274            "texture_3d" => TextureType::Sampled3D(Self::sampled_type(tplt)?),
275            "texture_cube" => TextureType::SampledCube(Self::sampled_type(tplt)?),
276            "texture_cube_array" => TextureType::SampledCubeArray(Self::sampled_type(tplt)?),
277            "texture_multisampled_2d" => TextureType::Multisampled2D(Self::sampled_type(tplt)?),
278            "texture_storage_1d" => {
279                let (tex, acc) = Self::texel_access(tplt)?;
280                TextureType::Storage1D(tex, acc)
281            }
282            "texture_storage_2d" => {
283                let (tex, acc) = Self::texel_access(tplt)?;
284                TextureType::Storage2D(tex, acc)
285            }
286            "texture_storage_2d_array" => {
287                let (tex, acc) = Self::texel_access(tplt)?;
288                TextureType::Storage2DArray(tex, acc)
289            }
290            "texture_storage_3d" => {
291                let (tex, acc) = Self::texel_access(tplt)?;
292                TextureType::Storage3D(tex, acc)
293            }
294            #[cfg(feature = "naga-ext")]
295            "texture_1d_array" => TextureType::Sampled1DArray(Self::sampled_type(tplt)?),
296            #[cfg(feature = "naga-ext")]
297            "texture_storage_1d_array" => {
298                let (tex, acc) = Self::texel_access(tplt)?;
299                TextureType::Storage1DArray(tex, acc)
300            }
301            #[cfg(feature = "naga-ext")]
302            "texture_multisampled_2d_array" => {
303                TextureType::Multisampled2DArray(Self::sampled_type(tplt)?)
304            }
305            _ => return Err(E::Builtin("not a templated texture type")),
306        };
307        Ok(Self { ty })
308    }
309    fn sampled_type(tplt: &[TpltParam]) -> Result<SampledType, E> {
310        match tplt {
311            [TpltParam::Type(ty)] => ty.try_into(),
312            [_] => Err(Error::Builtin(
313                "texture sampled type must be `i32`, `u32` or `f32`",
314            )),
315            _ => Err(Error::Builtin(
316                "sampled texture types take a single template parameter",
317            )),
318        }
319    }
320    fn texel_access(tplt: &[TpltParam]) -> Result<(TexelFormat, AccessMode), E> {
321        match tplt {
322            [
323                TpltParam::Enumerant(Enumerant::TexelFormat(texel)),
324                TpltParam::Enumerant(Enumerant::AccessMode(access)),
325            ] => Ok((*texel, *access)),
326            _ => Err(Error::Builtin(
327                "storage texture types take two template parameters",
328            )),
329        }
330    }
331    pub fn ty(&self) -> TextureType {
332        self.ty.clone()
333    }
334}
335
336pub struct BitcastTemplate {
337    ty: Type,
338}
339
340impl BitcastTemplate {
341    pub fn parse(tplt: &[TpltParam]) -> Result<BitcastTemplate, E> {
342        let ty = match tplt {
343            [TpltParam::Type(ty)] => Ok(ty.clone()),
344            _ => Err(E::TemplateArgs("bitcast")),
345        }?;
346        if ty.is_numeric() || ty.is_vec() && ty.inner_ty().is_numeric() {
347            Ok(BitcastTemplate { ty })
348        } else {
349            Err(Error::Builtin(
350                "bitcast template type must be a numeric scalar or numeric vector",
351            ))
352        }
353    }
354    pub fn ty(&self) -> &Type {
355        &self.ty
356    }
357}
358
359#[cfg(feature = "naga-ext")]
360#[derive(Clone, Debug, PartialEq, Eq, Hash)]
361pub struct AccelerationStructureTags {
362    tags: Vec<AccelerationStructureTag>,
363}
364
365#[cfg(feature = "naga-ext")]
366impl AccelerationStructureTags {
367    pub fn new(tags: Vec<AccelerationStructureTag>) -> Self {
368        Self { tags }
369    }
370
371    pub fn parse(tplt: &[TpltParam]) -> Result<Self, E> {
372        let tags = tplt
373            .iter()
374            .map(|param| match param {
375                TpltParam::Enumerant(Enumerant::AccelerationStructureTag(tag)) => Ok(*tag),
376                TpltParam::Instance(_) => Err(Error::Builtin(
377                    "expected an acceleration structure tag, not an instance",
378                )),
379                TpltParam::Type(_) => Err(Error::Builtin(
380                    "expected an acceleration structure tag, not a type",
381                )),
382                _ => Err(Error::Builtin("unknown acceleration structure tag")),
383            })
384            .collect::<Result<Vec<_>, E>>()?;
385        Ok(Self { tags })
386    }
387
388    pub fn tags(&self) -> &[AccelerationStructureTag] {
389        &self.tags
390    }
391}