[Stack Switching] Improve cont.bind fuzzing

src/tools/fuzzing.h

@@ -217,6 +217,9 @@ class TranslateToFuzzReader {
   // All arrays that are mutable.
   std::vector<HeapType> mutableArrays;
 
+  // Mapping of signatures to the continuations they are used by.
+  std::unordered_map<HeapType, std::vector<HeapType>> sigConts;
+
   // All tags that are valid as exception tags (which cannot have results).
   std::vector<Tag*> exceptionTags;
 

src/tools/fuzzing/fuzzing.cpp

@@ -574,6 +574,7 @@ void TranslateToFuzzReader::setupHeapTypes() {
         break;
       case HeapTypeKind::Cont:
         interestingHeapSubTypes[cont].push_back(type);
+        sigConts[type.getContinuation().type].push_back(type);
         break;
       case HeapTypeKind::Basic:
         WASM_UNREACHABLE("unexpected kind");
@@ -5458,20 +5459,94 @@ Expression* TranslateToFuzzReader::makeBrOn(Type type) {
 }
 
 Expression* TranslateToFuzzReader::makeContBind(Type type) {
-  auto sig = type.getHeapType().getContinuation().type.getSignature();
-  // Add a single param to be bound. TODO: Add multiple, and look in
-  // interestingHeapTypes.
-  std::vector<Type> newParams;
-  for (auto t : sig.params) {
-    newParams.push_back(t);
-  }
-  auto newParam = getSingleConcreteType();
-  newParams.insert(newParams.begin(), newParam);
-  auto newSig = Signature(Type(newParams), sig.results);
-  auto newCont = Continuation(newSig);
-  auto newType = Type(newCont, NonNullable, Exact);
-  std::vector<Expression*> newArgs{make(newParam)};
-  return builder.makeContBind(type.getHeapType(), newArgs, make(newType));
+  // We must output a signature that corresponds to the type we were given.
+  auto outputSigType = type.getHeapType().getContinuation().type;
+  auto outputSig = outputSigType.getSignature();
+  auto numOutputParams = outputSig.params.size();
+
+  // Look for a compatible signature, one who we are a suffix of. For example,
+  // with params [x,y,z] we'd want a signature like [a,b,x,y,z] so that we can
+  // bind a and b. That will be the input signature, which is longer, and after
+  // cont.bind it becomes the output signature. Don't always do this, however.
+  std::optional<HeapType> inputSigType;
+  if (!oneIn(4)) {
+    auto& funcTypes = interestingHeapSubTypes[HeapTypes::func];
+    // Filter out signatures with incompatible results.
+    std::vector<HeapType> relevantFuncTypes;
+    for (auto funcType : funcTypes) {
+      if (funcType.getSignature().results == outputSig.results) {
+        relevantFuncTypes.push_back(funcType);
+      }
+    }
+    if (!relevantFuncTypes.empty()) {
+      int tries = fuzzParams->TRIES;
+      while (tries-- > 0) {
+        auto pickedSigType = pick(relevantFuncTypes);
+        auto pickedSig = pickedSigType.getSignature();
+        assert(pickedSig.results == outputSig.results);
+        auto numinputParams = pickedSig.params.size();
+        if (numinputParams < numOutputParams) {
+          // Too short.
+          continue;
+        }
+        // Ignoring the input params at the start, compare the tails.
+        auto numAddedParams = numinputParams - numOutputParams;
+        bool bad = false;
+        for (Index i = 0; i < numOutputParams; i++) {
+          if (!Type::isSubType(outputSig.params[i],
+                               pickedSig.params[numAddedParams + i])) {
+            bad = true;
+            break;
+          }
+        }
+        if (!bad) {
+          inputSigType = pickedSigType;
+          break;
+        }
+      }
+    }
+  }
+
+  Index numAddedParams;
+  if (inputSigType) {
+    numAddedParams =
+      inputSigType->getSignature().params.size() - numOutputParams;
+  } else {
+    // We failed to find a signature, either use the current one (binding no
+    // input params) or invent a input one, adding one param.
+    if (oneIn(2)) {
+      inputSigType = outputSigType;
+      numAddedParams = 0;
+    } else {
+      std::vector<Type> inputParams;
+      for (auto t : outputSig.params) {
+        inputParams.push_back(t);
+      }
+      auto inputParam = getSingleConcreteType();
+      inputParams.insert(inputParams.begin(), inputParam);
+      inputSigType = Signature(Type(inputParams), outputSig.results);
+      numAddedParams = 1;
+    }
+  }
+  auto inputSig = inputSigType->getSignature();
+
+  // Pick a continuation type for the signature. If existing continuations use
+  // it, usually pick one of them.
+  auto& inputSigConts = sigConts[*inputSigType];
+  HeapType inputCont;
+  if (!inputSigConts.empty() && !oneIn(5)) {
+    inputCont = pick(inputSigConts);
+  } else {
+    inputCont = Continuation(inputSig);
+  }
+  Type inputType = Type(inputCont, NonNullable, Exact);
+
+  // Generate the new args and the cont.bind.
+  std::vector<Expression*> newArgs;
+  for (Index i = 0; i < numAddedParams; i++) {
+    newArgs.push_back(make(inputSig.params[i]));
+  }
+  return builder.makeContBind(type.getHeapType(), newArgs, make(inputType));
 }
 
 bool TranslateToFuzzReader::maybeSignedGet(const Field& field) {