[ReorderFunctions] Count ref.func uses in bodies and globals

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 in function bodies, and in global initializers, as well.

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

Author: artemkulyk

CHANGELOG.md

@@ -22,6 +22,8 @@ Current Trunk
 - Reject non-natural alignment for atomic memory operations at parse time (#8962)
 - Note that fast-math mode can ignore the difference between negative zero and
   zero (like clang and gcc). (#9056)
+- reorder-functions now counts ref.func uses in function bodies and global
+  initializers. (#9160)
 
 v132
 ----

src/passes/ReorderFunctions.cpp

@@ -30,6 +30,7 @@
 #include <memory>
 
 #include <ir/element-utils.h>
+#include <ir/find_all.h>
 #include <pass.h>
 #include <wasm.h>
 
@@ -54,6 +55,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;
 };
@@ -82,7 +89,16 @@ struct ReorderFunctions : public Pass {
     }
     ElementUtils::iterAllElementFunctionNames(
       module, [&](Name name) { counts[name]++; });
-    // TODO: count all RefFunc as well
+    // find counts on ref.funcs in globals (those in element segments are
+    // counted above, and those in functions by the scanner)
+    for (auto& global : module->globals) {
+      if (!global->init) {
+        continue;
+      }
+      for (auto* refFunc : FindAll<RefFunc>(global->init).list) {
+        counts[refFunc->func]++;
+      }
+    }
     // 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,60 @@
+;; 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 global initializers. The expected order is $c (two
+;; ref.func uses and a call), $a (one ref.func use in a function and one in a
+;; global), then $b (a single call).
+
+(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)
+  )
+)