Use UniqueArena for types.

Ensure that each distinct type occurs only once in `Module::types`, so that we
can use `Eq` on `Type` or `TypeInner` for type equivalence, without being
confused by differing `Handle<Type>` values that point to identical types.

This removes a number of duplicate types from the ir snapshots.

Cargo.toml

@@ -20,6 +20,7 @@ bitflags = "1"
 bit-set = "0.5"
 codespan-reporting = { version = "0.11.0", optional = true }
 fxhash = "0.2"
+indexmap = "1.6" # 1.7 has MSRV 1.49
 log = "0.4"
 num-traits = "0.2"
 spirv = { version = "0.2", optional = true }
@@ -36,8 +37,8 @@ glsl-in = ["pp-rs"]
 glsl-validate = []
 glsl-out = ["petgraph"]
 msl-out = []
-serialize = ["serde"]
-deserialize = ["serde"]
+serialize = ["serde", "indexmap/serde-1"]
+deserialize = ["serde", "indexmap/serde-1"]
 spv-in = ["petgraph", "spirv"]
 spv-out = ["spirv"]
 wgsl-in = ["codespan-reporting", "hexf-parse"]

src/arena.rs

@@ -6,8 +6,11 @@ use std::{cmp::Ordering, fmt, hash, marker::PhantomData, num::NonZeroU32, ops};
 type Index = NonZeroU32;
 
 use crate::Span;
+use indexmap::set::IndexSet;
 
 /// A strongly typed reference to an arena item.
+///
+/// A `Handle` value can be used as an index into an [`Arena`] or [`UniqueArena`].
 #[cfg_attr(feature = "serialize", derive(serde::Serialize))]
 #[cfg_attr(feature = "deserialize", derive(serde::Deserialize))]
 #[cfg_attr(
@@ -81,7 +84,8 @@ impl<T> Handle<T> {
         use std::convert::TryFrom;
 
         let handle_index = u32::try_from(index + 1)
-            .and_then(Index::try_from)
+            .ok()
+            .and_then(Index::new)
             .expect("Failed to insert into UniqueArena. Handle overflows");
         Handle::new(handle_index)
     }
@@ -369,3 +373,192 @@ mod tests {
         assert!(arena[t1] != arena[t2]);
     }
 }
