1 files changed, 30 insertions, 9 deletions

diff --git a/src/tools/fuzzing/fuzzing.cpp b/src/tools/fuzzing/fuzzing.cpp
index bdaf71e83..ef7b3a5b7 100644
--- a/src/tools/fuzzing/fuzzing.cpp
+++ b/src/tools/fuzzing/fuzzing.cpp
@@ -16,6 +16,7 @@
 
 #include "tools/fuzzing.h"
 #include "ir/gc-type-utils.h"
+#include "ir/local-structural-dominance.h"
 #include "ir/module-utils.h"
 #include "ir/subtypes.h"
 #include "ir/type-updating.h"
@@ -559,15 +560,37 @@ void TranslateToFuzzReader::addHashMemorySupport() {
 }
 
 TranslateToFuzzReader::FunctionCreationContext::~FunctionCreationContext() {
+  // We must ensure non-nullable locals validate. Later down we'll run
+  // TypeUpdating::handleNonDefaultableLocals which will make them validate by
+  // turning them nullable + add ref.as_non_null to fix up types. That has the
+  // downside of making them trap at runtime, however, and also we lose the non-
+  // nullability in the type, so we prefer to do a manual fixup that avoids a
+  // trap, which we do by writing a non-nullable value into the local at the
+  // function entry.
+  // TODO: We could be more precise and use a LocalGraph here, at the cost of
+  //       doing more work.
+  LocalStructuralDominance info(
+    func, parent.wasm, LocalStructuralDominance::NonNullableOnly);
+  for (auto index : info.nonDominatingIndices) {
+    // Do not always do this, but with high probability, to reduce the amount of
+    // traps.
+    if (!parent.oneIn(5)) {
+      auto* value = parent.makeTrivial(func->getLocalType(index));
+      func->body = parent.builder.makeSequence(
+        parent.builder.makeLocalSet(index, value), func->body);
+    }
+  }
+
+  // Then, to handle remaining cases we did not just fix up, do the general
+  // fixup to ensure we validate.
+  TypeUpdating::handleNonDefaultableLocals(func, parent.wasm);
+
   if (HANG_LIMIT > 0) {
     parent.addHangLimitChecks(func);
   }
   assert(breakableStack.empty());
   assert(hangStack.empty());
   parent.funcContext = nullptr;
-
-  // We must ensure non-nullable locals validate.
-  TypeUpdating::handleNonDefaultableLocals(func, parent.wasm);
 }
 
 Expression* TranslateToFuzzReader::makeHangLimitCheck() {
@@ -621,11 +644,6 @@ Function* TranslateToFuzzReader::addFunction() {
   Index numVars = upToSquared(MAX_VARS);
   for (Index i = 0; i < numVars; i++) {
     auto type = getConcreteType();
-    if (!type.isDefaultable()) {
-      // We can't use a nondefaultable type as a var, as those must be
-      // initialized to some default value.
-      continue;
-    }
     funcContext->typeLocals[type].push_back(params.size() + func->vars.size());
     func->vars.push_back(type);
   }
@@ -1251,7 +1269,10 @@ Expression* TranslateToFuzzReader::makeTrivial(Type type) {
   } nester(*this);
 
   if (type.isConcrete()) {
-    if (oneIn(2) && funcContext) {
+    // If we have a function context, use a local half the time. Use a local
+    // less often if the local is non-nullable, however, as then we might be
+    // using it before it was set, which would trap.
+    if (funcContext && oneIn(type.isNonNullable() ? 10 : 2)) {
       return makeLocalGet(type);
     } else {
       return makeConst(type);