[wasmparser] Reduce the size of TypeId for better validation performance (#844)

Ekleog · Ekleog-NEAR · web-flow · commit 24d48083cea0 · 2022-12-09T09:36:41.000-06:00
* Reduce the size of TypeId for better validation performance

* Add a benchmark showcasing TypeId size reduction performance benefits

* Reduce TypeId size further to 16 bytes

Performance benefits are less obvious, but it gets us further away from
a regression

Co-authored-by: Léo Gaspard &lt;leo@near.org&gt;
diff --git a/crates/wasmparser/benches/lots-of-types.wasm b/crates/wasmparser/benches/lots-of-types.wasm
diff --git a/crates/wasmparser/src/validator/component.rs b/crates/wasmparser/src/validator/component.rs
@@ -114,12 +114,12 @@ impl ComponentState {
             check_max(current.type_count(), 1, MAX_WASM_TYPES, "types", offset)?;
         }
 
-        current.core_types.push(TypeId {
-            type_size: ty.type_size(),
-            index: types.len(),
-            type_index: Some(current.core_types.len()),
-            is_core: true,
-        });
+        current.core_types.push(TypeId::new(
+            ty.type_size(),
+            types.len(),
+            Some(current.core_types.len()),
+            true,
+        ));
         types.push(ty);
 
         Ok(())
@@ -142,12 +142,8 @@ impl ComponentState {
             exports: module.exports.clone(),
         });
 
-        self.core_modules.push(TypeId {
-            type_size: ty.type_size(),
-            index: types.len(),
-            type_index: None,
-            is_core: true,
-        });
+        self.core_modules
+            .push(TypeId::new(ty.type_size(), types.len(), None, true));
 
         types.push(ty);
 
@@ -213,12 +209,12 @@ impl ComponentState {
             check_max(current.type_count(), 1, MAX_WASM_TYPES, "types", offset)?;
         }
 
-        current.types.push(TypeId {
-            type_size: ty.type_size(),
-            index: types.len(),
-            type_index: Some(current.types.len()),
-            is_core: false,
-        });
+        current.types.push(TypeId::new(
+            ty.type_size(),
+            types.len(),
+            Some(current.types.len()),
+            false,
+        ));
         types.push(ty);
 
         Ok(())
@@ -387,12 +383,8 @@ impl ComponentState {
 
         let lowered_ty = Type::Func(info.into_func_type());
 
-        self.core_funcs.push(TypeId {
-            type_size: lowered_ty.type_size(),
-            index: types.len(),
-            type_index: None,
-            is_core: true,
-        });
+        self.core_funcs
+            .push(TypeId::new(lowered_ty.type_size(), types.len(), None, true));
 
         types.push(lowered_ty);
 
@@ -406,12 +398,8 @@ impl ComponentState {
             exports: mem::take(&mut component.exports),
         });
 
-        self.components.push(TypeId {
-            type_size: ty.type_size(),
-            index: types.len(),
-            type_index: None,
-            is_core: false,
-        });
+        self.components
+            .push(TypeId::new(ty.type_size(), types.len(), None, false));
 
         types.push(ty);
     }
@@ -1079,12 +1067,7 @@ impl ComponentState {
             kind: InstanceTypeKind::Instantiated(module_type_id),
         });
 
-        let id = TypeId {
-            type_size: ty.type_size(),
-            index: types.len(),
-            type_index: None,
-            is_core: true,
-        };
+        let id = TypeId::new(ty.type_size(), types.len(), None, true);
 
         types.push(ty);
 
@@ -1223,12 +1206,7 @@ impl ComponentState {
             kind: ComponentInstanceTypeKind::Instantiated(component_type_id),
         });
 
-        let id = TypeId {
-            type_size: ty.type_size(),
-            index: types.len(),
-            type_index: None,
-            is_core: false,
-        };
+        let id = TypeId::new(ty.type_size(), types.len(), None, false);
 
         types.push(ty);
 
@@ -1330,12 +1308,7 @@ impl ComponentState {
             kind: ComponentInstanceTypeKind::Exports(inst_exports),
         });
 
-        let id = TypeId {
-            type_size: ty.type_size(),
-            index: types.len(),
-            type_index: None,
-            is_core: false,
-        };
+        let id = TypeId::new(ty.type_size(), types.len(), None, false);
 
         types.push(ty);
 
@@ -1418,12 +1391,7 @@ impl ComponentState {
             kind: InstanceTypeKind::Exports(inst_exports),
         });
 
-        let id = TypeId {
-            type_size: ty.type_size(),
-            index: types.len(),
-            type_index: None,
-            is_core: true,
-        };
+        let id = TypeId::new(ty.type_size(), types.len(), None, true);
 
         types.push(ty);
 
