Bug 1731121 - wasm: Make TypeContext refcounted and shareable. r=lth

This commit makes TypeContext refcounted/shareable. This is so that
we can change RttValue's representation to be:
  1. RefPtr<TypeContext> typeContext;
  2. uint32_t typeIndex;
Effectively adding an explicit type context field to RttValue instead
of using an implicit global TypeContext, as it is now.

The big change here is that as TypeContext is now refcounted, it
needs to be heap-allocated, and cannot be a plain member of
ModuleEnvironment, as it is now. This adds some extra indirections.
If this is unacceptable, then I think we can find a workaround.

The next commit will implement the changes to RttValue.

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

js/src/wasm/AsmJS.cpp

@@ -1388,6 +1388,8 @@ class MOZ_STACK_CLASS ModuleValidatorShared {
   }
 
  protected:
+  [[nodiscard]] bool initModuleEnvironment() { return moduleEnv_.initTypes(0); }
+
   [[nodiscard]] bool addStandardLibraryMathInfo() {
     static constexpr struct {
       const char* name;
@@ -1898,24 +1900,24 @@ class MOZ_STACK_CLASS ModuleValidator : public ModuleValidatorShared {
  private:
   // Helpers:
   bool newSig(FuncType&& sig, uint32_t* sigIndex) {
-    if (moduleEnv_.types.length() >= MaxTypes) {
+    if (moduleEnv_.types->length() >= MaxTypes) {
       return failCurrentOffset("too many signatures");
     }
 
-    *sigIndex = moduleEnv_.types.length();
-    return moduleEnv_.types.append(std::move(sig)) &&
+    *sigIndex = moduleEnv_.types->length();
+    return moduleEnv_.types->append(std::move(sig)) &&
            moduleEnv_.typeIds.append(TypeIdDesc());
   }
   bool declareSig(FuncType&& sig, uint32_t* sigIndex) {
     SigSet::AddPtr p = sigSet_.lookupForAdd(sig);
     if (p) {
       *sigIndex = p->sigIndex();
-      MOZ_ASSERT(moduleEnv_.types.funcType(*sigIndex) == sig);
+      MOZ_ASSERT(moduleEnv_.types->funcType(*sigIndex) == sig);
       return true;
     }
 
     return newSig(std::move(sig), sigIndex) &&
-           sigSet_.add(p, HashableSig(*sigIndex, moduleEnv_.types));
+           sigSet_.add(p, HashableSig(*sigIndex, *moduleEnv_.types));
   }
 
  private:
@@ -1960,6 +1962,9 @@ class MOZ_STACK_CLASS ModuleValidator : public ModuleValidatorShared {
                              !parser_.pc_->sc()->hasExplicitUseStrict();
     asmJSMetadata_->source = do_AddRef(parser_.ss);
 
+    if (!initModuleEnvironment()) {
+      return false;
+    }
     return addStandardLibraryMathInfo();
   }
 
@@ -2058,7 +2063,7 @@ class MOZ_STACK_CLASS ModuleValidator : public ModuleValidatorShared {
       return false;
     }
 
-    return funcImportMap_.add(p, NamedSig(name, sigIndex, moduleEnv_.types),
+    return funcImportMap_.add(p, NamedSig(name, sigIndex, *moduleEnv_.types),
                               *importIndex);
   }
 
@@ -2089,15 +2094,15 @@ 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),
+          FuncDesc(&moduleEnv_.types->funcType(funcTypeIndex),
                    &moduleEnv_.typeIds[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),
+          FuncDesc(&moduleEnv_.types->funcType(funcTypeIndex),
                    &moduleEnv_.typeIds[funcTypeIndex], funcTypeIndex);
     }
     for (const Export& exp : moduleEnv_.exports) {
@@ -3993,7 +3998,7 @@ static bool CheckFunctionSignature(ModuleValidator<Unit>& m, ParseNode* usepn,
     return m.addFuncDef(name, usepn->pn_pos.begin, std::move(sig), func);
   }
 
-  const FuncType& existingSig = m.env().types.funcType(existing->sigIndex());
+  const FuncType& existingSig = m.env().types->funcType(existing->sigIndex());
 
   if (!CheckSignatureAgainstExisting(m, usepn, sig, existingSig)) {
     return false;
@@ -4067,7 +4072,7 @@ static bool CheckFuncPtrTableAgainstExisting(ModuleValidator<Unit>& m,
     }
 
     if (!CheckSignatureAgainstExisting(
-            m, usepn, sig, m.env().types.funcType(table.sigIndex()))) {
+            m, usepn, sig, m.env().types->funcType(table.sigIndex()))) {
       return false;
     }
 
@@ -6231,7 +6236,7 @@ static bool CheckFuncPtrTable(ModuleValidator<Unit>& m, ParseNode* decl) {
           elem, "function-pointer table's elements must be names of functions");
     }
 
-    const FuncType& funcSig = m.env().types.funcType(func->sigIndex());
+    const FuncType& funcSig = m.env().types->funcType(func->sigIndex());
     if (sig) {
       if (*sig != funcSig) {
         return m.fail(elem, "all functions in table must have same signature");

js/src/wasm/WasmBaselineCompile.cpp

@@ -4337,7 +4337,7 @@ bool BaseCompiler::emitCallIndirect() {
 
   sync();
 
-  const FuncType& funcType = moduleEnv_.types[funcTypeIndex].funcType();
+  const FuncType& funcType = (*moduleEnv_.types)[funcTypeIndex].funcType();
 
   // Stack: ... arg1 .. argn callee
 
@@ -5942,7 +5942,7 @@ bool BaseCompiler::emitStructNewWithRtt() {
     return true;
   }
 
-  const StructType& structType = moduleEnv_.types[typeIndex].structType();
+  const StructType& structType = (*moduleEnv_.types)[typeIndex].structType();
 
   // Allocate zeroed storage.  The parameter to StructNew is a rtt value that is
   // guaranteed to be at the top of the stack by validation.
@@ -6034,7 +6034,7 @@ bool BaseCompiler::emitStructGet(FieldExtension extension) {
     return true;
   }
 
-  const StructType& structType = moduleEnv_.types[typeIndex].structType();
+  const StructType& structType = (*moduleEnv_.types)[typeIndex].structType();
 
   RegRef rp = popRef();
 
@@ -6086,7 +6086,7 @@ bool BaseCompiler::emitStructSet() {
     return true;
   }
 
-  const StructType& structType = moduleEnv_.types[typeIndex].structType();
+  const StructType& structType = (*moduleEnv_.types)[typeIndex].structType();
   const StructField& structField = structType.fields_[fieldIndex];
 
   // Reserve this register early if we will need it so that it is not taken by
@@ -6153,7 +6153,7 @@ bool BaseCompiler::emitArrayNewWithRtt() {
     return true;
   }
 
-  const ArrayType& arrayType = moduleEnv_.types[typeIndex].arrayType();
+  const ArrayType& arrayType = (*moduleEnv_.types)[typeIndex].arrayType();
 
   // Allocate zeroed storage.  The parameter to ArrayNew is a rtt value and
   // length that are guaranteed to be at the top of the stack by validation.
@@ -6238,7 +6238,7 @@ bool BaseCompiler::emitArrayGet(FieldExtension extension) {
     return true;
   }
 
-  const ArrayType& arrayType = moduleEnv_.types[typeIndex].arrayType();
+  const ArrayType& arrayType = (*moduleEnv_.types)[typeIndex].arrayType();
 
   RegI32 index = popI32();
   RegRef rp = popRef();
@@ -6286,7 +6286,7 @@ bool BaseCompiler::emitArraySet() {
     return true;
   }
 
-  const ArrayType& arrayType = moduleEnv_.types[typeIndex].arrayType();
+  const ArrayType& arrayType = (*moduleEnv_.types)[typeIndex].arrayType();
 
   // Reserve this register early if we will need it so that it is not taken by
   // any register used in this function.
@@ -9680,7 +9680,7 @@ bool js::wasm::BaselineCompileFunctions(const ModuleEnvironment& moduleEnv,
     if (!locals.appendAll(moduleEnv.funcs[func.index].type->args())) {
       return false;
     }
-    if (!DecodeLocalEntries(d, moduleEnv.types, moduleEnv.features, &locals)) {
+    if (!DecodeLocalEntries(d, *moduleEnv.types, moduleEnv.features, &locals)) {
       return false;
     }
 

js/src/wasm/WasmContext.cpp

@@ -30,8 +30,8 @@ bool wasm::Context::ensureTypeContext(JSContext* cx) {
   }
   // Default options should be OK here
   FeatureOptions options;
-  typeContext = js::MakeUnique<TypeContext>(FeatureArgs::build(cx, options),
-                                            TypeDefVector());
+  typeContext =
+      js_new<TypeContext>(FeatureArgs::build(cx, options), TypeDefVector());
   return !!typeContext;
 }
 

js/src/wasm/WasmContext.h

@@ -48,7 +48,7 @@ class Context {
   [[nodiscard]] bool ensureTypeContext(JSContext* cx);
 
   // The global type context.
-  UniquePtr<TypeContext> typeContext;
+  MutableTypeContext typeContext;
 
   // about:memory reporting
 

js/src/wasm/WasmGenerator.cpp

@@ -275,9 +275,9 @@ bool ModuleGenerator::init(Metadata* maybeAsmJSMetadata) {
     // Copy type definitions to metadata that are required at runtime,
     // allocating global data so that codegen can find the type id's at
     // runtime.
-    for (uint32_t typeIndex = 0; typeIndex < moduleEnv_->types.length();
+    for (uint32_t typeIndex = 0; typeIndex < moduleEnv_->types->length();
          typeIndex++) {
-      const TypeDef& typeDef = moduleEnv_->types[typeIndex];
+      const TypeDef& typeDef = (*moduleEnv_->types)[typeIndex];
       TypeIdDesc& typeId = moduleEnv_->typeIds[typeIndex];
 
       if (TypeIdDesc::isGlobal(typeDef)) {
@@ -308,12 +308,12 @@ bool ModuleGenerator::init(Metadata* maybeAsmJSMetadata) {
     if (moduleEnv_->functionReferencesEnabled()) {
       // Do a linear pass to create a map from src index to dest index.
       RenumberVector renumbering;
-      if (!renumbering.reserve(moduleEnv_->types.length())) {
+      if (!renumbering.reserve(moduleEnv_->types->length())) {
         return false;
       }
       for (uint32_t srcIndex = 0, destIndex = 0;
-           srcIndex < moduleEnv_->types.length(); srcIndex++) {
-        const TypeDef& typeDef = moduleEnv_->types[srcIndex];
+           srcIndex < moduleEnv_->types->length(); srcIndex++) {
+        const TypeDef& typeDef = (*moduleEnv_->types)[srcIndex];
         if (!TypeIdDesc::isGlobal(typeDef)) {
           renumbering.infallibleAppend(UINT32_MAX);
           continue;

js/src/wasm/WasmIntrinsic.cpp

@@ -120,25 +120,30 @@ bool wasm::CompileIntrinsicModule(JSContext* cx,
   }
   moduleEnv.memory = Some(MemoryDesc(Limits(0, Nothing(), sharedMemory)));
 
+  // Initialize the type section
+  if (!moduleEnv.initTypes(ops.size())) {
+    return false;
+  }
+
   // Add (type (func (params ...))) for each intrinsic. The function types will
   // be deduplicated by the runtime
-  for (const IntrinsicOp& op : ops) {
+  for (uint32_t funcIndex = 0; funcIndex < ops.size(); funcIndex++) {
+    const IntrinsicOp& op = ops[funcIndex];
     const Intrinsic& intrinsic = Intrinsic::getFromOp(op);
 
     FuncType type;
-    if (!intrinsic.funcType(&type) ||
-        !moduleEnv.types.append(TypeDef(std::move(type))) ||
-        !moduleEnv.typeIds.append(TypeIdDesc())) {
+    if (!intrinsic.funcType(&type)) {
       ReportOutOfMemory(cx);
       return false;
     }
+    (*moduleEnv.types)[funcIndex] = TypeDef(std::move(type));
   }
 
   // Add (func (type $i)) declarations. Do this after all types have been added
   // as the function declaration metadata uses pointers into the type vectors
   // that must be stable.
   for (uint32_t funcIndex = 0; funcIndex < ops.size(); funcIndex++) {
-    FuncDesc decl(&moduleEnv.types[funcIndex].funcType(),
+    FuncDesc decl(&(*moduleEnv.types)[funcIndex].funcType(),
                   &moduleEnv.typeIds[funcIndex], funcIndex);
     if (!moduleEnv.funcs.append(decl)) {
       ReportOutOfMemory(cx);

js/src/wasm/WasmIonCompile.cpp

@@ -1605,7 +1605,7 @@ class FunctionCompiler {
       return true;
     }
 
-    const FuncType& funcType = moduleEnv_.types[funcTypeIndex].funcType();
+    const FuncType& funcType = (*moduleEnv_.types)[funcTypeIndex].funcType();
     const TypeIdDesc& funcTypeId = moduleEnv_.typeIds[funcTypeIndex];
 
     CalleeDesc callee;
@@ -2740,7 +2740,7 @@ static bool EmitCallIndirect(FunctionCompiler& f, bool oldStyle) {
     return true;
   }
 
-  const FuncType& funcType = f.moduleEnv().types[funcTypeIndex].funcType();
+  const FuncType& funcType = (*f.moduleEnv().types)[funcTypeIndex].funcType();
 
   CallCompileState call;
   if (!EmitCallArgs(f, funcType, args, &call)) {
@@ -5780,7 +5780,7 @@ bool wasm::IonCompileFunctions(const ModuleEnvironment& moduleEnv,
     if (!locals.appendAll(funcType.args())) {
       return false;
     }
-    if (!DecodeLocalEntries(d, moduleEnv.types, moduleEnv.features, &locals)) {
+    if (!DecodeLocalEntries(d, *moduleEnv.types, moduleEnv.features, &locals)) {
       return false;
     }
 

js/src/wasm/WasmJS.cpp

@@ -4195,21 +4195,22 @@ JSFunction* WasmFunctionCreate(JSContext* cx, HandleFunction func,
                                   OptimizedBackend::Ion, DebugEnabled::False);
   compilerEnv.computeParameters();
 
-  // Add the import for the function
-  FuncType funcType = FuncType(std::move(params), std::move(results));
-  TypeDef funcTypeDef = TypeDef(std::move(funcType));
-  if (!moduleEnv.types.append(std::move(funcTypeDef))) {
-    return nullptr;
-  }
-  if (!moduleEnv.typeIds.resize(1)) {
+  // Initialize the type section
+  if (!moduleEnv.initTypes(1)) {
     return nullptr;
   }
+  FuncType funcType = FuncType(std::move(params), std::move(results));
+  (*moduleEnv.types)[0] = TypeDef(std::move(funcType));
+
+  // Add an (import (func ...))
   FuncDesc funcDesc =
-      FuncDesc(&moduleEnv.types[0].funcType(), &moduleEnv.typeIds[0], 0);
+      FuncDesc(&(*moduleEnv.types)[0].funcType(), &moduleEnv.typeIds[0], 0);
   if (!moduleEnv.funcs.append(funcDesc) ||
       !moduleEnv.funcImportGlobalDataOffsets.resize(1)) {
     return nullptr;
   }
+
+  // Add an (export (func 0))
   moduleEnv.declareFuncExported(0, false, false);
 
   // We will be looking up and using the function in the future by index so the