Require RefFunc to have the proper type (WebAssembly#7376)

tlively · web-flow · commit cd3b26d6e2e2 · 2025-03-21T10:03:53.000-07:00
As a holdout from before GC was implemented, we previously allowed
RefFunc expressions to have type `funcref` rather than a specific
signature type matching that of the referenced function. Remove this
allowance and start requiring the types to be correct and precise to
eliminate the possibility of stale types inhibiting (or invalidating!)
optimizations.

Update various older passes to update the types of RefFuncs, including
those in tables, to keep their output passing validation. Also update
the kitchen sink example test to construct RefFunc expressions with the
correct type via the C API.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -21,6 +21,10 @@ Current Trunk
  - `string` is now a subtype of `ext` (rather than `any`). This allows better
    transformations for strings, like an inverse of StringLowering, but will
    error on codebases that depend on being able to pass strings into anyrefs.
+ - Require the type of RefFunc expressions to match the type of the referenced
+   function. It is no longer valid to type them as funcref in the IR.
+ - The C and JS APIs for creating RefFunc expressions now take a HeapType
+   instead of a Type.
 
 v122
 ----
diff --git a/src/binaryen-c.cpp b/src/binaryen-c.cpp
@@ -1606,12 +1606,11 @@ BinaryenExpressionRef BinaryenRefAs(BinaryenModuleRef module,
     Builder(*(Module*)module).makeRefAs(RefAsOp(op), (Expression*)value));
 }
 
-BinaryenExpressionRef
-BinaryenRefFunc(BinaryenModuleRef module, const char* func, BinaryenType type) {
-  // TODO: consider changing the C API to receive a heap type
-  Type type_(type);
+BinaryenExpressionRef BinaryenRefFunc(BinaryenModuleRef module,
+                                      const char* func,
+                                      BinaryenHeapType type) {
   return static_cast<Expression*>(
-    Builder(*(Module*)module).makeRefFunc(func, type_.getHeapType()));
+    Builder(*(Module*)module).makeRefFunc(func, HeapType(type)));
 }
 
 BinaryenExpressionRef BinaryenRefEq(BinaryenModuleRef module,
@@ -6259,8 +6258,12 @@ bool TypeBuilderBuildAndDispose(TypeBuilderRef builder,
   auto* B = (TypeBuilder*)builder;
   auto result = B->build();
   if (auto err = result.getError()) {
-    *errorIndex = err->index;
-    *errorReason = static_cast<TypeBuilderErrorReason>(err->reason);
+    if (errorIndex) {
+      *errorIndex = err->index;
+    }
+    if (errorReason) {
+      *errorReason = static_cast<TypeBuilderErrorReason>(err->reason);
+    }
     delete B;
     return false;
   }
diff --git a/src/binaryen-c.h b/src/binaryen-c.h
@@ -961,7 +961,7 @@ BINARYEN_API BinaryenExpressionRef BinaryenRefAs(BinaryenModuleRef module,
                                                  BinaryenExpressionRef value);
 BINARYEN_API BinaryenExpressionRef BinaryenRefFunc(BinaryenModuleRef module,
                                                    const char* func,
-                                                   BinaryenType type);
+                                                   BinaryenHeapType type);
 BINARYEN_API BinaryenExpressionRef BinaryenRefEq(BinaryenModuleRef module,
                                                  BinaryenExpressionRef left,
                                                  BinaryenExpressionRef right);
diff --git a/src/ir/element-utils.h b/src/ir/element-utils.h
@@ -42,7 +42,7 @@ template<typename T>
 inline void iterAllElementFunctionNames(const Module* wasm, T visitor) {
   for (auto& segment : wasm->elementSegments) {
     iterElementSegmentFunctionNames(
-      segment.get(), [&](Name& name, Index i) { visitor(name); });
+      segment.get(), [&](const Name& name, Index i) { visitor(name); });
   }
 }
 