@@ -1638,12 +1606,12 @@ impl ComponentState {
         let current = components.last_mut().unwrap();
         check_max(current.type_count(), 1, MAX_WASM_TYPES, "types", offset)?;
 
-        current.core_types.push(TypeId {
-            type_size: ty.type_size,
-            index: ty.index,
-            type_index: Some(current.core_types.len()),
-            is_core: true,
-        });
+        current.core_types.push(TypeId::new(
+            ty.type_size as usize,
+            ty.index,
+            Some(current.core_types.len()),
+            true,
+        ));
 
         Ok(())
     }
@@ -1655,12 +1623,12 @@ impl ComponentState {
         let current = components.last_mut().unwrap();
         check_max(current.type_count(), 1, MAX_WASM_TYPES, "types", offset)?;
 
-        current.types.push(TypeId {
-            type_size: ty.type_size,
-            index: ty.index,
-            type_index: Some(current.types.len()),
-            is_core: false,
-        });
+        current.types.push(TypeId::new(
+            ty.type_size as usize,
+            ty.index,
+            Some(current.types.len()),
+            false,
+        ));
 
         Ok(())
     }
diff --git a/crates/wasmparser/src/validator/core.rs b/crates/wasmparser/src/validator/core.rs
@@ -489,12 +489,12 @@ impl Module {
             check_max(self.types.len(), 1, MAX_WASM_TYPES, "types", offset)?;
         }
 
-        self.types.push(TypeId {
-            type_size: ty.type_size(),
-            index: types.len(),
-            type_index: Some(self.types.len()),
-            is_core: true,
-        });
+        self.types.push(TypeId::new(
+            ty.type_size(),
+            types.len(),
+            Some(self.types.len()),
+            true,
+        ));
         types.push(ty);
         Ok(())
     }
diff --git a/crates/wasmparser/src/validator/types.rs b/crates/wasmparser/src/validator/types.rs
@@ -324,21 +324,52 @@ fn push_primitive_wasm_types(ty: &PrimitiveValType, lowered_types: &mut LoweredT
 /// Represents a unique identifier for a type known to a [`crate::Validator`].
 #[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
 pub struct TypeId {
+    /// The index into the global list of types.
+    pub(crate) index: usize,
     /// The effective type size for the type.
     ///
     /// This is stored as part of the ID to avoid having to recurse through
     /// the global type list when calculating type sizes.
-    pub(crate) type_size: usize,
-    /// The index into the global list of types.
-    pub(crate) index: usize,
-    /// The original type index.
+    pub(crate) type_size: u32,
+    /// Metadata about the type hash, used to avoid two isomorphic types from
+    /// being considered as equal for Eq and Hash.
     ///
-    /// This will be `None` for implicitly defined types, e.g. types for
-    /// modules definitions, component definitions, instantiations, and function
-    /// lowerings.
-    pub(crate) type_index: Option<usize>,
-    /// Whether or not the type is a core type.
-    pub(crate) is_core: bool,
+    /// This contains:
+    /// - Bit (1 << 31): True iff the type is a core type
+    /// - Bit (1 << 30): False iff the type is implicitly defined, eg. types
+    ///   for module definitions, component definitions, instantiations and
+    ///   function lowerings.
+    /// - Other bits: if (1 << 62) is not true, then the type index
+    pub(crate) type_hash: u32,
+}
+
+impl TypeId {
+    pub(crate) fn new(
+        type_size: usize,
+        index: usize,
+        type_index: Option<usize>,
+        is_core: bool,
+    ) -> TypeId {
+        let mut type_hash = 0;
+        if is_core {
+            type_hash |= 1 << 31;
+        }
+        if let Some(type_index) = type_index {
+            type_hash |= 1 << 30;
+            let type_index: u32 = type_index
+                .try_into()
+                .expect("failed converting usize to u32");
+            if type_index & (0b11 << 30) != 0 {
+                panic!("type index bigger than 2**30");
+            }
+            type_hash |= type_index;
+        }
+        TypeId {
+            type_size: type_size.try_into().expect("too big type"),
+            index,
+            type_hash,
+        }
+    }
 }
 
 /// A unified type definition for validating WebAssembly modules and components.
@@ -520,7 +551,7 @@ impl ComponentValType {
     pub(crate) fn type_size(&self) -> usize {
         match self {
             Self::Primitive(_) => 1,
-            Self::Type(id) => id.type_size,
+            Self::Type(id) => id.type_size as usize,
         }
     }
 }
@@ -592,7 +623,7 @@ impl EntityType {
 
     pub(crate) fn type_size(&self) -> usize {
         match self {
-            Self::Func(id) | Self::Tag(id) => id.type_size,
+            Self::Func(id) | Self::Tag(id) => id.type_size as usize,
             Self::Table(_) | Self::Memory(_) | Self::Global(_) => 1,
         }
     }
@@ -800,7 +831,7 @@ impl ComponentEntityType {
             | Self::Func(ty)
             | Self::Type(ty)
             | Self::Instance(ty)
-            | Self::Component(ty) => ty.type_size,
+            | Self::Component(ty) => ty.type_size as usize,
             Self::Value(ty) => ty.type_size(),
         }
     }
