UPSTREAM: Add handle validation pass

commits.md

@@ -0,0 +1,7 @@
+- [ ] Split out `ExprFedDepError`
+- [ ] Combine handle validation pass errors at current implementation level
+- [ ] Split out the pass!
+	- [ ] squash! add `ExpressionTypeResolver::check`
+	- [ ] squash! `impl Index<Expression> for ExpressionTypeResolver`
+	- [ ] note: trade-offs were made, we're now doing multiple passes. Boo for
+		perf, yay for abstraction.

src/arena.rs

@@ -15,6 +15,15 @@ pub struct BadHandle {
     pub index: usize,
 }
 
+impl BadHandle {
+    fn new<T>(handle: Handle<T>) -> Self {
+        Self {
+            kind: std::any::type_name::<T>(),
+            index: handle.index(),
+        }
+    }
+}
+
 /// A strongly typed reference to an arena item.
 ///
 /// A `Handle` value can be used as an index into an [`Arena`] or [`UniqueArena`].
@@ -275,10 +284,9 @@ impl<T> Arena<T> {
     }
 
     pub fn try_get(&self, handle: Handle<T>) -> Result<&T, BadHandle> {
-        self.data.get(handle.index()).ok_or_else(|| BadHandle {
-            kind: std::any::type_name::<T>(),
-            index: handle.index(),
-        })
+        self.data
+            .get(handle.index())
+            .ok_or_else(|| BadHandle::new(handle))
     }
 
     /// Get a mutable reference to an element in the arena.
@@ -313,6 +321,12 @@ impl<T> Arena<T> {
             Span::default()
         }
     }
+
+    pub fn check_contains_handle(&self, handle: Handle<T>) -> Result<(), BadHandle> {
+        (handle.index() < self.data.len())
+            .then_some(())
+            .ok_or_else(|| BadHandle::new(handle))
+    }
 }
 
 #[cfg(feature = "deserialize")]
@@ -533,10 +547,15 @@ impl<T: Eq + hash::Hash> UniqueArena<T> {
 
     /// Return this arena's value at `handle`, if that is a valid handle.
     pub fn get_handle(&self, handle: Handle<T>) -> Result<&T, BadHandle> {
-        self.set.get_index(handle.index()).ok_or_else(|| BadHandle {
-            kind: std::any::type_name::<T>(),
-            index: handle.index(),
-        })
+        self.set
+            .get_index(handle.index())
+            .ok_or_else(|| BadHandle::new(handle))
+    }
+
+    pub fn check_contains_handle(&self, handle: Handle<T>) -> Result<(), BadHandle> {
+        (handle.index() < self.set.len())
+            .then_some(())
+            .ok_or_else(|| BadHandle::new(handle))
     }
 }
 

src/back/hlsl/conv.rs

@@ -40,12 +40,12 @@ impl crate::TypeInner {
         }
     }
 
