ExtractFunction: Do not always remove the memory and table (#3877)

Instead, run RemoveUnusedModuleElements, which does that sort of thing. That
is, this pass just "extracts" the function by turning all others into imports, and then
they should almost all be removable via RemoveUnusedModuleElements, depending
on whether they are used in the table or not, whether the extracted function calls
them, etc.

Without this, we would error if a function was in the table, and so this fixes #3876

4 files changed, 68 insertions, 19 deletions

diff --git a/src/passes/ExtractFunction.cpp b/src/passes/ExtractFunction.cpp
index ca05b6108..626ac1ee7 100644
--- a/src/passes/ExtractFunction.cpp
+++ b/src/passes/ExtractFunction.cpp
@@ -14,9 +14,11 @@
  * limitations under the License.
  */
 
-// Removes code from all functions but one, leaving a valid module
-// with (mostly) just the code you want to debug (function-parallel,
-// non-lto) passes on.
+// Removes code from all functions but one, leaving a valid module with (mostly)
+// just the code from that function, as best we can.
+//
+// This pass will run --remove-unused-module-elements automatically for you, in
+// order to remove as many things as possible.
 
 #include "pass.h"
 #include "wasm.h"
@@ -44,20 +46,20 @@ struct ExtractFunction : public Pass {
     if (!found) {
       Fatal() << "could not find the function to extract\n";
     }
-    // clear data
-    module->memory.segments.clear();
-    // TODO: if the extracted function contains a call_indirect, the referenced
-    // table should not be removed.
-    module->tables.clear();
-    // leave just an export for the thing we want
-    if (!module->getExportOrNull(name)) {
-      module->exports.clear();
-      auto* export_ = new Export;
-      export_->name = name;
-      export_->value = name;
-      export_->kind = ExternalKind::Function;
-      module->addExport(export_);
-    }
+
+    // Leave just one export, for the thing we want.
+    module->exports.clear();
+    auto* export_ = new Export;
+    export_->name = name;
+    export_->value = name;
+    export_->kind = ExternalKind::Function;
+    module->addExport(export_);
+
+    // Remove unneeded things.
+    PassRunner postRunner(runner);
+    postRunner.add("remove-unused-module-elements");
+    postRunner.setIsNested(true);
+    postRunner.run();
   }
 };
 
diff --git a/src/passes/RemoveUnusedModuleElements.cpp b/src/passes/RemoveUnusedModuleElements.cpp
index b5cf2635a..59d9a2a9b 100644
--- a/src/passes/RemoveUnusedModuleElements.cpp
+++ b/src/passes/RemoveUnusedModuleElements.cpp
@@ -229,6 +229,10 @@ struct RemoveUnusedModuleElements : public Pass {
              analyzer.reachable.count(
                ModuleElement(ModuleElementKind::Table, curr->name)) == 0;
     });
+    // TODO: After removing elements, we may be able to remove more things, and
+    //       should continue to work. (For example, after removing a reference
+    //       to a function from an element segment, we may be able to remove
+    //       that function, etc.)
 
     // Handle the memory
     if (!exportsMemory && !analyzer.usesMemory) {
diff --git a/test/passes/extract-function_pass-arg=extract@foo.txt b/test/passes/extract-function_pass-arg=extract@foo.txt
index 6726bb3a4..3caf274c3 100644
--- a/test/passes/extract-function_pass-arg=extract@foo.txt
+++ b/test/passes/extract-function_pass-arg=extract@foo.txt
@@ -1,9 +1,28 @@
 (module
  (type $none_=>_none (func))
  (import "env" "bar" (func $bar))
- (import "env" "other" (func $other))
  (export "foo" (func $foo))
  (func $foo
   (call $bar)
  )
 )
+(module
+ (type $none_=>_none (func))
+ (import "env" "other" (func $other))
+ (export "foo" (func $foo))
+ (func $foo
+  (nop)
+ )
+)
+(module
+ (type $none (func))
+ (import "env" "other" (func $other))
+ (table $t 10 funcref)
+ (elem $0 (i32.const 0) $other)
+ (export "foo" (func $foo))
+ (func $foo
+  (call_indirect (type $none)
+   (i32.const 10)
+  )
+ )
+)
diff --git a/test/passes/extract-function_pass-arg=extract@foo.wast b/test/passes/extract-function_pass-arg=extract@foo.wast
index 876910300..ab1d8b269 100644
--- a/test/passes/extract-function_pass-arg=extract@foo.wast
+++ b/test/passes/extract-function_pass-arg=extract@foo.wast
@@ -9,4 +9,28 @@
     (drop (i32.const 1))
   )
 )
-
+(module
+  ;; Use another function in the table, but the table is not used in the
+  ;; extracted function
+  (table $t 10 funcref)
+  (elem $0 (table $t) (i32.const 0) func $other)
+  (func $foo
+  )
+  (func $other
+    (drop (i32.const 1))
+  )
+)
+(module
+  ;; Use another function in the table, and the table *is* used. As a result,
+  ;; the table and its elements will remain. The called function, $other, will
+  ;; remain as an import that is placed in the table.
+  (type $none (func))
+  (table $t 10 funcref)
+  (elem $0 (table $t) (i32.const 0) func $other)
+  (func $foo
+   (call_indirect (type $none) (i32.const 10))
+  )
+  (func $other
+    (drop (i32.const 1))
+  )
+)