[wasm-split] Use a fresh table when reference types are enabled (#6726)

Rather than trying to trampoline primary-to-secondary calls through an
existing table, just create a fresh table for this purpose. This ensures
that modifications to the existing tables cannot interfere with
primary-to-secondary calls and conversely that loading the secondary
module cannot overwrite modifications to the tables.

4 files changed, 254 insertions, 104 deletions

diff --git a/src/ir/module-splitting.cpp b/src/ir/module-splitting.cpp
index 50809a9a4..22f8b301f 100644
--- a/src/ir/module-splitting.cpp
+++ b/src/ir/module-splitting.cpp
@@ -147,15 +147,23 @@ void TableSlotManager::addSlot(Name func, Slot slot) {
 }
 
 TableSlotManager::TableSlotManager(Module& module) : module(module) {
+  if (module.features.hasReferenceTypes()) {
+    // Just create a new table to manage all primary-to-secondary calls lazily.
+    // Do not re-use slots for functions that will already be in existing
+    // tables, since that is not correct in the face of table mutations.
+    // TODO: Reduce overhead by creating a separate table for each function type
+    // if WasmGC is enabled.
+    return;
+  }
+
   // TODO: Reject or handle passive element segments
-  // TODO: If reference types are enabled, just create a fresh table to make bad
-  // interactions with user code impossible.
   auto funcref = Type(HeapType::func, Nullable);
   auto it = std::find_if(
     module.tables.begin(),
     module.tables.end(),
     [&](std::unique_ptr<Table>& table) { return table->type == funcref; });
   if (it == module.tables.end()) {
+    // There is no indirect function table, so we will create one lazily.
     return;
   }
 
@@ -560,18 +568,15 @@ void ModuleSplitter::indirectReferencesToSecondaryFunctions() {
   } gatherer(*this);
   gatherer.walkModule(&primary);
 
-  // Find all RefFuncs in active elementSegments, which we can ignore: tables
-  // are the means by which we connect the modules, and are handled directly.
-  // Passive segments, however, are like RefFuncs in code, and we need to not
-  // ignore them here.
+  // Ignore references to secondary functions that occur in the active segment
+  // that will contain the imported placeholders. Indirect calls to table slots
+  // initialized by that segment will already go to the right place once the
+  // secondary module has been loaded and the table has been patched.
   std::unordered_set<RefFunc*> ignore;
-  for (auto& seg : primary.elementSegments) {
-    if (!seg->table.is()) {
-      continue;
-    }
-    for (auto* curr : seg->data) {
-      if (auto* refFunc = curr->dynCast<RefFunc>()) {
-        ignore.insert(refFunc);
+  if (tableManager.activeSegment) {
+    for (auto* expr : tableManager.activeSegment->data) {
+      if (auto* ref = expr->dynCast<RefFunc>()) {
+        ignore.insert(ref);
       }
     }
   }
diff --git a/test/example/module-splitting.txt b/test/example/module-splitting.txt
index b083e8967..69fabc816 100644
--- a/test/example/module-splitting.txt
+++ b/test/example/module-splitting.txt
@@ -399,14 +399,24 @@ After:
  (type $0 (func (param i32) (result i32)))
  (import "placeholder" "0" (func $placeholder_0 (type $0) (param i32) (result i32)))
  (table $table 1 funcref)
- (elem $0 (i32.const 0) $placeholder_0)
+ (table $0 1 funcref)
+ (elem $0 (table $table) (i32.const 0) func $trampoline_foo)
+ (elem $0_1 (table $0) (i32.const 0) func $placeholder_0)
  (export "%table" (table $table))
+ (export "%table_1" (table $0))
+ (func $trampoline_foo (type $0) (param $0 i32) (result i32)
+  (call_indirect $0 (type $0)
+   (local.get $0)
+   (i32.const 0)
+  )
+ )
 )
 Secondary:
 (module
  (type $0 (func (param i32) (result i32)))
+ (import "primary" "%table_1" (table $0 1 funcref))
  (import "primary" "%table" (table $table 1 funcref))
- (elem $0 (i32.const 0) $foo)
+ (elem $0 (table $0) (i32.const 0) func $foo)
  (func $foo (type $0) (param $0 i32) (result i32)
   (local.get $0)
  )
@@ -427,16 +437,25 @@ After:
 (module
  (type $0 (func (param i32) (result i32)))
  (import "placeholder" "0" (func $placeholder_0 (type $0) (param i32) (result i32)))
- (import "placeholder" "1" (func $placeholder_1 (type $0) (param i32) (result i32)))
  (table $table 2 funcref)
- (elem $0 (i32.const 0) $placeholder_0 $placeholder_1)
+ (table $0 1 funcref)
+ (elem $0 (table $table) (i32.const 0) func $trampoline_foo $trampoline_foo)
+ (elem $0_1 (table $0) (i32.const 0) func $placeholder_0)
  (export "%table" (table $table))
+ (export "%table_1" (table $0))
+ (func $trampoline_foo (type $0) (param $0 i32) (result i32)
+  (call_indirect $0 (type $0)
+   (local.get $0)
+   (i32.const 0)
+  )
+ )
 )
 Secondary:
 (module
  (type $0 (func (param i32) (result i32)))
+ (import "primary" "%table_1" (table $0 1 funcref))
  (import "primary" "%table" (table $table 2 funcref))
- (elem $0 (i32.const 0) $foo $foo)
+ (elem $0 (table $0) (i32.const 0) func $foo)
  (func $foo (type $0) (param $0 i32) (result i32)
   (local.get $0)
  )
@@ -457,23 +476,33 @@ Keeping: <none>
 After:
 (module
  (type $0 (func (param i32) (result i32)))
- (import "placeholder" "42" (func $placeholder_42 (type $0) (param i32) (result i32)))
+ (import "placeholder" "0" (func $placeholder_0 (type $0) (param i32) (result i32)))
  (table $table 1000 funcref)
- (elem $0 (i32.const 42) $placeholder_42)
+ (table $0 1 funcref)
+ (elem $0 (table $table) (i32.const 42) func $trampoline_foo)
+ (elem $0_1 (table $0) (i32.const 0) func $placeholder_0)
  (export "foo" (func $foo))
  (export "%table" (table $table))
+ (export "%table_2" (table $0))
  (func $foo (type $0) (param $0 i32) (result i32)
-  (call_indirect $table (type $0)
+  (call_indirect $0 (type $0)
+   (local.get $0)
+   (i32.const 0)
+  )
+ )
+ (func $trampoline_foo (type $0) (param $0 i32) (result i32)
+  (call_indirect $0 (type $0)
    (local.get $0)
-   (i32.const 42)
+   (i32.const 0)
   )
  )
 )
 Secondary:
 (module
  (type $0 (func (param i32) (result i32)))
+ (import "primary" "%table_2" (table $0 1 funcref))
  (import "primary" "%table" (table $table 1000 funcref))
- (elem $0 (i32.const 42) $foo)
+ (elem $0 (table $0) (i32.const 0) func $foo)
  (func $foo (type $0) (param $0 i32) (result i32)
   (local.get $0)
  )
@@ -498,23 +527,33 @@ After:
  (import "env" "base" (global $base i32))
  (import "placeholder" "0" (func $placeholder_0 (type $0) (param i32) (result i32)))
  (table $table 1000 funcref)
- (elem $0 (global.get $base) $placeholder_0)
+ (table $0 1 funcref)
+ (elem $0 (table $table) (global.get $base) func $trampoline_foo)
+ (elem $0_1 (table $0) (i32.const 0) func $placeholder_0)
  (export "foo" (func $foo))
  (export "%table" (table $table))
+ (export "%table_2" (table $0))
  (export "%global" (global $base))
  (func $foo (type $0) (param $0 i32) (result i32)
-  (call_indirect $table (type $0)
+  (call_indirect $0 (type $0)
+   (local.get $0)
+   (i32.const 0)
+  )
+ )
+ (func $trampoline_foo (type $0) (param $0 i32) (result i32)
+  (call_indirect $0 (type $0)
    (local.get $0)
-   (global.get $base)
+   (i32.const 0)
   )
  )
 )
 Secondary:
 (module
  (type $0 (func (param i32) (result i32)))
+ (import "primary" "%table_2" (table $0 1 funcref))
  (import "primary" "%table" (table $table 1000 funcref))
  (import "primary" "%global" (global $base i32))
- (elem $0 (global.get $base) $foo)
+ (elem $0 (table $0) (i32.const 0) func $foo)
  (func $foo (type $0) (param $0 i32) (result i32)
   (local.get $0)
  )
@@ -538,25 +577,34 @@ After:
  (type $0 (func (param i32) (result i32)))
  (import "env" "base" (global $base i32))
  (import "placeholder" "0" (func $placeholder_0 (type $0) (param i32) (result i32)))
- (import "placeholder" "1" (func $placeholder_1 (type $0) (param i32) (result i32)))
  (table $table 1000 funcref)
- (elem $0 (global.get $base) $placeholder_0 $placeholder_1)
+ (table $0 1 funcref)
+ (elem $0 (table $table) (global.get $base) func $trampoline_foo $trampoline_foo)
+ (elem $0_1 (table $0) (i32.const 0) func $placeholder_0)
  (export "foo" (func $foo))
  (export "%table" (table $table))
+ (export "%table_2" (table $0))
  (export "%global" (global $base))
  (func $foo (type $0) (param $0 i32) (result i32)
-  (call_indirect $table (type $0)
+  (call_indirect $0 (type $0)
    (local.get $0)
-   (global.get $base)
+   (i32.const 0)
+  )
+ )
+ (func $trampoline_foo (type $0) (param $0 i32) (result i32)
+  (call_indirect $0 (type $0)
+   (local.get $0)
+   (i32.const 0)
   )
  )
 )
 Secondary:
 (module
  (type $0 (func (param i32) (result i32)))
+ (import "primary" "%table_2" (table $0 1 funcref))
  (import "primary" "%table" (table $table 1000 funcref))
  (import "primary" "%global" (global $base i32))
- (elem $0 (global.get $base) $foo $foo)
+ (elem $0 (table $0) (i32.const 0) func $foo)
  (func $foo (type $0) (param $0 i32) (result i32)
   (local.get $0)
  )
@@ -584,34 +632,38 @@ After:
  (type $0 (func (param i32) (result i32)))
  (type $1 (func))
  (import "env" "base" (global $base i32))
- (import "placeholder" "1" (func $placeholder_1 (type $0) (param i32) (result i32)))
+ (import "placeholder" "0" (func $placeholder_0 (type $0) (param i32) (result i32)))
  (table $table 1000 funcref)
- (elem $0 (global.get $base) $null $placeholder_1)
+ (table $0 1 funcref)
+ (elem $0 (table $table) (global.get $base) func $null $trampoline_foo)
+ (elem $0_1 (table $0) (i32.const 0) func $placeholder_0)
  (export "foo" (func $foo))
- (export "%null" (func $null))
  (export "%table" (table $table))
+ (export "%table_2" (table $0))
  (export "%global" (global $base))
  (func $null (type $1)
   (nop)
  )
  (func $foo (type $0) (param $0 i32) (result i32)
-  (call_indirect $table (type $0)
+  (call_indirect $0 (type $0)
+   (local.get $0)
+   (i32.const 0)
+  )
+ )
+ (func $trampoline_foo (type $0) (param $0 i32) (result i32)
+  (call_indirect $0 (type $0)
    (local.get $0)
-   (i32.add
-    (global.get $base)
-    (i32.const 1)
-   )
+   (i32.const 0)
   )
  )
 )
 Secondary:
 (module
  (type $0 (func (param i32) (result i32)))
- (type $1 (func))
+ (import "primary" "%table_2" (table $0 1 funcref))
  (import "primary" "%table" (table $table 1000 funcref))
  (import "primary" "%global" (global $base i32))
- (import "primary" "%null" (func $null (type $1)))
- (elem $0 (global.get $base) $null $foo)
+ (elem $0 (table $0) (i32.const 0) func $foo)
  (func $foo (type $0) (param $0 i32) (result i32)
   (local.get $0)
  )
@@ -799,23 +851,36 @@ After:
 (module
  (type $0 (func))
  (import "placeholder" "0" (func $placeholder_0 (type $0)))
- (import "placeholder" "2" (func $placeholder_2 (type $0)))
+ (import "placeholder" "1" (func $placeholder_1 (type $0)))
  (table $table 4 funcref)
- (elem $0 (i32.const 0) $placeholder_0 $bar $placeholder_2 $quux)
+ (table $0 2 funcref)
+ (elem $0 (table $table) (i32.const 0) func $trampoline_foo $bar $trampoline_baz $quux)
+ (elem $0_1 (table $0) (i32.const 0) func $placeholder_0 $placeholder_1)
  (export "%table" (table $table))
+ (export "%table_1" (table $0))
  (func $bar (type $0)
   (nop)
  )
  (func $quux (type $0)
   (nop)
  )
+ (func $trampoline_foo (type $0)
+  (call_indirect $0 (type $0)
+   (i32.const 0)
+  )
+ )
+ (func $trampoline_baz (type $0)
+  (call_indirect $0 (type $0)
+   (i32.const 1)
+  )
+ )
 )
 Secondary:
 (module
  (type $0 (func))
+ (import "primary" "%table_1" (table $0 2 funcref))
  (import "primary" "%table" (table $table 4 funcref))
- (elem $0 (i32.const 0) $foo)
- (elem $1 (i32.const 2) $baz)
+ (elem $0 (table $0) (i32.const 0) func $foo $baz)
  (func $baz (type $0)
   (nop)
  )
@@ -850,11 +915,13 @@ After:
  (type $0 (func))
  (import "env" "base" (global $base i32))
  (import "placeholder" "0" (func $placeholder_0 (type $0)))
- (import "placeholder" "2" (func $placeholder_2 (type $0)))
+ (import "placeholder" "1" (func $placeholder_1 (type $0)))
  (table $table 4 funcref)
- (elem $0 (global.get $base) $placeholder_0 $bar $placeholder_2 $quux)
- (export "%bar" (func $bar))
+ (table $0 2 funcref)
+ (elem $0 (table $table) (global.get $base) func $trampoline_foo $bar $trampoline_baz $quux)
+ (elem $0_1 (table $0) (i32.const 0) func $placeholder_0 $placeholder_1)
  (export "%table" (table $table))
+ (export "%table_1" (table $0))
  (export "%global" (global $base))
  (func $bar (type $0)
   (nop)
@@ -862,14 +929,24 @@ After:
  (func $quux (type $0)
   (nop)
  )
+ (func $trampoline_foo (type $0)
+  (call_indirect $0 (type $0)
+   (i32.const 0)
+  )
+ )
+ (func $trampoline_baz (type $0)
+  (call_indirect $0 (type $0)
+   (i32.const 1)
+  )
+ )
 )
 Secondary:
 (module
  (type $0 (func))
+ (import "primary" "%table_1" (table $0 2 funcref))
  (import "primary" "%table" (table $table 4 funcref))
  (import "primary" "%global" (global $base i32))
- (import "primary" "%bar" (func $bar (type $0)))
- (elem $0 (global.get $base) $foo $bar $baz)
+ (elem $0 (table $0) (i32.const 0) func $foo $baz)
  (func $baz (type $0)
   (nop)
  )
@@ -903,20 +980,38 @@ After:
  (type $0 (func))
  (import "placeholder" "0" (func $placeholder_0 (type $0)))
  (import "placeholder" "1" (func $placeholder_1 (type $0)))
- (import "placeholder" "3" (func $placeholder_3 (type $0)))
+ (import "placeholder" "2" (func $placeholder_2 (type $0)))
  (table $table 4 funcref)
- (elem $0 (i32.const 0) $placeholder_0 $placeholder_1 $baz $placeholder_3)
+ (table $0 3 funcref)
+ (elem $0 (table $table) (i32.const 0) func $trampoline_foo $trampoline_bar $baz $trampoline_quux)
+ (elem $0_1 (table $0) (i32.const 0) func $placeholder_0 $placeholder_1 $placeholder_2)
  (export "%table" (table $table))
+ (export "%table_1" (table $0))
  (func $baz (type $0)
   (nop)
  )
+ (func $trampoline_foo (type $0)
+  (call_indirect $0 (type $0)
+   (i32.const 0)
+  )
+ )
+ (func $trampoline_bar (type $0)
+  (call_indirect $0 (type $0)
+   (i32.const 1)
+  )
+ )
+ (func $trampoline_quux (type $0)
+  (call_indirect $0 (type $0)
+   (i32.const 2)
+  )
+ )
 )
 Secondary:
 (module
  (type $0 (func))
+ (import "primary" "%table_1" (table $0 3 funcref))
  (import "primary" "%table" (table $table 4 funcref))
- (elem $0 (i32.const 0) $foo $bar)
- (elem $1 (i32.const 3) $quux)
+ (elem $0 (table $0) (i32.const 0) func $foo $bar $quux)
  (func $bar (type $0)
   (nop)
  )
@@ -955,23 +1050,40 @@ After:
  (import "env" "base" (global $base i32))
  (import "placeholder" "0" (func $placeholder_0 (type $0)))
  (import "placeholder" "1" (func $placeholder_1 (type $0)))
- (import "placeholder" "3" (func $placeholder_3 (type $0)))
+ (import "placeholder" "2" (func $placeholder_2 (type $0)))
  (table $table 4 funcref)
- (elem $0 (global.get $base) $placeholder_0 $placeholder_1 $baz $placeholder_3)
- (export "%baz" (func $baz))
+ (table $0 3 funcref)
+ (elem $0 (table $table) (global.get $base) func $trampoline_foo $trampoline_bar $baz $trampoline_quux)
+ (elem $0_1 (table $0) (i32.const 0) func $placeholder_0 $placeholder_1 $placeholder_2)
  (export "%table" (table $table))
+ (export "%table_1" (table $0))
  (export "%global" (global $base))
  (func $baz (type $0)
   (nop)
  )
+ (func $trampoline_foo (type $0)
+  (call_indirect $0 (type $0)
+   (i32.const 0)
+  )
+ )
+ (func $trampoline_bar (type $0)
+  (call_indirect $0 (type $0)
+   (i32.const 1)
+  )
+ )
+ (func $trampoline_quux (type $0)
+  (call_indirect $0 (type $0)
+   (i32.const 2)
+  )
+ )
 )
 Secondary:
 (module
  (type $0 (func))
+ (import "primary" "%table_1" (table $0 3 funcref))
  (import "primary" "%table" (table $table 4 funcref))
  (import "primary" "%global" (global $base i32))
- (import "primary" "%baz" (func $baz (type $0)))
- (elem $0 (global.get $base) $foo $bar $baz $quux)
+ (elem $0 (table $0) (i32.const 0) func $foo $bar $quux)
  (func $bar (type $0)
   (nop)
  )
@@ -1002,23 +1114,32 @@ After:
 (module
  (type $0 (func))
  (import "env" "base" (global $base i32))
- (import "placeholder" "1" (func $placeholder_1 (type $0)))
+ (import "placeholder" "0" (func $placeholder_0 (type $0)))
  (table $table 2 funcref)
- (elem $0 (global.get $base) $foo $placeholder_1)
+ (table $0 1 funcref)
+ (elem $0 (table $table) (global.get $base) func $foo $trampoline_bar)
+ (elem $0_1 (table $0) (i32.const 0) func $placeholder_0)
  (export "%foo" (func $foo))
  (export "%table" (table $table))
+ (export "%table_2" (table $0))
  (export "%global" (global $base))
  (func $foo (type $0)
   (nop)
  )
+ (func $trampoline_bar (type $0)
+  (call_indirect $0 (type $0)
+   (i32.const 0)
+  )
+ )
 )
 Secondary:
 (module
  (type $0 (func))
+ (import "primary" "%table_2" (table $0 1 funcref))
  (import "primary" "%table" (table $table 2 funcref))
  (import "primary" "%global" (global $base i32))
  (import "primary" "%foo" (func $foo (type $0)))
- (elem $0 (global.get $base) $foo $bar)
+ (elem $0 (table $0) (i32.const 0) func $bar)
  (func $bar (type $0)
   (call $foo)
  )
@@ -1045,24 +1166,28 @@ Keeping: foo
 After:
 (module
  (type $0 (func (param i32) (result i32)))
- (import "placeholder" "1" (func $placeholder_1 (type $0) (param i32) (result i32)))
- (table $table 2 2 funcref)
- (elem $0 (i32.const 0) $foo $placeholder_1)
+ (import "placeholder" "0" (func $placeholder_0 (type $0) (param i32) (result i32)))
+ (table $table 1 1 funcref)
+ (table $0 1 funcref)
+ (elem $0 (table $table) (i32.const 0) func $foo)
+ (elem $0_1 (table $0) (i32.const 0) func $placeholder_0)
  (export "%foo" (func $foo))
  (export "%table" (table $table))
+ (export "%table_2" (table $0))
  (func $foo (type $0) (param $0 i32) (result i32)
-  (call_indirect $table (type $0)
+  (call_indirect $0 (type $0)
+   (i32.const 0)
    (i32.const 0)
-   (i32.const 1)
   )
  )
 )
 Secondary:
 (module
  (type $0 (func (param i32) (result i32)))
- (import "primary" "%table" (table $table 2 2 funcref))
+ (import "primary" "%table_2" (table $0 1 funcref))
+ (import "primary" "%table" (table $table 1 1 funcref))
  (import "primary" "%foo" (func $foo (type $0) (param i32) (result i32)))
- (elem $0 (i32.const 1) $bar)
+ (elem $0 (table $0) (i32.const 0) func $bar)
  (func $bar (type $0) (param $0 i32) (result i32)
   (call $foo
    (i32.const 1)
diff --git a/test/lit/wasm-split/passive.wast b/test/lit/wasm-split/passive.wast
index a72dcada5..322288201 100644
--- a/test/lit/wasm-split/passive.wast
+++ b/test/lit/wasm-split/passive.wast
@@ -1,31 +1,31 @@
 ;; NOTE: Assertions have been generated by update_lit_checks.py --all-items and should not be edited.
 
-;; RUN: wasm-split %s --split-funcs=second-in-table -g -o1 %t.1.wasm -o2 %t.2.wasm -all
-;; RUN: wasm-dis %t.1.wasm | filecheck %s --check-prefix PRIMARY
-;; RUN: wasm-dis %t.2.wasm | filecheck %s --check-prefix SECONDARY
+;; RUN: wasm-split %s -all --split-funcs=second-in-table -g -o1 %t.1.wasm -o2 %t.2.wasm
+;; RUN: wasm-dis -all %t.1.wasm | filecheck %s --check-prefix PRIMARY
+;; RUN: wasm-dis -all %t.2.wasm | filecheck %s --check-prefix SECONDARY
 
 (module
  (type $func-array (array (mut funcref)))
 
  ;; PRIMARY:      (type $0 (func))
 
- ;; PRIMARY:      (import "placeholder" "0" (func $placeholder_0))
+ ;; PRIMARY:      (import "placeholder" "0" (func $placeholder_0 (type $0)))
 
  ;; PRIMARY:      (table $table 3 funcref)
  (table $table 3 funcref)
 
- ;; Workaround for https://github.com/WebAssembly/binaryen/issues/6572 - we
- ;; error without an active segment.
- (elem (i32.const 0))
-
- ;; PRIMARY:      (elem $0 (i32.const 0) $placeholder_0)
+ ;; PRIMARY:      (table $1 1 funcref)
 
  ;; PRIMARY:      (elem $passive func $in-table $1)
  (elem $passive func $in-table $second-in-table)
 
+ ;; PRIMARY:      (elem $1 (table $1) (i32.const 0) func $placeholder_0)
+
  ;; PRIMARY:      (export "table" (table $table))
 
- ;; PRIMARY:      (func $in-table
+ ;; PRIMARY:      (export "table_1" (table $1))
+
+ ;; PRIMARY:      (func $in-table (type $0)
  ;; PRIMARY-NEXT:  (nop)
  ;; PRIMARY-NEXT: )
  (func $in-table
@@ -35,11 +35,13 @@
 
  ;; SECONDARY:      (type $0 (func))
 
+ ;; SECONDARY:      (import "primary" "table_1" (table $timport$0 1 funcref))
+
  ;; SECONDARY:      (import "primary" "table" (table $table 3 funcref))
 
- ;; SECONDARY:      (elem $0 (i32.const 0) $second-in-table)
+ ;; SECONDARY:      (elem $0 (table $timport$0) (i32.const 0) func $second-in-table)
 
- ;; SECONDARY:      (func $second-in-table
+ ;; SECONDARY:      (func $second-in-table (type $0)
  ;; SECONDARY-NEXT:  (nop)
  ;; SECONDARY-NEXT: )
  (func $second-in-table
@@ -47,8 +49,8 @@
   ;; handle it by adding a trampoline from the segment as a new function "$1".
  )
 )
-;; PRIMARY:      (func $1
-;; PRIMARY-NEXT:  (call_indirect (type $0)
+;; PRIMARY:      (func $1 (type $0)
+;; PRIMARY-NEXT:  (call_indirect $1 (type $0)
 ;; PRIMARY-NEXT:   (i32.const 0)
 ;; PRIMARY-NEXT:  )
 ;; PRIMARY-NEXT: )
diff --git a/test/lit/wasm-split/ref.func.wast b/test/lit/wasm-split/ref.func.wast
index f7832dcf4..39cce5dea 100644
--- a/test/lit/wasm-split/ref.func.wast
+++ b/test/lit/wasm-split/ref.func.wast
@@ -1,8 +1,8 @@
 ;; NOTE: Assertions have been generated by update_lit_checks.py --all-items and should not be edited.
 
-;; RUN: wasm-split %s --split-funcs=second,second-in-table -g -o1 %t.1.wasm -o2 %t.2.wasm -all
-;; RUN: wasm-dis %t.1.wasm | filecheck %s --check-prefix PRIMARY
-;; RUN: wasm-dis %t.2.wasm | filecheck %s --check-prefix SECONDARY
+;; RUN: wasm-split %s -all --split-funcs=second,second-in-table -g -o1 %t.1.wasm -o2 %t.2.wasm
+;; RUN: wasm-dis -all %t.1.wasm | filecheck %s --check-prefix PRIMARY
+;; RUN: wasm-dis -all %t.2.wasm | filecheck %s --check-prefix SECONDARY
 
 ;; Test that we handle ref.func operations properly as we split out $second.
 ;; ref.funcs that refer to the other module must be fixed up to refer to
@@ -10,33 +10,41 @@
 (module
  ;; PRIMARY:      (type $0 (func))
 
- ;; PRIMARY:      (import "placeholder" "1" (func $placeholder_1))
+ ;; PRIMARY:      (import "placeholder" "0" (func $placeholder_0 (type $0)))
 
- ;; PRIMARY:      (import "placeholder" "2" (func $placeholder_2))
+ ;; PRIMARY:      (import "placeholder" "1" (func $placeholder_1 (type $0)))
 
  ;; PRIMARY:      (global $glob1 (ref func) (ref.func $prime))
 
  ;; PRIMARY:      (global $glob2 (ref func) (ref.func $2))
 
- ;; PRIMARY:      (table $table 3 3 funcref)
+ ;; PRIMARY:      (table $table 1 1 funcref)
  (table $table 1 1 funcref)
 
  (global $glob1 (ref func) (ref.func $prime))
 
  (global $glob2 (ref func) (ref.func $second))
 
- ;; PRIMARY:      (elem $elem (i32.const 0) $in-table $placeholder_1 $placeholder_2)
+ ;; PRIMARY:      (table $1 2 funcref)
+
+ ;; PRIMARY:      (elem $elem (table $table) (i32.const 0) func $in-table $3)
  (elem $elem (i32.const 0) $in-table $second-in-table)
 
+ ;; PRIMARY:      (elem $1 (table $1) (i32.const 0) func $placeholder_0 $placeholder_1)
+
+ ;; PRIMARY:      (elem declare func $2 $prime)
+
  ;; PRIMARY:      (export "prime" (func $prime))
 
  ;; PRIMARY:      (export "table" (table $table))
 
+ ;; PRIMARY:      (export "table_2" (table $1))
+
  ;; PRIMARY:      (export "global" (global $glob1))
 
- ;; PRIMARY:      (export "global_3" (global $glob2))
+ ;; PRIMARY:      (export "global_4" (global $glob2))
 
- ;; PRIMARY:      (func $prime
+ ;; PRIMARY:      (func $prime (type $0)
  ;; PRIMARY-NEXT:  (drop
  ;; PRIMARY-NEXT:   (ref.func $prime)
  ;; PRIMARY-NEXT:  )
@@ -55,17 +63,21 @@
 
  ;; SECONDARY:      (type $0 (func))
 
- ;; SECONDARY:      (import "primary" "table" (table $table 3 3 funcref))
+ ;; SECONDARY:      (import "primary" "table_2" (table $timport$0 2 funcref))
+
+ ;; SECONDARY:      (import "primary" "table" (table $table 1 1 funcref))
 
  ;; SECONDARY:      (import "primary" "global" (global $glob1 (ref func)))
 
- ;; SECONDARY:      (import "primary" "global_3" (global $glob2 (ref func)))
+ ;; SECONDARY:      (import "primary" "global_4" (global $glob2 (ref func)))
 
- ;; SECONDARY:      (import "primary" "prime" (func $prime))
+ ;; SECONDARY:      (import "primary" "prime" (func $prime (type $0)))
 
- ;; SECONDARY:      (elem $0 (i32.const 1) $second-in-table $second)
+ ;; SECONDARY:      (elem $0 (table $timport$0) (i32.const 0) func $second $second-in-table)
 
- ;; SECONDARY:      (func $second
+ ;; SECONDARY:      (elem declare func $prime)
+
+ ;; SECONDARY:      (func $second (type $0)
  ;; SECONDARY-NEXT:  (drop
  ;; SECONDARY-NEXT:   (ref.func $prime)
  ;; SECONDARY-NEXT:  )
@@ -82,7 +94,7 @@
   )
  )
 
- ;; PRIMARY:      (func $in-table
+ ;; PRIMARY:      (func $in-table (type $0)
  ;; PRIMARY-NEXT:  (nop)
  ;; PRIMARY-NEXT: )
  (func $in-table
@@ -91,7 +103,7 @@
   ;; table is a list of ref.funcs.
  )
 
- ;; SECONDARY:      (func $second-in-table
+ ;; SECONDARY:      (func $second-in-table (type $0)
  ;; SECONDARY-NEXT:  (nop)
  ;; SECONDARY-NEXT: )
  (func $second-in-table
@@ -99,8 +111,14 @@
   ;; (but we will get a placeholder, as all split-out functions do).
  )
 )
-;; PRIMARY:      (func $2
-;; PRIMARY-NEXT:  (call_indirect (type $0)
-;; PRIMARY-NEXT:   (i32.const 2)
+;; PRIMARY:      (func $2 (type $0)
+;; PRIMARY-NEXT:  (call_indirect $1 (type $0)
+;; PRIMARY-NEXT:   (i32.const 0)
+;; PRIMARY-NEXT:  )
+;; PRIMARY-NEXT: )
+
+;; PRIMARY:      (func $3 (type $0)
+;; PRIMARY-NEXT:  (call_indirect $1 (type $0)
+;; PRIMARY-NEXT:   (i32.const 1)
 ;; PRIMARY-NEXT:  )
 ;; PRIMARY-NEXT: )