Refactor the TypeTables type (#3971)

* Remove duplicate `TypeTables` type

This was once needed historically but it is no longer needed.

* Make the internals of `TypeTables` private

Instead of reaching internally for the `wasm_signatures` map an `Index`
implementation now exists to indirect accesses through the type of the
index being accessed. For the component model this table of types will
grow a number of other tables and this'll assist in consuming sites not
having to worry so much about which map they're reaching into.

crates/cranelift/src/compiler.rs

@@ -256,7 +256,7 @@ impl wasmtime_environ::Compiler for Compiler {
         let compiled_trampolines = translation
             .exported_signatures
             .iter()
-            .map(|i| self.host_to_wasm_trampoline(&types.wasm_signatures[*i]))
+            .map(|i| self.host_to_wasm_trampoline(&types[*i]))
             .collect::<Result<Vec<_>, _>>()?;
 
         let mut func_starts = Vec::with_capacity(funcs.len());

crates/cranelift/src/func_environ.rs

@@ -1486,7 +1486,7 @@ impl<'module_environment> cranelift_wasm::FuncEnvironment for FuncEnvironment<'m
         index: TypeIndex,
     ) -> WasmResult<ir::SigRef> {
         let index = self.module.types[index].unwrap_function();
-        let sig = crate::indirect_signature(self.isa, &self.types.wasm_signatures[index]);
+        let sig = crate::indirect_signature(self.isa, &self.types[index]);
         Ok(func.import_signature(sig))
     }
 

crates/cranelift/src/lib.rs

@@ -226,7 +226,7 @@ fn func_signature(
         _ => wasmtime_call_conv(isa),
     };
     let mut sig = blank_sig(isa, call_conv);
-    push_types(isa, &mut sig, &types.wasm_signatures[func.signature]);
+    push_types(isa, &mut sig, &types[func.signature]);
     return sig;
 }
 

crates/environ/src/module.rs

@@ -7,7 +7,7 @@ use serde::{Deserialize, Serialize};
 use std::collections::BTreeMap;
 use std::convert::TryFrom;
 use std::mem;
-use std::ops::Range;
+use std::ops::{Index, Range};
 use wasmtime_types::*;
 
 /// Implemenation styles for WebAssembly linear memory.
