[ReorderFunctions] Count ref.func uses

ReorderFunctions ordered functions using calls, plus exports, start and
element segments. It did not count ref.func, so a function that is only
used through a reference could be sorted as if it were unused. Count
ref.func uses in function bodies and in module code (global initializers,
element segments, table initializers).

The size effect is small, but the counts now match how functions are
actually used.

Author: artemkulyk

src/passes/ReorderFunctions.cpp

@@ -29,7 +29,6 @@
 
 #include <memory>
 
-#include <ir/element-utils.h>
 #include <pass.h>
 #include <wasm.h>
 
@@ -54,6 +53,12 @@ struct CallCountScanner : public WalkerPass<PostWalker<CallCountScanner>> {
     (*counts)[curr->target]++;
   }
 
+  void visitRefFunc(RefFunc* curr) {
+    // can't add a new element in parallel
+    assert(counts->count(curr->func) > 0);
+    (*counts)[curr->func]++;
+  }
+
 private:
   NameCountMap* counts;
 };
@@ -69,8 +74,10 @@ struct ReorderFunctions : public Pass {
     for (auto& func : module->functions) {
       counts.try_emplace(func->name, 0);
     }
-    // find counts on function calls
+    // find counts on function calls and ref.funcs in function bodies
     CallCountScanner(&counts).run(getPassRunner(), module);
+    // find counts on ref.funcs in module code: globals, element segments, etc.
+    CallCountScanner(&counts).walkModuleCode(module);
     // find counts on global usages
     if (module->start.is()) {
       counts[module->start]++;
@@ -80,9 +87,6 @@ struct ReorderFunctions : public Pass {
         counts[*curr->getInternalName()]++;
       }
     }
-    ElementUtils::iterAllElementFunctionNames(
-      module, [&](Name name) { counts[name]++; });
-    // TODO: count all RefFunc as well
     // TODO: count the declaration section as well, which adds another mention
     // sort
     std::sort(module->functions.begin(),

test/lit/passes/reorder-functions.wast

@@ -0,0 +1,59 @@
+;; NOTE: Assertions have been generated by update_lit_checks.py --all-items and should not be edited.
+;; RUN: foreach %s %t wasm-opt --all-features --reorder-functions -S -o - | filecheck %s
+
+;; Check that ref.func references count towards a function's usage, both in
+;; function bodies and in module code (global initializers and element
+;; segments). The expected order is $c, then $a, then $b.
+
+(module
+  ;; CHECK:      (type $0 (func (result i32)))
+
+  ;; CHECK:      (global $g funcref (ref.func $a))
+
+  ;; CHECK:      (elem declare func $a $c)
+
+  ;; CHECK:      (func $c (type $0) (result i32)
+  ;; CHECK-NEXT:  (i32.const 3)
+  ;; CHECK-NEXT: )
+
+  ;; CHECK:      (func $a (type $0) (result i32)
+  ;; CHECK-NEXT:  (i32.const 1)
+  ;; CHECK-NEXT: )
+  (func $a (result i32)
+    (i32.const 1)
+  )
+  ;; CHECK:      (func $b (type $0) (result i32)
+  ;; CHECK-NEXT:  (i32.const 2)
+  ;; CHECK-NEXT: )
+  (func $b (result i32)
+    (i32.const 2)
+  )
+  (func $c (result i32)
+    (i32.const 3)
+  )
+  (global $g funcref
+    (ref.func $a)
+  )
+  ;; CHECK:      (func $use (type $0) (result i32)
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (ref.func $a)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (ref.func $c)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (ref.func $c)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (call $c)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (call $b)
+  ;; CHECK-NEXT: )
+  (func $use (result i32)
+    (drop (ref.func $a))
+    (drop (ref.func $c))
+    (drop (ref.func $c))
+    (drop (call $c))
+    (call $b)
+  )
+)