1 files changed, 23 insertions, 50 deletions

diff --git a/src/passes/RemoveUnusedModuleElements.cpp b/src/passes/RemoveUnusedModuleElements.cpp
index 6cd050da9..069137ff6 100644
--- a/src/passes/RemoveUnusedModuleElements.cpp
+++ b/src/passes/RemoveUnusedModuleElements.cpp
@@ -25,6 +25,7 @@
 
 #include "wasm.h"
 #include "pass.h"
+#include "ir/module-utils.h"
 #include "ir/utils.h"
 #include "asm_v_wasm.h"
 
@@ -63,15 +64,15 @@ struct ReachabilityAnalyzer : public PostWalker<ReachabilityAnalyzer> {
         reachable.insert(curr);
         if (curr.first == ModuleElementKind::Function) {
           // if not an import, walk it
-          auto* func = module->getFunctionOrNull(curr.second);
-          if (func) {
+          auto* func = module->getFunction(curr.second);
+          if (!func->imported()) {
             walk(func->body);
           }
         } else {
           // if not imported, it has an init expression we need to walk
-          auto* glob = module->getGlobalOrNull(curr.second);
-          if (glob) {
-            walk(glob->init);
+          auto* global = module->getGlobal(curr.second);
+          if (!global->imported()) {
+            walk(global->init);
           }
         }
       }
@@ -83,11 +84,6 @@ struct ReachabilityAnalyzer : public PostWalker<ReachabilityAnalyzer> {
       queue.emplace_back(ModuleElementKind::Function, curr->target);
     }
   }
-  void visitCallImport(CallImport* curr) {
-    if (reachable.count(ModuleElement(ModuleElementKind::Function, curr->target)) == 0) {
-      queue.emplace_back(ModuleElementKind::Function, curr->target);
-    }
-  }
   void visitCallIndirect(CallIndirect* curr) {
     usesTable = true;
   }
@@ -131,19 +127,17 @@ struct ReachabilityAnalyzer : public PostWalker<ReachabilityAnalyzer> {
 // Finds function type usage
 
 struct FunctionTypeAnalyzer : public PostWalker<FunctionTypeAnalyzer> {
-  std::vector<Import*> functionImports;
+  std::vector<Function*> functionImports;
   std::vector<Function*> functions;
   std::vector<CallIndirect*> indirectCalls;
 
-  void visitImport(Import* curr) {
-    if (curr->kind == ExternalKind::Function && curr->functionType.is()) {
-      functionImports.push_back(curr);
-    }
-  }
-
   void visitFunction(Function* curr) {
     if (curr->type.is()) {
-      functions.push_back(curr);
+      if (curr->imported()) {
+        functionImports.push_back(curr);
+      } else {
+        functions.push_back(curr);
+      }
     }
   }
 
@@ -176,9 +170,9 @@ struct RemoveUnusedModuleElements : public Pass {
     }
     // If told to, root all the functions
     if (rootAllFunctions) {
-      for (auto& func : module->functions) {
+      ModuleUtils::iterDefinedFunctions(*module, [&](Function* func) {
         roots.emplace_back(ModuleElementKind::Function, func->name);
-      }
+      });
     }
     // Exports are roots.
     bool exportsMemory = false;
@@ -194,15 +188,14 @@ struct RemoveUnusedModuleElements : public Pass {
         exportsTable = true;
       }
     }
-    // Check for special imports are roots.
+    // Check for special imports, which are roots.
     bool importsMemory = false;
     bool importsTable = false;
-    for (auto& curr : module->imports) {
-      if (curr->kind == ExternalKind::Memory) {
-        importsMemory = true;
-      } else if (curr->kind == ExternalKind::Table) {
-        importsTable = true;
-      }
+    if (module->memory.imported()) {
+      importsMemory = true;
+    }
+    if (module->table.imported()) {
+      importsTable = true;
     }
     // For now, all functions that can be called indirectly are marked as roots.
     for (auto& segment : module->table.segments) {
@@ -225,17 +218,6 @@ struct RemoveUnusedModuleElements : public Pass {
         return analyzer.reachable.count(ModuleElement(ModuleElementKind::Global, curr->name)) == 0;
       }), v.end());
     }
-    {
-      auto& v = module->imports;
-      v.erase(std::remove_if(v.begin(), v.end(), [&](const std::unique_ptr<Import>& curr) {
-        if (curr->kind == ExternalKind::Function) {
-          return analyzer.reachable.count(ModuleElement(ModuleElementKind::Function, curr->name)) == 0;
-        } else if (curr->kind == ExternalKind::Global) {
-          return analyzer.reachable.count(ModuleElement(ModuleElementKind::Global, curr->name)) == 0;
-        }
-        return false;
-      }), v.end());
-    }
     module->updateMaps();
     // Handle the memory and table
     if (!exportsMemory && !analyzer.usesMemory) {
@@ -245,10 +227,9 @@ struct RemoveUnusedModuleElements : public Pass {
       }
       if (module->memory.segments.empty()) {
         module->memory.exists = false;
-        module->memory.imported = false;
+        module->memory.module = module->memory.base = Name();
         module->memory.initial = 0;
         module->memory.max = 0;
-        removeImport(ExternalKind::Memory, module);
       }
     }
     if (!exportsTable && !analyzer.usesTable) {
@@ -258,21 +239,13 @@ struct RemoveUnusedModuleElements : public Pass {
       }
       if (module->table.segments.empty()) {
         module->table.exists = false;
-        module->table.imported = false;
+        module->table.module = module->table.base = Name();
         module->table.initial = 0;
         module->table.max = 0;
-        removeImport(ExternalKind::Table, module);
       }
     }
   }
 
-  void removeImport(ExternalKind kind, Module* module) {
-    auto& v = module->imports;
-    v.erase(std::remove_if(v.begin(), v.end(), [&](const std::unique_ptr<Import>& curr) {
-      return curr->kind == kind;
-    }), v.end());
-  }
-
   void optimizeFunctionTypes(Module* module) {
     FunctionTypeAnalyzer analyzer;
     analyzer.walkModule(module);
@@ -294,7 +267,7 @@ struct RemoveUnusedModuleElements : public Pass {
     };
     // canonicalize all uses of function types
     for (auto* import : analyzer.functionImports) {
-      import->functionType = canonicalize(import->functionType);
+      import->type = canonicalize(import->type);
     }
     for (auto* func : analyzer.functions) {
       func->type = canonicalize(func->type);