Bug 1546138 - Wasm: Implement 'ref.func'. r=bbouvier

This commit implements the 'ref.func' instruction by emitting an instance call
to WasmInstanceObject::getExportedFunction.

The referenced function must be used in an element segment to validate.
See [1] for more details.

[1] WebAssembly/reference-types#31

Differential Revision: https://phabricator.services.mozilla.com/D40586

js/src/jit-test/tests/wasm/gc/ref-func.js

@@ -0,0 +1,95 @@
+// |jit-test| skip-if: !wasmReftypesEnabled()
+
+// 'ref.func' parses, validates and returns a non-null value
+wasmFullPass(`
+	(module
+		(elem declared $run)
+		(func $run (result i32)
+			ref.func $run
+			ref.is_null
+		)
+		(export "run" $run)
+	)
+`, 0);
+
+// function returning reference to itself
+{
+	let {f1} = wasmEvalText(`
+		(module
+			(elem declared $f1)
+			(func $f1 (result funcref) ref.func $f1)
+			(export "f1" $f1)
+		)
+	`).exports;
+	assertEq(f1(), f1);
+}
+
+// function returning reference to a different function
+{
+	let {f1, f2} = wasmEvalText(`
+		(module
+			(elem declared $f1)
+			(func $f1)
+			(func $f2 (result funcref) ref.func $f1)
+			(export "f1" $f1)
+			(export "f2" $f2)
+		)
+	`).exports;
+	assertEq(f2(), f1);
+}
+
+// function returning reference to function in a different module
+{
+	let i1 = wasmEvalText(`
+		(module
+			(elem declared $f1)
+			(func $f1)
+			(export "f1" $f1)
+		)
+	`);
+	let i2 = wasmEvalText(`
+		(module
+			(import $f1 "" "f1" (func))
+			(elem declared $f1)
+			(func $f2 (result funcref) ref.func $f1)
+			(export "f1" $f1)
+			(export "f2" $f2)
+		)
+	`, {"": i1.exports});
+
+	let f1 = i1.exports.f1;
+	let f2 = i2.exports.f2;
+	assertEq(f2(), f1);
+}
+
+// function index must be valid
+assertErrorMessage(() => {
+	wasmEvalText(`
+		(module
+			(func (result funcref) ref.func 10)
+		)
+	`);
+}, WebAssembly.CompileError, /function index out of range/);
+
+function validFuncRefText(forwardDeclare) {
+	return wasmEvalText(`
+		(module
+			(table 1 funcref)
+			(func $test (result funcref) ref.func $referenced)
+			(func $referenced)
+			${forwardDeclare}
+		)
+	`);
+}
+
+// referenced function must be forward declared somehow
+assertErrorMessage(() => validFuncRefText(''), WebAssembly.CompileError, /function index is not in an element segment/);
+
+// referenced function can be forward declared via segments
+assertEq(validFuncRefText('(elem 0 (i32.const 0) $referenced)') instanceof WebAssembly.Instance, true);
+assertEq(validFuncRefText('(elem passive $referenced)') instanceof WebAssembly.Instance, true);
+assertEq(validFuncRefText('(elem declared $referenced)') instanceof WebAssembly.Instance, true);
+
+// referenced function cannot be forward declared via start section or export
+assertErrorMessage(() => validFuncRefText('(start $referenced)'), WebAssembly.CompileError, /function index is not in an element segment/);
+assertErrorMessage(() => validFuncRefText('(export "referenced" $referenced)'), WebAssembly.CompileError, /function index is not in an element segment/);

js/src/wasm/WasmAST.h

