Skip to main content

wgsl_types/
inst.rs

1//! [`Instance`]s of WGSL [`Type`]s.
2
3use itertools::Itertools;
4use std::{
5    cell::{Ref, RefCell, RefMut},
6    ops::Index,
7    rc::Rc,
8};
9
10use crate::{
11    Error, f16,
12    syntax::{AccessMode, AddressSpace},
13    ty::{StructType, Ty, Type},
14};
15
16type E = Error;
17
18/// Path to a memory view of an instance.
19///
20/// See [Instance::view]
21#[derive(Clone, Debug, PartialEq, Eq)]
22pub enum MemView {
23    /// View the whole instance.
24    Whole,
25    /// Access a `struct` member or `vec` component.
26    Member(String, Box<MemView>),
27    /// Access an `array`, `vec` or `mat` component.
28    Index(usize, Box<MemView>),
29}
30
31impl MemView {
32    pub fn append_member(&mut self, comp: String) {
33        match self {
34            MemView::Whole => *self = MemView::Member(comp, Box::new(MemView::Whole)),
35            MemView::Member(_, v) | MemView::Index(_, v) => v.append_member(comp),
36        }
37    }
38    pub fn append_index(&mut self, index: usize) {
39        match self {
40            MemView::Whole => *self = MemView::Index(index, Box::new(MemView::Whole)),
41            MemView::Member(_, v) | MemView::Index(_, v) => v.append_index(index),
42        }
43    }
44}
45
46/// Instance of a plain type.
47///
48/// Reference: <https://www.w3.org/TR/WGSL/#plain-types-section>
49#[derive(Clone, Debug, PartialEq)]
50pub enum Instance {
51    Literal(LiteralInstance),
52    Struct(StructInstance),
53    Array(ArrayInstance),
54    Vec(VecInstance),
55    Mat(MatInstance),
56    Ptr(PtrInstance),
57    Ref(RefInstance),
58    Atomic(AtomicInstance),
59    /// This variant is used to refer to currently unreachable instances, including
60    /// instances of *opaque* types, such as textures and samplers, and instances that
61    /// cannot be computed yet.
62    Opaque(Type),
63}
64
65impl Instance {
66    pub fn unwrap_literal(self) -> LiteralInstance {
67        match self {
68            Instance::Literal(field_0) => field_0,
69            val => panic!("called `Instance::unwrap_literal()` on a `{val}` value"),
70        }
71    }
72    pub fn unwrap_literal_ref(&self) -> &LiteralInstance {
73        match self {
74            Instance::Literal(field_0) => field_0,
75            val => panic!("called `Instance::unwrap_literal_ref()` on a `{val}` value"),
76        }
77    }
78    pub fn unwrap_vec(self) -> VecInstance {
79        match self {
80            Instance::Vec(field_0) => field_0,
81            val => panic!("called `Instance::unwrap_vec()` on a `{val}` value"),
82        }
83    }
84    pub fn unwrap_vec_ref(&self) -> &VecInstance {
85        match self {
86            Instance::Vec(field_0) => field_0,
87            val => panic!("called `Instance::unwrap_vec_ref()` on a `{val}` value"),
88        }
89    }
90    pub fn unwrap_vec_mut(&mut self) -> &mut VecInstance {
91        match self {
92            Instance::Vec(field_0) => field_0,
93            val => panic!("called `Instance::unwrap_vec_mut()` on a `{val}` value"),
94        }
95    }
96}
97
98macro_rules! from_enum {
99    ($target_enum:ident :: $field:ident ( $from:ident )) => {
100        impl From<$from> for $target_enum {
101            fn from(value: $from) -> Self {
102                $target_enum::$field(value)
103            }
104        }
105    };
106}
107
108from_enum!(Instance::Literal(LiteralInstance));
109from_enum!(Instance::Struct(StructInstance));
110from_enum!(Instance::Array(ArrayInstance));
111from_enum!(Instance::Vec(VecInstance));
112from_enum!(Instance::Mat(MatInstance));
113from_enum!(Instance::Ptr(PtrInstance));
114from_enum!(Instance::Ref(RefInstance));
115from_enum!(Instance::Atomic(AtomicInstance));
116from_enum!(Instance::Opaque(Type));
117
118// Transitive `From` implementations.
119// They have to be implemented manually unfortunately.
120
121macro_rules! impl_transitive_from {
122    ($from:ident => $middle:ident => $into:ident) => {
123        impl From<$from> for $into {
124            fn from(value: $from) -> Self {
125                $into::from($middle::from(value))
126            }
127        }
128    };
129}
130
131impl_transitive_from!(bool => LiteralInstance => Instance);
132impl_transitive_from!(i64 => LiteralInstance => Instance);
133impl_transitive_from!(f64 => LiteralInstance => Instance);
134impl_transitive_from!(i32 => LiteralInstance => Instance);
135impl_transitive_from!(u32 => LiteralInstance => Instance);
136impl_transitive_from!(f32 => LiteralInstance => Instance);
137
138impl Instance {
139    /// Get an instance representing a memory view.
140    ///
141    /// There are two ways to create a memory view:
142    /// * Accessing a `struct` component (`struct.member`)
143    /// * Indexing an `array`, `vec`, or `mat` (`arr[n]`)
144    ///
145    /// Reference: <https://www.w3.org/TR/WGSL/#memory-views>
146    pub fn view(&self, view: &MemView) -> Result<&Instance, E> {
147        match view {
148            MemView::Whole => Ok(self),
149            MemView::Member(m, v) => match self {
150                Instance::Struct(s) => {
151                    let inst = s.member(m).ok_or_else(|| E::Component(s.ty(), m.clone()))?;
152                    inst.view(v)
153                }
154                _ => Err(E::Component(self.ty(), m.clone())),
155            },
156            MemView::Index(i, view) => match self {
157                Instance::Array(a) => {
158                    let inst = a
159                        .components
160                        .get(*i)
161                        .ok_or(E::OutOfBounds(*i, a.ty(), a.n()))?;
162                    inst.view(view)
163                }
164                Instance::Vec(v) => {
165                    let inst = v
166                        .components
167                        .get(*i)
168                        .ok_or(E::OutOfBounds(*i, v.ty(), v.n()))?;
169                    inst.view(view)
170                }
171                Instance::Mat(m) => {
172                    let inst = m
173                        .components
174                        .get(*i)
175                        .ok_or(E::OutOfBounds(*i, m.ty(), m.c()))?;
176                    inst.view(view)
177                }
178                _ => Err(E::NotIndexable(self.ty())),
179            },
180        }
181    }
182
183    /// Get an instance representing a memory view.
184    ///
185    /// See [Self::view]
186    pub fn view_mut(&mut self, view: &MemView) -> Result<&mut Instance, E> {
187        let ty = self.ty();
188        match view {
189            MemView::Whole => Ok(self),
190            MemView::Member(m, v) => match self {
191                Instance::Struct(s) => {
192                    let inst = s.member_mut(m).ok_or_else(|| E::Component(ty, m.clone()))?;
193                    inst.view_mut(v)
194                }
195                _ => Err(E::Component(ty, m.clone())),
196            },
197            MemView::Index(i, view) => match self {
198                Instance::Array(a) => {
199                    let n = a.n();
200                    let inst = a.components.get_mut(*i).ok_or(E::OutOfBounds(*i, ty, n))?;
201                    inst.view_mut(view)
202                }
203                Instance::Vec(v) => {
204                    let n = v.n();
205                    let inst = v.components.get_mut(*i).ok_or(E::OutOfBounds(*i, ty, n))?;
206                    inst.view_mut(view)
207                }
208                Instance::Mat(m) => {
209                    let c = m.c();
210                    let inst = m.components.get_mut(*i).ok_or(E::OutOfBounds(*i, ty, c))?;
211                    inst.view_mut(view)
212                }
213                _ => Err(E::NotIndexable(ty)),
214            },
215        }
216    }
217
218    /// Mutate the instance.
219    ///
220    /// This is the operation performed by the assignment operator.
221    pub fn write(&mut self, value: Instance) -> Result<Instance, E> {
222        if value.ty() != self.ty() {
223            return Err(E::WriteRefType(value.ty(), self.ty()));
224        }
225        let old = std::mem::replace(self, value);
226        Ok(old)
227    }
228}
229
230/// Instance of a numeric literal type.
231#[derive(Clone, Copy, Debug, PartialEq)]
232pub enum LiteralInstance {
233    Bool(bool),
234    AbstractInt(i64),
235    AbstractFloat(f64),
236    I32(i32),
237    U32(u32),
238    F32(f32),
239    F16(f16),
240    #[cfg(feature = "naga-ext")]
241    I64(i64), // identity if representable
242    #[cfg(feature = "naga-ext")]
243    U64(u64), // reinterpretation of bits
244    #[cfg(feature = "naga-ext")]
245    F64(f64),
246}
247
248from_enum!(LiteralInstance::Bool(bool));
249from_enum!(LiteralInstance::AbstractInt(i64));
250from_enum!(LiteralInstance::AbstractFloat(f64));
251from_enum!(LiteralInstance::I32(i32));
252from_enum!(LiteralInstance::U32(u32));
253from_enum!(LiteralInstance::F32(f32));
254from_enum!(LiteralInstance::F16(f16));
255
256impl LiteralInstance {
257    pub fn unwrap_bool(self) -> bool {
258        match self {
259            LiteralInstance::Bool(field_0) => field_0,
260            val => panic!("called `LiteralInstance::unwrap_bool()` on a `{val}` value"),
261        }
262    }
263
264    pub fn unwrap_abstract_int(self) -> i64 {
265        match self {
266            LiteralInstance::AbstractInt(field_0) => field_0,
267            val => panic!("called `LiteralInstance::unwrap_abstract_int()` on a `{val}` value"),
268        }
269    }
270    pub fn unwrap_abstract_float(self) -> f64 {
271        match self {
272            LiteralInstance::AbstractFloat(field_0) => field_0,
273            val => panic!("called `LiteralInstance::unwrap_abstract_float()` on a `{val}` value"),
274        }
275    }
276    pub fn unwrap_i32(self) -> i32 {
277        match self {
278            LiteralInstance::I32(field_0) => field_0,
279            val => panic!("called `LiteralInstance::unwrap_i32()` on a `{val}` value"),
280        }
281    }
282    pub fn unwrap_u32(self) -> u32 {
283        match self {
284            LiteralInstance::U32(field_0) => field_0,
285            val => panic!("called `LiteralInstance::unwrap_u32()` on a `{val}` value"),
286        }
287    }
288    pub fn unwrap_f32(self) -> f32 {
289        match self {
290            LiteralInstance::F32(field_0) => field_0,
291            val => panic!("called `LiteralInstance::unwrap_f32()` on a `{val}` value"),
292        }
293    }
294    pub fn unwrap_f16(self) -> f16 {
295        match self {
296            LiteralInstance::F16(field_0) => field_0,
297            val => panic!("called `LiteralInstance::unwrap_f16()` on a `{val}` value"),
298        }
299    }
300    #[cfg(feature = "naga-ext")]
301    pub fn unwrap_i64(self) -> i64 {
302        match self {
303            LiteralInstance::I64(field_0) => field_0,
304            val => panic!("called `LiteralInstance::unwrap_i64()` on a `{val}` value"),
305        }
306    }
307    #[cfg(feature = "naga-ext")]
308    pub fn unwrap_u64(self) -> u64 {
309        match self {
310            LiteralInstance::U64(field_0) => field_0,
311            val => panic!("called `LiteralInstance::unwrap_u64()` on a `{val}` value"),
312        }
313    }
314    #[cfg(feature = "naga-ext")]
315    pub fn unwrap_f64(self) -> f64 {
316        match self {
317            LiteralInstance::F64(field_0) => field_0,
318            val => panic!("called `LiteralInstance::unwrap_f64()` on a `{val}` value"),
319        }
320    }
321}
322
323/// Instance of a `struct` type.
324///
325/// Reference: <https://www.w3.org/TR/WGSL/#struct-types>
326#[derive(Clone, Debug, PartialEq)]
327pub struct StructInstance {
328    pub ty: StructType,
329    pub members: Vec<Instance>,
330}
331
332impl StructInstance {
333    /// Create a `struct` instance.
334    ///
335    /// # Panics
336    /// * if there is not the right number of members
337    /// * if the members are not of the right type
338    pub fn new(ty: StructType, members: Vec<Instance>) -> Self {
339        assert_eq!(ty.members.len(), members.len());
340        for (m, m_ty) in members.iter().zip(&ty.members) {
341            assert_eq!(m_ty.ty, m.ty());
342        }
343
344        Self { ty, members }
345    }
346    /// Get a `struct` member value by name.
347    pub fn member(&self, name: &str) -> Option<&Instance> {
348        self.members
349            .iter()
350            .zip(&self.ty.members)
351            .find_map(|(inst, m_ty)| (m_ty.name == name).then_some(inst))
352    }
353    /// Get a `struct` member value by name.
354    pub fn member_mut(&mut self, name: &str) -> Option<&mut Instance> {
355        self.members
356            .iter_mut()
357            .zip(&self.ty.members)
358            .find_map(|(inst, m_ty)| (m_ty.name == name).then_some(inst))
359    }
360    // pub fn iter_members(&self) -> impl Iterator<Item = &(String, Instance)> {
361    //     self.members.iter()
362    // }
363}
364
365/// Instance of an `array<T, N>` type.
366///
367/// Reference: <https://www.w3.org/TR/WGSL/#array-types>
368#[derive(Clone, Debug, PartialEq, Default)]
369pub struct ArrayInstance {
370    components: Vec<Instance>,
371    pub runtime_sized: bool,
372}
373
374impl ArrayInstance {
375    /// Construct an `array`.
376    ///
377    /// # Panics
378    /// * if the components is empty
379    /// * if the components are not all the same type
380    pub fn new(components: Vec<Instance>, runtime_sized: bool) -> Self {
381        assert!(!components.is_empty());
382        assert!(components.iter().map(|c| c.ty()).all_equal());
383        Self {
384            components,
385            runtime_sized,
386        }
387    }
388    /// The element count.
389    pub fn n(&self) -> usize {
390        self.components.len()
391    }
392    /// Get an element by index.
393    pub fn get(&self, i: usize) -> Option<&Instance> {
394        self.components.get(i)
395    }
396    /// Get an element by index.
397    pub fn get_mut(&mut self, i: usize) -> Option<&mut Instance> {
398        self.components.get_mut(i)
399    }
400    pub fn iter(&self) -> impl Iterator<Item = &Instance> {
401        self.components.iter()
402    }
403    pub fn iter_mut(&mut self) -> impl Iterator<Item = &mut Instance> {
404        self.components.iter_mut()
405    }
406    pub fn as_slice(&self) -> &[Instance] {
407        self.components.as_slice()
408    }
409}
410
411impl IntoIterator for ArrayInstance {
412    type Item = Instance;
413    type IntoIter = <Vec<Instance> as IntoIterator>::IntoIter;
414
415    fn into_iter(self) -> Self::IntoIter {
416        self.components.into_iter()
417    }
418}
419
420/// Instance of a `vecN<T>` type.
421///
422/// Reference: <https://www.w3.org/TR/WGSL/#vector-types>
423#[derive(Clone, Debug, PartialEq)]
424pub struct VecInstance {
425    components: ArrayInstance,
426}
427
428impl VecInstance {
429    /// Construct a `vec`.
430    ///
431    /// # Panics
432    /// * if the components length is not [2, 3, 4]
433    /// * if the components are not all the same type
434    /// * if the type is not a scalar
435    pub fn new(components: Vec<Instance>) -> Self {
436        assert!((2..=4).contains(&components.len()));
437        let components = ArrayInstance::new(components, false);
438        assert!(components.inner_ty().is_scalar());
439        Self { components }
440    }
441    /// The component count.
442    pub fn n(&self) -> usize {
443        self.components.n()
444    }
445    /// Get a component by index.
446    pub fn get(&self, i: usize) -> Option<&Instance> {
447        self.components.get(i)
448    }
449    /// Get a component by index.
450    pub fn get_mut(&mut self, i: usize) -> Option<&mut Instance> {
451        self.components.get_mut(i)
452    }
453    pub fn iter(&self) -> impl Iterator<Item = &Instance> {
454        self.components.iter()
455    }
456    pub fn iter_mut(&mut self) -> impl Iterator<Item = &mut Instance> {
457        self.components.iter_mut()
458    }
459    pub fn as_slice(&self) -> &[Instance] {
460        self.components.as_slice()
461    }
462}
463
464impl IntoIterator for VecInstance {
465    type Item = Instance;
466    type IntoIter = <ArrayInstance as IntoIterator>::IntoIter;
467
468    fn into_iter(self) -> Self::IntoIter {
469        self.components.into_iter()
470    }
471}
472
473impl Index<usize> for VecInstance {
474    type Output = Instance;
475
476    fn index(&self, index: usize) -> &Self::Output {
477        self.get(index).unwrap()
478    }
479}
480
481impl<T: Into<Instance>> From<[T; 2]> for VecInstance {
482    fn from(components: [T; 2]) -> Self {
483        Self::new(components.map(Into::into).to_vec())
484    }
485}
486impl<T: Into<Instance>> From<[T; 3]> for VecInstance {
487    fn from(components: [T; 3]) -> Self {
488        Self::new(components.map(Into::into).to_vec())
489    }
490}
491impl<T: Into<Instance>> From<[T; 4]> for VecInstance {
492    fn from(components: [T; 4]) -> Self {
493        Self::new(components.map(Into::into).to_vec())
494    }
495}
496
497/// Instance of a `matCxR<T>` type.
498///
499/// Reference: <https://www.w3.org/TR/WGSL/#matrix-types>
500#[derive(Clone, Debug, PartialEq)]
501pub struct MatInstance {
502    /// Column vectors of the matrix
503    components: Vec<Instance>,
504}
505
506impl MatInstance {
507    /// Construct a `mat` from column vectors.
508    ///
509    /// # Panics
510    /// * if the number of columns is not [2, 3, 4]
511    /// * if the columns don't have the same number of rows
512    /// * if the number of rows is not [2, 3, 4]
513    /// * if the elements don't have the same type
514    /// * if the type is not a scalar
515    pub fn from_cols(components: Vec<Instance>) -> Self {
516        assert!((2..=4).contains(&components.len()));
517        assert!(
518            components
519                .iter()
520                .map(|c| c.unwrap_vec_ref().n())
521                .all_equal(),
522            "MatInstance columns must have the same number for rows"
523        );
524        assert!(
525            components.iter().map(|c| c.ty()).all_equal(),
526            "MatInstance columns must have the same type"
527        );
528        Self { components }
529    }
530
531    /// The row count.
532    pub fn r(&self) -> usize {
533        self.components.first().unwrap().unwrap_vec_ref().n()
534    }
535    /// The column count.
536    pub fn c(&self) -> usize {
537        self.components.len()
538    }
539    /// Get a column vector.
540    pub fn col(&self, i: usize) -> Option<&Instance> {
541        self.components.get(i)
542    }
543    /// Get a column vector.
544    pub fn col_mut(&mut self, i: usize) -> Option<&mut Instance> {
545        self.components.get_mut(i)
546    }
547    /// Get a component.
548    pub fn get(&self, col: usize, row: usize) -> Option<&Instance> {
549        self.col(col).and_then(|v| v.unwrap_vec_ref().get(row))
550    }
551    /// Get a component.
552    pub fn get_mut(&mut self, i: usize, j: usize) -> Option<&mut Instance> {
553        self.col_mut(i).and_then(|v| v.unwrap_vec_mut().get_mut(j))
554    }
555    pub fn iter_cols(&self) -> impl Iterator<Item = &Instance> {
556        self.components.iter()
557    }
558    pub fn iter_cols_mut(&mut self) -> impl Iterator<Item = &mut Instance> {
559        self.components.iter_mut()
560    }
561    pub fn iter(&self) -> impl Iterator<Item = &Instance> {
562        self.components
563            .iter()
564            .flat_map(|v| v.unwrap_vec_ref().iter())
565    }
566    pub fn iter_mut(&mut self) -> impl Iterator<Item = &mut Instance> {
567        self.components
568            .iter_mut()
569            .flat_map(|v| v.unwrap_vec_mut().iter_mut())
570    }
571}
572impl IntoIterator for MatInstance {
573    type Item = Instance;
574    type IntoIter = <Vec<Instance> as IntoIterator>::IntoIter;
575
576    fn into_iter(self) -> Self::IntoIter {
577        self.components.into_iter()
578    }
579}
580
581/// Instance of a `ptr<AS,T,AM>` type.
582///
583/// Reference: <https://www.w3.org/TR/WGSL/#ref-ptr-types>
584#[derive(Clone, Debug, PartialEq)]
585pub struct PtrInstance {
586    pub ptr: RefInstance,
587}
588
589impl From<RefInstance> for PtrInstance {
590    fn from(r: RefInstance) -> Self {
591        Self { ptr: r }
592    }
593}
594
595/// Instance of a `ref<AS,T,AM>` type.
596///
597/// Reference: <https://www.w3.org/TR/WGSL/#ref-ptr-types>
598#[derive(Clone, Debug, PartialEq)]
599pub struct RefInstance {
600    /// Inner type
601    pub ty: Type,
602    pub space: AddressSpace,
603    pub access: AccessMode,
604    pub view: MemView,
605    pub ptr: Rc<RefCell<Instance>>,
606}
607
608impl RefInstance {
609    pub fn new(inst: Instance, space: AddressSpace, access: AccessMode) -> Self {
610        let ty = inst.ty();
611        Self {
612            ty,
613            space,
614            access,
615            view: MemView::Whole,
616            ptr: Rc::new(RefCell::new(inst)),
617        }
618    }
619}
620
621impl From<PtrInstance> for RefInstance {
622    fn from(p: PtrInstance) -> Self {
623        p.ptr
624    }
625}
626
627impl RefInstance {
628    /// Get a reference to a `struct` or `vec` member.
629    pub fn view_member(&self, comp: String) -> Result<Self, E> {
630        if !self.access.is_read() {
631            return Err(E::NotRead);
632        }
633        let mut view = self.view.clone();
634        view.append_member(comp);
635        let ty = self.ptr.borrow().view(&view)?.ty();
636        Ok(Self {
637            ty,
638            space: self.space,
639            access: self.access,
640            view,
641            ptr: self.ptr.clone(),
642        })
643    }
644    /// Get a reference to an `array`, `vec` or `mat` component.
645    pub fn view_index(&self, index: usize) -> Result<Self, E> {
646        if !self.access.is_read() {
647            return Err(E::NotRead);
648        }
649        let mut view = self.view.clone();
650        view.append_index(index);
651        let ty = self.ptr.borrow().view(&view)?.ty();
652        Ok(Self {
653            ty,
654            space: self.space,
655            access: self.access,
656            view,
657            ptr: self.ptr.clone(),
658        })
659    }
660
661    pub fn read<'a>(&'a self) -> Result<Ref<'a, Instance>, E> {
662        if !self.access.is_read() {
663            return Err(E::NotRead);
664        }
665        Ok(Ref::<'a, Instance>::map(self.ptr.borrow(), |r| {
666            r.view(&self.view).expect("invalid reference")
667        }))
668    }
669
670    pub fn write(&self, value: Instance) -> Result<(), E> {
671        if !self.access.is_write() {
672            return Err(E::NotWrite);
673        }
674        if value.ty() != self.ty {
675            return Err(E::WriteRefType(value.ty(), self.ty.clone()));
676        }
677        let mut r = self.ptr.borrow_mut();
678        let view = r.view_mut(&self.view).expect("invalid reference");
679        assert!(view.ty() == value.ty());
680        let _ = std::mem::replace(view, value);
681        Ok(())
682    }
683
684    pub fn read_write<'a>(&'a self) -> Result<RefMut<'a, Instance>, E> {
685        if !self.access.is_write() {
686            return Err(E::NotReadWrite);
687        }
688        Ok(RefMut::<'a, Instance>::map(self.ptr.borrow_mut(), |r| {
689            r.view_mut(&self.view).expect("invalid reference")
690        }))
691    }
692}
693
694/// `atomic<T>` Instance.
695///
696/// Reference: <https://www.w3.org/TR/WGSL/#atomic-types>
697#[derive(Clone, Debug, PartialEq)]
698pub struct AtomicInstance {
699    content: Box<Instance>,
700}
701
702impl AtomicInstance {
703    /// # Panics
704    /// * if the instance is not an i32 or u32
705    pub fn new(inst: Instance) -> Self {
706        assert!(matches!(inst.ty(), Type::I32 | Type::U32));
707        Self {
708            content: inst.into(),
709        }
710    }
711
712    pub fn inner(&self) -> &Instance {
713        &self.content
714    }
715
716    pub fn inner_mut(&mut self) -> &mut Instance {
717        &mut self.content
718    }
719}