Bug 1790626 - wasm: Move 'immediate' type representation to FuncType …

…and only use with call_indirect. r=yury

TypeIdDesc describes how to load the 'type id' for any type. Right now the
'type id' is an immediate integer for small function types, a *FuncType for
larger function types, and *RttValue for everything else.

The immediate integer case is an optimization for signature checks in
call_indirect. We can simplify our code by adding an 'immediateType_'
field to `FuncType` which stores an alternate representation of the
function type (if any) that can be used with call_indirect. Then
`TypeIdDesc`/`TypeDefWithId` are not needed during instantiation
anymore and we can just pass `TypeDefVector` from compilation to
runtime.

As a drive-by fix, some code for setting up the global data are of
Instance was simplified to allocate whole contiguous areas, instead
of looping, in some cases.

Some uses of TypeIdDesc still remain after this commit, but they are
only in codegen for figuring out how to emit signature checks. The
class is renamed and simplified to CallIndirectId to represent this.

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

js/src/jit/CodeGenerator.cpp

@@ -12814,7 +12814,7 @@ static bool CreateStackMapFromLSafepoint(LSafepoint& safepoint,
 }
 
 bool CodeGenerator::generateWasm(
-    wasm::TypeIdDesc funcTypeId, wasm::BytecodeOffset trapOffset,
+    wasm::CallIndirectId callIndirectId, wasm::BytecodeOffset trapOffset,
     const wasm::ArgTypeVector& argTypes, const RegisterOffsets& trapExitLayout,
     size_t trapExitLayoutNumWords, wasm::FuncOffsets* offsets,
     wasm::StackMaps* stackMaps, wasm::Decoder* decoder) {
@@ -12824,7 +12824,8 @@ bool CodeGenerator::generateWasm(
 
   size_t nInboundStackArgBytes = StackArgAreaSizeUnaligned(argTypes);
 
-  wasm::GenerateFunctionPrologue(masm, funcTypeId, mozilla::Nothing(), offsets);
+  wasm::GenerateFunctionPrologue(masm, callIndirectId, mozilla::Nothing(),
+                                 offsets);
 
   MOZ_ASSERT(masm.framePushed() == 0);
 

js/src/jit/CodeGenerator.h

@@ -106,7 +106,7 @@ class CodeGenerator final : public CodeGeneratorSpecific {
 
   [[nodiscard]] bool generate();
   [[nodiscard]] bool generateWasm(
-      wasm::TypeIdDesc funcTypeId, wasm::BytecodeOffset trapOffset,
+      wasm::CallIndirectId callIndirectId, wasm::BytecodeOffset trapOffset,
       const wasm::ArgTypeVector& argTys, const RegisterOffsets& trapExitLayout,
       size_t trapExitLayoutNumWords, wasm::FuncOffsets* offsets,
       wasm::StackMaps* stackMaps, wasm::Decoder* decoder);

js/src/jit/MacroAssembler.cpp

@@ -3958,15 +3958,15 @@ void MacroAssembler::wasmCallIndirect(const wasm::CallSiteDesc& desc,
 
   // Write the functype-id into the ABI functype-id register.
 
-  const wasm::TypeIdDesc funcTypeId = callee.wasmTableSigId();
-  switch (funcTypeId.kind()) {
-    case wasm::TypeIdDescKind::Global:
-      loadWasmGlobalPtr(funcTypeId.globalDataOffset(), WasmTableCallSigReg);
+  const wasm::CallIndirectId callIndirectId = callee.wasmTableSigId();
+  switch (callIndirectId.kind()) {
+    case wasm::CallIndirectIdKind::Global:
+      loadWasmGlobalPtr(callIndirectId.globalDataOffset(), WasmTableCallSigReg);
       break;
-    case wasm::TypeIdDescKind::Immediate:
-      move32(Imm32(funcTypeId.immediate()), WasmTableCallSigReg);
+    case wasm::CallIndirectIdKind::Immediate:
+      move32(Imm32(callIndirectId.immediate()), WasmTableCallSigReg);
       break;
-    case wasm::TypeIdDescKind::None:
+    case wasm::CallIndirectIdKind::None:
       break;
   }
 

js/src/wasm/AsmJS.cpp

@@ -1941,8 +1941,7 @@ class MOZ_STACK_CLASS ModuleValidator : public ModuleValidatorShared {
     }
 
     *sigIndex = moduleEnv_.types->length();
-    return moduleEnv_.types->append(std::move(sig)) &&
-           moduleEnv_.typeIds.append(TypeIdDesc());
+    return moduleEnv_.types->append(std::move(sig));
   }
   bool declareSig(FuncType&& sig, uint32_t* sigIndex) {
     SigSet::AddPtr p = sigSet_.lookupForAdd(sig);
@@ -2130,16 +2129,14 @@ class MOZ_STACK_CLASS ModuleValidator : public ModuleValidatorShared {
       uint32_t funcTypeIndex = r.front().key().sigIndex();
       MOZ_ASSERT(!moduleEnv_.funcs[funcIndex].type);
       moduleEnv_.funcs[funcIndex] =
-          FuncDesc(&moduleEnv_.types->funcType(funcTypeIndex),
-                   &moduleEnv_.typeIds[funcTypeIndex], funcTypeIndex);
+          FuncDesc(&moduleEnv_.types->funcType(funcTypeIndex), funcTypeIndex);
     }
     for (const Func& func : funcDefs_) {
       uint32_t funcIndex = funcImportMap_.count() + func.funcDefIndex();
       uint32_t funcTypeIndex = func.sigIndex();
       MOZ_ASSERT(!moduleEnv_.funcs[funcIndex].type);
       moduleEnv_.funcs[funcIndex] =
-          FuncDesc(&moduleEnv_.types->funcType(funcTypeIndex),
-                   &moduleEnv_.typeIds[funcTypeIndex], funcTypeIndex);
+          FuncDesc(&moduleEnv_.types->funcType(funcTypeIndex), funcTypeIndex);
     }
     for (const Export& exp : moduleEnv_.exports) {
       if (exp.kind() != DefinitionKind::Function) {
@@ -2150,10 +2147,7 @@ class MOZ_STACK_CLASS ModuleValidator : public ModuleValidatorShared {
                                      /* canRefFunc */ false);
     }
 
-    if (!moduleEnv_.funcImportGlobalDataOffsets.resize(
-            funcImportMap_.count())) {
-      return nullptr;
-    }
+    moduleEnv_.numFuncImports = funcImportMap_.count();
 
     MOZ_ASSERT(asmJSMetadata_->asmJSFuncNames.empty());
     if (!asmJSMetadata_->asmJSFuncNames.resize(funcImportMap_.count())) {
@@ -6185,10 +6179,11 @@ static bool CheckFunction(ModuleValidator<Unit>& m) {
     }
   }
 
+  FuncType sig(std::move(args), std::move(results));
+
   ModuleValidatorShared::Func* func = nullptr;
-  if (!CheckFunctionSignature(m, funNode,
-                              FuncType(std::move(args), std::move(results)),
-                              FunctionName(funNode), &func)) {
+  if (!CheckFunctionSignature(m, funNode, std::move(sig), FunctionName(funNode),
+                              &func)) {
     return false;
   }
 

js/src/wasm/WasmBCClass-inl.h

@@ -29,10 +29,6 @@ const FuncType& BaseCompiler::funcType() const {
   return *moduleEnv_.funcs[func_.index].type;
 }
 
-const TypeIdDesc& BaseCompiler::funcTypeId() const {
-  return *moduleEnv_.funcs[func_.index].typeId;
-}
-
 bool BaseCompiler::usesMemory() const { return moduleEnv_.usesMemory(); }
 
 bool BaseCompiler::usesSharedMemory() const {

js/src/wasm/WasmBCClass.h

@@ -324,7 +324,6 @@ struct BaseCompiler final {
   // WasmBaselineObject-inl.h.
 
   inline const FuncType& funcType() const;
-  inline const TypeIdDesc& funcTypeId() const;
   inline bool usesMemory() const;
   inline bool usesSharedMemory() const;
   inline bool isMem32() const;

js/src/wasm/WasmBaselineCompile.cpp

@@ -378,7 +378,7 @@ bool BaseCompiler::beginFunction() {
   }
 
   GenerateFunctionPrologue(
-      masm, *moduleEnv_.funcs[func_.index].typeId,
+      masm, CallIndirectId::forFunc(moduleEnv_, func_.index),
       compilerEnv_.mode() == CompileMode::Tier1 ? Some(func_.index) : Nothing(),
       &offsets_);
 
@@ -1574,15 +1574,16 @@ bool BaseCompiler::callIndirect(uint32_t funcTypeIndex, uint32_t tableIndex,
                                 const Stk& indexVal, const FunctionCall& call,
                                 CodeOffset* fastCallOffset,
                                 CodeOffset* slowCallOffset) {
-  const TypeIdDesc& funcTypeId = moduleEnv_.typeIds[funcTypeIndex];
-  MOZ_ASSERT(funcTypeId.kind() != TypeIdDescKind::None);
+  CallIndirectId callIndirectId =
+      CallIndirectId::forFuncType(moduleEnv_, funcTypeIndex);
+  MOZ_ASSERT(callIndirectId.kind() != CallIndirectIdKind::None);
 
   const TableDesc& table = moduleEnv_.tables[tableIndex];
 
   loadI32(indexVal, RegI32(WasmTableCallIndexReg));
 
   CallSiteDesc desc(bytecodeOffset(), CallSiteDesc::Indirect);
-  CalleeDesc callee = CalleeDesc::wasmTable(table, funcTypeId);
+  CalleeDesc callee = CalleeDesc::wasmTable(table, callIndirectId);
   OutOfLineCode* oob = addOutOfLineCode(
       new (alloc_) OutOfLineAbortingTrap(Trap::OutOfBounds, bytecodeOffset()));
   if (!oob) {
@@ -4706,7 +4707,7 @@ bool BaseCompiler::emitCall() {
 
   CodeOffset raOffset;
   if (import) {
-    raOffset = callImport(moduleEnv_.funcImportGlobalDataOffsets[funcIndex],
+    raOffset = callImport(moduleEnv_.offsetOfFuncImportInstanceData(funcIndex),
                           baselineCall);
   } else {
     raOffset = callDefinition(funcIndex, baselineCall);
@@ -6305,12 +6306,11 @@ void BaseCompiler::emitBarrieredClear(RegPtr valueAddr) {
 #ifdef ENABLE_WASM_GC
 
 void BaseCompiler::emitGcCanon(uint32_t typeIndex) {
-  const TypeIdDesc& typeId = moduleEnv_.typeIds[typeIndex];
   RegRef rp = needRef();
 #  ifndef RABALDR_PIN_INSTANCE
   fr.loadInstancePtr(InstanceReg);
 #  endif
-  masm.loadWasmGlobalPtr(typeId.globalDataOffset(), rp);
+  masm.loadWasmGlobalPtr(moduleEnv_.offsetOfTypeId(typeIndex), rp);
   pushRef(rp);
 }
 

js/src/wasm/WasmCode.cpp

@@ -705,7 +705,6 @@ void LazyStubTier::addSizeOfMisc(MallocSizeOf mallocSizeOf, size_t* code,
 
 size_t Metadata::sizeOfExcludingThis(MallocSizeOf mallocSizeOf) const {
   return SizeOfVectorExcludingThis(types, mallocSizeOf) +
-         typeIds.sizeOfExcludingThis(mallocSizeOf) +
          globals.sizeOfExcludingThis(mallocSizeOf) +
          tables.sizeOfExcludingThis(mallocSizeOf) +
          tags.sizeOfExcludingThis(mallocSizeOf) +

js/src/wasm/WasmCode.h

@@ -375,7 +375,7 @@ using ModuleHash = uint8_t[8];
 
 struct Metadata : public ShareableBase<Metadata>, public MetadataCacheablePod {
   TypeDefVector types;
-  TypeIdDescVector typeIds;
+  uint32_t typeIdsOffsetStart;
   GlobalDescVector globals;
   TableDescVector tables;
   TagDescVector tags;

js/src/wasm/WasmCodegenTypes.cpp

@@ -21,6 +21,7 @@
 #include "wasm/WasmExprType.h"
 #include "wasm/WasmStubs.h"
 #include "wasm/WasmTypeDef.h"
+#include "wasm/WasmValidate.h"
 #include "wasm/WasmValue.h"
 
 using mozilla::MakeEnumeratedRange;
@@ -153,6 +154,28 @@ const CodeRange* wasm::LookupInSorted(const CodeRangeVector& codeRanges,
   return &codeRanges[match];
 }
 
+CallIndirectId CallIndirectId::forFunc(const ModuleEnvironment& moduleEnv,
+                                       uint32_t funcIndex) {
+  return CallIndirectId::forFuncType(moduleEnv,
+                                     moduleEnv.funcs[funcIndex].typeIndex);
+}
+
+CallIndirectId CallIndirectId::forFuncType(const ModuleEnvironment& moduleEnv,
+                                           uint32_t funcTypeIndex) {
+  // asm.js tables are homogenous and don't require a signature check
+  if (moduleEnv.isAsmJS()) {
+    return CallIndirectId();
+  }
+
+  const FuncType& funcType = moduleEnv.types->type(funcTypeIndex).funcType();
+  if (funcType.hasImmediateTypeId()) {
+    return CallIndirectId(CallIndirectIdKind::Immediate,
+                          funcType.immediateTypeId());
+  }
+  return CallIndirectId(CallIndirectIdKind::Global,
+                        moduleEnv.offsetOfTypeId(funcTypeIndex));
+}
+
 CalleeDesc CalleeDesc::function(uint32_t funcIndex) {
   CalleeDesc c;
   c.which_ = Func;
@@ -165,13 +188,14 @@ CalleeDesc CalleeDesc::import(uint32_t globalDataOffset) {
   c.u.import.globalDataOffset_ = globalDataOffset;
   return c;
 }
-CalleeDesc CalleeDesc::wasmTable(const TableDesc& desc, TypeIdDesc funcTypeId) {
+CalleeDesc CalleeDesc::wasmTable(const TableDesc& desc,
+                                 CallIndirectId callIndirectId) {
   CalleeDesc c;
   c.which_ = WasmTable;
   c.u.table.globalDataOffset_ = desc.globalDataOffset;
   c.u.table.minLength_ = desc.initialLength;
   c.u.table.maxLength_ = desc.maximumLength;
-  c.u.table.funcTypeId_ = funcTypeId;
+  c.u.table.callIndirectId_ = callIndirectId;
   return c;
 }
 CalleeDesc CalleeDesc::asmJSTable(const TableDesc& desc) {

js/src/wasm/WasmCodegenTypes.h

@@ -44,6 +44,7 @@ namespace wasm {
 
 using mozilla::EnumeratedArray;
 
+struct ModuleEnvironment;
 struct TableDesc;
 struct V128;
 
@@ -614,6 +615,42 @@ struct TryNote {
 WASM_DECLARE_CACHEABLE_POD(TryNote);
 WASM_DECLARE_POD_VECTOR(TryNote, TryNoteVector)
 
+// CallIndirectId describes how to compile a call_indirect and matching
+// signature check in the function prologue for a given function type.
+
+enum class CallIndirectIdKind { None, Immediate, Global };
+
+class CallIndirectId {
+  CallIndirectIdKind kind_;
+  size_t bits_;
+
+  CallIndirectId(CallIndirectIdKind kind, size_t bits)
+      : kind_(kind), bits_(bits) {}
+
+ public:
+  CallIndirectId() : kind_(CallIndirectIdKind::None), bits_(0) {}
+
+  // Get the CallIndirectId for a function in a specific module.
+  static CallIndirectId forFunc(const ModuleEnvironment& moduleEnv,
+                                uint32_t funcIndex);
+
+  // Get the CallIndirectId for a function type in a specific module.
+  static CallIndirectId forFuncType(const ModuleEnvironment& moduleEnv,
+                                    uint32_t funcTypeIndex);
+
+  CallIndirectIdKind kind() const { return kind_; }
+  bool isGlobal() const { return kind_ == CallIndirectIdKind::Global; }
+
+  uint32_t immediate() const {
+    MOZ_ASSERT(kind_ == CallIndirectIdKind::Immediate);
+    return bits_;
+  }
+  uint32_t globalDataOffset() const {
+    MOZ_ASSERT(kind_ == CallIndirectIdKind::Global);
+    return bits_;
+  }
+};
+
 // CalleeDesc describes how to compile one of the variety of asm.js/wasm calls.
 // This is hoisted into WasmCodegenTypes.h for sharing between Ion and Baseline.
 
@@ -658,20 +695,17 @@ class CalleeDesc {
       uint32_t globalDataOffset_;
       uint32_t minLength_;
       Maybe<uint32_t> maxLength_;
-      TypeIdDesc funcTypeId_;
+      CallIndirectId callIndirectId_;
     } table;
     SymbolicAddress builtin_;
   } u;
 
-  WASM_CHECK_CACHEABLE_POD(which_, u.funcIndex_, u.import.globalDataOffset_,
-                           u.table.globalDataOffset_, u.table.minLength_,
-                           u.table.maxLength_, u.table.funcTypeId_, u.builtin_);
-
  public:
   CalleeDesc() = default;
   static CalleeDesc function(uint32_t funcIndex);
   static CalleeDesc import(uint32_t globalDataOffset);
-  static CalleeDesc wasmTable(const TableDesc& desc, TypeIdDesc funcTypeId);
+  static CalleeDesc wasmTable(const TableDesc& desc,
+                              CallIndirectId callIndirectId);
   static CalleeDesc asmJSTable(const TableDesc& desc);
   static CalleeDesc builtin(SymbolicAddress callee);
   static CalleeDesc builtinInstanceMethod(SymbolicAddress callee);
@@ -694,9 +728,9 @@ class CalleeDesc {
     MOZ_ASSERT(isTable());
     return u.table.globalDataOffset_ + offsetof(TableInstanceData, elements);
   }
-  TypeIdDesc wasmTableSigId() const {
+  CallIndirectId wasmTableSigId() const {
     MOZ_ASSERT(which_ == WasmTable);
-    return u.table.funcTypeId_;
+    return u.table.callIndirectId_;
   }
   uint32_t wasmTableMinLength() const {
     MOZ_ASSERT(which_ == WasmTable);
@@ -713,8 +747,6 @@ class CalleeDesc {
   bool isFuncRef() const { return which_ == FuncRef; }
 };
 
-WASM_DECLARE_CACHEABLE_POD(CalleeDesc);
-
 }  // namespace wasm
 }  // namespace js
 

js/src/wasm/WasmCompile.cpp

@@ -674,7 +674,7 @@ static bool DecodeCodeSection(const ModuleEnvironment& env, DecoderT& d,
   }
 
   for (uint32_t funcDefIndex = 0; funcDefIndex < numFuncDefs; funcDefIndex++) {
-    if (!DecodeFunctionBody(d, mg, env.numFuncImports() + funcDefIndex)) {
+    if (!DecodeFunctionBody(d, mg, env.numFuncImports + funcDefIndex)) {
       return false;
     }
   }

js/src/wasm/WasmConstants.h

@@ -112,8 +112,6 @@ static constexpr TypeCode AbstractReferenceTypeCode = TypeCode::ExternRef;
 
 static constexpr TypeCode AbstractReferenceTypeIndexCode = TypeCode::Ref;
 
-enum class TypeIdDescKind { None, Immediate, Global };
-
 // A wasm::Trap represents a wasm-defined trap that can occur during execution
 // which triggers a WebAssembly.RuntimeError. Generated code may jump to a Trap
 // symbolically, passing the bytecode offset to report as the trap offset. The