Use fast hashing functions with IndexMap and IndexSet everywhere.

src/arena.rs

@@ -5,8 +5,7 @@ use std::{cmp::Ordering, fmt, hash, marker::PhantomData, num::NonZeroU32, ops};
 /// the same size and representation as `Handle<T>`.
 type Index = NonZeroU32;
 
-use crate::Span;
-use indexmap::set::IndexSet;
+use crate::{FastIndexSet, Span};
 
 #[derive(Clone, Copy, Debug, thiserror::Error, PartialEq)]
 #[error("Handle {index} of {kind} is either not present, or inaccessible yet")]
@@ -516,11 +515,11 @@ mod tests {
 /// `UniqueArena` is `HashSet`-like.
 #[cfg_attr(feature = "clone", derive(Clone))]
 pub struct UniqueArena<T> {
-    set: IndexSet<T>,
+    set: FastIndexSet<T>,
 
     /// Spans for the elements, indexed by handle.
     ///
-    /// The length of this vector is always equal to `set.len()`. `IndexSet`
+    /// The length of this vector is always equal to `set.len()`. `FastIndexSet`
     /// 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.
@@ -532,7 +531,7 @@ impl<T> UniqueArena<T> {
     /// Create a new arena with no initial capacity allocated.
     pub fn new() -> Self {
         UniqueArena {
-            set: IndexSet::new(),
+            set: FastIndexSet::default(),
             #[cfg(feature = "span")]
             span_info: Vec::new(),
         }
@@ -741,7 +740,7 @@ where
     where
         D: serde::Deserializer<'de>,
     {
-        let set = IndexSet::deserialize(deserializer)?;
+        let set = FastIndexSet::deserialize(deserializer)?;
         #[cfg(feature = "span")]
         let span_info = std::iter::repeat(Span::default()).take(set.len()).collect();
 

src/compact/mod.rs

@@ -95,9 +95,9 @@ pub fn compact(module: &mut crate::Module) {
 
     // Drop unused types from the type arena.
     //
-    // `IndexSet`s don't have an underlying Vec<T> that we can steal, compact in
-    // place, and then rebuild the `IndexSet` from. So we have to rebuild the
-    // type arena from scratch.
+    // `FastIndexSet`s don't have an underlying Vec<T> that we can
+    // steal, compact in place, and then rebuild the `FastIndexSet`
+    // from. So we have to rebuild the type arena from scratch.
     log::trace!("compacting types");
     let mut new_types = arena::UniqueArena::new();
     for (old_handle, mut ty, span) in module.types.drain_all() {

src/front/spv/mod.rs

@@ -40,7 +40,7 @@ use function::*;
 use crate::{
     arena::{Arena, Handle, UniqueArena},
     proc::{Alignment, Layouter},
-    FastHashMap, FastHashSet,
+    FastHashMap, FastHashSet, FastIndexMap,
 };
 
 use num_traits::cast::FromPrimitive;
@@ -596,11 +596,7 @@ pub struct Frontend<I> {
     /// use that target block id.
     ///
     /// Used to preserve allocations between instruction parsing.
-    switch_cases: indexmap::IndexMap<
-        spirv::Word,
-        (BodyIndex, Vec<i32>),
-        std::hash::BuildHasherDefault<rustc_hash::FxHasher>,
-    >,
+    switch_cases: FastIndexMap<spirv::Word, (BodyIndex, Vec<i32>)>,
 
     /// Tracks access to gl_PerVertex's builtins, it is used to cull unused builtins since initializing those can
     /// affect performance and the mere presence of some of these builtins might cause backends to error since they
@@ -641,7 +637,7 @@ impl<I: Iterator<Item = u32>> Frontend<I> {
             dummy_functions: Arena::new(),
             function_call_graph: GraphMap::new(),
             options: options.clone(),
-            switch_cases: indexmap::IndexMap::default(),
+            switch_cases: FastIndexMap::default(),
             gl_per_vertex_builtin_access: FastHashSet::default(),
         }
     }

src/front/wgsl/lower/mod.rs

@@ -6,8 +6,7 @@ use crate::front::wgsl::parse::number::Number;
 use crate::front::wgsl::parse::{ast, conv};
 use crate::front::{Emitter, Typifier};
 use crate::proc::{ensure_block_returns, Alignment, Layouter, ResolveContext, TypeResolution};
-use crate::{Arena, FastHashMap, Handle, Span};
-use indexmap::IndexMap;
+use crate::{Arena, FastHashMap, FastIndexMap, Handle, Span};
 
 mod construction;
 
@@ -101,7 +100,7 @@ pub struct StatementContext<'source, 'temp, 'out> {
     naga_expressions: &'out mut Arena<crate::Expression>,
     /// Stores the names of expressions that are assigned in `let` statement
     /// Also stores the spans of the names, for use in errors.
-    named_expressions: &'out mut IndexMap<Handle<crate::Expression>, (String, Span)>,
+    named_expressions: &'out mut FastIndexMap<Handle<crate::Expression>, (String, Span)>,
     arguments: &'out [crate::FunctionArgument],
     module: &'out mut crate::Module,
 }
@@ -915,7 +914,7 @@ impl<'source, 'temp> Lowerer<'source, 'temp> {
         let mut local_table = FastHashMap::default();
         let mut local_variables = Arena::new();
         let mut expressions = Arena::new();
-        let mut named_expressions = IndexMap::default();
+        let mut named_expressions = FastIndexMap::default();
 
         let arguments = f
             .arguments

src/lib.rs

@@ -310,12 +310,13 @@ pub type FastHashSet<K> = rustc_hash::FxHashSet<K>;
 pub type FastIndexSet<K> =
     indexmap::IndexSet<K, std::hash::BuildHasherDefault<rustc_hash::FxHasher>>;
 
+/// Insertion-order-preserving hash map (`IndexMap<K, V>`), but with the same
+/// hasher as `FastHashMap<K, V>` (faster but not resilient to DoS attacks).
+pub type FastIndexMap<K, V> =
+    indexmap::IndexMap<K, V, std::hash::BuildHasherDefault<rustc_hash::FxHasher>>;
+
 /// Map of expressions that have associated variable names
-pub(crate) type NamedExpressions = indexmap::IndexMap<
-    Handle<Expression>,
-    String,
-    std::hash::BuildHasherDefault<rustc_hash::FxHasher>,
->;
+pub(crate) type NamedExpressions = FastIndexMap<Handle<Expression>, String>;
 
 /// Early fragment tests.
 ///
@@ -1995,7 +1996,7 @@ pub struct SpecialTypes {
     ///
     /// Call [`Module::generate_predeclared_type`] to populate this if
     /// needed and return the handle.
-    pub predeclared_types: indexmap::IndexMap<PredeclaredType, Handle<Type>>,
+    pub predeclared_types: FastIndexMap<PredeclaredType, Handle<Type>>,
 }
 
 /// Shader module.