From 1f8e6dd5b2442ecabed625f94cae485737f8693e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Alexander=20Thei=C3=9Fen?= <alex.theissen@me.com>
Date: Tue, 13 Oct 2020 10:42:36 +0200
Subject: [PATCH] Don't generate duplicate thunks

Previously, functions that appear in multiple places
(exported, start function, table) would generate a thunk
for each place they appear in. Those additional thunks are
identical and only only one of them would be referenced.
Main offender are tables with redundant entries.

This commit eliminates those duplicate thunks without adding
any additional overhead.
---
 src/stack_height/thunk.rs                 | 20 +++--------------
 tests/expectations/stack-height/start.wat | 20 ++---------------
 tests/expectations/stack-height/table.wat | 26 ++++-------------------
 3 files changed, 9 insertions(+), 57 deletions(-)

diff --git a/src/stack_height/thunk.rs b/src/stack_height/thunk.rs
index e3edb30..2bad3fe 100644
--- a/src/stack_height/thunk.rs
+++ b/src/stack_height/thunk.rs
@@ -13,7 +13,6 @@ struct Thunk {
 	signature: FunctionType,
 	// Index in function space of this thunk.
 	idx: Option<u32>,
-	original_func_idx: u32,
 	callee_stack_cost: u32,
 }
 
@@ -23,9 +22,6 @@ pub(crate) fn generate_thunks(
 ) -> Result<elements::Module, Error> {
 	// First, we need to collect all function indices that should be replaced by thunks
 
-	// Function indices which needs to generate thunks.
-	let mut need_thunks: Vec<u32> = Vec::new();
-
 	let mut replacement_map: Map<u32, Thunk> = {
 		let exports = module
 			.export_section()
@@ -57,12 +53,10 @@ pub(crate) fn generate_thunks(
 
 			// Don't generate a thunk if stack_cost of a callee is zero.
 			if callee_stack_cost != 0 {
-				need_thunks.push(func_idx);
 				replacement_map.insert(func_idx, Thunk {
 					signature: resolve_func_type(func_idx, &module)?.clone(),
 					idx: None,
 					callee_stack_cost,
-					original_func_idx: func_idx,
 				});
 			}
 		}
@@ -76,17 +70,9 @@ pub(crate) fn generate_thunks(
 	let mut next_func_idx = module.functions_space() as u32;
 
 	let mut mbuilder = builder::from_module(module);
-	for func_idx in need_thunks {
-		let mut thunk = replacement_map
-			.get_mut(&func_idx)
-			.expect(
-				"`func_idx` should come from `need_thunks`;
-				`need_thunks` is populated with the same items that in `replacement_map`;
-				qed"
-			);
-
+	for (func_idx, thunk) in replacement_map.iter_mut() {
 		let instrumented_call = instrument_call!(
-			thunk.original_func_idx as u32,
+			*func_idx,
 			thunk.callee_stack_cost as i32,
 			ctx.stack_height_global_idx(),
 			ctx.stack_limit()
@@ -113,7 +99,7 @@ pub(crate) fn generate_thunks(
 				// Signature of the thunk should match the original function signature.
 				.signature()
 					.with_params(thunk.signature.params().to_vec())
-					.with_return_type(thunk.signature.return_type().clone())
+					.with_return_type(thunk.signature.return_type())
 					.build()
 				.body()
 					.with_instructions(elements::Instructions::new(
diff --git a/tests/expectations/stack-height/start.wat b/tests/expectations/stack-height/start.wat
index aa38a2e..7a3e905 100644
--- a/tests/expectations/stack-height/start.wat
+++ b/tests/expectations/stack-height/start.wat
@@ -22,23 +22,7 @@
     i32.const 1
     i32.sub
     global.set 0)
-  (func (;4;) (type 1)
-    global.get 0
-    i32.const 1
-    i32.add
-    global.set 0
-    global.get 0
-    i32.const 1024
-    i32.gt_u
-    if  ;; label = @1
-      unreachable
-    end
-    call 1
-    global.get 0
-    i32.const 1
-    i32.sub
-    global.set 0)
   (global (;0;) (mut i32) (i32.const 0))
-  (export "exported_start" (func 4))
+  (export "exported_start" (func 3))
   (export "call" (func 2))
-  (start 4))
+  (start 3))
diff --git a/tests/expectations/stack-height/table.wat b/tests/expectations/stack-height/table.wat
index 710debc..3b05d18 100644
--- a/tests/expectations/stack-height/table.wat
+++ b/tests/expectations/stack-height/table.wat
@@ -26,25 +26,7 @@
     local.get 0
     local.get 1
     i32.add)
-  (func (;3;) (type 2) (param i32 i32) (result i32)
-    local.get 0
-    local.get 1
-    global.get 0
-    i32.const 2
-    i32.add
-    global.set 0
-    global.get 0
-    i32.const 1024
-    i32.gt_u
-    if  ;; label = @1
-      unreachable
-    end
-    call 2
-    global.get 0
-    i32.const 2
-    i32.sub
-    global.set 0)
-  (func (;4;) (type 1) (param i32)
+  (func (;3;) (type 1) (param i32)
     local.get 0
     global.get 0
     i32.const 2
@@ -61,7 +43,7 @@
     i32.const 2
     i32.sub
     global.set 0)
-  (func (;5;) (type 2) (param i32 i32) (result i32)
+  (func (;4;) (type 2) (param i32 i32) (result i32)
     local.get 0
     local.get 1
     global.get 0
@@ -81,5 +63,5 @@
     global.set 0)
   (table (;0;) 10 funcref)
   (global (;0;) (mut i32) (i32.const 0))
-  (export "i32.add" (func 5))
-  (elem (;0;) (i32.const 0) func 0 4 5))
+  (export "i32.add" (func 4))
+  (elem (;0;) (i32.const 0) func 0 3 4))