Support for SetVReg and add SetGPR.

These changes are useful for debugging and are required for being able
to change all GC roots.

MIPS is untested.

Change-Id: I2ba055de64264098bffe869a4fb192d0975f1c8f
diff --git a/src/constants_mips.h b/src/constants_mips.h
index 32fa158..08661e9 100644
--- a/src/constants_mips.h
+++ b/src/constants_mips.h
@@ -27,14 +27,14 @@
 
 enum Register {
   ZERO =  0,
-  AT   =  1,
-  V0   =  2,
+  AT   =  1,  // Assembler temporary.
+  V0   =  2,  // Values.
   V1   =  3,
-  A0   =  4,
+  A0   =  4,  // Arguments.
   A1   =  5,
   A2   =  6,
   A3   =  7,
-  T0   =  8,
+  T0   =  8,  // Temporaries.
   T1   =  9,
   T2   = 10,
   T3   = 11,
@@ -42,7 +42,7 @@
   T5   = 13,
   T6   = 14,
   T7   = 15,
-  S0   = 16,
+  S0   = 16,  // Saved values.
   S1   = 17,
   S2   = 18,
   S3   = 19,
@@ -50,14 +50,14 @@
   S5   = 21,
   S6   = 22,
   S7   = 23,
-  T8   = 24,
+  T8   = 24,  // More temporaries.
   T9   = 25,
-  K0   = 26,
+  K0   = 26,  // Reserved for trap handler.
   K1   = 27,
-  GP   = 28,
-  SP   = 29,
-  FP   = 30,
-  RA   = 31,
+  GP   = 28,  // Global pointer.
+  SP   = 29,  // Stack pointer.
+  FP   = 30,  // Saved value/frame pointer.
+  RA   = 31,  // Return address.
   kNumberOfCoreRegisters = 32,
   kNoRegister = -1  // Signals an illegal register.
 };
