[Wasm GC] Global Refining pass (#4344)

Fairly simple, this uses the existing infrastructure to find opportunities
to refine the type of a global variable. This a common pattern in j2wasm
for example, where a global begins as a null of $java.lang.Object (the
least specific type) but it is in practice always assigned an object of
some specific type.

6 files changed, 265 insertions, 0 deletions

diff --git a/src/passes/CMakeLists.txt b/src/passes/CMakeLists.txt
index 25c9cb6b3..3313637ba 100644
--- a/src/passes/CMakeLists.txt
+++ b/src/passes/CMakeLists.txt
@@ -30,6 +30,7 @@ set(passes_SOURCES
   Flatten.cpp
   FuncCastEmulation.cpp
   GenerateDynCalls.cpp
+  GlobalRefining.cpp
   GlobalTypeOptimization.cpp
   Heap2Local.cpp
   I64ToI32Lowering.cpp
diff --git a/src/passes/GlobalRefining.cpp b/src/passes/GlobalRefining.cpp
new file mode 100644
index 000000000..98848566b
--- /dev/null
+++ b/src/passes/GlobalRefining.cpp
@@ -0,0 +1,139 @@
+/*
+ * Copyright 2021 WebAssembly Community Group participants
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+//
+// Apply more specific subtypes to global variables where possible.
+//
+
+#include "ir/find_all.h"
+#include "ir/lubs.h"
+#include "ir/module-utils.h"
+#include "ir/utils.h"
+#include "pass.h"
+#include "wasm-type.h"
+#include "wasm.h"
+
+using namespace std;
+
+namespace wasm {
+
+namespace {
+
+struct GlobalRefining : public Pass {
+  void run(PassRunner* runner, Module* module) override {
+    if (getTypeSystem() != TypeSystem::Nominal) {
+      Fatal() << "GlobalRefining requires nominal typing";
+    }
+
+    // First, find all the global.sets.
+
+    struct GlobalInfo {
+      std::vector<GlobalSet*> sets;
+    };
+
+    ModuleUtils::ParallelFunctionAnalysis<GlobalInfo> analysis(
+      *module, [&](Function* func, GlobalInfo& info) {
+        if (func->imported()) {
+          return;
+        }
+        info.sets = std::move(FindAll<GlobalSet>(func->body).list);
+      });
+
+    // A map of globals to the lub for that global.
+    std::unordered_map<Name, LUBFinder> lubs;
+
+    // Combine all the information we gathered and compute lubs.
+    for (auto& [func, info] : analysis.map) {
+      for (auto* set : info.sets) {
+        lubs[set->name].noteUpdatableExpression(set->value);
+      }
+    }
+
+    bool optimized = false;
+
+    for (auto& global : module->globals) {
+      if (global->imported()) {
+        continue;
+      }
+
+      auto& lub = lubs[global->name];
+
+      // Note the initial value.
+      lub.noteUpdatableExpression(global->init);
+
+      // The initial value cannot be unreachable, but it might be null, and all
+      // other values might be too. In that case, we've noted nothing useful
+      // and we can move on.
+      if (!lub.noted()) {
+        continue;
+      }
+
+      auto oldType = global->type;
+      auto newType = lub.getBestPossible();
+      if (newType != oldType) {
+        // We found an improvement!
+        assert(Type::isSubType(newType, oldType));
+        global->type = newType;
+        lub.updateNulls();
+        optimized = true;
+      }
+    }
+
+    if (!optimized) {
+      return;
+    }
+
+    // Update function contents for their new parameter types: global.gets must
+    // now return the new type for any globals that we modified.
+    struct GetUpdater : public WalkerPass<PostWalker<GetUpdater>> {
+      bool isFunctionParallel() override { return true; }
+
+      GlobalRefining& parent;
+      Module& wasm;
+
+      GetUpdater(GlobalRefining& parent, Module& wasm)
+        : parent(parent), wasm(wasm) {}
+
+      GetUpdater* create() override { return new GetUpdater(parent, wasm); }
+
+      // If we modify anything in a function then we must refinalize so that
+      // types propagate outwards.
+      bool modified = false;
+
+      void visitGlobalGet(GlobalGet* curr) {
+        auto oldType = curr->type;
+        auto newType = wasm.getGlobal(curr->name)->type;
+        if (newType != oldType) {
+          curr->type = newType;
+          modified = true;
+        }
+      }
+
+      void visitFunction(Function* curr) {
+        if (modified) {
+          ReFinalize().walkFunctionInModule(curr, &wasm);
+        }
+      }
+    };
+    GetUpdater(*this, *module).run(runner, module);
+  }
+};
+
+} // anonymous namespace
+
+Pass* createGlobalRefiningPass() { return new GlobalRefining(); }
+
+} // namespace wasm
diff --git a/src/passes/pass.cpp b/src/passes/pass.cpp
index 49615d81b..d20a5edcd 100644
--- a/src/passes/pass.cpp
+++ b/src/passes/pass.cpp
@@ -155,6 +155,8 @@ void PassRegistry::registerPasses() {
   registerPass(
     "generate-stack-ir", "generate Stack IR", createGenerateStackIRPass);
   registerPass(
+    "global-refining", "refine the types of globals", createGlobalRefiningPass);
+  registerPass(
     "gto", "globally optimize GC types", createGlobalTypeOptimizationPass);
   registerPass("type-refining",
                "apply more specific subtypes to type fields where possible",
@@ -528,6 +530,7 @@ void PassRunner::addDefaultGlobalOptimizationPrePasses() {
   if (wasm->features.hasGC() && getTypeSystem() == TypeSystem::Nominal &&
       options.optimizeLevel >= 2) {
     addIfNoDWARFIssues("type-refining");
+    addIfNoDWARFIssues("global-refining");
     // Global type optimization can remove fields that are not needed, which can
     // remove ref.funcs that were once assigned to vtables but are no longer
     // needed, which can allow more code to be removed globally. After those,
diff --git a/src/passes/passes.h b/src/passes/passes.h
index e8e8ecc78..5a67b1d82 100644
--- a/src/passes/passes.h
+++ b/src/passes/passes.h
@@ -51,6 +51,7 @@ Pass* createFunctionMetricsPass();
 Pass* createGenerateDynCallsPass();
 Pass* createGenerateI64DynCallsPass();
 Pass* createGenerateStackIRPass();
+Pass* createGlobalRefiningPass();
 Pass* createGlobalTypeOptimizationPass();
 Pass* createHeap2LocalPass();
 Pass* createI64ToI32LoweringPass();
diff --git a/test/lit/help/optimization-opts.test b/test/lit/help/optimization-opts.test
index af23314de..6667c45bf 100644
--- a/test/lit/help/optimization-opts.test
+++ b/test/lit/help/optimization-opts.test
@@ -275,6 +275,8 @@
 ;; CHECK-NEXT:
 ;; CHECK-NEXT:   --generate-stack-ir                           generate Stack IR
 ;; CHECK-NEXT:
+;; CHECK-NEXT:   --global-refining                             refine the types of globals
+;; CHECK-NEXT:
 ;; CHECK-NEXT:   --gto                                         globally optimize GC types
 ;; CHECK-NEXT:
 ;; CHECK-NEXT:   --heap2local                                  replace GC allocations with
diff --git a/test/lit/passes/global-refining.wast b/test/lit/passes/global-refining.wast
new file mode 100644
index 000000000..633d88e5f
--- /dev/null
+++ b/test/lit/passes/global-refining.wast
@@ -0,0 +1,119 @@
+;; NOTE: Assertions have been generated by update_lit_checks.py --all-items and should not be edited.
+;; RUN: foreach %s %t wasm-opt --nominal --global-refining -all -S -o - | filecheck %s
+
+(module
+  ;; Globals with no assignments aside from their initial values. The first is
+  ;; a null, so we have nothing concrete to improve with (though we could use
+  ;; the type of the null perhaps, TODO). The second is a ref.func which lets
+  ;; us refine.
+  ;; CHECK:      (type $none_=>_none (func_subtype func))
+
+  ;; CHECK:      (global $func-null-init (mut anyref) (ref.null func))
+  (global $func-null-init (mut anyref) (ref.null func))
+  ;; CHECK:      (global $func-func-init (mut (ref $none_=>_none)) (ref.func $foo))
+  (global $func-func-init (mut anyref) (ref.func $foo))
+  ;; CHECK:      (func $foo (type $none_=>_none)
+  ;; CHECK-NEXT:  (nop)
+  ;; CHECK-NEXT: )
+  (func $foo)
+)
+
+(module
+  ;; Globals with later assignments of null. The global with a function in its
+  ;; init will update the null to allow it to refine.
+
+  ;; CHECK:      (type $none_=>_none (func_subtype func))
+
+  ;; CHECK:      (global $func-null-init (mut anyref) (ref.null func))
+  (global $func-null-init (mut anyref) (ref.null func))
+  ;; CHECK:      (global $func-func-init (mut (ref null $none_=>_none)) (ref.func $foo))
+  (global $func-func-init (mut anyref) (ref.func $foo))
+
+  ;; CHECK:      (func $foo (type $none_=>_none)
+  ;; CHECK-NEXT:  (global.set $func-null-init
+  ;; CHECK-NEXT:   (ref.null any)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (global.set $func-func-init
+  ;; CHECK-NEXT:   (ref.null $none_=>_none)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $foo
+   (global.set $func-null-init (ref.null any))
+   (global.set $func-func-init (ref.null any))
+  )
+)
+
+(module
+  ;; Globals with later assignments of something non-null. Both can be refined,
+  ;; and the one with a non-null initial value can even become non-nullable.
+
+  ;; CHECK:      (type $none_=>_none (func_subtype func))
+
+  ;; CHECK:      (global $func-null-init (mut (ref null $none_=>_none)) (ref.null $none_=>_none))
+  (global $func-null-init (mut anyref) (ref.null func))
+  ;; CHECK:      (global $func-func-init (mut (ref $none_=>_none)) (ref.func $foo))
+  (global $func-func-init (mut anyref) (ref.func $foo))
+
+  ;; CHECK:      (elem declare func $foo)
+
+  ;; CHECK:      (func $foo (type $none_=>_none)
+  ;; CHECK-NEXT:  (global.set $func-null-init
+  ;; CHECK-NEXT:   (ref.func $foo)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (global.set $func-func-init
+  ;; CHECK-NEXT:   (ref.func $foo)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $foo
+   (global.set $func-null-init (ref.func $foo))
+   (global.set $func-func-init (ref.func $foo))
+  )
+)
+
+(module
+  ;; A global with multiple later assignments. The refined type is more
+  ;; specific than the original, but less than each of the non-null values.
+
+  ;; CHECK:      (type $none_=>_none (func_subtype func))
+
+  ;; CHECK:      (type $i32_=>_none (func_subtype (param i32) func))
+
+  ;; CHECK:      (global $global (mut funcref) (ref.null func))
+  (global $global (mut anyref) (ref.null any))
+
+  ;; CHECK:      (elem declare func $bar $foo)
+
+  ;; CHECK:      (func $foo (type $none_=>_none)
+  ;; CHECK-NEXT:  (global.set $global
+  ;; CHECK-NEXT:   (ref.func $foo)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (global.set $global
+  ;; CHECK-NEXT:   (ref.func $bar)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (global.set $global
+  ;; CHECK-NEXT:   (ref.null func)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (global.set $global
+  ;; CHECK-NEXT:   (ref.null func)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (global.set $global
+  ;; CHECK-NEXT:   (ref.null func)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $foo
+   (global.set $global (ref.func $foo))
+   (global.set $global (ref.func $bar))
+   (global.set $global (ref.null func))
+   ;; These nulls will be updated.
+   (global.set $global (ref.null eq))
+   (global.set $global (ref.null data))
+  )
+
+  ;; CHECK:      (func $bar (type $i32_=>_none) (param $x i32)
+  ;; CHECK-NEXT:  (nop)
+  ;; CHECK-NEXT: )
+  (func $bar (param $x i32)
+    ;; A function with a different signature, whose reference is also assigned
+    ;; to the global.
+  )
+)