Fuzzer: Update the typeLocals data structure before mutation (#6537)

Rather than compute the map of type to locals of that type once, at the
start, also update it when relevant, as we can add more locals in some
cases. This allows us to local.get from those late-added locals too.

2 files changed, 20 insertions, 2 deletions

diff --git a/src/tools/fuzzing.h b/src/tools/fuzzing.h
index 4aef9c372..ea9598f85 100644
--- a/src/tools/fuzzing.h
+++ b/src/tools/fuzzing.h
@@ -150,6 +150,14 @@ private:
     }
 
     ~FunctionCreationContext();
+
+    // Fill in the typeLocals data structure.
+    void computeTypeLocals() {
+      typeLocals.clear();
+      for (Index i = 0; i < func->getNumLocals(); i++) {
+        typeLocals[func->getLocalType(i)].push_back(i);
+      }
+    }
   };
 
   FunctionCreationContext* funcContext = nullptr;
diff --git a/src/tools/fuzzing/fuzzing.cpp b/src/tools/fuzzing/fuzzing.cpp
index 0120dbef3..a35a94d71 100644
--- a/src/tools/fuzzing/fuzzing.cpp
+++ b/src/tools/fuzzing/fuzzing.cpp
@@ -695,7 +695,6 @@ Function* TranslateToFuzzReader::addFunction() {
   params.reserve(numParams);
   for (Index i = 0; i < numParams; i++) {
     auto type = getSingleConcreteType();
-    funcContext->typeLocals[type].push_back(params.size());
     params.push_back(type);
   }
   auto paramType = Type(params);
@@ -704,9 +703,9 @@ Function* TranslateToFuzzReader::addFunction() {
   Index numVars = upToSquared(MAX_VARS);
   for (Index i = 0; i < numVars; i++) {
     auto type = getConcreteType();
-    funcContext->typeLocals[type].push_back(params.size() + func->vars.size());
     func->vars.push_back(type);
   }
+  context.computeTypeLocals();
   // with small chance, make the body unreachable
   auto bodyType = func->getResults();
   if (oneIn(10)) {
@@ -722,6 +721,14 @@ Function* TranslateToFuzzReader::addFunction() {
   // may end up breaking them. TODO: do them after the fact, like with the
   // hang limit checks.
   if (allowOOB) {
+    // Notice the locals and their types again, as more may have been added
+    // during generation of the body. We want to be able to local.get from those
+    // as well.
+    // TODO: We could also add a "localize" phase here to stash even more things
+    //       in locals, so that they can be reused. But we would need to be
+    //       careful with non-nullable locals (which error if used before being
+    //       set, or trap if we make them nullable, both of which are bad).
+    context.computeTypeLocals();
     // Recombinations create duplicate code patterns.
     recombine(func);
     // Mutations add random small changes, which can subtly break duplicate
@@ -733,6 +740,7 @@ Function* TranslateToFuzzReader::addFunction() {
     // after.
     fixAfterChanges(func);
   }
+
   // Add hang limit checks after all other operations on the function body.
   wasm.addFunction(func);
   // Export some functions, but not all (to allow inlining etc.). Try to export
@@ -1194,7 +1202,9 @@ void TranslateToFuzzReader::modifyInitialFunctions() {
     // Optionally, fuzz the function contents.
     if (upTo(RESOLUTION) >= chance) {
       dropToLog(func);
+      // Notice params as well as any locals generated above.
       // TODO add some locals? and the rest of addFunction's operations?
+      context.computeTypeLocals();
       // TODO: if we add OOB checks after creation, then we can do it on
       //       initial contents too, and it may be nice to *not* run these
       //       passes, like we don't run them on new functions. But, we may