+
+/// An arena whose elements are guaranteed to be unique.
+///
+/// A `UniqueArena` holds a set of unique values of type `T`, each with an
+/// associated [`Span`]. Inserting a value returns a `Handle<T>`, which can be
+/// used to index the `UniqueArena` and obtain shared access to the `T` element.
+/// Access via a `Handle` is an array lookup - no hash lookup is necessary.
+///
+/// The element type must implement `Eq` and `Hash`. Insertions of equivalent
+/// elements, according to `Eq`, all return the same `Handle`.
+///
+/// Once inserted, elements may not be mutated.
+///
+/// `UniqueArena` is similar to [`Arena`]: If `Arena` is vector-like,
+/// `UniqueArena` is `HashSet`-like.
+pub struct UniqueArena<T> {
+    set: IndexSet<T>,
+
+    /// Spans for the elements, indexed by handle.
+    ///
+    /// The length of this vector is always equal to `set.len()`. `IndexSet`
+    /// promises that its elements "are indexed in a compact range, without
+    /// holes in the range 0..set.len()", so we can always use the indices
+    /// returned by insertion as indices into this vector.
+    #[cfg(feature = "span")]
+    span_info: Vec<Span>,
+}
+
+impl<T> UniqueArena<T> {
+    /// Create a new arena with no initial capacity allocated.
+    pub fn new() -> Self {
+        UniqueArena {
+            set: IndexSet::new(),
+            #[cfg(feature = "span")]
+            span_info: Vec::new(),
+        }
+    }
+
+    /// Return the current number of items stored in this arena.
+    pub fn len(&self) -> usize {
+        self.set.len()
+    }
+
+    /// Return `true` if the arena contains no elements.
+    pub fn is_empty(&self) -> bool {
+        self.set.is_empty()
+    }
+
+    /// Clears the arena, keeping all allocations.
+    pub fn clear(&mut self) {
+        self.set.clear();
+        #[cfg(feature = "span")]
+        self.span_info.clear();
+    }
+
+    /// Return the span associated with `handle`.
+    ///
+    /// If a value has been inserted multiple times, the span returned is the
+    /// one provided with the first insertion.
+    ///
+    /// If the `span` feature is not enabled, always return `Span::default`.
+    /// This can be detected with [`Span::is_defined`].
+    pub fn get_span(&self, handle: Handle<T>) -> Span {
+        #[cfg(feature = "span")]
+        {
+            *self
+                .span_info
+                .get(handle.index())
+                .unwrap_or(&Span::default())
+        }
+        #[cfg(not(feature = "span"))]
+        {
+            let _ = handle;
+            Span::default()
+        }
+    }
+}
+
+impl<T: Eq + hash::Hash> UniqueArena<T> {
+    /// Returns an iterator over the items stored in this arena, returning both
+    /// the item's handle and a reference to it.
+    pub fn iter(&self) -> impl DoubleEndedIterator<Item = (Handle<T>, &T)> {
+        self.set.iter().enumerate().map(|(i, v)| {
+            let position = i + 1;
+            let index = unsafe { Index::new_unchecked(position as u32) };
+            (Handle::new(index), v)
+        })
+    }
+
+    /// Insert a new value into the arena.
+    ///
+    /// Return a [`Handle<T>`], which can be used to index this arena to get a
+    /// shared reference to the element.
+    ///
+    /// If this arena already contains an element that is `Eq` to `value`,
+    /// return a `Handle` to the existing element, and drop `value`.
+    ///
+    /// When the `span` feature is enabled, if `value` is inserted into the
+    /// arena, associate `span` with it. An element's span can be retrieved with
+    /// the [`get_span`] method.
+    ///
+    /// [`Handle<T>`]: Handle
+    /// [`get_span`]: UniqueArena::get_span
+    pub fn fetch_or_append(&mut self, value: T, span: Span) -> Handle<T> {
+        let (index, added) = self.set.insert_full(value);
+
+        #[cfg(feature = "span")]
+        {
+            if added {
+                debug_assert!(index == self.span_info.len());
+                self.span_info.push(span);
+            }
+
+            debug_assert!(self.set.len() == self.span_info.len());
+        }
+
+        Handle::from_usize(index)
+    }
+
+    /// Return this arena's handle for `value`, if present.
+    ///
+    /// If this arena already contains an element equal to `value`,
+    /// return its handle. Otherwise, return `None`.
+    pub fn get(&self, value: &T) -> Option<Handle<T>> {
+        self.set
+            .get_index_of(value)
+            .map(|index| unsafe { Handle::from_usize_unchecked(index) })
+    }
+
+    /// Return this arena's value at `handle`, if that is a valid handle.
+    pub fn try_get(&self, handle: Handle<T>) -> Option<&T> {
+        self.set.get_index(handle.index())
+    }
+}
+
+impl<T> Default for UniqueArena<T> {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl<T: fmt::Debug + Eq + hash::Hash> fmt::Debug for UniqueArena<T> {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        f.debug_map().entries(self.iter()).finish()
+    }
+}
+
+impl<T> ops::Index<Handle<T>> for UniqueArena<T> {
+    type Output = T;
+    fn index(&self, handle: Handle<T>) -> &T {
+        &self.set[handle.index()]
+    }
+}
+
+#[cfg(feature = "serialize")]
+impl<T> serde::Serialize for UniqueArena<T>
+where
+    T: Eq + hash::Hash,
+    T: serde::Serialize,
+{
+    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
+    where
+        S: serde::Serializer,
+    {
+        self.set.serialize(serializer)
+    }
+}
+
+#[cfg(feature = "deserialize")]
+impl<'de, T> serde::Deserialize<'de> for UniqueArena<T>
+where
+    T: Eq + hash::Hash,
+    T: serde::Deserialize<'de>,
+{
+    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
+    where
+        D: serde::Deserializer<'de>,
+    {
+        let set = IndexSet::deserialize(deserializer)?;
+        #[cfg(feature = "span")]
+        let span_info = std::iter::repeat(Span::default()).take(set.len()).collect();
+
+        Ok(Self {
+            set,
+            #[cfg(feature = "span")]
+            span_info,
+        })
+    }
+}

src/back/msl/writer.rs

@@ -22,7 +22,7 @@ const ATOMIC_REFERENCE: &str = "&";
 
 struct TypeContext<'a> {
     handle: Handle<crate::Type>,
-    arena: &'a crate::Arena<crate::Type>,
+    arena: &'a crate::UniqueArena<crate::Type>,
     names: &'a FastHashMap<NameKey, String>,
     access: crate::StorageAccess,
     first_time: bool,
@@ -354,7 +354,7 @@ fn should_pack_struct_member(
     }
 }
 
