From 654156714c4787e5e4322ebfe2e78a16e47342e6 Mon Sep 17 00:00:00 2001
From: Alex Crichton <alex@alexcrichton.com>
Date: Thu, 25 Mar 2021 18:44:22 -0500
Subject: [PATCH] Deduplicate function signatures in wasm modules (#2772)

Currently wasmtime will generate a `SignatureIndex`-per-type in the
module itself, even if the module itself declares the same type multiple
times. To make matters worse if the same type is declared across
multiple modules used in a module-linking-using-module then the
signature will be recorded each time it's declared.

This commit adds a simple map to module translation to deduplicate these
function types. This should improve the performance of module-linking
graphs where the same function type may be declared in a number of
modules. For modules that don't use module linking this adds an extra
map that's not used too often, but the time spent managing it should be
dwarfed by other compile tasks.
---
 crates/environ/src/module_environ.rs | 26 ++++++++++++++++----------
 1 file changed, 16 insertions(+), 10 deletions(-)

diff --git a/crates/environ/src/module_environ.rs b/crates/environ/src/module_environ.rs
index 2f53bb4af838..57e2f6bb1e7e 100644
--- a/crates/environ/src/module_environ.rs
+++ b/crates/environ/src/module_environ.rs
@@ -41,6 +41,8 @@ pub struct ModuleEnvironment<'data> {
     /// Intern'd types for this entire translation, shared by all modules.
     types: TypeTables,
 
+    interned_func_types: HashMap<WasmFuncType, SignatureIndex>,
+
     // Various bits and pieces of configuration
     features: WasmFeatures,
     target_config: TargetFrontendConfig,
@@ -147,6 +149,7 @@ impl<'data> ModuleEnvironment<'data> {
             tunables: tunables.clone(),
             features: *features,
             first_module: true,
+            interned_func_types: Default::default(),
         }
     }
 
@@ -378,16 +381,19 @@ impl<'data> cranelift_wasm::ModuleEnvironment<'data> for ModuleEnvironment<'data
     }
 
     fn declare_type_func(&mut self, wasm: WasmFuncType, sig: ir::Signature) -> WasmResult<()> {
-        let sig = translate_signature(sig, self.pointer_type());
-
-        // FIXME(#2469): Signatures should be deduplicated in these two tables
-        // since `SignatureIndex` is already a index space separate from the
-        // module's index space. Note that this may get more urgent with
-        // module-linking modules where types are more likely to get repeated
-        // (across modules).
-        let sig_index = self.types.native_signatures.push(sig);
-        let sig_index2 = self.types.wasm_signatures.push(wasm);
-        debug_assert_eq!(sig_index, sig_index2);
+        // Deduplicate wasm function signatures through `interned_func_types`,
+        // which also deduplicates across wasm modules with module linking.
+        let sig_index = match self.interned_func_types.get(&wasm) {
+            Some(idx) => *idx,
+            None => {
+                let sig = translate_signature(sig, self.pointer_type());
+                let sig_index = self.types.native_signatures.push(sig);
+                let sig_index2 = self.types.wasm_signatures.push(wasm.clone());
+                debug_assert_eq!(sig_index, sig_index2);
+                self.interned_func_types.insert(wasm, sig_index);
+                sig_index
+            }
+        };
         self.result
             .module
             .types