[Turbofan] ContextRef::SerializeContextChain should take a depth param

We can save memory by only serializing a context chain to a
*required* depth if we know it.

Bug: v8:7790
Change-Id: I97d21f8cd7b56b26fddd95e00a26d5e520d96170
Reviewed-on: https://chromium-review.googlesource.com/c/v8/v8/+/1678358
Reviewed-by: Georg Neis <neis@chromium.org>
Commit-Queue: Michael Stanton <mvstanton@chromium.org>
Cr-Commit-Position: refs/heads/master@{#62795}

src/compiler/heap-refs.h

@@ -326,15 +326,14 @@ class ContextRef : public HeapObjectRef {
   using HeapObjectRef::HeapObjectRef;
   Handle<Context> object() const;
 
-  void SerializeContextChain();
+  // {previous} decrements {depth} by 1 for each previous link successfully
+  // followed. If {depth} != 0 on function return, then it only got
+  // partway to the desired depth. If {serialize} is true, then
+  // {previous} will cache its findings.
+  ContextRef previous(size_t* depth, bool serialize = false) const;
 
-  // {previous} decrements n by 1 for each previous link successfully
-  // followed. If depth != 0 on function return, then it only got
-  // partway to the desired depth.
-  ContextRef previous(size_t* depth) const;
-
-  void SerializeSlot(int index);
-  base::Optional<ObjectRef> get(int index) const;
+  // Only returns a value if the index is valid for this ContextRef.
+  base::Optional<ObjectRef> get(int index, bool serialize = false) const;
 
   // We only serialize the ScopeInfo if certain Promise
   // builtins are called.

src/compiler/js-heap-broker.cc

@@ -576,21 +576,15 @@ class ContextData : public HeapObjectData {
  public:
   ContextData(JSHeapBroker* broker, ObjectData** storage,
               Handle<Context> object);
-  void SerializeContextChain(JSHeapBroker* broker);
 
-  ContextData* previous() const {
-    return previous_;
-  }
+  // {previous} will return the closest valid context possible to desired
+  // {depth}, decrementing {depth} for each previous link successfully followed.
+  // If {serialize} is true, it will serialize contexts along the way.
+  ContextData* previous(JSHeapBroker* broker, size_t* depth, bool serialize);
 
-  void SerializeSlot(JSHeapBroker* broker, int index);
-
-  ObjectData* GetSlot(int index) {
-    auto search = slots_.find(index);
-    if (search != slots_.end()) {
-      return search->second;
-    }
-    return nullptr;
-  }
+  // Returns nullptr if the slot index isn't valid or wasn't serialized
+  // (unless {serialize} is true).
+  ObjectData* GetSlot(JSHeapBroker* broker, int index, bool serialize);
 
  private:
   ZoneMap<int, ObjectData*> slots_;
@@ -601,25 +595,46 @@ ContextData::ContextData(JSHeapBroker* broker, ObjectData** storage,
                          Handle<Context> object)
     : HeapObjectData(broker, storage, object), slots_(broker->zone()) {}
 
-void ContextData::SerializeContextChain(JSHeapBroker* broker) {
-  if (previous_ != nullptr) return;
+ContextData* ContextData::previous(JSHeapBroker* broker, size_t* depth,
+                                   bool serialize) {
+  if (*depth == 0) return this;
 
+  if (serialize && previous_ == nullptr) {
+    TraceScope tracer(broker, this, "ContextData::previous");
     Handle<Context> context = Handle<Context>::cast(object());
-  // Context::previous DCHECK-fails when called on the native context.
-  if (!context->IsNativeContext()) {
-    TraceScope tracer(broker, this, "ContextData::SerializeContextChain");
-    previous_ = broker->GetOrCreateData(context->previous())->AsContext();
-    previous_->SerializeContextChain(broker);
+    Context prev = context->previous();
+    if (!prev.is_null()) {
+      previous_ = broker->GetOrCreateData(prev)->AsContext();
     }
+  }
+
+  if (previous_ != nullptr) {
+    *depth = *depth - 1;
+    return previous_->previous(broker, depth, serialize);
+  }
+  return this;
 }
 
-void ContextData::SerializeSlot(JSHeapBroker* broker, int index) {
-  TraceScope tracer(broker, this, "ContextData::SerializeSlot");
-  TRACE(broker, "Serializing context slot " << index);
+ObjectData* ContextData::GetSlot(JSHeapBroker* broker, int index,
+                                 bool serialize) {
+  CHECK_GE(index, 0);
+  auto search = slots_.find(index);
+  if (search != slots_.end()) {
+    return search->second;
+  }
+
+  if (serialize) {
     Handle<Context> context = Handle<Context>::cast(object());
-  CHECK(index >= 0 && index < context->length());
+    if (index < context->length()) {
+      TraceScope tracer(broker, this, "ContextData::GetSlot");
+      TRACE(broker, "Serializing context slot " << index);
       ObjectData* odata = broker->GetOrCreateData(context->get(index));
       slots_.insert(std::make_pair(index, odata));
+      return odata;
+    }
+  }
+
+  return nullptr;
 }
 
 class NativeContextData : public ContextData {
@@ -2067,35 +2082,31 @@ bool ObjectRef::equals(const ObjectRef& other) const {
 
 Isolate* ObjectRef::isolate() const { return broker()->isolate(); }
 
-ContextRef ContextRef::previous(size_t* depth) const {
+ContextRef ContextRef::previous(size_t* depth, bool serialize) const {
   DCHECK_NOT_NULL(depth);
-  DCHECK_GE(*depth, 0);
   if (broker()->mode() == JSHeapBroker::kDisabled) {
     AllowHandleAllocation handle_allocation;
     AllowHandleDereference handle_dereference;
     Context current = *object();
-    while (*depth != 0) {
+    while (*depth != 0 && !current.previous().is_null()) {
       current = current.previous();
       (*depth)--;
     }
     return ContextRef(broker(), handle(current, broker()->isolate()));
   }
   ContextData* current = this->data()->AsContext();
-  while (*depth != 0 && current->previous() != nullptr) {
-    current = current->previous();
-    (*depth)--;
-  }
-  return ContextRef(broker(), current);
+  return ContextRef(broker(), current->previous(broker(), depth, serialize));
 }
 
-base::Optional<ObjectRef> ContextRef::get(int index) const {
+base::Optional<ObjectRef> ContextRef::get(int index, bool serialize) const {
   if (broker()->mode() == JSHeapBroker::kDisabled) {
     AllowHandleAllocation handle_allocation;
     AllowHandleDereference handle_dereference;
     Handle<Object> value(object()->get(index), broker()->isolate());
     return ObjectRef(broker(), value);
   }
-  ObjectData* optional_slot = data()->AsContext()->GetSlot(index);
+  ObjectData* optional_slot =
+      data()->AsContext()->GetSlot(broker(), index, serialize);
   if (optional_slot != nullptr) {
     return ObjectRef(broker(), optional_slot);
   }
@@ -3718,18 +3729,6 @@ void SourceTextModuleRef::Serialize() {
   data()->AsSourceTextModule()->Serialize(broker());
 }
 
-void ContextRef::SerializeContextChain() {
-  if (broker()->mode() == JSHeapBroker::kDisabled) return;
-  CHECK_EQ(broker()->mode(), JSHeapBroker::kSerializing);
-  data()->AsContext()->SerializeContextChain(broker());
-}
-
-void ContextRef::SerializeSlot(int index) {
-  if (broker()->mode() == JSHeapBroker::kDisabled) return;
-  CHECK_EQ(broker()->mode(), JSHeapBroker::kSerializing);
-  data()->AsContext()->SerializeSlot(broker(), index);
-}
-
 void NativeContextRef::Serialize() {
   if (broker()->mode() == JSHeapBroker::kDisabled) return;
   CHECK_EQ(broker()->mode(), JSHeapBroker::kSerializing);
@@ -4053,7 +4052,7 @@ GlobalAccessFeedback const* JSHeapBroker::ProcessFeedbackForGlobalAccess(
     }
     ContextRef context_ref(this, context);
     if (immutable) {
-      context_ref.SerializeSlot(context_slot_index);
+      context_ref.get(context_slot_index, true);
     }
     return new (zone())
         GlobalAccessFeedback(context_ref, context_slot_index, immutable);

src/compiler/js-heap-copy-reducer.cc

@@ -31,7 +31,6 @@ Reduction JSHeapCopyReducer::Reduce(Node* node) {
       if (object.IsJSFunction()) object.AsJSFunction().Serialize();
       if (object.IsJSObject()) object.AsJSObject().SerializeObjectCreateMap();
       if (object.IsSourceTextModule()) object.AsSourceTextModule().Serialize();
-      if (object.IsContext()) object.AsContext().SerializeContextChain();
       break;
     }
     case IrOpcode::kJSCreateArray: {

src/compiler/serializer-for-background-compilation.cc

@@ -1106,13 +1106,11 @@ void SerializerForBackgroundCompilation::VisitPopContext(
 void SerializerForBackgroundCompilation::ProcessImmutableLoad(
     ContextRef& context_ref, int slot, ContextProcessingMode mode) {
   DCHECK(mode == kSerializeSlot || mode == kSerializeSlotAndAddToAccumulator);
-  context_ref.SerializeSlot(slot);
+  base::Optional<ObjectRef> slot_value = context_ref.get(slot, true);
 
   // Also, put the object into the constant hints for the accumulator.
-  if (mode == kSerializeSlotAndAddToAccumulator) {
-    Handle<Object> slot_value(context_ref.object()->get(slot),
-                              broker()->isolate());
-    environment()->accumulator_hints().AddConstant(slot_value);
+  if (mode == kSerializeSlotAndAddToAccumulator && slot_value.has_value()) {
+    environment()->accumulator_hints().AddConstant(slot_value.value().object());
   }
 }
 
@@ -1128,31 +1126,23 @@ void SerializerForBackgroundCompilation::ProcessContextAccess(
     if (x->IsContext()) {
       // Walk this context to the given depth and serialize the slot found.
       ContextRef context_ref(broker(), x);
-      // TODO(mvstanton): Add a depth parameter to SerializeContextChain.
-      context_ref.SerializeContextChain();
-      if (mode != kIgnoreSlot) {
       size_t remaining_depth = depth;
-        context_ref = context_ref.previous(&remaining_depth);
-        if (remaining_depth == 0) {
+      context_ref = context_ref.previous(&remaining_depth, true);
+      if (remaining_depth == 0 && mode != kIgnoreSlot) {
         ProcessImmutableLoad(context_ref, slot, mode);
       }
     }
   }
-  }
   for (auto x : context_hints.virtual_contexts()) {
     if (x.distance <= static_cast<unsigned int>(depth)) {
       ContextRef context_ref(broker(), x.context);
-      // TODO(mvstanton): Add a depth parameter to SerializeContextChain.
-      context_ref.SerializeContextChain();
-      if (mode != kIgnoreSlot) {
       size_t remaining_depth = depth - x.distance;
-        context_ref = context_ref.previous(&remaining_depth);
-        if (remaining_depth == 0) {
+      context_ref = context_ref.previous(&remaining_depth, true);
+      if (remaining_depth == 0 && mode != kIgnoreSlot) {
         ProcessImmutableLoad(context_ref, slot, mode);
       }
     }
   }
-  }
 }
 
 void SerializerForBackgroundCompilation::VisitLdaContextSlot(
@@ -1437,7 +1427,7 @@ void SerializerForBackgroundCompilation::VisitCallJSRuntime(
   // BytecodeGraphBuilder::VisitCallJSRuntime needs the {runtime_index}
   // slot in the native context to be serialized.
   const int runtime_index = iterator->GetNativeContextIndexOperand(0);
-  broker()->native_context().SerializeSlot(runtime_index);
+  broker()->native_context().get(runtime_index, true);
 }
 
 Hints SerializerForBackgroundCompilation::RunChildSerializer(