-    pub(super) fn try_size_hlsl(
+    pub(super) fn size_hlsl(
         &self,
         types: &crate::UniqueArena<crate::Type>,
         constants: &crate::Arena<crate::Constant>,
-    ) -> Result<u32, crate::arena::BadHandle> {
-        Ok(match *self {
+    ) -> u32 {
+        match *self {
             Self::Matrix {
                 columns,
                 rows,
@@ -58,17 +58,16 @@ impl crate::TypeInner {
             Self::Array { base, size, stride } => {
                 let count = match size {
                     crate::ArraySize::Constant(handle) => {
-                        let constant = constants.try_get(handle)?;
-                        constant.to_array_length().unwrap_or(1)
+                        constants[handle].to_array_length().unwrap_or(1)
                     }
                     // A dynamically-sized array has to have at least one element
                     crate::ArraySize::Dynamic => 1,
                 };
-                let last_el_size = types[base].inner.try_size_hlsl(types, constants)?;
+                let last_el_size = types[base].inner.size_hlsl(types, constants);
                 ((count - 1) * stride) + last_el_size
             }
-            _ => self.try_size(constants)?,
-        })
+            _ => self.size(constants),
+        }
     }
 
     /// Used to generate the name of the wrapped type constructor

src/back/hlsl/writer.rs

@@ -805,6 +805,10 @@ impl<'a, W: fmt::Write> super::Writer<'a, W> {
     ///
     /// # Notes
     /// Ends in a newline
+    ///
+    /// # Panics
+    ///
+    /// This function may panic when given a malformed IR module.
     fn write_struct(
         &mut self,
         module: &Module,
@@ -829,10 +833,7 @@ impl<'a, W: fmt::Write> super::Writer<'a, W> {
                 }
             }
             let ty_inner = &module.types[member.ty].inner;
-            last_offset = member.offset
-                + ty_inner
-                    .try_size_hlsl(&module.types, &module.constants)
-                    .unwrap();
+            last_offset = member.offset + ty_inner.size_hlsl(&module.types, &module.constants);
 
             // The indentation is only for readability
             write!(self.out, "{}", back::INDENT)?;

src/back/wgsl/writer.rs

@@ -1609,15 +1609,7 @@ impl<W: Write> Writer<W> {
                     Function::Regular(fun_name) => {
                         write!(self.out, "{}(", fun_name)?;
                         self.write_expr(module, arg, func_ctx)?;
-                        if let Some(arg) = arg1 {
-                            write!(self.out, ", ")?;
-                            self.write_expr(module, arg, func_ctx)?;
-                        }
-                        if let Some(arg) = arg2 {
-                            write!(self.out, ", ")?;
-                            self.write_expr(module, arg, func_ctx)?;
-                        }
-                        if let Some(arg) = arg3 {
+                        for arg in IntoIterator::into_iter([arg1, arg2, arg3]).flatten() {
                             write!(self.out, ", ")?;
                             self.write_expr(module, arg, func_ctx)?;
                         }

src/proc/layouter.rs

@@ -1,4 +1,7 @@
-use crate::arena::{Arena, BadHandle, Handle, UniqueArena};
+use crate::{
+    arena::{Arena, BadHandle, Handle, UniqueArena},
+    WithSpan,
+};
 use std::{fmt::Display, num::NonZeroU32, ops};
 
 /// A newtype struct where its only valid values are powers of 2
@@ -130,8 +133,6 @@ pub enum LayoutErrorInner {
     InvalidStructMemberType(u32, Handle<crate::Type>),
     #[error("Type width must be a power of two")]
     NonPowerOfTwoWidth,
-    #[error("Array size is a bad handle")]
-    BadHandle(#[from] BadHandle),
 }
 
 #[derive(Clone, Copy, Debug, PartialEq, thiserror::Error)]
@@ -153,6 +154,20 @@ impl Layouter {
         self.layouts.clear();
     }
 
+    pub fn validate_handles(
+        &mut self,
+        types: &UniqueArena<crate::Type>,
+        constants: &Arena<crate::Constant>,
+    ) -> Result<(), WithSpan<BadHandle>> {
+        types
+            .iter()
+            .skip(self.layouts.len())
+            .try_for_each(|(_handle, ty)| {
+                // TODO: track the type we fail on?
+                ty.inner.validate_handles(constants)
+            })
+    }
+
     /// Extend this `Layouter` with layouts for any new entries in `types`.
     ///
     /// Ensure that every type in `types` has a corresponding [TypeLayout] in
@@ -166,7 +181,13 @@ impl Layouter {
     /// end can call this function at any time, passing its current type and
     /// constant arenas, and then assume that layouts are available for all
     /// types.
+    ///
+    /// # Panics
+    ///
+    /// If `types` contains invalid [`Handle`]s to `constants`, then this function will panic. You
+    /// can check for this condition by calling [`Self::validate_handles`].
     #[allow(clippy::or_fun_call)]
+    #[track_caller]
     pub fn update(
         &mut self,
         types: &UniqueArena<crate::Type>,
@@ -175,10 +196,8 @@ impl Layouter {
         use crate::TypeInner as Ti;
 
         for (ty_handle, ty) in types.iter().skip(self.layouts.len()) {
-            let size = ty
-                .inner
-                .try_size(constants)
-                .map_err(|error| LayoutErrorInner::BadHandle(error).with(ty_handle))?;
+            // phase-id: layouter
+            let size = ty.inner.size(constants);
             let layout = match ty.inner {
                 Ti::Scalar { width, .. } | Ti::Atomic { width, .. } => {
                     let alignment = Alignment::new(width as u32)
@@ -255,3 +274,52 @@ impl Layouter {
         Ok(())
     }
 }
+
+// #[derive(Debug)]
+// struct HandlesValidated<T>(T);
+
+// impl<T> HandlesValidated<T> {
+//     pub const fn new(t: T) -> Self {
+//         Self(t)
+//     }
+// }
+
+// impl<'a, T> Clone for HandlesValidated<&'a T> {
+//     fn clone(&self) -> Self {
+//         let Self(ref_) = self;
+//         Self(ref_)
+//     }
+// }
+
+// impl<'a, T> Copy for HandlesValidated<&'a T> {}
+
+// impl<'a, T> AsRef<T> for HandlesValidated<&'a T> {
+//     fn as_ref(&self) -> &T {
+//         todo!()
+//     }
+// }
+
+// impl<'a, T> Deref for HandlesValidated<&'a T> {
+//     type Target = T;
+
+//     fn deref(&self) -> &Self::Target {
+//         let Self(inner) = self;
+//         inner
+//     }
+// }
+
+// impl<'a, T> Deref for HandlesValidated<&'a mut T> {
+//     type Target = T;
+
+//     fn deref(&self) -> &Self::Target {
+//         let Self(inner) = self;
+//         inner
+//     }
+// }
+
+// impl<'a, T> DerefMut for HandlesValidated<&'a mut T> {
+//     fn deref_mut(&mut self) -> &mut Self::Target {
+//         let Self(inner) = self;
+//         inner
+//     }
+// }

src/proc/mod.rs

@@ -16,6 +16,8 @@ pub use namer::{EntryPointIndex, NameKey, Namer};
 pub use terminator::ensure_block_returns;
 pub use typifier::{ResolveContext, ResolveError, TypeResolution};
 
+use crate::{arena::BadHandle, WithSpan};
+
 impl From<super::StorageFormat> for super::ScalarKind {
     fn from(format: super::StorageFormat) -> Self {
         use super::{ScalarKind as Sk, StorageFormat as Sf};
@@ -97,11 +99,32 @@ impl super::TypeInner {
         }
     }
 
-    pub fn try_size(
+    pub fn validate_handles(
         &self,
         constants: &super::Arena<super::Constant>,
-    ) -> Result<u32, crate::arena::BadHandle> {
-        Ok(match *self {
+    ) -> Result<(), WithSpan<BadHandle>> {
+        match self {
+            &Self::Array {
+                base: _,
+                size: super::ArraySize::Constant(handle),
+                stride: _,
+            } => constants
+                .try_get(handle)
+                .map(|_| ())
+                .map_err(|e| WithSpan::new(e).with_handle(handle, constants)),
+            _ => Ok(()),
+        }
+    }
+
+    /// Get the size of this type.
+    ///
+    /// # Panics
+    ///
+    /// Panics if the `constants` doesn't contain a referenced handle. This may not happen in
+    /// a properly validated IR module. You can check for this condition with
+    /// [`Self::validate_handles`].
+    pub fn size(&self, constants: &super::Arena<super::Constant>) -> u32 {
+        match *self {
             Self::Scalar { kind: _, width } | Self::Atomic { kind: _, width } => width as u32,
             Self::Vector {
                 size,
@@ -122,8 +145,7 @@ impl super::TypeInner {
             } => {
                 let count = match size {
                     super::ArraySize::Constant(handle) => {
-                        let constant = constants.try_get(handle)?;
-                        constant.to_array_length().unwrap_or(1)
+                        constants[handle].to_array_length().unwrap_or(1)
                     }
                     // A dynamically-sized array has to have at least one element
                     super::ArraySize::Dynamic => 1,
@@ -132,13 +154,7 @@ impl super::TypeInner {
             }
             Self::Struct { span, .. } => span,
             Self::Image { .. } | Self::Sampler { .. } | Self::BindingArray { .. } => 0,
-        })
-    }
-
-    /// Get the size of this type. Panics if the `constants` doesn't contain
-    /// a referenced handle. This may not happen in a properly validated IR module.
-    pub fn size(&self, constants: &super::Arena<super::Constant>) -> u32 {
-        self.try_size(constants).unwrap()
+        }
     }
 
     /// Return the canonical form of `self`, or `None` if it's already in