[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

Author: kripken

CHANGELOG.md

@@ -14,6 +14,7 @@ full changeset diff at the end of each section.
 
 Current Trunk
 -------------
+- Add support for non-nullable locals in wasm GC.
 - Add a new flag to Directize, `--pass-arg=directize-initial-contents-immutable`
   which indicates the initial table contents are immutable. That is the case for
   LLVM, for example, and it allows us to optimize more indirect calls to direct

README.md

@@ -117,6 +117,25 @@ There are a few differences between Binaryen IR and the WebAssembly language:
     `(elem declare func $..)`. Binaryen will emit that data when necessary, but
     it does not represent it in IR. That is, IR can be worked on without needing
     to think about declaring function references.
+  * Binaryen IR allows non-nullable locals in the form that the wasm spec does,
+    (which was historically nicknamed "1a"), in which a `local.get` must be
+    structurally dominated by a `local.set` in order to validate (that ensures
+    we do not read the default value of null). Despite being aligned with the
+    wasm spec, there are some minor details that you may notice:
+    * A nameless `Block` in Binaryen IR does not interfere with validation.
+      Nameless blocks are never emitted into the binary format (we just emit
+      their contents), so we ignore them for purposes of non-nullable locals. As
+      a result, if you read wasm text emitted by Binaryen then you may see what
+      seems to be code that should not validate per the spec (and may not
+      validate in wasm text parsers), but that difference will not exist in the
+      binary format (binaries emitted by Binaryen will always work everywhere,
+      aside for bugs of course).
+    * The Binaryen pass runner will automatically fix up validation after each
+      pass (finding things that do not validate and fixing them up, usually by
+      demoting a local to be nullable). As a result you do not need to worry
+      much about this when writing Binaryen passes. For more details see the
+      `requiresNonNullableLocalFixups()` hook in `pass.h` and the
+      `LocalStructuralDominance` class.
 
 As a result, you might notice that round-trip conversions (wasm => Binaryen IR
 => wasm) change code a little in some corner cases.

scripts/test/shared.py

@@ -262,7 +262,8 @@ def has_shell_timeout():
     '--experimental-wasm-gc',
     '--experimental-wasm-typed-funcref',
     '--experimental-wasm-memory64',
-    '--experimental-wasm-extended-const'
+    '--experimental-wasm-extended-const',
+    '--experimental-wasm-nn-locals',
 ]
 
 # external tools

src/ir/CMakeLists.txt

@@ -12,6 +12,7 @@ set(ir_SOURCES
   possible-contents.cpp
   properties.cpp
   LocalGraph.cpp
+  LocalStructuralDominance.cpp
   ReFinalize.cpp
   stack-utils.cpp
   table-utils.cpp

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

src/ir/eh-utils.cpp

@@ -17,7 +17,6 @@
 #include "ir/eh-utils.h"
 #include "ir/branch-utils.h"
 #include "ir/find_all.h"
-#include "ir/type-updating.h"
 
 namespace wasm {
 
@@ -157,9 +156,6 @@ void handleBlockNestedPops(Function* func, Module& wasm) {
   for (auto* try_ : trys.list) {
     handleBlockNestedPop(try_, func, wasm);
   }
-  // Pops we handled can be of non-defaultable types, so we may have created
-  // non-nullable type locals. Fix them.
-  TypeUpdating::handleNonDefaultableLocals(func, wasm);
 }
 
 Pop* findPop(Expression* expr) {

src/ir/iteration.h

@@ -76,6 +76,7 @@ template<class Specific> class AbstractChildIterator {
 public:
   // The vector of children in the order emitted by wasm-delegations-fields
   // (which is in reverse execution order).
+  // TODO: rename this "reverseChildren"?
   SmallVector<Expression**, 4> children;
 
   AbstractChildIterator(Expression* parent) {

src/ir/linear-execution.h

@@ -49,7 +49,7 @@ struct LinearExecutionWalker : public PostWalker<SubType, VisitorType> {
 
     switch (curr->_id) {
       case Expression::Id::InvalidId:
-        abort();
+        WASM_UNREACHABLE("bad id");
       case Expression::Id::BlockId: {
         self->pushTask(SubType::doVisitBlock, currp);
         if (curr->cast<Block>()->name.is()) {