From 22dc70ea8d2d5b89dd8c39b169d8d75a8002aa5e Mon Sep 17 00:00:00 2001
From: Nick Fitzgerald <fitzgen@gmail.com>
Date: Wed, 18 Oct 2023 11:24:16 -0700
Subject: [PATCH] Use defined function index when guarding against recursive
 inlining

We were mistakenly using the regular function index, which is the same when
there are zero imported functions, so our existing test passed.
---
 src/optimize.rs       | 12 +++++-------
 tests/all/optimize.rs | 38 ++++++++++++++++++++++++++++++++++++++
 2 files changed, 43 insertions(+), 7 deletions(-)

diff --git a/src/optimize.rs b/src/optimize.rs
index ff8c26f..ccfc9d8 100644
--- a/src/optimize.rs
+++ b/src/optimize.rs
@@ -786,24 +786,22 @@ impl Optimizer {
 
         // ... and if that function index is not already on our winlining stack
         // (i.e. we aren't in a recursive inlining chain)...
-        if on_stack.contains(&callee_func_index) {
+        let defined_func_index = callee_func_index
+            .checked_sub(context.num_imported_funcs)
+            .unwrap();
+        if on_stack.contains(&defined_func_index) {
             return Ok(None);
         }
 
         // ... and if that function has the correct type (if this is not true
         // then either the profile is bogus/mismatched or every time the call
         // site was executed it trapped)...
-        if context.funcs[usize::try_from(callee_func_index - context.num_imported_funcs).unwrap()]
-            != type_index
-        {
+        if context.funcs[usize::try_from(defined_func_index).unwrap()] != type_index {
             return Ok(None);
         }
 
         // ... then we can winline this callee and push it onto our stack!
 
-        let defined_func_index = callee_func_index
-            .checked_sub(context.num_imported_funcs)
-            .unwrap();
         let locals_delta = *num_locals;
         let func_body = context.func_bodies[usize::try_from(defined_func_index).unwrap()].clone();
         let ops = func_body
diff --git a/tests/all/optimize.rs b/tests/all/optimize.rs
index 1c8a80c..10f48d5 100644
--- a/tests/all/optimize.rs
+++ b/tests/all/optimize.rs
@@ -590,6 +590,44 @@ fn dont_inline_direct_recursion() -> Result<()> {
     )
 }
 
+#[test]
+fn dont_inline_direct_recursion_with_imports() -> Result<()> {
+    assert_optimize(
+        Optimizer::new().min_total_calls(100),
+        &[&[(1, 100)]],
+        r#"
+(module
+  (type (func (param i32) (result i32)))
+
+  (import "foo" "bar" (func (type 0)))
+
+  (func (type 0)
+    local.get 0
+    local.get 0
+    call_indirect (type 0)
+  )
+
+  (table 100 100 funcref)
+  (elem (i32.const 0) funcref (ref.func 1))
+)
+        "#,
+        r#"
+(module
+  (type (;0;) (func (param i32) (result i32)))
+  (import "foo" "bar" (func (;0;) (type 0)))
+  (func (;1;) (type 0) (param i32) (result i32)
+    (local i32)
+    local.get 0
+    local.get 0
+    call_indirect (type 0)
+  )
+  (table (;0;) 100 100 funcref)
+  (elem (;0;) (i32.const 0) funcref (ref.func 1))
+)
+        "#,
+    )
+}
+
 #[test]
 fn multiple_funcref_tables() -> Result<()> {
     assert_optimize(