[Wasm GC] Support non-nullable locals in the "1a" form (#4959)

An overview of this is in the README in the diff here (conveniently, it is near the
top of the diff). Basically, we fix up nn locals after each pass, by default. This keeps
things easy to reason about - what validates is what is valid wasm - but there are
some minor nuances as mentioned there, in particular, we ignore nameless blocks
(which are commonly added by various passes; ignoring them means we can keep
more locals non-nullable).

The key addition here is LocalStructuralDominance which checks which local
indexes have the "structural dominance" property of 1a, that is, that each get has
a set in its block or an outer block that precedes it. I optimized that function quite
a lot to reduce the overhead of running that logic after each pass. The overhead
is something like 2% on J2Wasm and 0% on Dart (0%, because in this mode we
shrink code size, so there is less work actually, and it balances out).

Since we run fixups after each pass, this PR removes logic to manually call the
fixup code from various places we used to call it (like eh-utils and various passes).

Various passes are now marked as requiresNonNullableLocalFixups => false.
That lets us skip running the fixups after them, which we normally do automatically.
This helps avoid overhead. Most passes still need the fixups, though - any pass
that adds a local, or a named block, or moves code around, likely does.

This removes a hack in SimplifyLocals that is no longer needed. Before we
worked to avoid moving a set into a try, as it might not validate. Now, we just do it
and let fixups happen automatically if they need to: in the common code they
probably don't, so the extra complexity seems not worth it.

Also removes a hack from StackIR. That hack tried to avoid roundtrip adding a
nondefaultable local. But we have the logic to fix that up now, and opts will
likely keep it non-nullable as well.

Various tests end up updated here because now a local can be non-nullable -
previous fixups are no longer needed.

Note that this doesn't remove the gc-nn-locals feature. That has been useful for
testing, and may still be useful in the future - it basically just allows nn locals in
all positions (that can't read the null default value at the entry). We can consider
removing it separately.

Fixes #4824

1 files changed, 230 insertions, 0 deletions

diff --git a/src/ir/LocalStructuralDominance.cpp b/src/ir/LocalStructuralDominance.cpp
new file mode 100644
index 000000000..cfb75e006
--- /dev/null
+++ b/src/ir/LocalStructuralDominance.cpp
@@ -0,0 +1,230 @@
+/*
+ * Copyright 2022 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.
+ */
+
+#include "ir/iteration.h"
+#include "ir/local-structural-dominance.h"
+#include "support/small_vector.h"
+
+namespace wasm {
+
+LocalStructuralDominance::LocalStructuralDominance(Function* func,
+                                                   Module& wasm,
+                                                   Mode mode) {
+  if (!wasm.features.hasReferenceTypes()) {
+    // No references, so nothing to look at.
+    return;
+  }
+
+  bool hasRefVar = false;
+  for (auto var : func->vars) {
+    if (var.isRef()) {
+      hasRefVar = true;
+      break;
+    }
+  }
+  if (!hasRefVar) {
+    return;
+  }
+
+  if (mode == NonNullableOnly) {
+    bool hasNonNullableVar = false;
+    for (auto var : func->vars) {
+      // Check if we have any non-nullable vars at all.
+      if (var.isNonNullable()) {
+        hasNonNullableVar = true;
+        break;
+      }
+    }
+    if (!hasNonNullableVar) {
+      return;
+    }
+  }
+
+  struct Scanner : public PostWalker<Scanner> {
+    std::set<Index>& nonDominatingIndices;
+
+    // The locals that have been set, and so at the current time, they
+    // structurally dominate.
+    std::vector<bool> localsSet;
+
+    Scanner(Function* func, Mode mode, std::set<Index>& nonDominatingIndices)
+      : nonDominatingIndices(nonDominatingIndices) {
+      localsSet.resize(func->getNumLocals());
+
+      // Parameters always dominate.
+      for (Index i = 0; i < func->getNumParams(); i++) {
+        localsSet[i] = true;
+      }
+
+      for (Index i = func->getNumParams(); i < func->getNumLocals(); i++) {
+        auto type = func->getLocalType(i);
+        // Mark locals we don't need to care about as "set". We never do any
+        // work for such a local.
+        if (!type.isRef() || (mode == NonNullableOnly && type.isNullable())) {
+          localsSet[i] = true;
+        }
+      }
+
+      // Note that we do not need to start a scope for the function body.
+      // Logically there is a scope there, but there is no code after it, so
+      // there is nothing to clean up when that scope exits, so we may as well
+      // not even create a scope. Just start walking the body now.
+      walk(func->body);
+    }
+
+    using Locals = SmallVector<Index, 5>;
+
+    // When we exit a control flow scope, we must undo the locals that it set.
+    std::vector<Locals> cleanupStack;
+
+    static void doBeginScope(Scanner* self, Expression** currp) {
+      self->cleanupStack.emplace_back();
+    }
+
+    static void doEndScope(Scanner* self, Expression** currp) {
+      for (auto index : self->cleanupStack.back()) {
+        assert(self->localsSet[index]);
+        self->localsSet[index] = false;
+      }
+      self->cleanupStack.pop_back();
+    }
+
+    static void doLocalSet(Scanner* self, Expression** currp) {
+      auto index = (*currp)->cast<LocalSet>()->index;
+      if (!self->localsSet[index]) {
+        // This local is now set until the end of this scope.
+        self->localsSet[index] = true;
+        // If we are not in the topmost scope, note this for later cleanup.
+        if (!self->cleanupStack.empty()) {
+          self->cleanupStack.back().push_back(index);
+        }
+      }
+    }
+
+    static void scan(Scanner* self, Expression** currp) {
+      // Use a loop to avoid recursing on the last child - we can just go
+      // straight into a loop iteration for it.
+      while (1) {
+        Expression* curr = *currp;
+
+        switch (curr->_id) {
+          case Expression::Id::InvalidId:
+            WASM_UNREACHABLE("bad id");
+
+          // local.get can just be visited immediately, as it has no children.
+          case Expression::Id::LocalGetId: {
+            auto index = curr->cast<LocalGet>()->index;
+            if (!self->localsSet[index]) {
+              self->nonDominatingIndices.insert(index);
+            }
+            return;
+          }
+          case Expression::Id::LocalSetId: {
+            auto* set = curr->cast<LocalSet>();
+            if (!self->localsSet[set->index]) {
+              self->pushTask(doLocalSet, currp);
+            }
+            // Immediately continue in the loop.
+            currp = &set->value;
+            continue;
+          }
+
+          // Control flow structures.
+          case Expression::Id::BlockId: {
+            auto* block = curr->cast<Block>();
+            // Blocks with no name are never emitted in the binary format, so do
+            // not create a scope for them.
+            if (block->name.is()) {
+              self->pushTask(Scanner::doEndScope, currp);
+            }
+            auto& list = block->list;
+            for (int i = int(list.size()) - 1; i >= 0; i--) {
+              self->pushTask(Scanner::scan, &list[i]);
+            }
+            if (block->name.is()) {
+              // Just call the task immediately.
+              doBeginScope(self, currp);
+            }
+            return;
+          }
+          case Expression::Id::IfId: {
+            if (curr->cast<If>()->ifFalse) {
+              self->pushTask(Scanner::doEndScope, currp);
+              self->maybePushTask(Scanner::scan, &curr->cast<If>()->ifFalse);
+              self->pushTask(Scanner::doBeginScope, currp);
+            }
+            self->pushTask(Scanner::doEndScope, currp);
+            self->pushTask(Scanner::scan, &curr->cast<If>()->ifTrue);
+            self->pushTask(Scanner::doBeginScope, currp);
+            // Immediately continue in the loop.
+            currp = &curr->cast<If>()->condition;
+            continue;
+          }
+          case Expression::Id::LoopId: {
+            self->pushTask(Scanner::doEndScope, currp);
+            // Just call the task immediately.
+            doBeginScope(self, currp);
+            // Immediately continue in the loop.
+            currp = &curr->cast<Loop>()->body;
+            continue;
+          }
+          case Expression::Id::TryId: {
+            auto& list = curr->cast<Try>()->catchBodies;
+            for (int i = int(list.size()) - 1; i >= 0; i--) {
+              self->pushTask(Scanner::doEndScope, currp);
+              self->pushTask(Scanner::scan, &list[i]);
+              self->pushTask(Scanner::doBeginScope, currp);
+            }
+            self->pushTask(Scanner::doEndScope, currp);
+            // Just call the task immediately.
+            doBeginScope(self, currp);
+            // Immediately continue in the loop.
+            currp = &curr->cast<Try>()->body;
+            continue;
+          }
+
+          default: {
+            // Control flow structures have been handled. This is an expression,
+            // which we scan normally.
+            assert(!Properties::isControlFlowStructure(curr));
+            PostWalker<Scanner>::scan(self, currp);
+            return;
+          }
+        }
+      }
+    }
+
+    // Only local.set needs to be visited.
+    void pushTask(TaskFunc func, Expression** currp) {
+      // Visits to anything but a set can be ignored, so only very specific
+      // tasks need to actually be pushed here. In particular, we don't want to
+      // push tasks to call doVisit* when those callbacks do nothing.
+      if (func == scan || func == doLocalSet || func == doBeginScope ||
+          func == doEndScope) {
+        PostWalker<Scanner>::pushTask(func, currp);
+      }
+    }
+    void maybePushTask(TaskFunc func, Expression** currp) {
+      if (*currp) {
+        pushTask(func, currp);
+      }
+    }
+  };
+
+  Scanner(func, mode, nonDominatingIndices);
+}
+
+} // namespace wasm