Add a extract-function-index pass

scripts/foreach.py

@@ -0,0 +1,46 @@
+#!/usr/bin/env python3
+
+# Copyright 2021 WebAssembly Community Group participants
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import sys
+import subprocess
+
+from test import support
+
+
+# Usage: foreach infile tempfile cmd...
+#
+# Split 'infile', which contains multiple text modules, into separate temp files
+# containing one text module each and named `tempfile`.0, `tempfile`.1, etc. Run
+# `cmd` with the current temp file appended to it on all the temp files in
+# sequence. Exit with code 0 only if all of the subprocesses exited with code 0.
+def main():
+    infile = sys.argv[1]
+    tempfile = sys.argv[2]
+    cmd = sys.argv[3:]
+    returncode = 0
+    for i, (module, asserts) in enumerate(support.split_wast(infile)):
+        tempname = tempfile + '.' + str(i)
+        with open(tempname, 'w') as temp:
+            print(module, file=temp)
+        new_cmd = cmd + [tempname]
+        result = subprocess.run(new_cmd)
+        if result.returncode != 0:
+            returncode = result.returncode
+    sys.exit(returncode)
+
+
+if __name__ == '__main__':
+    main()

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

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",

src/passes/passes.h

@@ -42,6 +42,7 @@ Pass* createDuplicateImportEliminationPass();
 Pass* createDuplicateFunctionEliminationPass();
 Pass* createEmitTargetFeaturesPass();
 Pass* createExtractFunctionPass();
+Pass* createExtractFunctionIndexPass();
 Pass* createFlattenPass();
 Pass* createFuncCastEmulationPass();
 Pass* createFullPrinterPass();

test/lit/lit.cfg.py

@@ -15,7 +15,8 @@
     tool = tool_file[:-4] if tool_file.endswith('.exe') else tool_file
     config.substitutions.append((tool, tool_path))
 
-# Also make the `not` command available
-not_file = config.binaryen_src_root + '/scripts/not.py'
-python = sys.executable.replace('\\', '/')
-config.substitutions.append(('not', python + ' ' + not_file))
+# Also make the `not` and `foreach` commands available
+for tool in ('not', 'foreach'):
+    tool_file = config.binaryen_src_root + '/scripts/' + tool + '.py'
+    python = sys.executable.replace('\\', '/')
+    config.substitutions.append((tool, python + ' ' + tool_file))

test/lit/passes/extract-function.wast

@@ -0,0 +1,72 @@
+;; RUN: foreach %s %t wasm-opt --extract-function=foo -S -o - | filecheck %s
+;; RUN: foreach %s %t wasm-opt --extract-function --pass-arg=extract-function@foo -S -o - | filecheck %s
+;; RUN: foreach %s %t wasm-opt --extract-function-index=0 -S -o - | filecheck %s
+;; RUN: foreach %s %t wasm-opt --extract-function-index --pass-arg=extract-function-index@0 -S -o - | filecheck %s
+
+;; CHECK:      (module
+;; CHECK-NEXT:  (type $none_=>_none (func))
+;; CHECK-NEXT:  (import "env" "bar" (func $bar))
+;; CHECK-NEXT:  (export "foo" (func $foo))
+;; CHECK-NEXT:  (func $foo
+;; CHECK-NEXT:   (call $bar)
+;; CHECK-NEXT:  )
+;; CHECK-NEXT: )
+(module
+  (func $foo
+    (call $bar)
+  )
+  (func $bar
+    (call $foo)
+  )
+  (func $other
+    (drop (i32.const 1))
+  )
+)
+
+;; CHECK:      (module
+;; CHECK-NEXT:  (type $none_=>_none (func))
+;; CHECK-NEXT:  (import "env" "other" (func $other))
+;; CHECK-NEXT:  (export "foo" (func $foo))
+;; CHECK-NEXT:  (func $foo
+;; CHECK-NEXT:   (nop)
+;; CHECK-NEXT:  )
+;; CHECK-NEXT: )
+(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
+    (nop)
+  )
+  (func $other
+    (drop (i32.const 1))
+  )
+)
+
+;; CHECK:      (module
+;; CHECK-NEXT:  (type $none (func))
+;; CHECK-NEXT:  (import "env" "other" (func $other))
+;; CHECK-NEXT:  (table $t 10 funcref)
+;; CHECK-NEXT:  (elem $0 (i32.const 0) $other)
+;; CHECK-NEXT:  (export "foo" (func $foo))
+;; CHECK-NEXT:  (func $foo
+;; CHECK-NEXT:   (call_indirect (type $none)
+;; CHECK-NEXT:    (i32.const 10)
+;; CHECK-NEXT:   )
+;; CHECK-NEXT:  )
+;; CHECK-NEXT: )
+(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))
+  )
+)