[Wasm GC] Support non-nullable locals in the "1a" form

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,231 @@
+/*
+ * 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_set.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()) {