diff --git a/src/oat/runtime/arm/context_arm.cc b/src/oat/runtime/arm/context_arm.cc
index 057f41e..5bd4b3d 100644
--- a/src/oat/runtime/arm/context_arm.cc
+++ b/src/oat/runtime/arm/context_arm.cc
@@ -21,16 +21,20 @@
 namespace art {
 namespace arm {
 
-ArmContext::ArmContext() {
-#ifndef NDEBUG
-  // Initialize registers with easy to spot debug values
-  for (int i = 0; i < 16; i++) {
-    gprs_[i] = kBadGprBase + i;
+static const uint32_t gZero = 0;
+
+void ArmContext::Reset() {
+  for (size_t i = 0; i < kNumberOfCoreRegisters; i++) {
+    gprs_[i] = NULL;
   }
-  for (int i = 0; i < 32; i++) {
-    fprs_[i] = kBadFprBase + i;
+  for (size_t i = 0; i < kNumberOfSRegisters; i++) {
+    fprs_[i] = NULL;
   }
-#endif
+  gprs_[SP] = &sp_;
+  gprs_[PC] = &pc_;
+  // Initialize registers with easy to spot debug values.
+  sp_ = ArmContext::kBadGprBase + SP;
+  pc_ = ArmContext::kBadGprBase + PC;
 }
 
 void ArmContext::FillCalleeSaves(const StackVisitor& fr) {
@@ -41,40 +45,55 @@
   size_t fp_spill_count = __builtin_popcount(fp_core_spills);
   size_t frame_size = method->GetFrameSizeInBytes();
   if (spill_count > 0) {
-    // Lowest number spill is furthest away, walk registers and fill into context
+    // Lowest number spill is farthest away, walk registers and fill into context
     int j = 1;
-    for (int i = 0; i < 16; i++) {
+    for (size_t i = 0; i < kNumberOfCoreRegisters; i++) {
       if (((core_spills >> i) & 1) != 0) {
-        gprs_[i] = fr.LoadCalleeSave(spill_count - j, frame_size);
+        gprs_[i] = fr.CalleeSaveAddress(spill_count - j, frame_size);
         j++;
       }
     }
   }
   if (fp_spill_count > 0) {
-    // Lowest number spill is furthest away, walk registers and fill into context
+    // Lowest number spill is farthest away, walk registers and fill into context
     int j = 1;
-    for (int i = 0; i < 32; i++) {
+    for (size_t i = 0; i < kNumberOfSRegisters; i++) {
       if (((fp_core_spills >> i) & 1) != 0) {
-        fprs_[i] = fr.LoadCalleeSave(spill_count + fp_spill_count - j, frame_size);
+        fprs_[i] = fr.CalleeSaveAddress(spill_count + fp_spill_count - j, frame_size);
         j++;
       }
     }
   }
 }
 
+void ArmContext::SetGPR(uint32_t reg, uintptr_t value) {
+  CHECK_LT(reg, kNumberOfCoreRegisters);
+  CHECK_NE(gprs_[reg], &gZero); // Can't overwrite this static value since they are never reset.
+  CHECK(gprs_[reg] != NULL);
+  *gprs_[reg] = value;
+}
+
 void ArmContext::SmashCallerSaves() {
-  gprs_[0] = 0; // This needs to be 0 because we want a null/zero return value.
-  gprs_[1] = kBadGprBase + 1;
-  gprs_[2] = kBadGprBase + 2;
-  gprs_[3] = kBadGprBase + 3;
-  gprs_[IP] = kBadGprBase + IP;
-  gprs_[LR] = kBadGprBase + LR;
+  // This needs to be 0 because we want a null/zero return value.
+  gprs_[R0] = const_cast<uint32_t*>(&gZero);
+  gprs_[R1] = const_cast<uint32_t*>(&gZero);
+  gprs_[R2] = NULL;
+  gprs_[R3] = NULL;
 }
 
 extern "C" void art_do_long_jump(uint32_t*, uint32_t*);
 
 void ArmContext::DoLongJump() {
-  art_do_long_jump(&gprs_[0], &fprs_[S0]);
+  uintptr_t gprs[16];
+  uint32_t fprs[32];
+  for (size_t i = 0; i < kNumberOfCoreRegisters; ++i) {
+    gprs[i] = gprs_[i] != NULL ? *gprs_[i] : ArmContext::kBadGprBase + i;
+  }
+  for (size_t i = 0; i < kNumberOfSRegisters; ++i) {
+    fprs[i] = fprs_[i] != NULL ? *fprs_[i] : ArmContext::kBadGprBase + i;
+  }
+  DCHECK_EQ(reinterpret_cast<uintptr_t>(Thread::Current()), gprs[TR]);
+  art_do_long_jump(gprs, fprs);
 }
 
 }  // namespace arm
diff --git a/src/oat/runtime/arm/context_arm.h b/src/oat/runtime/arm/context_arm.h
index 6f42cc3..a2b9ebe 100644
--- a/src/oat/runtime/arm/context_arm.h
+++ b/src/oat/runtime/arm/context_arm.h
@@ -17,6 +17,7 @@
 #ifndef ART_SRC_OAT_RUNTIME_ARM_CONTEXT_ARM_H_
 #define ART_SRC_OAT_RUNTIME_ARM_CONTEXT_ARM_H_
 
+#include "locks.h"
 #include "constants_arm.h"
 #include "oat/runtime/context.h"
 
@@ -25,31 +26,39 @@
 
 class ArmContext : public Context {
  public:
-  ArmContext();
+  ArmContext() {
+    Reset();
+  }
+
   virtual ~ArmContext() {}
 
+  virtual void Reset();
+
   virtual void FillCalleeSaves(const StackVisitor& fr);
 
   virtual void SetSP(uintptr_t new_sp) {
-    gprs_[SP] = new_sp;
+    SetGPR(SP, new_sp);
   }
 
   virtual void SetPC(uintptr_t new_pc) {
-    gprs_[PC] = new_pc;
+    SetGPR(PC, new_pc);
   }
 
   virtual uintptr_t GetGPR(uint32_t reg) {
-    CHECK_GE(reg, 0u);
-    CHECK_LT(reg, 16u);
-    return gprs_[reg];
+    CHECK_LT(reg, kNumberOfCoreRegisters);
+    return *gprs_[reg];
   }
 
+  virtual void SetGPR(uint32_t reg, uintptr_t value);
   virtual void SmashCallerSaves();
   virtual void DoLongJump();
 
  private:
-  uintptr_t gprs_[16];
-  uint32_t fprs_[32];
+  // Pointers to register locations, initialized to NULL or the specific registers below.
+  uintptr_t* gprs_[kNumberOfCoreRegisters];
+  uint32_t* fprs_[kNumberOfSRegisters];
+  // Hold values for sp and pc if they are not located within a stack frame.
+  uintptr_t sp_, pc_;
 };
 
 }  // namespace arm
diff --git a/src/oat/runtime/context.h b/src/oat/runtime/context.h
index 317030f..895abf9 100644
--- a/src/oat/runtime/context.h
+++ b/src/oat/runtime/context.h
@@ -24,7 +24,8 @@
 
 class StackVisitor;
 
-// Representation of a thread's context on the executing machine
+// Representation of a thread's context on the executing machine, used to implement long jumps in
+// the quick stack frame layout.
 class Context {
  public:
   // Creates a context for the running architecture
@@ -32,6 +33,9 @@
 
   virtual ~Context() {}
 
+  // Re-initializes the registers for context re-use.
+  virtual void Reset() = 0;
+
   // Read values from callee saves in the given frame. The frame also holds
   // the method that holds the layout.
   virtual void FillCalleeSaves(const StackVisitor& fr) = 0;
@@ -45,12 +49,16 @@
   // Read the given GPR
   virtual uintptr_t GetGPR(uint32_t reg) = 0;
 
+  // Set the given GPR.
+  virtual void SetGPR(uint32_t reg, uintptr_t value) = 0;
+
   // Smash the caller save registers. If we're throwing, we don't want to return bogus values.
   virtual void SmashCallerSaves() = 0;
 
   // Switch execution of the executing context to this context
   virtual void DoLongJump() = 0;
 
+ protected:
   enum {
     kBadGprBase = 0xebad6070,
     kBadFprBase = 0xebad8070,
diff --git a/src/oat/runtime/mips/context_mips.cc b/src/oat/runtime/mips/context_mips.cc
index dc13c63..0c2f915 100644
--- a/src/oat/runtime/mips/context_mips.cc
+++ b/src/oat/runtime/mips/context_mips.cc
@@ -21,16 +21,20 @@
 namespace art {
 namespace mips {
 
-MipsContext::MipsContext() {
-#ifndef NDEBUG
+static const uint32_t gZero = 0;
+
+void MipsContext::Reset() {
+  for (size_t i = 0; i < kNumberOfCoreRegisters; i++) {
+    gprs_[i] = NULL;
+  }
+  for (size_t i = 0; i < kNumberOfFRegisters; i++) {
+    fprs_[i] = NULL;
+  }
+  gprs_[SP] = &sp_;
+  gprs_[RA] = &ra_;
   // Initialize registers with easy to spot debug values.
-  for (int i = 0; i < 32; i++) {
-    gprs_[i] = kBadGprBase + i;
-  }
-  for (int i = 0; i < 32; i++) {
-    fprs_[i] = kBadGprBase + i;
-  }
-#endif
+  sp_ = MipsContext::kBadGprBase + SP;
+  ra_ = MipsContext::kBadGprBase + RA;
 }
 
 void MipsContext::FillCalleeSaves(const StackVisitor& fr) {
@@ -41,39 +45,55 @@
   size_t fp_spill_count = __builtin_popcount(fp_core_spills);
   size_t frame_size = method->GetFrameSizeInBytes();
   if (spill_count > 0) {
-    // Lowest number spill is furthest away, walk registers and fill into context.
+    // Lowest number spill is farthest away, walk registers and fill into context.
     int j = 1;
-    for (int i = 0; i < 32; i++) {
+    for (size_t i = 0; i < kNumberOfCoreRegisters; i++) {
       if (((core_spills >> i) & 1) != 0) {
-        gprs_[i] = fr.LoadCalleeSave(spill_count - j, frame_size);
+        gprs_[i] = fr.CalleeSaveAddress(spill_count - j, frame_size);
         j++;
       }
     }
   }
   if (fp_spill_count > 0) {
-    // Lowest number spill is furthest away, walk registers and fill into context.
+    // Lowest number spill is farthest away, walk registers and fill into context.
     int j = 1;
-    for (int i = 0; i < 32; i++) {
+    for (size_t i = 0; i < kNumberOfFRegisters; i++) {
       if (((fp_core_spills >> i) & 1) != 0) {
-        fprs_[i] = fr.LoadCalleeSave(spill_count + fp_spill_count - j, frame_size);
+        fprs_[i] = fr.CalleeSaveAddress(spill_count + fp_spill_count - j, frame_size);
         j++;
       }
     }
   }
 }
 
+void MipsContext::SetGPR(uint32_t reg, uintptr_t value) {
+  CHECK_LT(reg, kNumberOfCoreRegisters);
+  CHECK_NE(gprs_[reg], &gZero); // Can't overwrite this static value since they are never reset.
+  CHECK(gprs_[reg] != NULL);
+  *gprs_[reg] = value;
+}
+
 void MipsContext::SmashCallerSaves() {
-  gprs_[V0] = 0; // This needs to be 0 because we want a null/zero return value.
-  gprs_[V1] = 0; // This needs to be 0 because we want a null/zero return value.
-  gprs_[A1] = kBadGprBase + A1;
-  gprs_[A2] = kBadGprBase + A2;
-  gprs_[A3] = kBadGprBase + A3;
+  // This needs to be 0 because we want a null/zero return value.
+  gprs_[V0] = const_cast<uint32_t*>(&gZero);
+  gprs_[V1] = const_cast<uint32_t*>(&gZero);
+  gprs_[A1] = NULL;
+  gprs_[A2] = NULL;
+  gprs_[A3] = NULL;
 }
 
 extern "C" void art_do_long_jump(uint32_t*, uint32_t*);
 
 void MipsContext::DoLongJump() {
-  art_do_long_jump(&gprs_[ZERO], &fprs_[F0]);
+  uintptr_t gprs[kNumberOfCoreRegisters];
+  uint32_t fprs[kNumberOfFRegisters];
+  for (size_t i = 0; i < kNumberOfCoreRegisters; ++i) {
+    gprs[i] = gprs_[i] != NULL ? *gprs_[i] : MipsContext::kBadGprBase + i;
+  }
+  for (size_t i = 0; i < kNumberOfFRegisters; ++i) {
+    fprs[i] = fprs_[i] != NULL ? *fprs_[i] : MipsContext::kBadGprBase + i;
+  }
+  art_do_long_jump(gprs, fprs);
 }
 
 }  // namespace mips
diff --git a/src/oat/runtime/mips/context_mips.h b/src/oat/runtime/mips/context_mips.h
index 1a86ca3..d4944a6 100644
--- a/src/oat/runtime/mips/context_mips.h
+++ b/src/oat/runtime/mips/context_mips.h
@@ -25,32 +25,38 @@
 
 class MipsContext : public Context {
  public:
-  MipsContext();
+  MipsContext() {
+    Reset();
+  }
   virtual ~MipsContext() {}
 
-  // No callee saves on mips
+  virtual void Reset();
+
   virtual void FillCalleeSaves(const StackVisitor& fr);
 
   virtual void SetSP(uintptr_t new_sp) {
-    gprs_[SP] = new_sp;
+    SetGPR(SP, new_sp);
   }
 
   virtual void SetPC(uintptr_t new_pc) {
-    gprs_[RA] = new_pc;
+    SetGPR(RA, new_pc);
   }
 
   virtual uintptr_t GetGPR(uint32_t reg) {
-    CHECK_GE(reg, 0u);
-    CHECK_LT(reg, 32u);
+    CHECK_LT(reg, kNumberOfCoreRegisters);
     return gprs_[reg];
   }
 
+  virtual void SetGPR(uint32_t reg, uintptr_t value);
   virtual void SmashCallerSaves();
   virtual void DoLongJump();
 
  private:
-  uintptr_t gprs_[32];
-  uint32_t fprs_[32];
+  // Pointers to registers in the stack, initialized to NULL except for the special cases below.
+  uintptr_t* gprs_[kNumberOfCoreRegisters];
+  uint32_t* fprs_[kNumberOfFRegisters];
+  // Hold values for sp and ra (return address) if they are not located within a stack frame.
+  uintptr_t sp_, ra_;
 };
 }  // namespace mips
 }  // namespace art
diff --git a/src/oat/runtime/x86/context_x86.cc b/src/oat/runtime/x86/context_x86.cc
index 4ff2283..4efdf81 100644
--- a/src/oat/runtime/x86/context_x86.cc
+++ b/src/oat/runtime/x86/context_x86.cc
@@ -21,14 +21,16 @@
 namespace art {
 namespace x86 {
 
-X86Context::X86Context() {
-#ifndef NDEBUG
-  // Initialize registers with easy to spot debug values.
-  for (int i = 0; i < 8; i++) {
-    gprs_[i] = kBadGprBase + i;
+static const uint32_t gZero = 0;
+
+void X86Context::Reset() {
+  for (int i = 0; i < kNumberOfCpuRegisters; i++) {
+    gprs_[i] = NULL;
   }
-  eip_ = 0xEBAD601F;
-#endif
+  gprs_[ESP] = &esp_;
+  // Initialize registers with easy to spot debug values.
+  esp_ = X86Context::kBadGprBase + ESP;
+  eip_ = X86Context::kBadGprBase + kNumberOfCpuRegisters;
 }
 
 void X86Context::FillCalleeSaves(const StackVisitor& fr) {
@@ -38,11 +40,11 @@
   DCHECK_EQ(method->GetFpSpillMask(), 0u);
   size_t frame_size = method->GetFrameSizeInBytes();
   if (spill_count > 0) {
-    // Lowest number spill is furthest away, walk registers and fill into context.
+    // Lowest number spill is farthest away, walk registers and fill into context.
     int j = 2;  // Offset j to skip return address spill.
-    for (int i = 0; i < 8; i++) {
+    for (int i = 0; i < kNumberOfCpuRegisters; i++) {
       if (((core_spills >> i) & 1) != 0) {
-        gprs_[i] = fr.LoadCalleeSave(spill_count - j, frame_size);
+        gprs_[i] = fr.CalleeSaveAddress(spill_count - j, frame_size);
         j++;
       }
     }
@@ -50,37 +52,40 @@
 }
 
 void X86Context::SmashCallerSaves() {
-  gprs_[EAX] = 0; // This needs to be 0 because we want a null/zero return value.
-  gprs_[ECX] = kBadGprBase + ECX;
-  gprs_[EDX] = kBadGprBase + EDX;
-  gprs_[EBX] = kBadGprBase + EBX;
+  // This needs to be 0 because we want a null/zero return value.
+  gprs_[EAX] = const_cast<uint32_t*>(&gZero);
+  gprs_[EDX] = const_cast<uint32_t*>(&gZero);
+  gprs_[ECX] = NULL;
+  gprs_[EBX] = NULL;
+}
+
+void X86Context::SetGPR(uint32_t reg, uintptr_t value){
+  CHECK_LT(reg, kNumberOfCpuRegisters);
+  CHECK_NE(gprs_[reg], &gZero);
+  CHECK(gprs_[reg] != NULL);
+  *gprs_[reg] = value;
 }
 
 void X86Context::DoLongJump() {
 #if defined(__i386__)
-  // We push all the registers using memory-memory pushes, we then pop-all to get the registers
-  // set up, we then pop esp which will move us down the stack to the delivery address. At the frame
-  // where the exception will be delivered, we push EIP so that the return will take us to the
-  // correct delivery instruction.
-  gprs_[ESP] -= 4;
-  *(reinterpret_cast<uintptr_t*>(gprs_[ESP])) = eip_;
+  // Array of GPR values, filled from the context backward for the long jump pop. We add a slot at
+  // the top for the stack pointer that doesn't get popped in a pop-all.
+  volatile uintptr_t gprs[kNumberOfCpuRegisters + 1];
+  for (size_t i = 0; i < kNumberOfCpuRegisters; ++i) {
+    gprs[kNumberOfCpuRegisters - i - 1] = gprs_[i] != NULL ? *gprs_[i] : X86Context::kBadGprBase + i;
+  }
+  // We want to load the stack pointer one slot below so that the ret will pop eip.
+  uintptr_t esp = gprs[kNumberOfCpuRegisters - ESP - 1] - kWordSize;
+  gprs[kNumberOfCpuRegisters] = esp;
+  *(reinterpret_cast<uintptr_t*>(esp)) = eip_;
   __asm__ __volatile__(
-      "pushl %4\n\t"
-      "pushl %0\n\t"
-      "pushl %1\n\t"
-      "pushl %2\n\t"
-      "pushl %3\n\t"
-      "pushl %4\n\t"
-      "pushl %5\n\t"
-      "pushl %6\n\t"
-      "pushl %7\n\t"
-      "popal\n\t"
-      "popl %%esp\n\t"
-      "ret\n\t"
-      :  //output
-      : "g"(gprs_[EAX]), "g"(gprs_[ECX]), "g"(gprs_[EDX]), "g"(gprs_[EBX]),
-        "g"(gprs_[ESP]), "g"(gprs_[EBP]), "g"(gprs_[ESI]), "g"(gprs_[EDI])
-      :);  // clobber
+      "movl %0, %%esp\n\t"  // ESP points to gprs.
+      "popal\n\t"           // Load all registers except ESP and EIP with values in gprs.
+      "popl %%esp\n\t"      // Load stack pointer.
+      "ret\n\t"             // From higher in the stack pop eip.
+      :  // output.
+      : "g"(&gprs[0])  // input.
+      :);  // clobber.
 #else
     UNIMPLEMENTED(FATAL);
 #endif
diff --git a/src/oat/runtime/x86/context_x86.h b/src/oat/runtime/x86/context_x86.h
index 3d6b1d9..ad49f8d 100644
--- a/src/oat/runtime/x86/context_x86.h
+++ b/src/oat/runtime/x86/context_x86.h
@@ -25,14 +25,17 @@
 
 class X86Context : public Context {
  public:
-  X86Context();
+  X86Context() {
+    Reset();
+  }
   virtual ~X86Context() {}
 
-  // No callee saves on X86
+  virtual void Reset();
+
   virtual void FillCalleeSaves(const StackVisitor& fr);
 
   virtual void SetSP(uintptr_t new_sp) {
-    gprs_[ESP] = new_sp;
+    SetGPR(ESP, new_sp);
   }
 
   virtual void SetPC(uintptr_t new_pc) {
@@ -40,17 +43,23 @@
   }
 
   virtual uintptr_t GetGPR(uint32_t reg) {
-    CHECK_GE(reg, 0u);
-    CHECK_LT(reg, 8u);
-    return gprs_[reg];
+    CHECK_LT(reg, kNumberOfCpuRegisters);
+    return *gprs_[reg];
   }
 
+  virtual void SetGPR(uint32_t reg, uintptr_t value);
+
   virtual void SmashCallerSaves();
   virtual void DoLongJump();
 
  private:
-  uintptr_t gprs_[8];
-  uintptr_t eip_;
+  // Pointers to register locations, floating point registers are all caller save. Values are
+  // initialized to NULL or the special registers below.
+  uintptr_t* gprs_[kNumberOfCpuRegisters];
+  // Hold values for esp and eip if they are not located within a stack frame. EIP is somewhat
+  // special in that it cannot be encoded normally as a register operand to an instruction (except
+  // in 64bit addressing modes).
+  uintptr_t esp_, eip_;
 };
 }  // namespace x86
 }  // namespace art
diff --git a/src/stack.cc b/src/stack.cc
index be6fe45..ed44df9 100644
--- a/src/stack.cc
+++ b/src/stack.cc
@@ -99,16 +99,20 @@
     uint32_t vmap_offset;
     // TODO: IsInContext stops before spotting floating point registers.
     if (vmap_table.IsInContext(vreg, vmap_offset, kind)) {
-      UNIMPLEMENTED(FATAL);
+      bool is_float = (kind == kFloatVReg) || (kind == kDoubleLoVReg) || (kind == kDoubleHiVReg);
+      uint32_t spill_mask = is_float ? m->GetFpSpillMask() : m->GetCoreSpillMask();
+      const uint32_t reg = vmap_table.ComputeRegister(spill_mask, vmap_offset, kReferenceVReg);
+      SetGPR(reg, new_value);
+    } else {
+      const DexFile::CodeItem* code_item = MethodHelper(m).GetCodeItem();
+      DCHECK(code_item != NULL) << PrettyMethod(m); // Can't be NULL or how would we compile its instructions?
+      uint32_t core_spills = m->GetCoreSpillMask();
+      uint32_t fp_spills = m->GetFpSpillMask();
+      size_t frame_size = m->GetFrameSizeInBytes();
+      int offset = GetVRegOffset(code_item, core_spills, fp_spills, frame_size, vreg);
+      byte* vreg_addr = reinterpret_cast<byte*>(GetCurrentQuickFrame()) + offset;
+      *reinterpret_cast<uint32_t*>(vreg_addr) = new_value;
     }
-    const DexFile::CodeItem* code_item = MethodHelper(m).GetCodeItem();
-    DCHECK(code_item != NULL) << PrettyMethod(m); // Can't be NULL or how would we compile its instructions?
-    uint32_t core_spills = m->GetCoreSpillMask();
-    uint32_t fp_spills = m->GetFpSpillMask();
-    size_t frame_size = m->GetFrameSizeInBytes();
-    int offset = GetVRegOffset(code_item, core_spills, fp_spills, frame_size, vreg);
-    byte* vreg_addr = reinterpret_cast<byte*>(GetCurrentQuickFrame()) + offset;
-    *reinterpret_cast<uint32_t*>(vreg_addr) = new_value;
   } else {
     return cur_shadow_frame_->SetVReg(vreg, new_value);
   }
@@ -119,6 +123,11 @@
   return context_->GetGPR(reg);
 }
 
+void StackVisitor::SetGPR(uint32_t reg, uintptr_t value) {
+  DCHECK (cur_quick_frame_ != NULL) << "This is a quick frame routine";
+  context_->SetGPR(reg, value);
+}
+
 uintptr_t StackVisitor::GetReturnPc() const {
   AbstractMethod** sp = GetCurrentQuickFrame();
   DCHECK(sp != NULL);
diff --git a/src/stack.h b/src/stack.h
index ce84807..ecfa846 100644
--- a/src/stack.h
+++ b/src/stack.h
@@ -51,9 +51,10 @@
   kUndefined,
 };
 
-// ShadowFrame has 3 possible layouts: portable (VRegs & references overlap),
-// interpreter (VRegs and separate references array), JNI (just VRegs, but where
-// VRegs really => references).
+// ShadowFrame has 3 possible layouts:
+//  - portable - a unified array of VRegs and references. Precise references need GC maps.
+//  - interpreter - separate VRegs and reference arrays. References are in the reference array.
+//  - JNI - just VRegs, but where every VReg holds a reference.
 class ShadowFrame {
  public:
   // Create ShadowFrame for interpreter.
@@ -77,7 +78,6 @@
   }
 
   void SetNumberOfVRegs(uint32_t number_of_vregs) {
-    DCHECK(number_of_vregs < kHasReferenceArray);
     number_of_vregs_ = number_of_vregs | (number_of_vregs_ & kHasReferenceArray);
   }
 
@@ -207,18 +207,21 @@
   ShadowFrame(uint32_t num_vregs, ShadowFrame* link, AbstractMethod* method, uint32_t dex_pc,
               bool has_reference_array)
       : number_of_vregs_(num_vregs), link_(link), method_(method), dex_pc_(dex_pc) {
+    CHECK_LT(num_vregs, static_cast<uint32_t>(kHasReferenceArray));
     if (has_reference_array) {
       number_of_vregs_ |= kHasReferenceArray;
       for (size_t i = 0; i < num_vregs; ++i) {
         SetVRegReference(i, NULL);
       }
-    }
-    for (size_t i = 0; i < num_vregs; ++i) {
-      SetVReg(i, 0);
+    } else {
+      for (size_t i = 0; i < num_vregs; ++i) {
+        SetVReg(i, 0);
+      }
     }
   }
 
   Object* const* References() const {
+    DCHECK(HasReferenceArray());
     const uint32_t* vreg_end = &vregs_[NumberOfVRegs()];
     return reinterpret_cast<Object* const*>(vreg_end);
   }
@@ -362,16 +365,15 @@
 
   size_t GetNativePcOffset() const SHARED_LOCKS_REQUIRED(Locks::mutator_lock_);
 
-  uintptr_t LoadCalleeSave(int num, size_t frame_size) const {
+  uintptr_t* CalleeSaveAddress(int num, size_t frame_size) const {
     // Callee saves are held at the top of the frame
-    AbstractMethod* method = GetMethod();
-    DCHECK(method != NULL);
+    DCHECK(GetMethod() != NULL);
     byte* save_addr =
         reinterpret_cast<byte*>(cur_quick_frame_) + frame_size - ((num + 1) * kPointerSize);
 #if defined(__i386__)
     save_addr -= kPointerSize;  // account for return address
 #endif
-    return *reinterpret_cast<uintptr_t*>(save_addr);
+    return reinterpret_cast<uintptr_t*>(save_addr);
   }
 
   // Returns the height of the stack in the managed stack frames, including transitions.
@@ -398,6 +400,7 @@
       SHARED_LOCKS_REQUIRED(Locks::mutator_lock_);
 
   uintptr_t GetGPR(uint32_t reg) const;
+  void SetGPR(uint32_t reg, uintptr_t value);
 
   uint32_t GetVReg(AbstractMethod** cur_quick_frame, const DexFile::CodeItem* code_item,
                    uint32_t core_spills, uint32_t fp_spills, size_t frame_size,
diff --git a/src/thread.cc b/src/thread.cc
index 7490d2a..5745c9b 100644
--- a/src/thread.cc
+++ b/src/thread.cc
@@ -1793,6 +1793,7 @@
     result = Context::Create();
   } else {
     long_jump_context_ = NULL;  // Avoid context being shared.
+    result->Reset();
   }
   return result;
 }