Bug 1847017 part 5 - Fix check for metadata hook for code shared across realms. r=jonco

In `MacroAssembler::checkAllocatorState` we were checking if the realm has a
metadata hook, but that doesn't work if code is shared across realms.

This patch adds a counter to the zone so that we can check if any realm has a
metadata hook. We now also only check this for object allocations so that string
and bigint allocations can still use the fast path.

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

js/src/gc/Zone.cpp

@@ -192,6 +192,8 @@ JS::Zone::Zone(JSRuntime* rt, Kind kind)
 Zone::~Zone() {
   MOZ_ASSERT_IF(regExps_.ref(), regExps().empty());
 
+  MOZ_ASSERT(numRealmsWithAllocMetadataBuilder_ == 0);
+
   DebugAPI::deleteDebugScriptMap(debugScriptMap);
   js_delete(finalizationObservers_.ref().release());
 

js/src/gc/Zone.h

@@ -287,6 +287,10 @@ class Zone : public js::ZoneAllocator, public js::gc::GraphNodeBase<JS::Zone> {
 
   js::MainThreadOrGCTaskOrIonCompileData<js::jit::JitZone*> jitZone_;
 
+  // Number of realms in this zone that have a non-null object allocation
+  // metadata builder.
+  js::MainThreadOrIonCompileData<size_t> numRealmsWithAllocMetadataBuilder_{0};
+
   // Last time at which JIT code was discarded for this zone. This is only set
   // when JitScripts and Baseline code are discarded as well.
   js::MainThreadData<mozilla::TimeStamp> lastDiscardedCodeTime_;
@@ -441,6 +445,17 @@ class Zone : public js::ZoneAllocator, public js::gc::GraphNodeBase<JS::Zone> {
 
   bool ensureJitZoneExists(JSContext* cx) { return !!getJitZone(cx); }
 
+  void incNumRealmsWithAllocMetadataBuilder() {
+    numRealmsWithAllocMetadataBuilder_++;
+  }
+  void decNumRealmsWithAllocMetadataBuilder() {
+    MOZ_ASSERT(numRealmsWithAllocMetadataBuilder_ > 0);
+    numRealmsWithAllocMetadataBuilder_--;
+  }
+  bool hasRealmWithAllocMetadataBuilder() const {
+    return numRealmsWithAllocMetadataBuilder_ > 0;
+  }
+
   void prepareForCompacting();
 
   void traceRootsInMajorGC(JSTracer* trc);

js/src/jit-test/tests/basic/metadata-hook-regexp-result.js

@@ -0,0 +1,20 @@
+var g = newGlobal({sameCompartmentAs: this});
+g.evaluate(`enableShellAllocationMetadataBuilder()`);
+
+function f() {
+    // Ensure a match stub is created for the zone.
+    var re = /abc.+/;
+    for (var i = 0; i < 100; i++) {
+        assertEq(re.exec("..abcd").index, 2);
+    }
+    // Allocated match result objects in the realm with the metadata hook
+    // must have metadata.
+    g.evaluate(`
+    var re = /abc.+/;
+    for (var i = 0; i < 100; i++) {
+        var obj = re.exec("..abcd");
+        assertEq(getAllocationMetadata(obj).stack.length > 0, true);
+    }
+    `)
+}
+f();

js/src/jit/CompileWrappers.cpp

@@ -200,8 +200,8 @@ const uint32_t* CompileRealm::addressOfGlobalWriteBarriered() {
   return &realm()->globalWriteBarriered;
 }
 
-bool CompileRealm::hasAllocationMetadataBuilder() {
-  return realm()->hasAllocationMetadataBuilder();
+bool CompileZone::hasRealmWithAllocMetadataBuilder() {
+  return zone()->hasRealmWithAllocMetadataBuilder();
 }
 
 JitCompileOptions::JitCompileOptions()

js/src/jit/CompileWrappers.h

@@ -123,6 +123,8 @@ class CompileZone {
 
   gc::AllocSite* catchAllAllocSite(JS::TraceKind traceKind,
                                    gc::CatchAllAllocSite siteKind);
+
+  bool hasRealmWithAllocMetadataBuilder();
 };
 
 class CompileRealm {
@@ -141,8 +143,6 @@ class CompileRealm {
 
   const GlobalObject* maybeGlobal();
   const uint32_t* addressOfGlobalWriteBarriered();
-
-  bool hasAllocationMetadataBuilder();
 };
 
 class JitCompileOptions {

js/src/jit/MacroAssembler.cpp

@@ -264,7 +264,8 @@ template void MacroAssembler::loadFromTypedBigIntArray(Scalar::Type arrayType,
 
 // Inlined version of gc::CheckAllocatorState that checks the bare essentials
 // and bails for anything that cannot be handled with our jit allocators.
-void MacroAssembler::checkAllocatorState(Label* fail) {
+void MacroAssembler::checkAllocatorState(Register temp, gc::AllocKind allocKind,
+                                         Label* fail) {
   // Don't execute the inline path if GC probes are built in.
 #ifdef JS_GC_PROBES
   jump(fail);
@@ -277,11 +278,16 @@ void MacroAssembler::checkAllocatorState(Label* fail) {
            fail);
 #endif
 
-  // Don't execute the inline path if the realm has an object metadata callback,
-  // as the metadata to use for the object may vary between executions of the
-  // op.
-  if (realm()->hasAllocationMetadataBuilder()) {
-    jump(fail);
+  // If the zone has a realm with an object allocation metadata hook, emit a
+  // guard for this. Note that IC stubs and some other trampolines can be shared
+  // across realms, so we don't bake in a realm pointer.
+  if (gc::IsObjectAllocKind(allocKind) &&
+      realm()->zone()->hasRealmWithAllocMetadataBuilder()) {
+    loadJSContext(temp);
+    loadPtr(Address(temp, JSContext::offsetOfRealm()), temp);
+    branchPtr(Assembler::NotEqual,
+              Address(temp, Realm::offsetOfAllocationMetadataBuilder()),
+              ImmWord(0), fail);
   }
 }
 
@@ -420,7 +426,7 @@ void MacroAssembler::allocateObject(Register result, Register temp,
                                     const AllocSiteInput& allocSite) {
   MOZ_ASSERT(gc::IsObjectAllocKind(allocKind));
 
-  checkAllocatorState(fail);
+  checkAllocatorState(temp, allocKind, fail);
 
   if (shouldNurseryAllocate(allocKind, initialHeap)) {
     MOZ_ASSERT(initialHeap == gc::Heap::Default);
@@ -685,7 +691,7 @@ void MacroAssembler::allocateString(Register result, Register temp,
   MOZ_ASSERT(allocKind == gc::AllocKind::STRING ||
              allocKind == gc::AllocKind::FAT_INLINE_STRING);
 
-  checkAllocatorState(fail);
+  checkAllocatorState(temp, allocKind, fail);
 
   if (shouldNurseryAllocate(allocKind, initialHeap)) {
     MOZ_ASSERT(initialHeap == gc::Heap::Default);
@@ -708,14 +714,16 @@ void MacroAssembler::newGCFatInlineString(Register result, Register temp,
 
 void MacroAssembler::newGCBigInt(Register result, Register temp,
                                  gc::Heap initialHeap, Label* fail) {
-  checkAllocatorState(fail);
+  constexpr gc::AllocKind allocKind = gc::AllocKind::BIGINT;
 
-  if (shouldNurseryAllocate(gc::AllocKind::BIGINT, initialHeap)) {
+  checkAllocatorState(temp, allocKind, fail);
+
+  if (shouldNurseryAllocate(allocKind, initialHeap)) {
     MOZ_ASSERT(initialHeap == gc::Heap::Default);
     return nurseryAllocateBigInt(result, temp, fail);
   }
 
-  freeListAllocate(result, temp, gc::AllocKind::BIGINT, fail);
+  freeListAllocate(result, temp, allocKind, fail);
 }
 
 void MacroAssembler::copySlotsFromTemplate(

js/src/jit/MacroAssembler.h

@@ -5189,7 +5189,7 @@ class MacroAssembler : public MacroAssemblerSpecific {
 
   // Inline allocation.
  private:
-  void checkAllocatorState(Label* fail);
+  void checkAllocatorState(Register temp, gc::AllocKind allocKind, Label* fail);
   bool shouldNurseryAllocate(gc::AllocKind allocKind, gc::Heap initialHeap);
   void nurseryAllocateObject(
       Register result, Register temp, gc::AllocKind allocKind,

js/src/vm/HelperThreads.cpp

@@ -314,9 +314,6 @@ static JSRuntime* GetSelectorRuntime(const CompilationSelector& selector) {
     JSRuntime* operator()(JSScript* script) {
       return script->runtimeFromMainThread();
     }
-    JSRuntime* operator()(Realm* realm) {
-      return realm->runtimeFromMainThread();
-    }
     JSRuntime* operator()(Zone* zone) { return zone->runtimeFromMainThread(); }
     JSRuntime* operator()(ZonesInState zbs) { return zbs.runtime; }
     JSRuntime* operator()(JSRuntime* runtime) { return runtime; }
@@ -328,7 +325,6 @@ static JSRuntime* GetSelectorRuntime(const CompilationSelector& selector) {
 static bool JitDataStructuresExist(const CompilationSelector& selector) {
   struct Matcher {
     bool operator()(JSScript* script) { return !!script->zone()->jitZone(); }
-    bool operator()(Realm* realm) { return !!realm->zone()->jitZone(); }
     bool operator()(Zone* zone) { return !!zone->jitZone(); }
     bool operator()(ZonesInState zbs) { return zbs.runtime->hasJitRuntime(); }
     bool operator()(JSRuntime* runtime) { return runtime->hasJitRuntime(); }
@@ -343,7 +339,6 @@ static bool IonCompileTaskMatches(const CompilationSelector& selector,
     jit::IonCompileTask* task_;
 
     bool operator()(JSScript* script) { return script == task_->script(); }
-    bool operator()(Realm* realm) { return realm == task_->script()->realm(); }
     bool operator()(Zone* zone) {
       return zone == task_->script()->zoneFromAnyThread();
     }

js/src/vm/HelperThreads.h

@@ -157,8 +157,8 @@ struct ZonesInState {
   JS::shadow::Zone::GCState state;
 };
 
-using CompilationSelector = mozilla::Variant<JSScript*, JS::Realm*, JS::Zone*,
-                                             ZonesInState, JSRuntime*>;
+using CompilationSelector =
+    mozilla::Variant<JSScript*, JS::Zone*, ZonesInState, JSRuntime*>;
 
 /*
  * Cancel scheduled or in progress Ion compilations.
@@ -169,10 +169,6 @@ inline void CancelOffThreadIonCompile(JSScript* script) {
   CancelOffThreadIonCompile(CompilationSelector(script));
 }
 
-inline void CancelOffThreadIonCompile(JS::Realm* realm) {
-  CancelOffThreadIonCompile(CompilationSelector(realm));
-}
-
 inline void CancelOffThreadIonCompile(JS::Zone* zone) {
   CancelOffThreadIonCompile(CompilationSelector(zone));
 }

js/src/vm/Realm.cpp

@@ -66,6 +66,10 @@ Realm::~Realm() {
     runtime_->lcovOutput().writeLCovResult(*lcovRealm_);
   }
 
+  if (allocationMetadataBuilder_) {
+    forgetAllocationMetadataBuilder();
+  }
+
   MOZ_ASSERT(runtime_->numRealms > 0);
   runtime_->numRealms--;
 }
@@ -338,18 +342,31 @@ void Realm::setAllocationMetadataBuilder(
     const js::AllocationMetadataBuilder* builder) {
   // Clear any jitcode in the runtime, which behaves differently depending on
   // whether there is a creation callback.
-  ReleaseAllJITCode(runtime_->gcContext());
+  if (bool(allocationMetadataBuilder_) != bool(builder)) {
+    ReleaseAllJITCode(runtime_->gcContext());
+    if (builder) {
+      zone()->incNumRealmsWithAllocMetadataBuilder();
+    } else {
+      zone()->decNumRealmsWithAllocMetadataBuilder();
+    }
+  }
 
   allocationMetadataBuilder_ = builder;
 }
 
 void Realm::forgetAllocationMetadataBuilder() {
+  if (!allocationMetadataBuilder_) {
+    return;
+  }
+
   // Unlike setAllocationMetadataBuilder, we don't have to discard all JIT
   // code here (code is still valid, just a bit slower because it doesn't do
   // inline GC allocations when a metadata builder is present), but we do want
-  // to cancel off-thread Ion compilations to avoid races when Ion calls
-  // hasAllocationMetadataBuilder off-thread.
-  CancelOffThreadIonCompile(this);
+  // to cancel off-thread Ion compilations to avoid races when Ion accesses
+  // numRealmsWithAllocMetadataBuilder_ off-thread.
+  CancelOffThreadIonCompile(zone());
+
+  zone()->decNumRealmsWithAllocMetadataBuilder();
 
   allocationMetadataBuilder_ = nullptr;
 }

js/src/vm/Realm.h

@@ -702,6 +702,9 @@ class JS::Realm : public JS::shadow::Realm {
   static constexpr size_t offsetOfCompartment() {
     return offsetof(JS::Realm, compartment_);
   }
+  static constexpr size_t offsetOfAllocationMetadataBuilder() {
+    return offsetof(JS::Realm, allocationMetadataBuilder_);
+  }
   static constexpr size_t offsetOfDebugModeBits() {
     return offsetof(JS::Realm, debugModeBits_);
   }