Simplify and fix heap type counting (#5110)

Annotations on array.get and array.set were not being counted and the code could
generally be simplified since `count` already ignores types that don't need to
be counted.

1 files changed, 12 insertions, 15 deletions

diff --git a/src/ir/module-utils.cpp b/src/ir/module-utils.cpp
index 84cd66c60..5044d303c 100644
--- a/src/ir/module-utils.cpp
+++ b/src/ir/module-utils.cpp
@@ -54,17 +54,15 @@ struct CodeScanner
     if (auto* call = curr->dynCast<CallIndirect>()) {
       counts.note(call->heapType);
     } else if (auto* call = curr->dynCast<CallRef>()) {
-      if (call->isReturn && call->target->type.isFunction()) {
-        counts.note(call->target->type);
-      }
+      counts.note(call->target->type);
     } else if (curr->is<RefNull>()) {
       counts.note(curr->type);
-    } else if (auto* make = curr->dynCast<StructNew>()) {
-      handleMake(make);
-    } else if (auto* make = curr->dynCast<ArrayNew>()) {
-      handleMake(make);
-    } else if (auto* make = curr->dynCast<ArrayInit>()) {
-      handleMake(make);
+    } else if (curr->is<StructNew>()) {
+      counts.note(curr->type);
+    } else if (curr->is<ArrayNew>()) {
+      counts.note(curr->type);
+    } else if (curr->is<ArrayInit>()) {
+      counts.note(curr->type);
     } else if (auto* cast = curr->dynCast<RefCast>()) {
       counts.note(cast->intendedType);
     } else if (auto* cast = curr->dynCast<RefTest>()) {
@@ -77,6 +75,11 @@ struct CodeScanner
       counts.note(get->ref->type);
     } else if (auto* set = curr->dynCast<StructSet>()) {
       counts.note(set->ref->type);
+    } else if (auto* get = curr->dynCast<ArrayGet>()) {
+      counts.note(get->ref->type);
+    } else if (auto* set = curr->dynCast<ArraySet>()) {
+      counts.note(set->ref->type);
+
     } else if (Properties::isControlFlowStructure(curr)) {
       if (curr->type.isTuple()) {
         // TODO: Allow control flow to have input types as well
@@ -86,12 +89,6 @@ struct CodeScanner
       }
     }
   }
-
-  template<typename T> void handleMake(T* curr) {
-    if (curr->type != Type::unreachable) {
-      counts.note(curr->type.getHeapType());
-    }
-  }
 };
 
 Counts getHeapTypeCounts(Module& wasm) {