-fn needs_array_length(ty: Handle<crate::Type>, arena: &crate::Arena<crate::Type>) -> bool {
+fn needs_array_length(ty: Handle<crate::Type>, arena: &crate::UniqueArena<crate::Type>) -> bool {
     if let crate::TypeInner::Struct { ref members, .. } = arena[ty].inner {
         if let Some(member) = members.last() {
             if let crate::TypeInner::Array {

src/back/spv/helpers.rs

@@ -1,4 +1,4 @@
-use crate::{Arena, Handle};
+use crate::{Handle, UniqueArena};
 use spirv::Word;
 
 pub(super) fn bytes_to_words(bytes: &[u8]) -> Vec<Word> {
@@ -35,7 +35,7 @@ pub(super) fn map_storage_class(class: crate::StorageClass) -> spirv::StorageCla
 pub(super) fn contains_builtin(
     binding: Option<&crate::Binding>,
     ty: Handle<crate::Type>,
-    arena: &Arena<crate::Type>,
+    arena: &UniqueArena<crate::Type>,
     built_in: crate::BuiltIn,
 ) -> bool {
     if let Some(&crate::Binding::BuiltIn(bi)) = binding {

src/back/spv/writer.rs

@@ -6,7 +6,7 @@ use super::{
     PipelineOptions, ResultMember, Writer, WriterFlags, BITS_PER_BYTE,
 };
 use crate::{
-    arena::{Arena, Handle},
+    arena::{Handle, UniqueArena},
     proc::TypeResolution,
     valid::{FunctionInfo, ModuleInfo},
 };
@@ -195,7 +195,7 @@ impl Writer {
 
     pub(super) fn get_pointer_id(
         &mut self,
-        arena: &Arena<crate::Type>,
+        arena: &UniqueArena<crate::Type>,
         handle: Handle<crate::Type>,
         class: spirv::StorageClass,
     ) -> Result<Word, Error> {
@@ -766,7 +766,7 @@ impl Writer {
 
     fn write_type_declaration_arena(
         &mut self,
-        arena: &Arena<crate::Type>,
+        arena: &UniqueArena<crate::Type>,
         handle: Handle<crate::Type>,
     ) -> Result<Word, Error> {
         let ty = &arena[handle];

src/front/glsl/constants.rs

@@ -1,12 +1,12 @@
 use crate::{
-    arena::{Arena, Handle},
+    arena::{Arena, Handle, UniqueArena},
     BinaryOperator, Constant, ConstantInner, Expression, ScalarKind, ScalarValue, Type, TypeInner,
     UnaryOperator,
 };
 
 #[derive(Debug)]
 pub struct ConstantSolver<'a> {
-    pub types: &'a mut Arena<Type>,
+    pub types: &'a mut UniqueArena<Type>,
     pub expressions: &'a Arena<Expression>,
     pub constants: &'a mut Arena<Constant>,
 }
@@ -531,18 +531,18 @@ mod tests {
 
     use crate::{
         Arena, Constant, ConstantInner, Expression, ScalarKind, ScalarValue, Type, TypeInner,
-        UnaryOperator, VectorSize,
+        UnaryOperator, UniqueArena, VectorSize,
     };
 
     use super::ConstantSolver;
 
     #[test]
     fn unary_op() {
-        let mut types = Arena::new();
+        let mut types = UniqueArena::new();
         let mut expressions = Arena::new();
         let mut constants = Arena::new();
 
-        let vec_ty = types.append(
+        let vec_ty = types.fetch_or_append(
             Type {
                 name: None,
                 inner: TypeInner::Vector {
@@ -698,7 +698,7 @@ mod tests {
         );
 
         let mut solver = ConstantSolver {
-            types: &mut Arena::new(),
+            types: &mut UniqueArena::new(),
             expressions: &expressions,
             constants: &mut constants,
         };
@@ -716,11 +716,11 @@ mod tests {
 
     #[test]
     fn access() {
-        let mut types = Arena::new();
+        let mut types = UniqueArena::new();
         let mut expressions = Arena::new();
         let mut constants = Arena::new();
 
-        let matrix_ty = types.append(
+        let matrix_ty = types.fetch_or_append(
             Type {
                 name: None,
                 inner: TypeInner::Matrix {
@@ -732,7 +732,7 @@ mod tests {
             Default::default(),
         );
 
-        let vec_ty = types.append(
+        let vec_ty = types.fetch_or_append(
             Type {
                 name: None,
                 inner: TypeInner::Vector {

src/front/glsl/functions.rs

@@ -900,7 +900,7 @@ impl Parser {
         }
 
         let (ty, value) = if !components.is_empty() {
-            let ty = self.module.types.append(
+            let ty = self.module.types.fetch_or_append(
                 Type {
                     name: None,
                     inner: TypeInner::Struct {

src/front/glsl/offset.rs

@@ -13,7 +13,7 @@ use super::{
     error::{Error, ErrorKind},
     Span,
 };
-use crate::{front::align_up, Arena, Constant, Handle, Type, TypeInner};
+use crate::{front::align_up, Arena, Constant, Handle, Type, TypeInner, UniqueArena};
 
 /// Struct with information needed for defining a struct member.
 ///
@@ -39,7 +39,7 @@ pub fn calculate_offset(
     mut ty: Handle<Type>,
     meta: Span,
     layout: StructLayout,
-    types: &mut Arena<Type>,
+    types: &mut UniqueArena<Type>,
     constants: &Arena<Constant>,
     errors: &mut Vec<Error>,
 ) -> TypeAlignSpan {

src/front/glsl/parser/declarations.rs

@@ -497,7 +497,7 @@ impl<'source> ParsingContext<'source> {
         let span = self.parse_struct_declaration_list(parser, &mut members, layout)?;
         self.expect(parser, TokenValue::RightBrace)?;
 
-        let mut ty = parser.module.types.append(
+        let mut ty = parser.module.types.fetch_or_append(
             Type {
                 name: Some(ty_name),
                 inner: TypeInner::Struct {