@@ -417,6 +417,7 @@ enum class AstExprKind {
   MemorySize,
   Nop,
   Pop,
+  RefFunc,
   RefNull,
   Return,
   SetGlobal,
@@ -1576,6 +1577,17 @@ class AstExtraConversionOperator final : public AstExpr {
   AstExpr* operand() const { return operand_; }
 };
 
+class AstRefFunc final : public AstExpr {
+  AstRef func_;
+
+ public:
+  static const AstExprKind Kind = AstExprKind::RefFunc;
+  explicit AstRefFunc(AstRef func)
+      : AstExpr(Kind, ExprType::FuncRef), func_(func) {}
+
+  AstRef& func() { return func_; }
+};
+
 class AstRefNull final : public AstExpr {
  public:
   static const AstExprKind Kind = AstExprKind::RefNull;

js/src/wasm/WasmBaselineCompile.cpp

@@ -6842,6 +6842,7 @@ class BaseCompiler final : public BaseCompilerInterface {
   MOZ_MUST_USE bool emitMemoryGrow();
   MOZ_MUST_USE bool emitMemorySize();
 
+  MOZ_MUST_USE bool emitRefFunc();
   MOZ_MUST_USE bool emitRefNull();
   void emitRefIsNull();
 
@@ -9858,6 +9859,20 @@ bool BaseCompiler::emitMemorySize() {
   return emitInstanceCall(lineOrBytecode, SASigMemorySize);
 }
 
+bool BaseCompiler::emitRefFunc() {
+  uint32_t lineOrBytecode = readCallSiteLineOrBytecode();
+  uint32_t funcIndex;
+  if (!iter_.readRefFunc(&funcIndex)) {
+    return false;
+  }
+  if (deadCode_) {
+    return true;
+  }
+
+  pushI32(funcIndex);
+  return emitInstanceCall(lineOrBytecode, SASigFuncRef);
+}
+
 bool BaseCompiler::emitRefNull() {
   if (!iter_.readRefNull()) {
     return false;
@@ -11487,6 +11502,9 @@ bool BaseCompiler::emitBody() {
             emitComparison(emitCompareRef, ValType::AnyRef, Assembler::Equal));
 #endif
 #ifdef ENABLE_WASM_REFTYPES
+      case uint16_t(Op::RefFunc):
+        CHECK_NEXT(emitRefFunc());
+        break;
       case uint16_t(Op::RefNull):
         CHECK_NEXT(emitRefNull());
         break;

js/src/wasm/WasmBuiltins.cpp

@@ -172,6 +172,8 @@ const SymbolicAddressSignature SASigTableSet = {SymbolicAddress::TableSet,
                                                 {_PTR, _I32, _RoN, _I32, _END}};
 const SymbolicAddressSignature SASigTableSize = {
     SymbolicAddress::TableSize, _I32, _Infallible, 2, {_PTR, _I32, _END}};
+const SymbolicAddressSignature SASigFuncRef = {
+    SymbolicAddress::FuncRef, _RoN, _FailOnInvalidRef, 2, {_PTR, _I32, _END}};
 const SymbolicAddressSignature SASigPostBarrier = {
     SymbolicAddress::PostBarrier, _VOID, _Infallible, 2, {_PTR, _PTR, _END}};
 const SymbolicAddressSignature SASigPostBarrierFiltering = {
@@ -836,6 +838,9 @@ void* wasm::AddressOf(SymbolicAddress imm, ABIFunctionType* abiType) {
     case SymbolicAddress::TableSize:
       *abiType = Args_General2;
       return FuncCast(Instance::tableSize, *abiType);
+    case SymbolicAddress::FuncRef:
+      *abiType = Args_General2;
+      return FuncCast(Instance::funcRef, *abiType);
     case SymbolicAddress::PostBarrier:
       *abiType = Args_General2;
       return FuncCast(Instance::postBarrier, *abiType);
@@ -957,6 +962,7 @@ bool wasm::NeedsBuiltinThunk(SymbolicAddress sym) {
     case SymbolicAddress::TableInit:
     case SymbolicAddress::TableSet:
     case SymbolicAddress::TableSize:
+    case SymbolicAddress::FuncRef:
     case SymbolicAddress::PostBarrier:
     case SymbolicAddress::PostBarrierFiltering:
     case SymbolicAddress::StructNew:

js/src/wasm/WasmBuiltins.h

@@ -64,6 +64,7 @@ extern const SymbolicAddressSignature SASigTableGrow;
 extern const SymbolicAddressSignature SASigTableInit;
 extern const SymbolicAddressSignature SASigTableSet;
 extern const SymbolicAddressSignature SASigTableSize;
+extern const SymbolicAddressSignature SASigFuncRef;
 extern const SymbolicAddressSignature SASigPostBarrier;
 extern const SymbolicAddressSignature SASigPostBarrierFiltering;
 extern const SymbolicAddressSignature SASigStructNew;

js/src/wasm/WasmFrameIter.cpp

@@ -1376,6 +1376,8 @@ static const char* ThunkedNativeToDescription(SymbolicAddress func) {
       return "call to native table.set function";
     case SymbolicAddress::TableSize:
       return "call to native table.size function";
+    case SymbolicAddress::FuncRef:
+      return "call to native func.ref function";
     case SymbolicAddress::PostBarrier:
       return "call to native GC postbarrier (in wasm)";
     case SymbolicAddress::PostBarrierFiltering:

js/src/wasm/WasmInstance.cpp

@@ -998,6 +998,33 @@ void Instance::initElems(uint32_t tableIndex, const ElemSegment& seg,
   return table.length();
 }
 
+/* static */ void* Instance::funcRef(Instance* instance, uint32_t funcIndex) {
+  MOZ_ASSERT(SASigFuncRef.failureMode == FailureMode::FailOnInvalidRef);
+  JSContext* cx = TlsContext.get();
+
+  Tier tier = instance->code().bestTier();
+  const MetadataTier& metadataTier = instance->metadata(tier);
+  const FuncImportVector& funcImports = metadataTier.funcImports;
+
+  // If this is an import, we need to recover the original wrapper function to
+  // maintain referential equality between a re-exported function and
+  // 'ref.func'. The imported function object is stable across tiers, which is
+  // what we want.
+  if (funcIndex < funcImports.length()) {
+    FuncImportTls& import = instance->funcImportTls(funcImports[funcIndex]);
+    return AnyRef::fromJSObject(import.fun).forCompiledCode();
+  }
+
+  RootedFunction fun(cx);
+  RootedWasmInstanceObject instanceObj(cx, instance->object());
+  if (WasmInstanceObject::getExportedFunction(cx, instanceObj, funcIndex,
+                                              &fun)) {
+    return AnyRef::fromJSObject(fun).forCompiledCode();
+  }
+
+  return AnyRef::invalid().forCompiledCode();
+}
+
 /* static */ void Instance::postBarrier(Instance* instance,
                                         gc::Cell** location) {
   MOZ_ASSERT(SASigPostBarrier.failureMode == FailureMode::Infallible);

js/src/wasm/WasmInstance.h

@@ -216,6 +216,7 @@ class Instance {
   static int32_t tableInit(Instance* instance, uint32_t dstOffset,
                            uint32_t srcOffset, uint32_t len, uint32_t segIndex,
                            uint32_t tableIndex);
+  static void* funcRef(Instance* instance, uint32_t funcIndex);
   static void postBarrier(Instance* instance, gc::Cell** location);
   static void postBarrierFiltering(Instance* instance, gc::Cell** location);
   static void* structNew(Instance* instance, uint32_t typeIndex);

js/src/wasm/WasmIonCompile.cpp

@@ -3312,6 +3312,47 @@ static bool EmitTableSize(FunctionCompiler& f) {
 #endif  // ENABLE_WASM_REFTYPES
 
 #ifdef ENABLE_WASM_REFTYPES
+static bool EmitRefFunc(FunctionCompiler& f) {
+  uint32_t funcIndex;
+  if (!f.iter().readRefFunc(&funcIndex)) {
+    return false;
+  }
+
+  if (f.inDeadCode()) {
+    return true;
+  }
+
+  uint32_t lineOrBytecode = f.readCallSiteLineOrBytecode();
+
+  const SymbolicAddressSignature& callee = SASigFuncRef;
+  CallCompileState args;
+  if (!f.passInstance(callee.argTypes[0], &args)) {
+    return false;
+  }
+
+  MDefinition* funcIndexArg = f.constant(Int32Value(funcIndex), MIRType::Int32);
+  if (!funcIndexArg) {
+    return false;
+  }
+  if (!f.passArg(funcIndexArg, callee.argTypes[1], &args)) {
+    return false;
+  }
+
+  if (!f.finishCall(&args)) {
+    return false;
+  }
+
+  // The return value here is either null, denoting an error, or a short-lived
+  // pointer to a location containing a possibly-null ref.
+  MDefinition* ret;
+  if (!f.builtinInstanceMethodCall(callee, lineOrBytecode, args, &ret)) {
+    return false;
+  }
+
+  f.iter().setResult(ret);
+  return true;
+}
+
 static bool EmitRefNull(FunctionCompiler& f) {
   if (!f.iter().readRefNull()) {
     return false;
@@ -3788,6 +3829,8 @@ static bool EmitBodyExprs(FunctionCompiler& f) {
                              MCompare::Compare_RefOrNull));
 #endif
 #ifdef ENABLE_WASM_REFTYPES
+      case uint16_t(Op::RefFunc):
+        CHECK(EmitRefFunc(f));
       case uint16_t(Op::RefNull):
         CHECK(EmitRefNull(f));
       case uint16_t(Op::RefIsNull):

js/src/wasm/WasmOpIter.h

@@ -429,6 +429,7 @@ class MOZ_STACK_CLASS OpIter : private Policy {
   MOZ_MUST_USE bool readI64Const(int64_t* i64);
   MOZ_MUST_USE bool readF32Const(float* f32);
   MOZ_MUST_USE bool readF64Const(double* f64);
+  MOZ_MUST_USE bool readRefFunc(uint32_t* funcTypeIndex);
   MOZ_MUST_USE bool readRefNull();
   MOZ_MUST_USE bool readCall(uint32_t* calleeIndex, ValueVector* argValues);
   MOZ_MUST_USE bool readCallIndirect(uint32_t* funcTypeIndex,
@@ -1472,6 +1473,22 @@ inline bool OpIter<Policy>::readF64Const(double* f64) {
   return readFixedF64(f64) && push(ValType::F64);
 }
 
+template <typename Policy>
+inline bool OpIter<Policy>::readRefFunc(uint32_t* funcTypeIndex) {
+  MOZ_ASSERT(Classify(op_) == OpKind::RefFunc);
+
+  if (!readVarU32(funcTypeIndex)) {
+    return fail("unable to read function index");
+  }
+  if (*funcTypeIndex >= env_.funcTypes.length()) {
+    return fail("function index out of range");
+  }
+  if (!env_.validForRefFunc.getBit(*funcTypeIndex)) {
+    return fail("function index is not in an element segment");
+  }
+  return push(StackType(ValType::FuncRef));
+}
+
 template <typename Policy>
 inline bool OpIter<Policy>::readRefNull() {
   MOZ_ASSERT(Classify(op_) == OpKind::RefNull);