3 files changed, 55 insertions, 28 deletions

diff --git a/src/passes/ExtractFunction.cpp b/src/passes/ExtractFunction.cpp
index ac3be9366..a219000c8 100644
--- a/src/passes/ExtractFunction.cpp
+++ b/src/passes/ExtractFunction.cpp
@@ -21,50 +21,73 @@
 // order to remove as many things as possible.
 
 #include "pass.h"
+#include "wasm-builder.h"
 #include "wasm.h"
 
 namespace wasm {
 
+static void extract(PassRunner* runner, Module* module, Name name) {
+  std::cerr << "extracting " << name << "\n";
+  bool found = false;
+  for (auto& func : module->functions) {
+    if (func->name != name) {
+      // Turn it into an import.
+      func->module = "env";
+      func->base = func->name;
+      func->vars.clear();
+      func->body = nullptr;
+    } else {
+      found = true;
+    }
+  }
+  if (!found) {
+    Fatal() << "could not find the function to extract\n";
+  }
+
+  // Leave just one export, for the thing we want.
+  module->exports.clear();
+  module->addExport(Builder::makeExport(name, name, ExternalKind::Function));
+
+  // Remove unneeded things.
+  PassRunner postRunner(runner);
+  postRunner.add("remove-unused-module-elements");
+  postRunner.setIsNested(true);
+  postRunner.run();
+}
+
 struct ExtractFunction : public Pass {
   void run(PassRunner* runner, Module* module) override {
     Name name = runner->options.getArgument(
       "extract-function",
       "ExtractFunction usage:  wasm-opt --extract-function=FUNCTION_NAME");
-    std::cerr << "extracting " << name << "\n";
-    bool found = false;
-    for (auto& func : module->functions) {
-      if (func->name != name) {
-        // Turn it into an import.
-        func->module = "env";
-        func->base = func->name;
-        func->vars.clear();
-        func->body = nullptr;
-      } else {
-        found = true;
+    extract(runner, module, name);
+  }
+};
+
+struct ExtractFunctionIndex : public Pass {
+  void run(PassRunner* runner, Module* module) override {
+    std::string index =
+      runner->options.getArgument("extract-function-index",
+                                  "ExtractFunctionIndex usage: wasm-opt "
+                                  "--extract-function-index=FUNCTION_INDEX");
+    for (char c : index) {
+      if (!std::isdigit(c)) {
+        Fatal() << "Expected numeric function index";
       }
     }
-    if (!found) {
-      Fatal() << "could not find the function to extract\n";
+    Index i = std::stoi(index);
+    if (i >= module->functions.size()) {
+      Fatal() << "Invalid function index";
     }
-
-    // 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();
+    // Assumes imports are at the beginning
+    Name name = module->functions[std::stoi(index)]->name;
+    extract(runner, module, name);
   }
 };
 
-// declare pass
+// declare passes
 
 Pass* createExtractFunctionPass() { return new ExtractFunction(); }
+Pass* createExtractFunctionIndexPass() { return new ExtractFunctionIndex(); }
 
 } // namespace wasm
diff --git a/src/passes/pass.cpp b/src/passes/pass.cpp
index 89b367272..10a688313 100644
--- a/src/passes/pass.cpp
+++ b/src/passes/pass.cpp
@@ -129,6 +129,9 @@ void PassRegistry::registerPasses() {
   registerPass("extract-function",
                "leaves just one function (useful for debugging)",
                createExtractFunctionPass);
+  registerPass("extract-function-index",
+               "leaves just one function selected by index",
+               createExtractFunctionIndexPass);
   registerPass(
     "flatten", "flattens out code, removing nesting", createFlattenPass);
   registerPass("fpcast-emu",
diff --git a/src/passes/passes.h b/src/passes/passes.h
index 4cc8040a4..5f265470d 100644
--- a/src/passes/passes.h
+++ b/src/passes/passes.h
@@ -42,6 +42,7 @@ Pass* createDuplicateImportEliminationPass();
 Pass* createDuplicateFunctionEliminationPass();
 Pass* createEmitTargetFeaturesPass();
 Pass* createExtractFunctionPass();
+Pass* createExtractFunctionIndexPass();
 Pass* createFlattenPass();
 Pass* createFuncCastEmulationPass();
 Pass* createFullPrinterPass();