[wasm] Introduce WasmContextCacheNodes to simplify WasmGraphBuilder interface.

This CL introduces a small struct to hold the {mem_start} and {mem_size}
node pointers that are managed in the function body decoder's SSA environment.
This struct insulates the function body decoder from further changes in
how context-specific information is represented in the compiler.

R=clemensh@chromium.org
CC=​mstarzinger@chromium.org

Bug: 
Change-Id: If17bef9fd2490ac11e4f3b3614f91333bb0b9528
Reviewed-on: https://chromium-review.googlesource.com/817282Reviewed-by: Clemens Hammacher <clemensh@chromium.org>
Commit-Queue: Ben L. Titzer <titzer@google.com>
Cr-Commit-Position: refs/heads/master@{#50004}

src/compiler/wasm-compiler.cc

@@ -122,10 +122,6 @@ Node* WasmGraphBuilder::Terminate(Node* effect, Node* control) {
   return terminate;
 }
 
-unsigned WasmGraphBuilder::InputCount(Node* node) {
-  return static_cast<unsigned>(node->InputCount());
-}
-
 bool WasmGraphBuilder::IsPhiWithMerge(Node* phi, Node* merge) {
   return phi && IrOpcode::IsPhiOpcode(phi->opcode()) &&
          NodeProperties::GetControlInput(phi) == merge;
@@ -3230,28 +3226,21 @@ void WasmGraphBuilder::BuildCWasmEntry(Address wasm_context_address) {
   }
 }
 
-// This function is used by WasmFullDecoder to create a node that loads the
-// {mem_start} variable from the WasmContext. It should not be used directly by
-// the WasmGraphBuilder. The WasmGraphBuilder should directly use {mem_start_},
-// which will always contain the correct node (stored in the SsaEnv).
-Node* WasmGraphBuilder::LoadMemStart() {
+void WasmGraphBuilder::InitContextCache(WasmContextCacheNodes* context_cache) {
   DCHECK_NOT_NULL(wasm_context_);
-  Node* mem_buffer = graph()->NewNode(
+  DCHECK_NOT_NULL(*control_);
+  DCHECK_NOT_NULL(*effect_);
+  // Load the memory start.
+  Node* mem_start = graph()->NewNode(
       jsgraph()->machine()->Load(MachineType::UintPtr()), wasm_context_,
       jsgraph()->Int32Constant(
           static_cast<int32_t>(offsetof(WasmContext, mem_start))),
       *effect_, *control_);
-  *effect_ = mem_buffer;
-  return mem_buffer;
-}
+  *effect_ = mem_start;
 
-// This function is used by WasmFullDecoder to create a node that loads the
-// {mem_size} variable from the WasmContext. It should not be used directly by
-// the WasmGraphBuilder. The WasmGraphBuilder should directly use {mem_size_},
-// which will always contain the correct node (stored in the SsaEnv).
-Node* WasmGraphBuilder::LoadMemSize() {
-  // Load mem_size from the WasmContext at runtime.
-  DCHECK_NOT_NULL(wasm_context_);
+  context_cache->mem_start = mem_start;
+
+  // Load the memory size.
   Node* mem_size = graph()->NewNode(
       jsgraph()->machine()->Load(MachineType::Uint32()), wasm_context_,
       jsgraph()->Int32Constant(
@@ -3262,7 +3251,69 @@ Node* WasmGraphBuilder::LoadMemSize() {
     mem_size = graph()->NewNode(jsgraph()->machine()->ChangeUint32ToUint64(),
                                 mem_size);
   }
-  return mem_size;
+
+  context_cache->mem_size = mem_size;
+}
+
+void WasmGraphBuilder::PrepareContextCacheForLoop(
+    WasmContextCacheNodes* context_cache, Node* control) {
+  // Introduce phis for mem_size and mem_start.
+  context_cache->mem_size =
+      Phi(MachineRepresentation::kWord32, 1, &context_cache->mem_size, control);
+  context_cache->mem_start = Phi(MachineType::PointerRepresentation(), 1,
+                                 &context_cache->mem_start, control);
+}
+
+void WasmGraphBuilder::NewContextCacheMerge(WasmContextCacheNodes* to,
+                                            WasmContextCacheNodes* from,
+                                            Node* merge) {
+  if (to->mem_size != from->mem_size) {
+    Node* vals[] = {to->mem_size, from->mem_size};
+    to->mem_size = Phi(MachineRepresentation::kWord32, 2, vals, merge);
+  }
+  if (to->mem_start != from->mem_start) {
+    Node* vals[] = {to->mem_start, from->mem_start};
+    to->mem_start = Phi(MachineType::PointerRepresentation(), 2, vals, merge);
+  }
+}
+
+void WasmGraphBuilder::MergeContextCacheInto(WasmContextCacheNodes* to,
+                                             WasmContextCacheNodes* from,
+                                             Node* merge) {
+  to->mem_size = CreateOrMergeIntoPhi(MachineRepresentation::kWord32, merge,
+                                      to->mem_size, from->mem_size);
+  to->mem_start = CreateOrMergeIntoPhi(MachineType::PointerRepresentation(),
+                                       merge, to->mem_start, from->mem_start);
+}
+
+Node* WasmGraphBuilder::CreateOrMergeIntoPhi(wasm::ValueType type, Node* merge,
+                                             Node* tnode, Node* fnode) {
+  if (IsPhiWithMerge(tnode, merge)) {
+    AppendToPhi(tnode, fnode);
+  } else if (tnode != fnode) {
+    uint32_t count = merge->InputCount();
+    Node** vals = Buffer(count);
+    for (uint32_t j = 0; j < count - 1; j++) vals[j] = tnode;
+    vals[count - 1] = fnode;
+    return Phi(type, count, vals, merge);
+  }
+  return tnode;
+}
+
+Node* WasmGraphBuilder::CreateOrMergeIntoEffectPhi(Node* merge, Node* tnode,
+                                                   Node* fnode) {
+  if (IsPhiWithMerge(tnode, merge)) {
+    AppendToPhi(tnode, fnode);
+  } else if (tnode != fnode) {
+    uint32_t count = merge->InputCount();
+    Node** effects = Buffer(count);
+    for (uint32_t j = 0; j < count - 1; j++) {
+      effects[j] = tnode;
+    }
+    effects[count - 1] = fnode;
+    tnode = EffectPhi(count, effects, merge);
+  }
+  return tnode;
 }
 
 void WasmGraphBuilder::GetGlobalBaseAndOffset(MachineType mem_type,
@@ -3299,17 +3350,20 @@ void WasmGraphBuilder::GetGlobalBaseAndOffset(MachineType mem_type,
 }
 
 Node* WasmGraphBuilder::MemBuffer(uint32_t offset) {
-  DCHECK_NOT_NULL(*mem_start_);
-  if (offset == 0) return *mem_start_;
-  return graph()->NewNode(jsgraph()->machine()->IntAdd(), *mem_start_,
+  DCHECK_NOT_NULL(context_cache_);
+  Node* mem_start = context_cache_->mem_start;
+  DCHECK_NOT_NULL(mem_start);
+  if (offset == 0) return mem_start;
+  return graph()->NewNode(jsgraph()->machine()->IntAdd(), mem_start,
                           jsgraph()->IntPtrConstant(offset));
 }
 
 Node* WasmGraphBuilder::CurrentMemoryPages() {
   // CurrentMemoryPages can not be called from asm.js.
   DCHECK_EQ(wasm::kWasmOrigin, env_->module->origin());
-  DCHECK_NOT_NULL(*mem_size_);
-  Node* mem_size = *mem_size_;
+  DCHECK_NOT_NULL(context_cache_);
+  Node* mem_size = context_cache_->mem_size;
+  DCHECK_NOT_NULL(mem_size);
   if (jsgraph()->machine()->Is64()) {
     mem_size = graph()->NewNode(jsgraph()->machine()->TruncateInt64ToInt32(),
                                 mem_size);
@@ -3435,7 +3489,9 @@ void WasmGraphBuilder::BoundsCheckMem(MachineType memtype, Node* index,
                                       uint32_t offset,
                                       wasm::WasmCodePosition position) {
   if (FLAG_wasm_no_bounds_checks) return;
-  DCHECK_NOT_NULL(*mem_size_);
+  DCHECK_NOT_NULL(context_cache_);
+  Node* mem_size = context_cache_->mem_size;
+  DCHECK_NOT_NULL(mem_size);
 
   uint32_t min_size = env_->module->initial_pages * wasm::WasmModule::kPageSize;
   uint32_t max_size =
@@ -3460,12 +3516,11 @@ void WasmGraphBuilder::BoundsCheckMem(MachineType memtype, Node* index,
     Node* cond;
     if (jsgraph()->machine()->Is32()) {
       cond = graph()->NewNode(jsgraph()->machine()->Uint32LessThanOrEqual(),
-                              jsgraph()->Int32Constant(end_offset), *mem_size_);
+                              jsgraph()->Int32Constant(end_offset), mem_size);
     } else {
       cond = graph()->NewNode(
           jsgraph()->machine()->Uint64LessThanOrEqual(),
-          jsgraph()->Int64Constant(static_cast<int64_t>(end_offset)),
-          *mem_size_);
+          jsgraph()->Int64Constant(static_cast<int64_t>(end_offset)), mem_size);
     }
     TrapIfFalse(wasm::kTrapMemOutOfBounds, cond, position);
   } else {
@@ -3485,11 +3540,11 @@ void WasmGraphBuilder::BoundsCheckMem(MachineType memtype, Node* index,
   Node* effective_size;
   if (jsgraph()->machine()->Is32()) {
     effective_size =
-        graph()->NewNode(jsgraph()->machine()->Int32Sub(), *mem_size_,
+        graph()->NewNode(jsgraph()->machine()->Int32Sub(), mem_size,
                          jsgraph()->Int32Constant(end_offset - 1));
   } else {
     effective_size = graph()->NewNode(
-        jsgraph()->machine()->Int64Sub(), *mem_size_,
+        jsgraph()->machine()->Int64Sub(), mem_size,
         jsgraph()->Int64Constant(static_cast<int64_t>(end_offset - 1)));
   }
 
@@ -3658,15 +3713,18 @@ Node* WasmGraphBuilder::StoreMem(MachineType memtype, Node* index,
 Node* WasmGraphBuilder::BuildAsmjsLoadMem(MachineType type, Node* index) {
   // TODO(turbofan): fold bounds checks for constant asm.js loads.
   // asm.js semantics use CheckedLoad (i.e. OOB reads return 0ish).
-  DCHECK_NOT_NULL(*mem_size_);
-  DCHECK_NOT_NULL(*mem_start_);
+  DCHECK_NOT_NULL(context_cache_);
+  Node* mem_start = context_cache_->mem_start;
+  Node* mem_size = context_cache_->mem_size;
+  DCHECK_NOT_NULL(mem_start);
+  DCHECK_NOT_NULL(mem_size);
   if (jsgraph()->machine()->Is64()) {
     index =
         graph()->NewNode(jsgraph()->machine()->ChangeUint32ToUint64(), index);
   }
   const Operator* op = jsgraph()->machine()->CheckedLoad(type);
   Node* load =
-      graph()->NewNode(op, *mem_start_, index, *mem_size_, *effect_, *control_);
+      graph()->NewNode(op, mem_start, index, mem_size, *effect_, *control_);
   *effect_ = load;
   return load;
 }
@@ -3675,16 +3733,19 @@ Node* WasmGraphBuilder::BuildAsmjsStoreMem(MachineType type, Node* index,
                                            Node* val) {
   // TODO(turbofan): fold bounds checks for constant asm.js stores.
   // asm.js semantics use CheckedStore (i.e. ignore OOB writes).
-  DCHECK_NOT_NULL(*mem_size_);
-  DCHECK_NOT_NULL(*mem_start_);
+  DCHECK_NOT_NULL(context_cache_);
+  Node* mem_start = context_cache_->mem_start;
+  Node* mem_size = context_cache_->mem_size;
+  DCHECK_NOT_NULL(mem_start);
+  DCHECK_NOT_NULL(mem_size);
   if (jsgraph()->machine()->Is64()) {
     index =
         graph()->NewNode(jsgraph()->machine()->ChangeUint32ToUint64(), index);
   }
   const Operator* op =
       jsgraph()->machine()->CheckedStore(type.representation());
-  Node* store = graph()->NewNode(op, *mem_start_, index, *mem_size_, val,
-                                 *effect_, *control_);
+  Node* store = graph()->NewNode(op, mem_start, index, mem_size, val, *effect_,
+                                 *control_);
   *effect_ = store;
   return val;
 }

src/compiler/wasm-compiler.h

@@ -209,6 +209,16 @@ enum CWasmEntryParameters {
 Handle<Code> CompileCWasmEntry(Isolate* isolate, wasm::FunctionSig* sig,
                                Address wasm_context_address);
 
+// Values from the {WasmContext} are cached between WASM-level function calls.
+// This struct allows the SSA environment handling this cache to be defined
+// and manipulated in wasm-compiler.{h,cc} instead of inside the WASM decoder.
+// (Note that currently, the globals base is immutable in a context, so not
+// cached here.)
+struct WasmContextCacheNodes {
+  Node* mem_start;
+  Node* mem_size;
+};
+
 // Abstracts details of building TurboFan graph nodes for wasm to separate
 // the wasm decoder from the internal details of TurboFan.
 typedef ZoneVector<Node*> NodeVector;
@@ -239,6 +249,9 @@ class WasmGraphBuilder {
   Node* Terminate(Node* effect, Node* control);
   Node* Merge(unsigned count, Node** controls);
   Node* Phi(wasm::ValueType type, unsigned count, Node** vals, Node* control);
+  Node* CreateOrMergeIntoPhi(wasm::ValueType type, Node* merge, Node* tnode,
+                             Node* fnode);
+  Node* CreateOrMergeIntoEffectPhi(Node* merge, Node* tnode, Node* fnode);
   Node* EffectPhi(unsigned count, Node** effects, Node* control);
   Node* NumberConstant(int32_t value);
   Node* Uint32Constant(uint32_t value);
@@ -259,7 +272,6 @@ class WasmGraphBuilder {
   Node* ConvertExceptionTagToRuntimeId(uint32_t tag);
   Node* GetExceptionRuntimeId();
   Node** GetExceptionValues(const wasm::WasmException* except_decl);
-  unsigned InputCount(Node* node);
   bool IsPhiWithMerge(Node* phi, Node* merge);
   bool ThrowsException(Node* node, Node** if_success, Node** if_exception);
   void AppendToMerge(Node* merge, Node* from);
@@ -362,9 +374,18 @@ class WasmGraphBuilder {
   void GetGlobalBaseAndOffset(MachineType mem_type, uint32_t offset,
                               Node** base_node, Node** offset_node);
 
-  void set_mem_size(Node** mem_size) { this->mem_size_ = mem_size; }
-
-  void set_mem_start(Node** mem_start) { this->mem_start_ = mem_start; }
+  // Utilities to manipulate sets of context cache nodes.
+  void InitContextCache(WasmContextCacheNodes* context_cache);
+  void PrepareContextCacheForLoop(WasmContextCacheNodes* context_cache,
+                                  Node* control);
+  void NewContextCacheMerge(WasmContextCacheNodes* to,
+                            WasmContextCacheNodes* from, Node* merge);
+  void MergeContextCacheInto(WasmContextCacheNodes* to,
+                             WasmContextCacheNodes* from, Node* merge);
+
+  void set_context_cache(WasmContextCacheNodes* context_cache) {
+    this->context_cache_ = context_cache;
+  }
 
   wasm::FunctionSig* GetFunctionSignature() { return sig_; }
 
@@ -413,8 +434,7 @@ class WasmGraphBuilder {
   NodeVector function_table_sizes_;
   Node** control_ = nullptr;
   Node** effect_ = nullptr;
-  Node** mem_size_ = nullptr;
-  Node** mem_start_ = nullptr;
+  WasmContextCacheNodes* context_cache_ = nullptr;
   Node* globals_start_ = nullptr;
   Node** cur_buffer_;
   size_t cur_bufsize_;

src/wasm/function-body-decoder-impl.h

@@ -716,9 +716,9 @@ class WasmDecoder : public Decoder {
         case kExprGrowMemory:
         case kExprCallFunction:
         case kExprCallIndirect:
-          // Add mem_size and mem_start to the assigned set.
-          assigned->Add(locals_count - 2);  // mem_size
-          assigned->Add(locals_count - 1);  // mem_start
+          // Add context cache nodes to the assigned set.
+          // TODO(titzer): make this more clear.
+          assigned->Add(locals_count - 1);
           length = OpcodeLength(decoder, pc);
           break;
         case kExprEnd:

src/wasm/function-body-decoder.cc

@@ -37,8 +37,7 @@ struct SsaEnv {
   State state;
   TFNode* control;
   TFNode* effect;
-  TFNode* mem_size;
-  TFNode* mem_start;
+  compiler::WasmContextCacheNodes context_cache;
   TFNode** locals;
 
   bool go() { return state >= kReached; }
@@ -47,8 +46,7 @@ struct SsaEnv {
     locals = nullptr;
     control = nullptr;
     effect = nullptr;
-    mem_size = nullptr;
-    mem_start = nullptr;
+    context_cache = {0};
   }
   void SetNotMerged() {
     if (state == kMerged) state = kReached;
@@ -94,8 +92,7 @@ class WasmGraphBuildingInterface {
     SsaEnv* ssa_env =
         reinterpret_cast<SsaEnv*>(decoder->zone()->New(sizeof(SsaEnv)));
     uint32_t num_locals = decoder->NumLocals();
-    // The '+ 2' here is to accommodate for mem_size and mem_start nodes.
-    uint32_t env_count = num_locals + 2;
+    uint32_t env_count = num_locals;
     size_t size = sizeof(TFNode*) * env_count;
     ssa_env->state = SsaEnv::kReached;
     ssa_env->locals =
@@ -128,20 +125,15 @@ class WasmGraphBuildingInterface {
     // Initialize effect and control before loading the context.
     builder_->set_effect_ptr(&ssa_env->effect);
     builder_->set_control_ptr(&ssa_env->control);
-    // Always load mem_size and mem_start from the WasmContext into the ssa.
     LoadContextIntoSsa(ssa_env);
     SetEnv(ssa_env);
   }
 
   // Reload the wasm context variables from the WasmContext structure attached
-  // to the memory object into the Ssa Environment. This does not automatically
-  // set the mem_size_ and mem_start_ pointers in WasmGraphBuilder.
+  // to the memory object into the Ssa Environment.
   void LoadContextIntoSsa(SsaEnv* ssa_env) {
     if (!ssa_env || !ssa_env->go()) return;
-    DCHECK_NOT_NULL(builder_->Effect());
-    DCHECK_NOT_NULL(builder_->Control());
-    ssa_env->mem_size = builder_->LoadMemSize();
-    ssa_env->mem_start = builder_->LoadMemStart();
+    builder_->InitContextCache(&ssa_env->context_cache);
   }
 
   void StartFunctionBody(Decoder* decoder, Control* block) {
@@ -368,7 +360,7 @@ class WasmGraphBuildingInterface {
 
   void GrowMemory(Decoder* decoder, const Value& value, Value* result) {
     result->node = BUILD(GrowMemory, value.node);
-    // Reload mem_size and mem_start after growing memory.
+    // Always reload the context cache after growing memory.
     LoadContextIntoSsa(ssa_env_);
   }
 
@@ -551,13 +543,10 @@ class WasmGraphBuildingInterface {
     }
 #endif
     ssa_env_ = env;
-    // TODO(wasm): Create a WasmEnv class with control, effect, mem_size and
-    // mem_start. SsaEnv can inherit from it. This way WasmEnv can be passed
-    // directly to WasmGraphBuilder instead of always copying four pointers.
+    // TODO(wasm): combine the control and effect pointers with context cache.
     builder_->set_control_ptr(&env->control);
     builder_->set_effect_ptr(&env->effect);
-    builder_->set_mem_size(&env->mem_size);
-    builder_->set_mem_start(&env->mem_start);
+    builder_->set_context_cache(&env->context_cache);
   }
 
   TFNode* CheckForException(Decoder* decoder, TFNode* node) {
@@ -587,8 +576,8 @@ class WasmGraphBuildingInterface {
     } else {
       DCHECK_EQ(SsaEnv::kMerged, try_info->catch_env->state);
       try_info->exception =
-          CreateOrMergeIntoPhi(kWasmI32, try_info->catch_env->control,
-                               try_info->exception, if_exception);
+          builder_->CreateOrMergeIntoPhi(kWasmI32, try_info->catch_env->control,
+                                         try_info->exception, if_exception);
     }
 
     SetEnv(success_env);
@@ -629,8 +618,8 @@ class WasmGraphBuildingInterface {
       DCHECK_NOT_NULL(val.node);
       DCHECK(val.type == old.type || val.type == kWasmVar);
       old.node = first ? val.node
-                       : CreateOrMergeIntoPhi(old.type, target->control,
-                                              old.node, val.node);
+                       : builder_->CreateOrMergeIntoPhi(
+                             old.type, target->control, old.node, val.node);
     }
   }
 
@@ -643,8 +632,7 @@ class WasmGraphBuildingInterface {
         to->locals = from->locals;
         to->control = from->control;
         to->effect = from->effect;
-        to->mem_size = from->mem_size;
-        to->mem_start = from->mem_start;
+        to->context_cache = from->context_cache;
         break;
       }
       case SsaEnv::kReached: {  // Create a new merge.
@@ -668,47 +656,26 @@ class WasmGraphBuildingInterface {
                 builder_->Phi(decoder->GetLocalType(i), 2, vals, merge);
           }
         }
-        // Merge mem_size and mem_start.
-        if (to->mem_size != from->mem_size) {
-          TFNode* vals[] = {to->mem_size, from->mem_size};
-          to->mem_size =
-              builder_->Phi(MachineRepresentation::kWord32, 2, vals, merge);
-        }
-        if (to->mem_start != from->mem_start) {
-          TFNode* vals[] = {to->mem_start, from->mem_start};
-          to->mem_start = builder_->Phi(MachineType::PointerRepresentation(), 2,
-                                        vals, merge);
-        }
+        // Start a new merge from the context cache.
+        builder_->NewContextCacheMerge(&to->context_cache, &from->context_cache,
+                                       merge);
         break;
       }
       case SsaEnv::kMerged: {
         TFNode* merge = to->control;
-        // Extend the existing merge.
+        // Extend the existing merge control node.
         builder_->AppendToMerge(merge, from->control);
         // Merge effects.
-        if (builder_->IsPhiWithMerge(to->effect, merge)) {
-          builder_->AppendToPhi(to->effect, from->effect);
-        } else if (to->effect != from->effect) {
-          uint32_t count = builder_->InputCount(merge);
-          TFNode** effects = builder_->Buffer(count);
-          for (uint32_t j = 0; j < count - 1; j++) {
-            effects[j] = to->effect;
-          }
-          effects[count - 1] = from->effect;
-          to->effect = builder_->EffectPhi(count, effects, merge);
-        }
+        to->effect = builder_->CreateOrMergeIntoEffectPhi(merge, to->effect,
+                                                          from->effect);
         // Merge locals.
         for (int i = decoder->NumLocals() - 1; i >= 0; i--) {
-          to->locals[i] = CreateOrMergeIntoPhi(decoder->GetLocalType(i), merge,
-                                               to->locals[i], from->locals[i]);
+          to->locals[i] = builder_->CreateOrMergeIntoPhi(
+              decoder->GetLocalType(i), merge, to->locals[i], from->locals[i]);
         }
-        // Merge mem_size and mem_start.
-        to->mem_size =
-            CreateOrMergeIntoPhi(MachineRepresentation::kWord32, merge,
-                                 to->mem_size, from->mem_size);
-        to->mem_start =
-            CreateOrMergeIntoPhi(MachineType::PointerRepresentation(), merge,
-                                 to->mem_start, from->mem_start);
+        // Merge the context caches.
+        builder_->MergeContextCacheInto(&to->context_cache,
+                                        &from->context_cache, merge);
         break;
       }
       default:
@@ -717,20 +684,6 @@ class WasmGraphBuildingInterface {
     return from->Kill();
   }
 
-  TFNode* CreateOrMergeIntoPhi(ValueType type, TFNode* merge, TFNode* tnode,
-                               TFNode* fnode) {
-    if (builder_->IsPhiWithMerge(tnode, merge)) {
-      builder_->AppendToPhi(tnode, fnode);
-    } else if (tnode != fnode) {
-      uint32_t count = builder_->InputCount(merge);
-      TFNode** vals = builder_->Buffer(count);
-      for (uint32_t j = 0; j < count - 1; j++) vals[j] = tnode;
-      vals[count - 1] = fnode;
-      return builder_->Phi(type, count, vals, merge);
-    }
-    return tnode;
-  }
-
   SsaEnv* PrepareForLoop(Decoder* decoder, SsaEnv* env) {
     if (!env->go()) return Split(decoder, env);
     env->state = SsaEnv::kMerged;
@@ -738,27 +691,21 @@ class WasmGraphBuildingInterface {
     env->control = builder_->Loop(env->control);
     env->effect = builder_->EffectPhi(1, &env->effect, env->control);
     builder_->Terminate(env->effect, env->control);
-    // The '+ 2' here is to be able to set mem_size and mem_start as assigned.
+    // The '+ 1' here is to be able to set the context cache as assigned.
     BitVector* assigned = WasmDecoder<validate>::AnalyzeLoopAssignment(
-        decoder, decoder->pc(), decoder->total_locals() + 2, decoder->zone());
+        decoder, decoder->pc(), decoder->total_locals() + 1, decoder->zone());
     if (decoder->failed()) return env;
     if (assigned != nullptr) {
       // Only introduce phis for variables assigned in this loop.
-      int mem_size_index = decoder->total_locals();
-      int mem_start_index = decoder->total_locals() + 1;
+      int context_cache_index = decoder->total_locals();
       for (int i = decoder->NumLocals() - 1; i >= 0; i--) {
         if (!assigned->Contains(i)) continue;
         env->locals[i] = builder_->Phi(decoder->GetLocalType(i), 1,
                                        &env->locals[i], env->control);
       }
-      // Introduce phis for mem_size and mem_start if necessary.
-      if (assigned->Contains(mem_size_index)) {
-        env->mem_size = builder_->Phi(MachineRepresentation::kWord32, 1,
-                                      &env->mem_size, env->control);
-      }
-      if (assigned->Contains(mem_start_index)) {
-        env->mem_start = builder_->Phi(MachineType::PointerRepresentation(), 1,
-                                       &env->mem_start, env->control);
+      // Introduce phis for context cache pointers if necessary.
+      if (assigned->Contains(context_cache_index)) {
+        builder_->PrepareContextCacheForLoop(&env->context_cache, env->control);
       }
 
       SsaEnv* loop_body_env = Split(decoder, env);
@@ -773,11 +720,8 @@ class WasmGraphBuildingInterface {
                                      &env->locals[i], env->control);
     }
 
-    // Conservatively introduce phis for mem_size and mem_start.
-    env->mem_size = builder_->Phi(MachineRepresentation::kWord32, 1,
-                                  &env->mem_size, env->control);
-    env->mem_start = builder_->Phi(MachineType::PointerRepresentation(), 1,
-                                   &env->mem_start, env->control);
+    // Conservatively introduce phis for context cache.
+    builder_->PrepareContextCacheForLoop(&env->context_cache, env->control);
 
     SsaEnv* loop_body_env = Split(decoder, env);
     builder_->StackCheck(decoder->position(), &loop_body_env->effect,
@@ -791,7 +735,7 @@ class WasmGraphBuildingInterface {
     SsaEnv* result =
         reinterpret_cast<SsaEnv*>(decoder->zone()->New(sizeof(SsaEnv)));
     // The '+ 2' here is to accommodate for mem_size and mem_start nodes.
-    size_t size = sizeof(TFNode*) * (decoder->NumLocals() + 2);
+    size_t size = sizeof(TFNode*) * (decoder->NumLocals());
     result->control = from->control;
     result->effect = from->effect;
 
@@ -801,13 +745,11 @@ class WasmGraphBuildingInterface {
           size > 0 ? reinterpret_cast<TFNode**>(decoder->zone()->New(size))
                    : nullptr;
       memcpy(result->locals, from->locals, size);
-      result->mem_size = from->mem_size;
-      result->mem_start = from->mem_start;
+      result->context_cache = from->context_cache;
     } else {
       result->state = SsaEnv::kUnreachable;
       result->locals = nullptr;
-      result->mem_size = nullptr;
-      result->mem_start = nullptr;
+      result->context_cache = {0};
     }
 
     return result;
@@ -823,8 +765,7 @@ class WasmGraphBuildingInterface {
     result->locals = from->locals;
     result->control = from->control;
     result->effect = from->effect;
-    result->mem_size = from->mem_size;
-    result->mem_start = from->mem_start;
+    result->context_cache = from->context_cache;
     from->Kill(SsaEnv::kUnreachable);
     return result;
   }
@@ -836,8 +777,7 @@ class WasmGraphBuildingInterface {
     result->control = nullptr;
     result->effect = nullptr;
     result->locals = nullptr;
-    result->mem_size = nullptr;
-    result->mem_start = nullptr;
+    result->context_cache = {0};
     return result;
   }