[TypedFunctionReferences] Add Typed Function References feature and use the types (#3388)

This adds the new feature and starts to use the new types where relevant. We
use them even without the feature being enabled, as we don't know the features
during wasm loading - but the hope is that given the type is a subtype, it should
all work out. In practice, if you print out the internal type you may see a typed
function reference-specific type for a ref.func for example, instead of a generic
funcref, but it should not affect anything else.

This PR does not support non-nullable types, that is, everything is nullable
for now. As suggested by @tlively this is simpler for now and leaves nullability
for later work (which will apparently require let or something else, and many
passes may need to be changed).

To allow this PR to work, we need to provide a type on creating a RefFunc. The
wasm-builder.h internal API is updated for this, as are the C and JS APIs,
which are breaking changes. cc @dcodeIO

We must also write and read function types properly. This PR improves
collectSignatures to find all the types, and also to sort them by the
dependencies between them (as we can't emit X in the binary if it depends
on Y, and Y has not been emitted - we need to give Y's index). This sorting
ends up changing a few test outputs.

InstrumentLocals support for printing function types that are not funcref
is disabled for now, until we figure out how to make that work and/or
decide if it's important enough to work on.

The fuzzer has various fixes to emit valid types for things (mostly
whitespace there). Also two drive-by fixes to call makeTrivial where it
should be (when we fail to create a specific node, we can't just try to make
another node, in theory it could infinitely recurse).

Binary writing changes here to replace calls to a standalone function to
write out a type with one that is called on the binary writer object itself,
which maintains a mapping of type indexes (getFunctionSignatureByIndex).

Author: kripken

CHANGELOG.md

@@ -15,6 +15,9 @@ full changeset diff at the end of each section.
 Current Trunk
 -------------
 
+- `RefFunc` C and JS API constructors (`BinaryenRefFunc` and `ref.func`
+  respectively) now take an extra `type` parameter, similar to `RefNull`. This
+  is necessary for typed function references support.
 - JS API functions for atomic notify/wait instructions are renamed.
   - `module.atomic.notify` -> `module.memory.atomic.notify`
   - `module.i32.atomic.wait` -> `module.memory.atomic.wait32`

src/binaryen-c.cpp

@@ -1187,9 +1187,11 @@ BinaryenExpressionRef BinaryenRefIsNull(BinaryenModuleRef module,
     Builder(*(Module*)module).makeRefIsNull((Expression*)value));
 }
 
-BinaryenExpressionRef BinaryenRefFunc(BinaryenModuleRef module,
-                                      const char* func) {
-  return static_cast<Expression*>(Builder(*(Module*)module).makeRefFunc(func));
+BinaryenExpressionRef
+BinaryenRefFunc(BinaryenModuleRef module, const char* func, BinaryenType type) {
+  Type type_(type);
+  return static_cast<Expression*>(
+    Builder(*(Module*)module).makeRefFunc(func, type_));
 }
 
 BinaryenExpressionRef BinaryenRefEq(BinaryenModuleRef module,

src/binaryen-c.h

@@ -792,7 +792,8 @@ BINARYEN_API BinaryenExpressionRef BinaryenRefNull(BinaryenModuleRef module,
 BINARYEN_API BinaryenExpressionRef
 BinaryenRefIsNull(BinaryenModuleRef module, BinaryenExpressionRef value);
 BINARYEN_API BinaryenExpressionRef BinaryenRefFunc(BinaryenModuleRef module,
-                                                   const char* func);
+                                                   const char* func,
+                                                   BinaryenType type);
 BINARYEN_API BinaryenExpressionRef BinaryenRefEq(BinaryenModuleRef module,
                                                  BinaryenExpressionRef left,
                                                  BinaryenExpressionRef right);

src/ir/ReFinalize.cpp

@@ -126,7 +126,11 @@ void ReFinalize::visitMemorySize(MemorySize* curr) { curr->finalize(); }
 void ReFinalize::visitMemoryGrow(MemoryGrow* curr) { curr->finalize(); }
 void ReFinalize::visitRefNull(RefNull* curr) { curr->finalize(); }
 void ReFinalize::visitRefIsNull(RefIsNull* curr) { curr->finalize(); }
-void ReFinalize::visitRefFunc(RefFunc* curr) { curr->finalize(); }
+void ReFinalize::visitRefFunc(RefFunc* curr) {
+  // TODO: should we look up the function and update the type from there? This
+  // could handle a change to the function's type, but is also not really what
+  // this class has been meant to do.
+}
 void ReFinalize::visitRefEq(RefEq* curr) { curr->finalize(); }
 void ReFinalize::visitTry(Try* curr) { curr->finalize(); }
 void ReFinalize::visitThrow(Throw* curr) { curr->finalize(); }

src/ir/module-utils.h

@@ -414,16 +414,29 @@ collectSignatures(Module& wasm,
       Counts& counts;
 
       TypeCounter(Counts& counts) : counts(counts) {}
+
       void visitExpression(Expression* curr) {
-        if (auto* call = curr->dynCast<CallIndirect>()) {
+        if (curr->is<RefNull>()) {
+          maybeNote(curr->type);
+        } else if (auto* call = curr->dynCast<CallIndirect>()) {
           counts[call->sig]++;
         } else if (Properties::isControlFlowStructure(curr)) {
-          // TODO: Allow control flow to have input types as well
+          maybeNote(curr->type);
           if (curr->type.isTuple()) {
+            // TODO: Allow control flow to have input types as well
             counts[Signature(Type::none, curr->type)]++;
           }
         }
       }
+
+      void maybeNote(Type type) {
+        if (type.isRef()) {
+          auto heapType = type.getHeapType();
+          if (heapType.isSignature()) {
+            counts[heapType.getSignature()]++;
+          }
+        }
+      }
     };
     TypeCounter(counts).walk(func->body);
   };
@@ -434,6 +447,14 @@ collectSignatures(Module& wasm,
   Counts counts;
   for (auto& curr : wasm.functions) {
     counts[curr->sig]++;
+    for (auto type : curr->vars) {
+      if (type.isRef()) {
+        auto heapType = type.getHeapType();
+        if (heapType.isSignature()) {
+          counts[heapType.getSignature()]++;
+        }
+      }
+    }
   }
   for (auto& curr : wasm.events) {
     counts[curr->sig]++;
@@ -444,10 +465,61 @@ collectSignatures(Module& wasm,
       counts[innerPair.first] += innerPair.second;
     }
   }
+
+  // TODO: recursively traverse each reference type, which may have a child type
+  //       this is itself a reference type.
+
+  // We must sort all the dependencies of a signature before it. For example,
+  // (func (param (ref (func)))) must appear after (func). To do that, find the
+  // depth of dependencies of each signature. For example, if A depends on B
+  // which depends on C, then A's depth is 2, B's is 1, and C's is 0 (assuming
+  // no other dependencies).
+  Counts depthOfDependencies;
+  std::unordered_map<Signature, std::unordered_set<Signature>> isDependencyOf;
+  // To calculate the depth of dependencies, we'll do a flow analysis, visiting
+  // each signature as we find out new things about it.
+  std::set<Signature> toVisit;
+  for (auto& pair : counts) {
+    auto sig = pair.first;
+    depthOfDependencies[sig] = 0;
+    toVisit.insert(sig);
+    for (Type type : {sig.params, sig.results}) {
+      for (auto element : type) {
+        if (element.isRef()) {
+          auto heapType = element.getHeapType();
+          if (heapType.isSignature()) {
+            isDependencyOf[heapType.getSignature()].insert(sig);
+          }
+        }
+      }
+    }
+  }
+  while (!toVisit.empty()) {
+    auto iter = toVisit.begin();
+    auto sig = *iter;
+    toVisit.erase(iter);
+    // Anything that depends on this has a depth of dependencies equal to this
+    // signature's, plus this signature itself.
+    auto newDepth = depthOfDependencies[sig] + 1;
+    if (newDepth > counts.size()) {
+      Fatal() << "Cyclic signatures detected, cannot sort them.";
+    }
+    for (auto& other : isDependencyOf[sig]) {
+      if (depthOfDependencies[other] < newDepth) {
+        // We found something new to propagate.
+        depthOfDependencies[other] = newDepth;
+        toVisit.insert(other);
+      }
+    }
+  }
+  // Sort by frequency and then simplicity, and also keeping every signature
+  // before things that depend on it.
   std::vector<std::pair<Signature, size_t>> sorted(counts.begin(),
                                                    counts.end());
   std::sort(sorted.begin(), sorted.end(), [&](auto a, auto b) {
-    // order by frequency then simplicity
+    if (depthOfDependencies[a.first] != depthOfDependencies[b.first]) {
+      return depthOfDependencies[a.first] < depthOfDependencies[b.first];
+    }
     if (a.second != b.second) {
       return a.second > b.second;
     }

src/js/binaryen.js-post.js

@@ -2112,8 +2112,8 @@ function wrapModule(module, self = {}) {
     'is_null'(value) {
       return Module['_BinaryenRefIsNull'](module, value);
     },
-    'func'(func) {
-      return preserveStack(() => Module['_BinaryenRefFunc'](module, strToStack(func)));
+    'func'(func, type) {
+      return preserveStack(() => Module['_BinaryenRefFunc'](module, strToStack(func), type));
     },
     'eq'(left, right) {
       return Module['_BinaryenRefEq'](module, left, right);

src/passes/InstrumentLocals.cpp

@@ -135,45 +135,48 @@ struct InstrumentLocals : public WalkerPass<PostWalker<InstrumentLocals>> {
     Builder builder(*getModule());
     Name import;
     auto type = curr->value->type;
-    if (type.isFunction()) {
-      import = set_funcref;
-    } else {
-      TODO_SINGLE_COMPOUND(curr->value->type);
-      switch (type.getBasic()) {
-        case Type::i32:
-          import = set_i32;
-          break;
-        case Type::i64:
-          return; // TODO
-        case Type::f32:
-          import = set_f32;
-          break;
-        case Type::f64:
-          import = set_f64;
-          break;
-        case Type::v128:
-          import = set_v128;
-          break;
-        case Type::externref:
-          import = set_externref;
-          break;
-        case Type::exnref:
-          import = set_exnref;
-          break;
-        case Type::anyref:
-          import = set_anyref;
-          break;
-        case Type::eqref:
-          import = set_eqref;
-          break;
-        case Type::i31ref:
-          import = set_i31ref;
-          break;
-        case Type::unreachable:
-          return; // nothing to do here
-        default:
-          WASM_UNREACHABLE("unexpected type");
-      }
+    if (type.isFunction() && type != Type::funcref) {
+      // FIXME: support typed function references
+      return;
+    }
+    TODO_SINGLE_COMPOUND(curr->value->type);
+    switch (type.getBasic()) {
+      case Type::i32:
+        import = set_i32;
+        break;
+      case Type::i64:
+        return; // TODO
+      case Type::f32:
+        import = set_f32;
+        break;
+      case Type::f64:
+        import = set_f64;
+        break;
+      case Type::v128:
+        import = set_v128;
+        break;
+      case Type::funcref:
+        import = set_funcref;
+        break;
+      case Type::externref:
+        import = set_externref;
+        break;
+      case Type::exnref:
+        import = set_exnref;
+        break;
+      case Type::anyref:
+        import = set_anyref;
+        break;
+      case Type::eqref:
+        import = set_eqref;
+        break;
+      case Type::i31ref:
+        import = set_i31ref;
+        break;
+      case Type::unreachable:
+        return; // nothing to do here
+      default:
+        WASM_UNREACHABLE("unexpected type");
     }
     curr->value = builder.makeCall(import,
                                    {builder.makeConst(int32_t(id++)),

src/tools/fuzzing.h

@@ -321,6 +321,10 @@ class TranslateToFuzzReader {
       }
       return Type(types);
     }
+    if (type.isFunction() && type != Type::funcref) {
+      // TODO: specific typed function references types.
+      return type;
+    }
     SmallVector<Type, 2> options;
     options.push_back(type); // includes itself
     TODO_SINGLE_COMPOUND(type);
@@ -653,6 +657,10 @@ class TranslateToFuzzReader {
     Index numVars = upToSquared(MAX_VARS);
     for (Index i = 0; i < numVars; i++) {
       auto type = getConcreteType();
+      if (type.isRef() && !type.isNullable()) {
+        // We can't use a nullable type as a var, which is null-initialized.
+        continue;
+      }
       funcContext->typeLocals[type].push_back(params.size() +
                                               func->vars.size());
       func->vars.push_back(type);
@@ -1371,7 +1379,6 @@ class TranslateToFuzzReader {
   }
 
   Expression* makeCall(Type type) {
-    // seems ok, go on
     int tries = TRIES;
     bool isReturn;
     while (tries-- > 0) {
@@ -1392,7 +1399,7 @@ class TranslateToFuzzReader {
       return builder.makeCall(target->name, args, type, isReturn);
     }
     // we failed to find something
-    return make(type);
+    return makeTrivial(type);
   }
 
   Expression* makeCallIndirect(Type type) {
@@ -1418,7 +1425,7 @@ class TranslateToFuzzReader {
         i = 0;
       }
       if (i == start) {
-        return make(type);
+        return makeTrivial(type);
       }
     }
     // with high probability, make sure the type is valid  otherwise, most are
@@ -2018,12 +2025,28 @@ class TranslateToFuzzReader {
         if (!wasm.functions.empty() && !oneIn(wasm.functions.size())) {
           target = pick(wasm.functions).get();
         }
-        return builder.makeRefFunc(target->name);
+        auto type = Type(HeapType(target->sig), /* nullable = */ true);
+        return builder.makeRefFunc(target->name, type);
       }
       if (type == Type::i31ref) {
         return builder.makeI31New(makeConst(Type::i32));
       }
-      return builder.makeRefNull(type);
+      if (oneIn(2) && type.isNullable()) {
+        return builder.makeRefNull(type);
+      }
+      // TODO: randomize the order
+      for (auto& func : wasm.functions) {
+        // FIXME: RefFunc type should be non-nullable, but we emit nullable
+        //        types for now.
+        if (type == Type(HeapType(func->sig), /* nullable = */ true)) {
+          return builder.makeRefFunc(func->name, type);
+        }
+      }
+      // We failed to find a function, so create a null reference if we can.
+      if (type.isNullable()) {
+        return builder.makeRefNull(type);
+      }
+      WASM_UNREACHABLE("un-handleable non-nullable type");
     }
     if (type.isTuple()) {
       std::vector<Expression*> operands;
@@ -2972,6 +2995,7 @@ class TranslateToFuzzReader {
              Type::anyref,
              Type::eqref,
              Type::i31ref));
+    // TODO: emit typed function references types
   }
 
   Type getSingleConcreteType() { return pick(getSingleConcreteTypes()); }
@@ -2997,12 +3021,24 @@ class TranslateToFuzzReader {
 
   Type getEqReferenceType() { return pick(getEqReferenceTypes()); }
 
+  Type getMVPType() {
+    return pick(items(FeatureOptions<Type>().add(
+      FeatureSet::MVP, Type::i32, Type::i64, Type::f32, Type::f64)));
+  }
+
   Type getTupleType() {
     std::vector<Type> elements;
-    size_t numElements = 2 + upTo(MAX_TUPLE_SIZE - 1);
-    elements.resize(numElements);
-    for (size_t i = 0; i < numElements; ++i) {
-      elements[i] = getSingleConcreteType();
+    size_t maxElements = 2 + upTo(MAX_TUPLE_SIZE - 1);
+    for (size_t i = 0; i < maxElements; ++i) {
+      auto type = getSingleConcreteType();
+      // Don't add a non-nullable type into a tuple, as currently we can't spill
+      // them into locals (that would require a "let").
+      if (!type.isNullable()) {
+        elements.push_back(type);
+      }
+    }
+    while (elements.size() < 2) {
+      elements.push_back(getMVPType());
     }
     return Type(elements);
   }

src/tools/tool-options.h

@@ -89,6 +89,8 @@ struct ToolOptions : public Options {
       .addFeature(FeatureSet::Multivalue, "multivalue functions")
       .addFeature(FeatureSet::GC, "garbage collection")
       .addFeature(FeatureSet::Memory64, "memory64")
+      .addFeature(FeatureSet::TypedFunctionReferences,
+                  "typed function references")
       .add("--no-validation",
            "-n",
            "Disables validation, assumes inputs are correct",