Bug 1452406 - Support calling RecompileInfo::shouldSweep under CodeGenerator::link instead of asserting. r=tcampbell

js/src/jit/CodeGenerator.cpp

@@ -10237,14 +10237,11 @@ CodeGenerator::link(JSContext* cx, CompilerConstraintList* constraints)
     // will trickle to jit::Compile() and return Method_Skipped.
     uint32_t warmUpCount = script->getWarmUpCount();
 
-    JitRuntime* jrt = cx->runtime()->jitRuntime();
+    IonCompilationId compilationId = cx->runtime()->jitRuntime()->nextCompilationId();
-    IonCompilationId compilationId = jrt->nextCompilationId();
+    cx->zone()->types.currentCompilationIdRef().emplace(compilationId);
-#ifdef DEBUG
+    auto resetCurrentId = mozilla::MakeScopeExit([cx] {
-    jrt->currentCompilationId().emplace(compilationId);
+        cx->zone()->types.currentCompilationIdRef().reset();
-    auto resetCurrentId = mozilla::MakeScopeExit([jrt] {
-        jrt->currentCompilationId().reset();
     });
-#endif
 
     // Record constraints. If an error occured, returns false and potentially
     // prevent future compilations. Otherwise, if an invalidation occured, then

js/src/jit/Ion.cpp

@@ -2870,7 +2870,7 @@ jit::Invalidate(TypeZone& types, FreeOp* fop,
         if (cancelOffThread)
             CancelOffThreadIonCompile(info.script());
 
-        IonScript* ionScript = info.maybeIonScriptToInvalidate();
+        IonScript* ionScript = info.maybeIonScriptToInvalidate(types);
         if (!ionScript)
             continue;
 
@@ -2897,7 +2897,7 @@ jit::Invalidate(TypeZone& types, FreeOp* fop,
     // IonScript will be immediately destroyed. Otherwise, it will be held live
     // until its last invalidated frame is destroyed.
     for (const RecompileInfo& info : invalid) {
-        IonScript* ionScript = info.maybeIonScriptToInvalidate();
+        IonScript* ionScript = info.maybeIonScriptToInvalidate(types);
         if (!ionScript)
             continue;
 
@@ -2919,7 +2919,7 @@ jit::Invalidate(TypeZone& types, FreeOp* fop,
 
     // Finally, null out script->ion for IonScripts that are still on the stack.
     for (const RecompileInfo& info : invalid) {
-        if (info.maybeIonScriptToInvalidate())
+        if (info.maybeIonScriptToInvalidate(types))
             ClearIonScriptAfterInvalidation(cx, info.script(), resetUses);
     }
 }

js/src/jit/JitCompartment.h

@@ -65,9 +65,6 @@ class JitRuntime
     ActiveThreadData<ExecutableAllocator> execAlloc_;
 
     ActiveThreadData<uint64_t> nextCompilationId_;
-#ifdef DEBUG
-    ActiveThreadData<mozilla::Maybe<IonCompilationId>> currentCompilationId_;
-#endif
 
     // Shared exception-handler tail.
     ExclusiveAccessLockWriteOnceData<uint32_t> exceptionTailOffset_;
@@ -190,11 +187,6 @@ class JitRuntime
     IonCompilationId nextCompilationId() {
         return IonCompilationId(nextCompilationId_++);
     }
-#ifdef DEBUG
-    mozilla::Maybe<IonCompilationId>& currentCompilationId() {
-        return currentCompilationId_.ref();
-    }
-#endif
 
     TrampolinePtr getVMWrapper(const VMFunction& f) const;
     JitCode* debugTrapHandler(JSContext* cx);

js/src/vm/TypeInference-inl.h

@@ -36,11 +36,13 @@ namespace js {
 /////////////////////////////////////////////////////////////////////
 
 jit::IonScript*
-RecompileInfo::maybeIonScriptToInvalidate() const
+RecompileInfo::maybeIonScriptToInvalidate(const TypeZone& zone) const
 {
+    MOZ_ASSERT(script_->zone() == zone.zone());
+
     // Make sure this is not called under CodeGenerator::link (before the
     // IonScript is created).
-    MOZ_ASSERT(!TlsContext.get()->runtime()->jitRuntime()->currentCompilationId().isSome());
+    MOZ_ASSERT_IF(zone.currentCompilationId(), zone.currentCompilationId().ref() != id_);
 
     if (!script_->hasIonScript() || script_->ionScript()->compilationId() != id_)
         return nullptr;
@@ -49,11 +51,19 @@ RecompileInfo::maybeIonScriptToInvalidate() const
 }
 
 inline bool
-RecompileInfo::shouldSweep()
+RecompileInfo::shouldSweep(const TypeZone& zone)
 {
     if (IsAboutToBeFinalizedUnbarriered(&script_))
         return true;
-    return maybeIonScriptToInvalidate() == nullptr;
+
+    MOZ_ASSERT(script_->zone() == zone.zone());
+
+    // Don't sweep if we're called under CodeGenerator::link, before the
+    // IonScript is created.
+    if (zone.currentCompilationId() && zone.currentCompilationId().ref() == id_)
+        return false;
+
+    return maybeIonScriptToInvalidate(zone) == nullptr;
 }
 
 /////////////////////////////////////////////////////////////////////

js/src/vm/TypeInference.cpp

@@ -1213,7 +1213,7 @@ class TypeCompilerConstraint : public TypeConstraint
     }
 
     bool sweep(TypeZone& zone, TypeConstraint** res) override {
-        if (data.shouldSweep() || compilation.shouldSweep())
+        if (data.shouldSweep() || compilation.shouldSweep(zone))
             return false;
         *res = zone.typeLifoAlloc().new_<TypeCompilerConstraint<T> >(compilation, data);
         return true;
@@ -1422,7 +1422,7 @@ bool
 js::FinishCompilation(JSContext* cx, HandleScript script, CompilerConstraintList* constraints,
                       IonCompilationId compilationId, bool* isValidOut)
 {
-    MOZ_ASSERT(*cx->runtime()->jitRuntime()->currentCompilationId() == compilationId);
+    MOZ_ASSERT(*cx->zone()->types.currentCompilationId() == compilationId);
 
     if (constraints->failed())
         return false;
@@ -4415,7 +4415,7 @@ JSScript::maybeSweepTypes(AutoClearTypeInferenceStateOnOOM* oom)
         RecompileInfoVector& inlinedCompilations = types_->inlinedCompilations();
         size_t dest = 0;
         for (size_t i = 0; i < inlinedCompilations.length(); i++) {
-            if (inlinedCompilations[i].shouldSweep())
+            if (inlinedCompilations[i].shouldSweep(types))
                 continue;
             inlinedCompilations[dest] = inlinedCompilations[i];
             dest++;
@@ -4485,6 +4485,7 @@ Zone::addSizeOfIncludingThis(mozilla::MallocSizeOf mallocSizeOf,
 TypeZone::TypeZone(Zone* zone)
   : zone_(zone),
     typeLifoAlloc_(zone->group(), (size_t) TYPE_LIFO_ALLOC_PRIMARY_CHUNK_SIZE),
+    currentCompilationId_(zone->group()),
     generation(zone->group(), 0),
     sweepTypeLifoAlloc(zone->group(), (size_t) TYPE_LIFO_ALLOC_PRIMARY_CHUNK_SIZE),
     sweepReleaseTypes(zone->group(), false),

js/src/vm/TypeInference.h

@@ -1139,9 +1139,9 @@ class RecompileInfo
         return script_;
     }
 
-    inline jit::IonScript* maybeIonScriptToInvalidate() const;
+    inline jit::IonScript* maybeIonScriptToInvalidate(const TypeZone& zone) const;
 
-    inline bool shouldSweep();
+    inline bool shouldSweep(const TypeZone& zone);
 
     bool operator==(const RecompileInfo& other) const {
         return script_== other.script_ && id_ == other.id_;
@@ -1340,6 +1340,9 @@ class TypeZone
     static const size_t TYPE_LIFO_ALLOC_PRIMARY_CHUNK_SIZE = 8 * 1024;
     ZoneGroupData<LifoAlloc> typeLifoAlloc_;
 
+    // Under CodeGenerator::link, the id of the current compilation.
+    ZoneGroupData<mozilla::Maybe<IonCompilationId>> currentCompilationId_;
+
     TypeZone(const TypeZone&) = delete;
     void operator=(const TypeZone&) = delete;
 
@@ -1388,6 +1391,13 @@ class TypeZone
         MOZ_RELEASE_ASSERT(sweepingTypes != sweeping);
         sweepingTypes = sweeping;
     }
+
+    mozilla::Maybe<IonCompilationId> currentCompilationId() const {
+        return currentCompilationId_.ref();
+    }
+    mozilla::Maybe<IonCompilationId>& currentCompilationIdRef() {
+        return currentCompilationId_.ref();
+    }
 };
 
 enum SpewChannel {