[wasm-metadce] Note ref.func connections + fix rooting of segment offsets (#5791)

1 files changed, 28 insertions, 13 deletions

diff --git a/src/tools/wasm-metadce.cpp b/src/tools/wasm-metadce.cpp
index 029d4aae8..e5e7d8706 100644
--- a/src/tools/wasm-metadce.cpp
+++ b/src/tools/wasm-metadce.cpp
@@ -195,6 +195,17 @@ struct MetaDCEGraph {
 
       void visitGlobalGet(GlobalGet* curr) { handleGlobal(curr->name); }
       void visitGlobalSet(GlobalSet* curr) { handleGlobal(curr->name); }
+      void visitRefFunc(RefFunc* curr) {
+        Name dceName;
+        if (!getModule()->getFunction(curr->func)->imported()) {
+          dceName = parent->functionToDCENode[curr->func];
+        } else {
+          dceName =
+            parent->importIdToDCENode[parent->getFunctionImportId(curr->func)];
+        }
+        assert(!parentDceName.isNull());
+        parent->nodes[parentDceName].reaches.push_back(dceName);
+      }
 
     private:
       MetaDCEGraph* parent;
@@ -209,7 +220,9 @@ struct MetaDCEGraph {
           // it's an import.
           dceName = parent->importIdToDCENode[parent->getGlobalImportId(name)];
         }
-        if (!parentDceName.isNull()) {
+        if (parentDceName.isNull()) {
+          parent->roots.insert(dceName);
+        } else {
           parent->nodes[parentDceName].reaches.push_back(dceName);
         }
       }
@@ -248,18 +261,8 @@ struct MetaDCEGraph {
         return std::make_unique<Scanner>(parent);
       }
 
-      void visitCall(Call* curr) {
-        if (!getModule()->getFunction(curr->target)->imported()) {
-          parent->nodes[parent->functionToDCENode[getFunction()->name]]
-            .reaches.push_back(parent->functionToDCENode[curr->target]);
-        } else {
-          assert(parent->functionToDCENode.count(getFunction()->name) > 0);
-          parent->nodes[parent->functionToDCENode[getFunction()->name]]
-            .reaches.push_back(
-              parent
-                ->importIdToDCENode[parent->getFunctionImportId(curr->target)]);
-        }
-      }
+      void visitCall(Call* curr) { handleFunction(curr->target); }
+      void visitRefFunc(RefFunc* curr) { handleFunction(curr->func); }
       void visitGlobalGet(GlobalGet* curr) { handleGlobal(curr->name); }
       void visitGlobalSet(GlobalSet* curr) { handleGlobal(curr->name); }
       void visitThrow(Throw* curr) { handleTag(curr->tag); }
@@ -272,6 +275,18 @@ struct MetaDCEGraph {
     private:
       MetaDCEGraph* parent;
 
+      void handleFunction(Name name) {
+        if (!getModule()->getFunction(name)->imported()) {
+          parent->nodes[parent->functionToDCENode[getFunction()->name]]
+            .reaches.push_back(parent->functionToDCENode[name]);
+        } else {
+          assert(parent->functionToDCENode.count(getFunction()->name) > 0);
+          parent->nodes[parent->functionToDCENode[getFunction()->name]]
+            .reaches.push_back(
+              parent->importIdToDCENode[parent->getFunctionImportId(name)]);
+        }
+      }
+
       void handleGlobal(Name name) {
         if (!getFunction()) {
           return; // non-function stuff (initializers) are handled separately