diff --git a/src/js/binaryen.js-post.js b/src/js/binaryen.js-post.js
@@ -3275,6 +3275,7 @@ Module['getFunctionInfo'] = function(func) {
     'name': UTF8ToString(Module['_BinaryenFunctionGetName'](func)),
     'module': UTF8ToString(Module['_BinaryenFunctionImportGetModule'](func)),
     'base': UTF8ToString(Module['_BinaryenFunctionImportGetBase'](func)),
+    'type': Module['_BinaryenFunctionGetType'](func),
     'params': Module['_BinaryenFunctionGetParams'](func),
     'results': Module['_BinaryenFunctionGetResults'](func),
     'vars': getAllNested(func, Module['_BinaryenFunctionGetNumVars'], Module['_BinaryenFunctionGetVar']),
@@ -4893,6 +4894,9 @@ Module['Function'] = (() => {
   Function['getName'] = function(func) {
     return UTF8ToString(Module['_BinaryenFunctionGetName'](func));
   };
+  Function['getType'] = function(func) {
+    return Module['_BinaryenFunctionGetType'](func);
+  }
   Function['getParams'] = function(func) {
     return Module['_BinaryenFunctionGetParams'](func);
   };
diff --git a/src/passes/FuncCastEmulation.cpp b/src/passes/FuncCastEmulation.cpp
@@ -162,25 +162,34 @@ struct FuncCastEmulation : public Pass {
     HeapType ABIType(
       Signature(Type(std::vector<Type>(numParams, Type::i64)), Type::i64));
     // Add a thunk for each function in the table, and do the call through it.
-    std::unordered_map<Name, Name> funcThunks;
-    ElementUtils::iterAllElementFunctionNames(module, [&](Name& name) {
-      auto iter = funcThunks.find(name);
-      if (iter == funcThunks.end()) {
-        auto thunk = makeThunk(name, module, numParams);
-        funcThunks[name] = thunk;
-        name = thunk;
-      } else {
-        name = iter->second;
+    std::unordered_map<Name, Function*> funcThunks;
+    for (auto& segment : module->elementSegments) {
+      if (!segment->type.isFunction()) {
+        continue;
       }
-    });
+      for (Index i = 0; i < segment->data.size(); ++i) {
+        auto* ref = segment->data[i]->dynCast<RefFunc>();
+        if (!ref) {
+          continue;
+        }
+        auto [iter, inserted] = funcThunks.insert({ref->func, nullptr});
+        if (inserted) {
+          iter->second = makeThunk(ref->func, module, numParams);
+        }
+        auto* thunk = iter->second;
+        ref->func = thunk->name;
+        // TODO: Make this exact.
+        ref->type = Type(thunk->type, NonNullable);
+      }
+    }
 
     // update call_indirects
     ParallelFuncCastEmulation(ABIType, numParams).run(getPassRunner(), module);
   }
 
 private:
   // Creates a thunk for a function, casting args and return value as needed.
-  Name makeThunk(Name name, Module* module, Index numParams) {
+  Function* makeThunk(Name name, Module* module, Index numParams) {
     Name thunk = std::string("byn$fpcast-emu$") + name.toString();
     if (module->getFunctionOrNull(thunk)) {
       Fatal() << "FuncCastEmulation::makeThunk seems a thunk name already in "
@@ -207,8 +216,7 @@ struct FuncCastEmulation : public Pass {
                            {}, // no vars
                            toABI(call, module));
     thunkFunc->hasExplicitName = true;
-    module->addFunction(std::move(thunkFunc));
-    return thunk;
+    return module->addFunction(std::move(thunkFunc));
   }
 };
 
diff --git a/src/passes/I64ToI32Lowering.cpp b/src/passes/I64ToI32Lowering.cpp
@@ -297,6 +297,39 @@ struct I64ToI32Lowering : public WalkerPass<PostWalker<I64ToI32Lowering>> {
       });
   }
 
+  void visitRefFunc(RefFunc* curr) {
+    auto sig = curr->type.getHeapType().getSignature();
+
+    auto lowerTypes = [](Type types) {
+      bool hasI64 = false;
+      for (auto t : types) {
+        if (t == Type::i64) {
+          hasI64 = true;
+          break;
+        }
+      }
+      if (!hasI64) {
+        return types;
+      }
+      std::vector<Type> newTypes;
+      for (auto t : types) {
+        if (t == Type::i64) {
+          newTypes.push_back(Type::i32);
+          newTypes.push_back(Type::i32);
+        } else {
+          newTypes.push_back(t);
+        }
+      }
+      return Type(newTypes);
+    };
+
+    auto newParams = lowerTypes(sig.params);
+    auto newResults = lowerTypes(sig.results);
+    if (newParams != sig.params || newResults != sig.results) {
+      curr->type = curr->type.with(HeapType(Signature(newParams, newResults)));
+    }
+  }
+
   void visitLocalGet(LocalGet* curr) {
     const auto mappedIndex = indexMap[curr->index];
     // Need to remap the local into the new naming scheme, regardless of
diff --git a/src/passes/LegalizeJSInterface.cpp b/src/passes/LegalizeJSInterface.cpp
@@ -108,57 +108,52 @@ struct LegalizeJSInterface : public Pass {
     // for each illegal import, we must call a legalized stub instead
     for (auto* im : originalFunctions) {
       if (im->imported() && isIllegal(im)) {
-        auto funcName = makeLegalStubForCalledImport(im, module);
-        illegalImportsToLegal[im->name] = funcName;
-        // we need to use the legalized version in the tables, as the import
-        // from JS is legal for JS. Our stub makes it look like a native wasm
-        // function.
-        ElementUtils::iterAllElementFunctionNames(module, [&](Name& name) {
-          if (name == im->name) {
-            name = funcName;
-          }
-        });
+        auto* func = makeLegalStubForCalledImport(im, module);
+        illegalImportsToLegal[im->name] = func;
       }
     }
 
     if (!illegalImportsToLegal.empty()) {
-      // fix up imports: call_import of an illegal must be turned to a call of a
-      // legal. the same must be done with ref.funcs.
+      // fix up imports: call of an illegal import must be turned to a call of a
+      // legal import. the same must be done with ref.funcs.
       struct Fixer : public WalkerPass<PostWalker<Fixer>> {
         bool isFunctionParallel() override { return true; }
 
         std::unique_ptr<Pass> create() override {
           return std::make_unique<Fixer>(illegalImportsToLegal);
         }
 
-        std::map<Name, Name>* illegalImportsToLegal;
+        std::unordered_map<Name, Function*>& illegalImportsToLegal;
 
-        Fixer(std::map<Name, Name>* illegalImportsToLegal)
+        Fixer(std::unordered_map<Name, Function*>& illegalImportsToLegal)
           : illegalImportsToLegal(illegalImportsToLegal) {}
 
         void visitCall(Call* curr) {
-          auto iter = illegalImportsToLegal->find(curr->target);
-          if (iter == illegalImportsToLegal->end()) {
+          auto iter = illegalImportsToLegal.find(curr->target);
+          if (iter == illegalImportsToLegal.end()) {
             return;
           }
 
-          replaceCurrent(
-            Builder(*getModule())
-              .makeCall(
-                iter->second, curr->operands, curr->type, curr->isReturn));
+          replaceCurrent(Builder(*getModule())
+                           .makeCall(iter->second->name,
+                                     curr->operands,
+                                     curr->type,
+                                     curr->isReturn));
         }
 
         void visitRefFunc(RefFunc* curr) {
-          auto iter = illegalImportsToLegal->find(curr->func);
-          if (iter == illegalImportsToLegal->end()) {
+          auto iter = illegalImportsToLegal.find(curr->func);
+          if (iter == illegalImportsToLegal.end()) {
             return;
           }
 
-          curr->func = iter->second;
+          curr->func = iter->second->name;
+          // TODO: Make this exact.
+          curr->type = Type(iter->second->type, NonNullable);
         }
       };
 
-      Fixer fixer(&illegalImportsToLegal);
+      Fixer fixer(illegalImportsToLegal);
       fixer.run(getPassRunner(), module);
       fixer.runOnModuleCode(getPassRunner(), module);
 
@@ -174,7 +169,7 @@ struct LegalizeJSInterface : public Pass {
 
 private:
   // map of illegal to legal names for imports
-  std::map<Name, Name> illegalImportsToLegal;
+  std::unordered_map<Name, Function*> illegalImportsToLegal;
   bool exportedHelpers = false;
   Function* getTempRet0 = nullptr;
   Function* setTempRet0 = nullptr;
@@ -274,7 +269,7 @@ struct LegalizeJSInterface : public Pass {
 
   // wasm calls the import, so it must call a stub that calls the actual legal
   // JS import
-  Name makeLegalStubForCalledImport(Function* im, Module* module) {
+  Function* makeLegalStubForCalledImport(Function* im, Module* module) {
     Builder builder(*module);
     auto legalIm = std::make_unique<Function>();
     legalIm->name = Name(std::string("legalimport$") + im->name.toString());
@@ -315,14 +310,14 @@ struct LegalizeJSInterface : public Pass {
     }
     legalIm->type = Signature(Type(params), call->type);
 
-    const auto& stubName = stub->name;
-    if (!module->getFunctionOrNull(stubName)) {
+    auto* stubPtr = stub.get();
+    if (!module->getFunctionOrNull(stub->name)) {
       module->addFunction(std::move(stub));
     }
     if (!module->getFunctionOrNull(legalIm->name)) {
       module->addFunction(std::move(legalIm));
     }
-    return stubName;
+    return stubPtr;
   }
 
   static Function*
diff --git a/src/passes/PrintCallGraph.cpp b/src/passes/PrintCallGraph.cpp
@@ -96,7 +96,7 @@ struct PrintCallGraph : public Pass {
     CallPrinter printer(module);
 
     // Indirect Targets
-    ElementUtils::iterAllElementFunctionNames(module, [&](Name& name) {
+    ElementUtils::iterAllElementFunctionNames(module, [&](Name name) {
       auto* func = module->getFunction(name);
       o << "  \"" << func->name << "\" [style=\"filled, rounded\"];\n";
     });
diff --git a/src/passes/RemoveImports.cpp b/src/passes/RemoveImports.cpp
@@ -51,7 +51,7 @@ struct RemoveImports : public WalkerPass<PostWalker<RemoveImports>> {
     // Do not remove names referenced in a table
     std::set<Name> indirectNames;
     ElementUtils::iterAllElementFunctionNames(
-      curr, [&](Name& name) { indirectNames.insert(name); });
+      curr, [&](Name name) { indirectNames.insert(name); });
     for (auto& name : names) {
       if (indirectNames.find(name) == indirectNames.end()) {
         curr->removeFunction(name);
diff --git a/src/passes/RemoveUnusedModuleElements.cpp b/src/passes/RemoveUnusedModuleElements.cpp
@@ -694,7 +694,7 @@ struct RemoveUnusedModuleElements : public Pass {
     // For now, all functions that can be called indirectly are marked as roots.
     // TODO: Compute this based on which ElementSegments are actually used,
     //       and which functions have a call_indirect of the proper type.
-    ElementUtils::iterAllElementFunctionNames(module, [&](Name& name) {
+    ElementUtils::iterAllElementFunctionNames(module, [&](Name name) {
       roots.emplace_back(ModuleElementKind::Function, name);
     });
 
diff --git a/src/passes/ReorderFunctions.cpp b/src/passes/ReorderFunctions.cpp
@@ -81,7 +81,7 @@ struct ReorderFunctions : public Pass {
       }
     }
     ElementUtils::iterAllElementFunctionNames(
-      module, [&](Name& name) { counts[name]++; });
+      module, [&](Name name) { counts[name]++; });
     // TODO: count all RefFunc as well
     // TODO: count the declaration section as well, which adds another mention
     // sort
diff --git a/src/wasm/wasm-validator.cpp b/src/wasm/wasm-validator.cpp
@@ -2364,23 +2364,21 @@ void FunctionValidator::visitRefFunc(RefFunc* curr) {
                     "ref.func should have a non-nullable reference type")) {
     return;
   }
+  if (!shouldBeTrue(curr->type.isSignature(),
+                    curr,
+                    "ref.func must have a function reference type")) {
+    return;
+  }
   if (!info.validateGlobally) {
     return;
   }
   auto* func = getModule()->getFunctionOrNull(curr->func);
-  shouldBeTrue(!!func, curr, "function argument of ref.func must exist");
-  shouldBeTrue(curr->type.isFunction(),
+  if (!shouldBeTrue(!!func, curr, "function argument of ref.func must exist")) {
+    return;
+  }
+  shouldBeTrue(func->type == curr->type.getHeapType(),
                curr,
-               "ref.func must have a function reference type");
-  shouldBeTrue(
-    !curr->type.isNullable(), curr, "ref.func must have non-nullable type");
-  // TODO: verify it also has a typed function references type, and the right
-  // one,
-  //   curr->type.getHeapType().getSignature()
-  // That is blocked on having the ability to create signature types in the C
-  // API (for now those users create the type with funcref). This also needs to
-  // be fixed in LegalizeJSInterface and FuncCastEmulation and other places that
-  // update function types.
+               "function reference type must match referenced function type");
 }
 
 void FunctionValidator::visitRefEq(RefEq* curr) {
diff --git a/src/wasm/wasm.cpp b/src/wasm/wasm.cpp
@@ -818,6 +818,7 @@ void RefIsNull::finalize() {
 void RefFunc::finalize() {
   // No-op. We assume that the full proper typed function type has been applied
   // previously.
+  assert(type.isSignature());
 }
 
 void RefFunc::finalize(Type type_) { type = type_; }
diff --git a/test/binaryen.js/expressions.js b/test/binaryen.js/expressions.js
@@ -1475,22 +1475,20 @@ console.log("# RefAs");
 console.log("# RefFunc");
 (function testRefFunc() {
   const module = new binaryen.Module();
+  module.addFunction("a", binaryen.none, binaryen.none, [], module.nop());
+  var type = binaryen.Function(module.getFunction("a")).type;
 
   var func = "a";
-  const theRefFunc = binaryen.RefFunc(module.ref.func(func, binaryen.funcref));
+  const theRefFunc = binaryen.RefFunc(module.ref.func(func, type));
   assert(theRefFunc instanceof binaryen.RefFunc);
   assert(theRefFunc instanceof binaryen.Expression);
   assert(theRefFunc.func === func);
-  // TODO: check the type. the type is (ref func), that is, a non-nullable func,
-  //       which differs from funcref. we don't have the ability to create such
-  //       a type in the C/JS APIs yet.
+  assert(theRefFunc.type === type);
 
   theRefFunc.func = func = "b";
   assert(theRefFunc.func === func);
-  theRefFunc.type = binaryen.f64;
   theRefFunc.finalize();
-  // TODO The type is a subtype of funcref, but we can't check that in the JS
-  //      API atm.
+  assert(theRefFunc.type === type);
 
   console.log(theRefFunc.toText());
   assert(
diff --git a/test/binaryen.js/kitchen-sink.js b/test/binaryen.js/kitchen-sink.js
diff --git a/test/binaryen.js/kitchen-sink.js.txt b/test/binaryen.js/kitchen-sink.js.txt
diff --git a/test/binaryen.js/validation_errors.js b/test/binaryen.js/validation_errors.js
diff --git a/test/example/c-api-kitchen-sink.c b/test/example/c-api-kitchen-sink.c
diff --git a/test/example/c-api-kitchen-sink.txt b/test/example/c-api-kitchen-sink.txt

Original file line number	Diff line number	Diff line change
`@@ -42,7 +42,7 @@ template<typename T>`
`42`	`42`	`inline void iterAllElementFunctionNames(const Module* wasm, T visitor) {`
`43`	`43`	`for (auto& segment : wasm->elementSegments) {`
`44`	`44`	`iterElementSegmentFunctionNames(`
`45`		`- segment.get(), [&](Name& name, Index i) { visitor(name); });`
	`45`	`+ segment.get(), [&](const Name& name, Index i) { visitor(name); });`
`46`	`46`	`}`
`47`	`47`	`}`
`48`	`48`