1 files changed, 5 insertions, 4 deletions
diff --git a/src/passes/opt-utils.h b/src/passes/opt-utils.h
index b333779f7..5f4ab545c 100644
--- a/src/passes/opt-utils.h
+++ b/src/passes/opt-utils.h
@@ -21,6 +21,7 @@
 #include <unordered_set>
 
 #include <ir/element-utils.h>
+#include <ir/module-utils.h>
 #include <pass.h>
 #include <wasm.h>
 
@@ -84,10 +85,10 @@ inline void replaceFunctions(PassRunner* runner,
       name = iter->second;
     }
   };
-  // replace direct calls
-  FunctionRefReplacer(maybeReplace).run(runner, &module);
-  // replace in table
-  ElementUtils::iterAllElementFunctionNames(&module, maybeReplace);
+  // replace direct calls in code both functions and module elements
+  FunctionRefReplacer replacer(maybeReplace);
+  replacer.run(runner, &module);
+  replacer.walkModuleCode(&module);
 
   // replace in start
   if (module.start.is()) {