[compiler] Generalize JSContextSpecialization.

src/compiler/js-context-specialization.cc

Index: src/compiler/js-context-specialization.cc
diff --git a/src/compiler/js-context-specialization.cc b/src/compiler/js-context-specialization.cc
index 2b12d3add6bf8ae6a6c30ffe6b9f7086ed76513e..9a2edc13e31ff1c830d829deb2ac2ae9aac99021 100644
--- a/src/compiler/js-context-specialization.cc
+++ b/src/compiler/js-context-specialization.cc
@@ -28,50 +28,81 @@ Reduction JSContextSpecialization::Reduce(Node* node) {
   return NoChange();
 }
 
+Reduction JSContextSpecialization::SimplifyJSLoadContext(Node* node,
+                                                         Node* new_context,
+                                                         size_t new_depth) {
+  DCHECK_EQ(IrOpcode::kJSLoadContext, node->opcode());
+  const ContextAccess& access = ContextAccessOf(node->op());
+  DCHECK_LE(new_depth, access.depth());
+
+  if (new_depth == access.depth() &&
+      new_context == NodeProperties::GetContextInput(node)) {
+    return NoChange();
+  }
 
-MaybeHandle<Context> JSContextSpecialization::GetSpecializationContext(
-    Node* node) {
-  DCHECK(node->opcode() == IrOpcode::kJSLoadContext ||
-         node->opcode() == IrOpcode::kJSStoreContext);
-  Node* const object = NodeProperties::GetContextInput(node);
-  return NodeProperties::GetSpecializationContext(object, context());
+  const Operator* op = jsgraph_->javascript()->LoadContext(
+      new_depth, access.index(), access.immutable());
+  NodeProperties::ReplaceContextInput(node, new_context);
+  NodeProperties::ChangeOp(node, op);
+  return Changed(node);
 }
 
+Reduction JSContextSpecialization::SimplifyJSStoreContext(Node* node,
+                                                          Node* new_context,
+                                                          size_t new_depth) {
+  DCHECK_EQ(IrOpcode::kJSStoreContext, node->opcode());
+  const ContextAccess& access = ContextAccessOf(node->op());
+  DCHECK_LE(new_depth, access.depth());
+
+  if (new_depth == access.depth() &&
+      new_context == NodeProperties::GetContextInput(node)) {
+    return NoChange();
+  }
+
+  const Operator* op =
+      jsgraph_->javascript()->StoreContext(new_depth, access.index());
+  NodeProperties::ReplaceContextInput(node, new_context);
+  NodeProperties::ChangeOp(node, op);
+  return Changed(node);
+}
 
 Reduction JSContextSpecialization::ReduceJSLoadContext(Node* node) {
   DCHECK_EQ(IrOpcode::kJSLoadContext, node->opcode());
 
-  // Get the specialization context from the node.
-  Handle<Context> context;
-  if (!GetSpecializationContext(node).ToHandle(&context)) return NoChange();
-
-  // Find the right parent context.
   const ContextAccess& access = ContextAccessOf(node->op());
-  for (size_t i = access.depth(); i > 0; --i) {
-    context = handle(context->previous(), isolate());
+  size_t depth = access.depth();
+
+  // First walk up the context chain in the graph as far as possible.
+  Node* outer = NodeProperties::GetOuterContext(node, &depth);
+
+  Handle<Context> concrete;
+  if (!NodeProperties::GetSpecializationContext(outer, context())
+           .ToHandle(&concrete)) {
+    // We do not have a concrete context object, so we can only partially reduce
+    // the load by folding-in the outer context node.
+    return SimplifyJSLoadContext(node, outer, depth);
+  }
+
+  // Now walk up the concrete context chain for the remaining depth.
+  for (; depth > 0; --depth) {
+    concrete = handle(concrete->previous(), isolate());
   }
 
-  // If the access itself is mutable, only fold-in the parent.
   if (!access.immutable()) {
-    // The access does not have to look up a parent, nothing to fold.
-    if (access.depth() == 0) {
-      return NoChange();
-    }
-    const Operator* op = jsgraph_->javascript()->LoadContext(
-        0, access.index(), access.immutable());
-    NodeProperties::ReplaceContextInput(node, jsgraph_->Constant(context));
-    NodeProperties::ChangeOp(node, op);
-    return Changed(node);
+    // We found the requested context object but since the context slot is
+    // mutable we can only partially reduce the load.
+    return SimplifyJSLoadContext(node, jsgraph()->Constant(concrete), depth);
   }
-  Handle<Object> value =
-      handle(context->get(static_cast<int>(access.index())), isolate());
 
   // Even though the context slot is immutable, the context might have escaped
   // before the function to which it belongs has initialized the slot.
-  // We must be conservative and check if the value in the slot is currently the
-  // hole or undefined. If it is neither of these, then it must be initialized.
+  // We must be conservative and check if the value in the slot is currently
+  // the hole or undefined. Only if it is neither of these, can we be sure that
+  // it won't change anymore.
+  Handle<Object> value(concrete->get(static_cast<int>(access.index())),
+                       isolate());
   if (value->IsUndefined(isolate()) || value->IsTheHole(isolate())) {
-    return NoChange();
+    return SimplifyJSLoadContext(node, jsgraph()->Constant(concrete), depth);
   }
 
   // Success. The context load can be replaced with the constant.
@@ -86,24 +117,27 @@ Reduction JSContextSpecialization::ReduceJSLoadContext(Node* node) {
 Reduction JSContextSpecialization::ReduceJSStoreContext(Node* node) {
   DCHECK_EQ(IrOpcode::kJSStoreContext, node->opcode());
 
-  // Get the specialization context from the node.
-  Handle<Context> context;
-  if (!GetSpecializationContext(node).ToHandle(&context)) return NoChange();
-
-  // The access does not have to look up a parent, nothing to fold.
   const ContextAccess& access = ContextAccessOf(node->op());
-  if (access.depth() == 0) {
-    return NoChange();
+  size_t depth = access.depth();
+
+  // First walk up the context chain in the graph until we reduce the depth to 0
+  // or hit a node that does not have a CreateXYZContext operator.
+  Node* outer = NodeProperties::GetOuterContext(node, &depth);
+
+  Handle<Context> concrete;
+  if (!NodeProperties::GetSpecializationContext(outer, context())
+           .ToHandle(&concrete)) {
+    // We do not have a concrete context object, so we can only partially reduce
+    // the load by folding-in the outer context node.
+    return SimplifyJSStoreContext(node, outer, depth);
   }
 
-  // Find the right parent context.
-  for (size_t i = access.depth(); i > 0; --i) {
-    context = handle(context->previous(), isolate());
+  // Now walk up the concrete context chain for the remaining depth.
+  for (; depth > 0; --depth) {
+    concrete = handle(concrete->previous(), isolate());
   }
 
-  NodeProperties::ReplaceContextInput(node, jsgraph_->Constant(context));
-  NodeProperties::ChangeOp(node, javascript()->StoreContext(0, access.index()));
-  return Changed(node);
+  return SimplifyJSStoreContext(node, jsgraph()->Constant(concrete), depth);
 }