Bug 1591598 - Support storing LazyScript/Scope in ScriptWarmUpData. r=jandem,jonco

Allow storing manually-barriered GC pointers in ScriptWarmUpData. This
updates the 'trace' method as needed. When switching types, the user must
first 'clear' the old type and then 'init' the new type. We continue to use
WarmUpCount(0) as the default safe state.

Depends on D55033

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

--HG--
extra : moz-landing-system : lando

js/src/jit/JitScript.cpp

@@ -149,9 +149,8 @@ bool JSScript::createJitScript(JSContext* cx) {
     return false;
   }
 
-  MOZ_ASSERT(!hasJitScript());
   prepareForDestruction.release();
-  warmUpData_.setJitScript(jitScript.release());
+  warmUpData_.initJitScript(jitScript.release());
   AddCellMemory(this, allocSize.value(), MemoryUse::JitScript);
 
   // We have a JitScript so we can set the script's jitCodeRaw_ pointer to the

js/src/vm/JSScript-inl.h

@@ -58,6 +58,28 @@ ScriptAndCounts::ScriptAndCounts(ScriptAndCounts&& sac)
 void SetFrameArgumentsObject(JSContext* cx, AbstractFramePtr frame,
                              HandleScript script, JSObject* argsobj);
 
+inline void ScriptWarmUpData::initEnclosingScript(LazyScript* enclosingScript) {
+  MOZ_ASSERT(data_ == ResetState());
+  setTaggedPtr<EnclosingScriptTag>(enclosingScript);
+  static_assert(std::is_base_of<gc::TenuredCell, LazyScript>::value,
+                "LazyScript must be TenuredCell to avoid post-barriers");
+}
+inline void ScriptWarmUpData::clearEnclosingScript() {
+  LazyScript::writeBarrierPre(toEnclosingScript());
+  data_ = ResetState();
+}
+
+inline void ScriptWarmUpData::initEnclosingScope(Scope* enclosingScope) {
+  MOZ_ASSERT(data_ == ResetState());
+  setTaggedPtr<EnclosingScopeTag>(enclosingScope);
+  static_assert(std::is_base_of<gc::TenuredCell, Scope>::value,
+                "Scope must be TenuredCell to avoid post-barriers");
+}
+inline void ScriptWarmUpData::clearEnclosingScope() {
+  Scope::writeBarrierPre(toEnclosingScope());
+  data_ = ResetState();
+}
+
 }  // namespace js
 
 inline JSFunction* JSScript::getFunction(size_t index) {

js/src/vm/JSScript.cpp

@@ -5273,12 +5273,32 @@ void RuntimeScriptData::markForCrossZone(JSContext* cx) {
 }
 
 void ScriptWarmUpData::trace(JSTracer* trc) {
-  if (isJitScript()) {
-    toJitScript()->trace(trc);
-    return;
-  }
+  uintptr_t tag = data_ & TagMask;
+  switch (tag) {
+    case EnclosingScriptTag: {
+      LazyScript* enclosingScript = toEnclosingScript();
+      TraceManuallyBarrieredEdge(trc, &enclosingScript, "enclosingScript");
+      setTaggedPtr<EnclosingScriptTag>(enclosingScript);
+      break;
+    }
 
-  MOZ_ASSERT(isWarmUpCount());
+    case EnclosingScopeTag: {
+      Scope* enclosingScope = toEnclosingScope();
+      TraceManuallyBarrieredEdge(trc, &enclosingScope, "enclosingScope");
+      setTaggedPtr<EnclosingScopeTag>(enclosingScope);
+      break;
+    }
+
+    case JitScriptTag: {
+      toJitScript()->trace(trc);
+      break;
+    }
+
+    default: {
+      MOZ_ASSERT(isWarmUpCount());
+      break;
+    }
+  }
 }
 
 void JSScript::traceChildren(JSTracer* trc) {

js/src/vm/JSScript.h

@@ -1397,8 +1397,16 @@ class ScriptSourceObject : public NativeObject {
 enum class GeneratorKind : bool { NotGenerator, Generator };
 enum class FunctionAsyncKind : bool { SyncFunction, AsyncFunction };
 
-// ScriptWarmUpData represents a pointer-sized field in JSScript that stores
-// one of the following:
+// ScriptWarmUpData represents a pointer-sized field in BaseScript that stores
+// one of the following using low-bit tags:
+//
+// * The enclosing LazyScript. This is only used while this script is lazy and
+//   its containing script is also lazy. This outer script must be compiled
+//   before the current script can in order to correctly build the scope chain.
+//
+// * The enclosing Scope. This is only used while this script is lazy and its
+//   containing script is compiled. This is the outer scope chain that will be
+//   used to compile this scipt.
 //
 // * The script's warm-up count. This is only used until the script has a
 //   JitScript. The Baseline Interpreter and JITs use the warm-up count stored
@@ -1407,8 +1415,10 @@ enum class FunctionAsyncKind : bool { SyncFunction, AsyncFunction };
 // * A pointer to the JitScript, when the script is warm enough for the Baseline
 //   Interpreter.
 //
-// Pointer tagging is used to distinguish those states.
 class ScriptWarmUpData {
+  uintptr_t data_ = ResetState();
+
+ private:
   static constexpr uintptr_t NumTagBits = 2;
   static constexpr uint32_t MaxWarmUpCount = UINT32_MAX >> NumTagBits;
 
@@ -1416,10 +1426,27 @@ class ScriptWarmUpData {
   // Public only for the JITs.
   static constexpr uintptr_t TagMask = (1 << NumTagBits) - 1;
   static constexpr uintptr_t JitScriptTag = 0;
-  static constexpr uintptr_t WarmUpCountTag = 1;
+  static constexpr uintptr_t EnclosingScriptTag = 1;
+  static constexpr uintptr_t EnclosingScopeTag = 2;
+  static constexpr uintptr_t WarmUpCountTag = 3;
 
  private:
-  uintptr_t data_ = 0 | WarmUpCountTag;
+  // A gc-safe value to clear to.
+  constexpr uintptr_t ResetState() { return 0 | WarmUpCountTag; }
+
+  template <uintptr_t Tag>
+  inline void setTaggedPtr(void* ptr) {
+    static_assert(Tag <= TagMask, "Tag must fit in TagMask");
+    MOZ_ASSERT((uintptr_t(ptr) & TagMask) == 0);
+    data_ = uintptr_t(ptr) | Tag;
+  }
+
+  template <typename T, uintptr_t Tag>
+  inline T getTaggedPtr() const {
+    static_assert(Tag <= TagMask, "Tag must fit in TagMask");
+    MOZ_ASSERT((data_ & TagMask) == Tag);
+    return reinterpret_cast<T>(data_ & ~TagMask);
+  }
 
   void setWarmUpCount(uint32_t count) {
     if (count > MaxWarmUpCount) {
@@ -1431,9 +1458,30 @@ class ScriptWarmUpData {
  public:
   void trace(JSTracer* trc);
 
+  bool isEnclosingScript() const {
+    return (data_ & TagMask) == EnclosingScriptTag;
+  }
+  bool isEnclosingScope() const {
+    return (data_ & TagMask) == EnclosingScopeTag;
+  }
   bool isWarmUpCount() const { return (data_ & TagMask) == WarmUpCountTag; }
   bool isJitScript() const { return (data_ & TagMask) == JitScriptTag; }
 
+  // NOTE: To change type safely, 'clear' the old tagged value and then 'init'
+  //       the new one. This will notify the GC appropriately.
+
+  LazyScript* toEnclosingScript() const {
+    return getTaggedPtr<LazyScript*, EnclosingScriptTag>();
+  }
+  inline void initEnclosingScript(LazyScript* enclosingScript);
+  inline void clearEnclosingScript();
+
+  Scope* toEnclosingScope() const {
+    return getTaggedPtr<Scope*, EnclosingScopeTag>();
+  }
+  inline void initEnclosingScope(Scope* enclosingScope);
+  inline void clearEnclosingScope();
+
   uint32_t toWarmUpCount() const {
     MOZ_ASSERT(isWarmUpCount());
     return data_ >> NumTagBits;
@@ -1448,20 +1496,17 @@ class ScriptWarmUpData {
   }
 
   jit::JitScript* toJitScript() const {
-    MOZ_ASSERT(isJitScript());
-    static_assert(JitScriptTag == 0, "Code depends on JitScriptTag being zero");
-    return reinterpret_cast<jit::JitScript*>(data_);
+    return getTaggedPtr<jit::JitScript*, JitScriptTag>();
   }
-  void setJitScript(jit::JitScript* jitScript) {
+  void initJitScript(jit::JitScript* jitScript) {
     MOZ_ASSERT(isWarmUpCount());
-    MOZ_ASSERT((uintptr_t(jitScript) & TagMask) == 0);
-    data_ = uintptr_t(jitScript) | JitScriptTag;
+    setTaggedPtr<JitScriptTag>(jitScript);
   }
   void clearJitScript() {
     MOZ_ASSERT(isJitScript());
-    setWarmUpCount(0);
+    data_ = ResetState();
   }
-};
+} JS_HAZ_GC_POINTER;
 
 static_assert(sizeof(ScriptWarmUpData) == sizeof(uintptr_t),
               "JIT code depends on ScriptWarmUpData being pointer-sized");