@@ -1089,7 +1089,23 @@ impl Module {
 #[derive(Default, Debug, Clone, Serialize, Deserialize)]
 #[allow(missing_docs)]
 pub struct TypeTables {
-    pub wasm_signatures: PrimaryMap<SignatureIndex, WasmFuncType>,
+    pub(crate) wasm_signatures: PrimaryMap<SignatureIndex, WasmFuncType>,
+}
+
+impl TypeTables {
+    /// Returns an iterator of all of the core wasm function signatures
+    /// registered in this instance.
+    pub fn wasm_signatures(&self) -> impl Iterator<Item = (SignatureIndex, &WasmFuncType)> {
+        self.wasm_signatures.iter()
+    }
+}
+
+impl Index<SignatureIndex> for TypeTables {
+    type Output = WasmFuncType;
+
+    fn index(&self, idx: SignatureIndex) -> &WasmFuncType {
+        &self.wasm_signatures[idx]
+    }
 }
 
 /// Type information about functions in a wasm module.

crates/environ/src/module_environ.rs

@@ -509,7 +509,7 @@ impl<'data> ModuleEnvironment<'data> {
 
                 if self.tunables.generate_native_debuginfo {
                     let sig_index = self.result.module.functions[func_index].signature;
-                    let sig = &self.types.wasm_signatures[sig_index];
+                    let sig = &self.types[sig_index];
                     let mut locals = Vec::new();
                     for pair in body.get_locals_reader()? {
                         locals.push(pair?);

crates/jit/src/instantiate.rs

@@ -17,7 +17,7 @@ use std::sync::Arc;
 use thiserror::Error;
 use wasmtime_environ::{
     CompileError, DefinedFuncIndex, FuncIndex, FunctionInfo, Module, ModuleTranslation, PrimaryMap,
-    SignatureIndex, StackMapInformation, Trampoline, Tunables, WasmFuncType, ELF_WASMTIME_ADDRMAP,
+    SignatureIndex, StackMapInformation, Trampoline, Tunables, ELF_WASMTIME_ADDRMAP,
     ELF_WASMTIME_TRAPS,
 };
 use wasmtime_runtime::{
@@ -358,14 +358,6 @@ pub fn mmap_vec_from_obj(obj: Object) -> Result<MmapVec> {
     }
 }
 
-/// This is intended to mirror the type tables in `wasmtime_environ`, except that
-/// it doesn't store the native signatures which are no longer needed past compilation.
-#[derive(Serialize, Deserialize)]
-#[allow(missing_docs)]
-pub struct TypeTables {
-    pub wasm_signatures: PrimaryMap<SignatureIndex, WasmFuncType>,
-}
-
 /// A compiled wasm module, ready to be instantiated.
 pub struct CompiledModule {
     wasm_data: Range<usize>,

crates/jit/src/lib.rs

@@ -30,7 +30,7 @@ mod unwind;
 pub use crate::code_memory::CodeMemory;
 pub use crate::instantiate::{
     finish_compile, mmap_vec_from_obj, subslice_range, CompiledModule, CompiledModuleInfo,
-    SetupError, SymbolizeContext, TypeTables,
+    SetupError, SymbolizeContext,
 };
 pub use demangling::*;
 pub use profiling::*;

crates/wasmtime/src/module.rs

@@ -13,9 +13,9 @@ use std::sync::Arc;
 use wasmparser::{Parser, ValidPayload, Validator};
 use wasmtime_environ::{
     DefinedFuncIndex, DefinedMemoryIndex, FunctionInfo, ModuleEnvironment, PrimaryMap,
-    SignatureIndex,
+    SignatureIndex, TypeTables,
 };
-use wasmtime_jit::{CompiledModule, CompiledModuleInfo, TypeTables};
+use wasmtime_jit::{CompiledModule, CompiledModuleInfo};
 use wasmtime_runtime::{
     CompiledModuleId, MemoryImage, MmapVec, ModuleMemoryImages, VMSharedSignatureIndex,
 };
@@ -418,13 +418,7 @@ impl Module {
         let (mmap, info) =
             wasmtime_jit::finish_compile(translation, obj, funcs, trampolines, tunables)?;
 
-        Ok((
+        Ok((mmap, Some(info), types))
-            mmap,
-            Some(info),
-            TypeTables {
-                wasm_signatures: types.wasm_signatures,
-            },
-        ))
     }
 
     /// Deserializes an in-memory compiled module previously created with
@@ -512,7 +506,7 @@ impl Module {
 
         let signatures = Arc::new(SignatureCollection::new_for_module(
             engine.signatures(),
-            &types.wasm_signatures,
+            &types,
             module.trampolines().map(|(idx, f, _)| (idx, f)),
         ));
 

crates/wasmtime/src/module/serialization.rs

@@ -48,8 +48,8 @@ use std::collections::BTreeMap;
 use std::path::Path;
 use std::str::FromStr;
 use std::sync::Arc;
-use wasmtime_environ::{FlagValue, Tunables};
+use wasmtime_environ::{FlagValue, Tunables, TypeTables};
-use wasmtime_jit::{subslice_range, CompiledModule, CompiledModuleInfo, TypeTables};
+use wasmtime_jit::{subslice_range, CompiledModule, CompiledModuleInfo};
 use wasmtime_runtime::MmapVec;
 
 const HEADER: &[u8] = b"\0wasmtime-aot";

crates/wasmtime/src/signatures.rs

@@ -6,7 +6,7 @@ use std::{
     sync::RwLock,
 };
 use std::{convert::TryFrom, sync::Arc};
-use wasmtime_environ::{PrimaryMap, SignatureIndex, WasmFuncType};
+use wasmtime_environ::{PrimaryMap, SignatureIndex, TypeTables, WasmFuncType};
 use wasmtime_runtime::{VMSharedSignatureIndex, VMTrampoline};
 
 /// Represents a collection of shared signatures.
@@ -27,14 +27,14 @@ impl SignatureCollection {
     /// and trampolines.
     pub fn new_for_module(
         registry: &SignatureRegistry,
-        signatures: &PrimaryMap<SignatureIndex, WasmFuncType>,
+        types: &TypeTables,
         trampolines: impl Iterator<Item = (SignatureIndex, VMTrampoline)>,
     ) -> Self {
         let (signatures, trampolines) = registry
             .0
             .write()
             .unwrap()
-            .register_for_module(signatures, trampolines);
+            .register_for_module(types, trampolines);
 
         Self {
             registry: registry.0.clone(),
@@ -89,7 +89,7 @@ struct SignatureRegistryInner {
 impl SignatureRegistryInner {
     fn register_for_module(
         &mut self,
-        signatures: &PrimaryMap<SignatureIndex, WasmFuncType>,
+        types: &TypeTables,
         trampolines: impl Iterator<Item = (SignatureIndex, VMTrampoline)>,
     ) -> (
         PrimaryMap<SignatureIndex, VMSharedSignatureIndex>,
@@ -98,8 +98,9 @@ impl SignatureRegistryInner {
         let mut sigs = PrimaryMap::default();
         let mut map = HashMap::default();
 
-        for (_, ty) in signatures.iter() {
+        for (idx, ty) in types.wasm_signatures() {
-            sigs.push(self.register(ty));
+            let b = sigs.push(self.register(ty));
+            assert_eq!(idx, b);
         }
 
         for (index, trampoline) in trampolines {

crates/wasmtime/src/types.rs

@@ -1,6 +1,5 @@
 use std::fmt;
-use wasmtime_environ::{EntityType, Global, Memory, Table, WasmFuncType, WasmType};
+use wasmtime_environ::{EntityType, Global, Memory, Table, TypeTables, WasmFuncType, WasmType};
-use wasmtime_jit::TypeTables;
 
 pub(crate) mod matching;
 
@@ -151,9 +150,7 @@ impl ExternType {
 
     pub(crate) fn from_wasmtime(types: &TypeTables, ty: &EntityType) -> ExternType {
         match ty {
-            EntityType::Function(idx) => {
+            EntityType::Function(idx) => FuncType::from_wasm_func_type(types[*idx].clone()).into(),
-                FuncType::from_wasm_func_type(types.wasm_signatures[*idx].clone()).into()
-            }
             EntityType::Global(ty) => GlobalType::from_wasmtime_global(ty).into(),
             EntityType::Memory(ty) => MemoryType::from_wasmtime_memory(ty).into(),
             EntityType::Table(ty) => TableType::from_wasmtime_table(ty).into(),

crates/wasmtime/src/types/matching.rs

@@ -2,8 +2,9 @@ use crate::linker::Definition;
 use crate::store::StoreOpaque;
 use crate::{signatures::SignatureCollection, Engine, Extern};
 use anyhow::{bail, Result};
-use wasmtime_environ::{EntityType, Global, Memory, SignatureIndex, Table, WasmFuncType, WasmType};
+use wasmtime_environ::{
-use wasmtime_jit::TypeTables;
+    EntityType, Global, Memory, SignatureIndex, Table, TypeTables, WasmFuncType, WasmType,
+};
 use wasmtime_runtime::VMSharedSignatureIndex;
 
 pub struct MatchCx<'a> {
@@ -120,7 +121,7 @@ impl MatchCx<'_> {
             return Ok(());
         }
         let msg = "function types incompatible";
-        let expected = &self.types.wasm_signatures[expected];
+        let expected = &self.types[expected];
         let actual = match self.engine.signatures().lookup_type(actual) {
             Some(ty) => ty,
             None => {