1use 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#[derive(Clone, Debug, PartialEq, Eq)]
22pub enum MemView {
23 Whole,
25 Member(String, Box<MemView>),
27 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#[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 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
118macro_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 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 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 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#[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), #[cfg(feature = "naga-ext")]
243 U64(u64), #[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#[derive(Clone, Debug, PartialEq)]
327pub struct StructInstance {
328 pub ty: StructType,
329 pub members: Vec<Instance>,
330}
331
332impl StructInstance {
333 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 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 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 }
364
365#[derive(Clone, Debug, PartialEq, Default)]
369pub struct ArrayInstance {
370 components: Vec<Instance>,
371 pub runtime_sized: bool,
372}
373
374impl ArrayInstance {
375 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 pub fn n(&self) -> usize {
390 self.components.len()
391 }
392 pub fn get(&self, i: usize) -> Option<&Instance> {
394 self.components.get(i)
395 }
396 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#[derive(Clone, Debug, PartialEq)]
424pub struct VecInstance {
425 components: ArrayInstance,
426}
427
428impl VecInstance {
429 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 pub fn n(&self) -> usize {
443 self.components.n()
444 }
445 pub fn get(&self, i: usize) -> Option<&Instance> {
447 self.components.get(i)
448 }
449 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#[derive(Clone, Debug, PartialEq)]
501pub struct MatInstance {
502 components: Vec<Instance>,
504}
505
506impl MatInstance {
507 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 pub fn r(&self) -> usize {
533 self.components.first().unwrap().unwrap_vec_ref().n()
534 }
535 pub fn c(&self) -> usize {
537 self.components.len()
538 }
539 pub fn col(&self, i: usize) -> Option<&Instance> {
541 self.components.get(i)
542 }
543 pub fn col_mut(&mut self, i: usize) -> Option<&mut Instance> {
545 self.components.get_mut(i)
546 }
547 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 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#[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#[derive(Clone, Debug, PartialEq)]
599pub struct RefInstance {
600 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 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 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#[derive(Clone, Debug, PartialEq)]
698pub struct AtomicInstance {
699 content: Box<Instance>,
700}
701
702impl AtomicInstance {
703 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}