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))
+  )
+)