diff --git a/src/ir/module-utils.cpp b/src/ir/module-utils.cpp index af8f3e9ff66..ed495c21267 100644 --- a/src/ir/module-utils.cpp +++ b/src/ir/module-utils.cpp @@ -1092,4 +1092,67 @@ IndexedHeapTypes getOptimizedIndexedHeapTypes(Module& wasm) { return indexedTypes; } +std::unordered_set getReferencedFuncs(Module& module) { + struct AddressedFuncsWalker : WalkerPass> { + // For each function, which functions are referenced in its body. + // The key for `nullptr` contains references that are not in a function + // (e.g. `elem` segments). + std::unordered_map>& + allReferencedFuncs; + // Points to `allReferencedFuncs`. + std::unordered_set* referencedFuncs = nullptr; + + AddressedFuncsWalker( + std::unordered_map>& + allReferencedFuncs) + : allReferencedFuncs(allReferencedFuncs), + referencedFuncs(&allReferencedFuncs.at(nullptr)) {} + + std::unique_ptr create() override { + return std::make_unique(allReferencedFuncs); + } + + bool isFunctionParallel() override { return true; } + + bool modifiesBinaryenIR() override { return false; } + + void doWalkFunction(Function* func) { + referencedFuncs = &allReferencedFuncs.at(func); + walk(func->body); + } + + void visitRefFunc(RefFunc* refFunc) { + referencedFuncs->insert(getModule()->getFunction(refFunc->func)); + } + }; + + std::unordered_map> + allReferencedFuncs; + allReferencedFuncs[nullptr]; + for (auto& func : module.functions) { + allReferencedFuncs[func.get()]; + } + + PassRunner runner(&module); + AddressedFuncsWalker walker(allReferencedFuncs); + walker.run(&runner, &module); + walker.runOnModuleCode(&runner, &module); + + std::unordered_set mergedReferencedFuncs; + for (auto& [_, referencedFuncs] : allReferencedFuncs) { + mergedReferencedFuncs.merge(referencedFuncs); + } + + for (const auto& export_ : module.exports) { + if (export_->kind != ExternalKind::Function) { + continue; + } + + mergedReferencedFuncs.insert( + module.getFunction(*export_->getInternalName())); + } + + return mergedReferencedFuncs; +} + } // namespace wasm::ModuleUtils diff --git a/src/ir/module-utils.h b/src/ir/module-utils.h index 70297c7cee2..ac5777fab81 100644 --- a/src/ir/module-utils.h +++ b/src/ir/module-utils.h @@ -505,6 +505,14 @@ struct IndexedHeapTypes { // frequency of use to minimize code size. IndexedHeapTypes getOptimizedIndexedHeapTypes(Module& wasm); +// Returns the set of functions that can potentially be the target of a call_ref +// or call_indirect. These are: +// * Functions that appear in a ref.func in the module. +// * Functions that are `elem declare`d (which is represented as ref.func in +// the IR). +// * Exported functions, since they can flow back to us as a reference. +std::unordered_set getReferencedFuncs(Module& module); + } // namespace wasm::ModuleUtils #endif // wasm_ir_module_h diff --git a/src/passes/GlobalEffects.cpp b/src/passes/GlobalEffects.cpp index e4e95a833f4..660b296b547 100644 --- a/src/passes/GlobalEffects.cpp +++ b/src/passes/GlobalEffects.cpp @@ -45,76 +45,6 @@ struct FuncInfo { std::unordered_set indirectCalledTypes; }; -// Only funcs that are referenced may be the target of an indirect call. A -// function is referenced if: -// - It appears in a ref.func expression (this includes `elem` statements -// because of how our IR is represented). -// - It's exported, because it may flow back to us as a reference. -// -// If a function doesn't meet any of these criteria, it can't be the target of -// an indirect call and we don't need to include its effects in indirect calls. -std::unordered_set getReferencedFuncs(Module& module, - PassRunner& passRunner) { - struct AddressedFuncsWalker : WalkerPass> { - // For each function, which functions are referenced in its body. - // The key for `nullptr` contains references that are not in a function - // (e.g. `elem` segments). - std::unordered_map>& - allReferencedFuncs; - // Points to `allReferencedFuncs`. - std::unordered_set* referencedFuncs = nullptr; - - AddressedFuncsWalker( - std::unordered_map>& - allReferencedFuncs) - : allReferencedFuncs(allReferencedFuncs), - referencedFuncs(&allReferencedFuncs[nullptr]) {} - - std::unique_ptr create() override { - return std::make_unique(allReferencedFuncs); - } - - bool isFunctionParallel() override { return true; } - - bool modifiesBinaryenIR() override { return false; } - - void doWalkFunction(Function* func) { - referencedFuncs = &allReferencedFuncs.at(func); - walk(func->body); - } - - void visitRefFunc(RefFunc* refFunc) { - referencedFuncs->insert(getModule()->getFunction(refFunc->func)); - } - }; - - std::unordered_map> - allReferencedFuncs; - for (auto& func : module.functions) { - allReferencedFuncs[func.get()]; - } - - AddressedFuncsWalker walker(allReferencedFuncs); - walker.run(&passRunner, &module); - walker.runOnModuleCode(&passRunner, &module); - - std::unordered_set mergedReferencedFuncs; - for (auto& [_, referencedFuncs] : allReferencedFuncs) { - mergedReferencedFuncs.merge(referencedFuncs); - } - - for (const auto& export_ : module.exports) { - if (export_->kind != ExternalKind::Function) { - continue; - } - - mergedReferencedFuncs.insert( - module.getFunction(*export_->getInternalName())); - } - - return mergedReferencedFuncs; -} - std::map analyzeFuncs(Module& module, const PassOptions& passOptions) { ModuleUtils::ParallelFunctionAnalysis analysis( @@ -456,7 +386,7 @@ struct GenerateGlobalEffects : public Pass { std::map funcInfos = analyzeFuncs(*module, getPassOptions()); - auto referencedFuncs = getReferencedFuncs(*module, *getPassRunner()); + auto referencedFuncs = ModuleUtils::getReferencedFuncs(*module); auto callGraph = buildCallGraph( *module, funcInfos, referencedFuncs, getPassOptions().worldMode); diff --git a/test/gtest/CMakeLists.txt b/test/gtest/CMakeLists.txt index 19052a46147..0ecf4bec7ff 100644 --- a/test/gtest/CMakeLists.txt +++ b/test/gtest/CMakeLists.txt @@ -29,6 +29,7 @@ set(unittest_SOURCES lattices.cpp leaves.cpp local-graph.cpp + module-utils.cpp possible-contents.cpp principal-type.cpp printing.cpp diff --git a/test/gtest/module-utils.cpp b/test/gtest/module-utils.cpp new file mode 100644 index 00000000000..6f546754c82 --- /dev/null +++ b/test/gtest/module-utils.cpp @@ -0,0 +1,76 @@ +/* + * Copyright 2026 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. + */ + +#include "gmock/gmock.h" +#include "gtest/gtest.h" + +#include "ir/module-utils.h" +#include "parser/wat-parser.h" +#include "wasm.h" + +using namespace wasm; +using namespace testing; + +namespace { + +class ModuleUtilsTest : public Test { +protected: + Module wasm; + + [[nodiscard]] bool parse(std::string_view text) { + wasm.features = FeatureSet::All; + auto parsed = WATParser::parseModule(wasm, text); + if (auto* err = parsed.getErr()) { + EXPECT_TRUE(false) << "Failed to parse WAT: " << err->msg; + return false; + } + return true; + } +}; + +TEST_F(ModuleUtilsTest, GetReferencedFuncs) { + auto moduleText = R"wasm( + (module + (import "env" "imported" (func $imported)) + + (table $table 1 funcref) + (elem (table $table) (i32.const 0) func $table_func) + + (func $table_func) + (func $ref_func) + (func $unreferenced) + + (func $exported + (drop (ref.func $ref_func)) + ) + (export "exported" (func $exported)) + ) + )wasm"; + + ASSERT_TRUE(parse(moduleText)); + + auto* tableFunc = wasm.getFunction("table_func"); + auto* refFunc = wasm.getFunction("ref_func"); + auto* exported = wasm.getFunction("exported"); + ASSERT_NE(tableFunc, nullptr); + ASSERT_NE(refFunc, nullptr); + ASSERT_NE(exported, nullptr); + + EXPECT_THAT(ModuleUtils::getReferencedFuncs(wasm), + UnorderedElementsAre(tableFunc, refFunc, exported)); +} + +} // namespace