Bug 1867193 - Reset more IC state when discarding IC stubs. r=iain

After discarding IC stubs, we would no longer support trial inlining, resulting
in potential performance cliffs.

With this patch we try to reset as much state as possible. The main exception is
that we preserve trial inlining data for call sites if the callee is active
on the stack and running in Baseline. In this case we also clone the IC stubs
for the call site.

For all other ICs we now reset the ICState and we also purge inactive inlined
ICScripts.

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

js/src/gc/Zone.cpp

@@ -455,8 +455,10 @@ void Zone::forceDiscardJitCode(JS::GCContext* gcx,
         }
 
         // If we did not release the JitScript, we need to purge IC stubs
-        // because the ICStubSpace will be purged below.
-        jitScript->purgeStubs(script);
+        // because the ICStubSpace will be purged below. Also purge all
+        // trial-inlined ICScripts that are not active on the stack.
+        jitScript->purgeInactiveICScripts();
+        jitScript->purgeStubs(script, newStubSpace);
 
         if (options.resetNurseryAllocSites ||
             options.resetPretenuredAllocSites) {

js/src/jit/BaselineBailouts.cpp

@@ -502,7 +502,8 @@ void BaselineStackBuilder::setNextCallee(
     JSFunction* nextCallee, TrialInliningState trialInliningState) {
   nextCallee_ = nextCallee;
 
-  if (trialInliningState == TrialInliningState::Inlined) {
+  if (trialInliningState == TrialInliningState::Inlined &&
+      !iter_.ionScript()->purgedICScripts()) {
     // Update icScript_ to point to the icScript of nextCallee
     const uint32_t pcOff = script_->pcToOffset(pc_);
     icScript_ = icScript_->findInlinedChild(pcOff);
@@ -511,8 +512,15 @@ void BaselineStackBuilder::setNextCallee(
     // just use the callee's own ICScript. We could still have the trial
     // inlined ICScript available, but we also could not if we transitioned
     // to TrialInliningState::Failure after being monomorphic inlined.
+    //
+    // Also use the callee's own ICScript if we purged callee ICScripts.
     icScript_ = nextCallee->nonLazyScript()->jitScript()->icScript();
   }
+
+  // Assert the ICScript matches nextCallee.
+  JSScript* calleeScript = nextCallee->nonLazyScript();
+  MOZ_RELEASE_ASSERT(icScript_->numICEntries() == calleeScript->numICEntries());
+  MOZ_RELEASE_ASSERT(icScript_->bytecodeSize() == calleeScript->length());
 }
 
 bool BaselineStackBuilder::done() {

js/src/jit/BaselineIC.h

@@ -291,7 +291,7 @@ class ICCacheIRStub final : public ICStub {
   void trace(JSTracer* trc);
   bool traceWeak(JSTracer* trc);
 
-  ICCacheIRStub* clone(JSContext* cx, ICStubSpace& newSpace);
+  ICCacheIRStub* clone(JSRuntime* rt, ICStubSpace& newSpace);
 
   // Returns true if this stub can call JS or VM code that can trigger a GC.
   bool makesGCCalls() const;

js/src/jit/CacheIRCompiler.cpp

@@ -1198,7 +1198,7 @@ void CacheIRWriter::copyStubData(uint8_t* dest) const {
   }
 }
 
-ICCacheIRStub* ICCacheIRStub::clone(JSContext* cx, ICStubSpace& newSpace) {
+ICCacheIRStub* ICCacheIRStub::clone(JSRuntime* rt, ICStubSpace& newSpace) {
   const CacheIRStubInfo* info = stubInfo();
   MOZ_ASSERT(info->makesGCCalls());
 
@@ -1217,7 +1217,7 @@ ICCacheIRStub* ICCacheIRStub::clone(JSContext* cx, ICStubSpace& newSpace) {
 
   // Because this can be called during sweeping when discarding JIT code, we
   // have to lock the store buffer
-  gc::AutoLockStoreBuffer lock(cx->runtime());
+  gc::AutoLockStoreBuffer lock(rt);
 
   uint32_t field = 0;
   while (true) {

js/src/jit/IonScript.h

@@ -106,6 +106,11 @@ class alignas(8) IonScript final : public TrailingArray {
   // Flag set if IonScript was compiled with profiling enabled.
   bool hasProfilingInstrumentation_ = false;
 
+  // If true, this IonScript was active on the stack when we discarded JIT code
+  // and inactive ICScripts. This means we should use the generic ICScripts for
+  // inlined functions when we bail out.
+  bool purgedICScripts_ = false;
+
   // Number of bytes this function reserves on the stack for slots spilled by
   // the register allocator.
   uint32_t localSlotsSize_ = 0;
@@ -355,6 +360,9 @@ class alignas(8) IonScript final : public TrailingArray {
     return hasProfilingInstrumentation_;
   }
 
+  bool purgedICScripts() const { return purgedICScripts_; }
+  void notePurgedICScripts() { purgedICScripts_ = true; }
+
   size_t sizeOfIncludingThis(mozilla::MallocSizeOf mallocSizeOf) const {
     return mallocSizeOf(this);
   }

js/src/jit/JSJitFrameIter.h

@@ -441,6 +441,8 @@ class SnapshotIterator {
   }
   inline BailoutKind bailoutKind() const { return snapshot_.bailoutKind(); }
 
+  IonScript* ionScript() const { return ionScript_; }
+
  public:
   // Read the next instruction available and get ready to either skip it or
   // evaluate it.

js/src/jit/JitScript.cpp

@@ -286,6 +286,30 @@ bool ICScript::hasInlinedChild(uint32_t pcOffset) {
   return false;
 }
 
+void ICScript::purgeInactiveICScripts() {
+  MOZ_ASSERT(inliningRoot());
+
+  if (!inlinedChildren_) {
+    return;
+  }
+
+  inlinedChildren_->eraseIf(
+      [](const CallSite& callsite) { return !callsite.callee_->active(); });
+
+  if (inlinedChildren_->empty()) {
+    inlinedChildren_.reset();
+    return;
+  }
+
+  // We have an active callee ICScript. This means the current ICScript must be
+  // active too.
+  MOZ_ASSERT(active());
+
+  for (CallSite& callsite : *inlinedChildren_) {
+    callsite.callee_->purgeInactiveICScripts();
+  }
+}
+
 void JitScript::resetWarmUpCount(uint32_t count) {
   icScript_.resetWarmUpCount(count);
   if (hasInliningRoot()) {
@@ -418,7 +442,25 @@ ICEntry* ICScript::interpreterICEntryFromPCOffset(uint32_t pcOffset) {
   return nullptr;
 }
 
-void JitScript::purgeStubs(JSScript* script) {
+void JitScript::purgeInactiveICScripts() {
+  if (!hasInliningRoot()) {
+    return;
+  }
+
+  icScript()->purgeInactiveICScripts();
+
+  inliningRoot()->purgeInactiveICScripts();
+  if (inliningRoot()->numInlinedScripts() == 0) {
+    inliningRoot_.reset();
+    icScript()->inliningRoot_ = nullptr;
+  } else {
+    // If a callee script is active on the stack, the root script must be active
+    // too.
+    MOZ_ASSERT(icScript()->active());
+  }
+}
+
+void JitScript::purgeStubs(JSScript* script, ICStubSpace& newStubSpace) {
   MOZ_ASSERT(script->jitScript() == this);
 
   Zone* zone = script->zone();
@@ -433,19 +475,53 @@ void JitScript::purgeStubs(JSScript* script) {
 
   JitSpew(JitSpew_BaselineIC, "Purging optimized stubs");
 
-  icScript()->purgeStubs(zone);
+  icScript()->purgeStubs(zone, newStubSpace);
   if (hasInliningRoot()) {
-    inliningRoot()->purgeStubs(zone);
+    inliningRoot()->purgeStubs(zone, newStubSpace);
   }
 
   notePurgedStubs();
 }
 
-void ICScript::purgeStubs(Zone* zone) {
+void ICScript::purgeStubs(Zone* zone, ICStubSpace& newStubSpace) {
   for (size_t i = 0; i < numICEntries(); i++) {
     ICEntry& entry = icEntry(i);
     ICFallbackStub* fallback = fallbackStub(i);
+
+    // If this is a trial inlining call site and the callee's ICScript hasn't
+    // been discarded, clone the IC chain instead of purging stubs. In this case
+    // both the current ICScript and the callee's inlined ICScript must be
+    // active on the stack.
+    //
+    // We can't purge the IC stubs in this case because it'd confuse trial
+    // inlining if we try to inline again later and we already have an ICScript
+    // for this call site.
+    if (fallback->trialInliningState() == TrialInliningState::Inlined &&
+        hasInlinedChild(fallback->pcOffset())) {
+      MOZ_ASSERT(active());
+      MOZ_ASSERT(findInlinedChild(fallback->pcOffset())->active());
+
+      JSRuntime* rt = zone->runtimeFromMainThread();
+      ICCacheIRStub* prev = nullptr;
+      ICStub* stub = entry.firstStub();
+      while (stub != fallback) {
+        ICCacheIRStub* clone = stub->toCacheIRStub()->clone(rt, newStubSpace);
+        if (prev) {
+          prev->setNext(clone);
+        } else {
+          entry.setFirstStub(clone);
+        }
+        MOZ_ASSERT(stub->toCacheIRStub()->next() == clone->next());
+        prev = clone;
+        stub = clone->next();
+      }
+      continue;
+    }
+
+    MOZ_ASSERT(!hasInlinedChild(fallback->pcOffset()));
+
     fallback->discardStubs(zone, &entry);
+    fallback->state().reset();
   }
 }
 
@@ -611,12 +687,17 @@ static void MarkActiveICScriptsAndCopyStubs(
     switch (frame.type()) {
       case FrameType::BaselineJS:
         frame.script()->jitScript()->icScript()->setActive();
+        // If the frame is using a trial-inlining ICScript, we have to preserve
+        // it too.
+        if (frame.baselineFrame()->icScript()->isInlined()) {
+          frame.baselineFrame()->icScript()->setActive();
+        }
         break;
       case FrameType::BaselineStub: {
         auto* layout = reinterpret_cast<BaselineStubFrameLayout*>(frame.fp());
         if (layout->maybeStubPtr() && !layout->maybeStubPtr()->isFallback()) {
           ICCacheIRStub* stub = layout->maybeStubPtr()->toCacheIRStub();
-          ICCacheIRStub* newStub = stub->clone(cx, newStubSpace);
+          ICCacheIRStub* newStub = stub->clone(cx->runtime(), newStubSpace);
           layout->setStubPtr(newStub);
         }
         break;
@@ -639,6 +720,9 @@ static void MarkActiveICScriptsAndCopyStubs(
              ++inlineIter) {
           inlineIter.script()->jitScript()->icScript()->setActive();
         }
+        // Because we're purging ICScripts, the bailout machinery should use
+        // the generic ICScript for inlined callees.
+        frame.ionScript()->notePurgedICScripts();
         break;
       }
       default:;

js/src/jit/JitScript.h

@@ -101,6 +101,15 @@ static IonScript* const IonCompilingScriptPtr =
  * BaselineScript and IonScript/WarpScript to which it
  * corresponds. They can be destroyed and recreated, and the ICScript
  * will remain valid.
+ *
+ * When we discard JIT code, we mark ICScripts that are active on the stack as
+ * active and then purge all of the inactive ICScripts. We also purge ICStubs,
+ * including the CallInlinedFunction stub at the trial inining call site, and
+ * reset the ICStates to allow trial inlining again later.
+ *
+ * If there's a BaselineFrame for an inlined ICScript, we'll preserve both this
+ * ICScript and the IC chain for the call site in the caller's ICScript.
+ * See ICScript::purgeStubs and ICScript::purgeInactiveICScripts.
  */
 
 class alignas(uintptr_t) ICScript final : public TrailingArray {
@@ -173,7 +182,9 @@ class alignas(uintptr_t) ICScript final : public TrailingArray {
   void removeInlinedChild(uint32_t pcOffset);
   bool hasInlinedChild(uint32_t pcOffset);
 
-  void purgeStubs(Zone* zone);
+  void purgeStubs(Zone* zone, ICStubSpace& newStubSpace);
+
+  void purgeInactiveICScripts();
 
   bool active() const { return active_; }
   void setActive() { active_ = true; }
@@ -429,7 +440,9 @@ class alignas(uintptr_t) JitScript final
 
   void trace(JSTracer* trc);
   void traceWeak(JSTracer* trc);
-  void purgeStubs(JSScript* script);
+  void purgeStubs(JSScript* script, ICStubSpace& newStubSpace);
+
+  void purgeInactiveICScripts();
 
   ICEntry& icEntryFromPCOffset(uint32_t pcOffset) {
     return icScript_.icEntryFromPCOffset(pcOffset);

js/src/jit/TrialInlining.cpp

@@ -6,6 +6,8 @@
 
 #include "jit/TrialInlining.h"
 
+#include "mozilla/DebugOnly.h"
+
 #include "jit/BaselineCacheIRCompiler.h"
 #include "jit/BaselineFrame.h"
 #include "jit/BaselineIC.h"
@@ -951,9 +953,9 @@ bool InliningRoot::traceWeak(JSTracer* trc) {
   return allSurvived;
 }
 
-void InliningRoot::purgeStubs(Zone* zone) {
+void InliningRoot::purgeStubs(Zone* zone, ICStubSpace& newStubSpace) {
   for (auto& inlinedScript : inlinedScripts_) {
-    inlinedScript->purgeStubs(zone);
+    inlinedScript->purgeStubs(zone, newStubSpace);
   }
 }
 
@@ -963,6 +965,24 @@ void InliningRoot::resetWarmUpCounts(uint32_t count) {
   }
 }
 
+void InliningRoot::purgeInactiveICScripts() {
+  mozilla::DebugOnly<uint32_t> totalSize = owningScript_->length();
+
+  for (auto& inlinedScript : inlinedScripts_) {
+    if (inlinedScript->active()) {
+      totalSize += inlinedScript->bytecodeSize();
+    } else {
+      MOZ_ASSERT(inlinedScript->bytecodeSize() < totalBytecodeSize_);
+      totalBytecodeSize_ -= inlinedScript->bytecodeSize();
+    }
+  }
+
+  MOZ_ASSERT(totalBytecodeSize_ == totalSize);
+
+  inlinedScripts_.eraseIf(
+      [](auto& inlinedScript) { return !inlinedScript->active(); });
+}
+
 #ifdef DEBUG
 bool InliningRoot::hasActiveICScript() const {
   for (auto& inlinedScript : inlinedScripts_) {

js/src/jit/TrialInlining.h

@@ -67,6 +67,7 @@ class ICCacheIRStub;
 class ICEntry;
 class ICFallbackStub;
 class ICScript;
+class ICStubSpace;
 
 /*
  * An InliningRoot is owned by a JitScript. In turn, it owns the set
@@ -86,9 +87,11 @@ class InliningRoot {
 
   uint32_t numInlinedScripts() const { return inlinedScripts_.length(); }
 
-  void purgeStubs(Zone* zone);
+  void purgeStubs(Zone* zone, ICStubSpace& newStubSpace);
   void resetWarmUpCounts(uint32_t count);
 
+  void purgeInactiveICScripts();
+
 #ifdef DEBUG
   bool hasActiveICScript() const;
 #endif