Support for JNI local reference cookie.

This also fixes a cross compilation bug in reseting the top of the
indirect reference table following a down call.

Change-Id: I40d913a6f86dadfe87b58d6d13a1ff3613f270ac
diff --git a/src/assembler.h b/src/assembler.h
index 7e14f3a..0779b58 100644
--- a/src/assembler.h
+++ b/src/assembler.h
@@ -334,6 +334,8 @@
   // Load routines
   virtual void Load(ManagedRegister dest, FrameOffset src, size_t size) = 0;
 
+  virtual void Load(ManagedRegister dest, ThreadOffset src, size_t size) = 0;
+
   virtual void LoadRef(ManagedRegister dest, FrameOffset  src) = 0;
 
   virtual void LoadRef(ManagedRegister dest, ManagedRegister base,
@@ -363,11 +365,18 @@
   virtual void Copy(FrameOffset dest, ManagedRegister src_base, Offset src_offset,
                     ManagedRegister scratch, size_t size) = 0;
 
+  virtual void Copy(ManagedRegister dest_base, Offset dest_offset, FrameOffset src,
+                    ManagedRegister scratch, size_t size) = 0;
+
   virtual void Copy(FrameOffset dest, FrameOffset src_base, Offset src_offset,
                     ManagedRegister scratch, size_t size) = 0;
 
-  virtual void Copy(ThreadOffset dest_base, Offset dest_offset, FrameOffset src,
-                    ManagedRegister scratch, ManagedRegister scratch2, size_t size) = 0;
+  virtual void Copy(ManagedRegister dest, Offset dest_offset,
+                    ManagedRegister src, Offset src_offset,
+                    ManagedRegister scratch, size_t size) = 0;
+
+  virtual void Copy(FrameOffset dest, Offset dest_offset, FrameOffset src, Offset src_offset,
+                    ManagedRegister scratch, size_t size) = 0;
 
   virtual void MemoryBarrier(ManagedRegister scratch) = 0;
 
diff --git a/src/assembler_arm.cc b/src/assembler_arm.cc
index 6c57bad..cd0c5f5 100644
--- a/src/assembler_arm.cc
+++ b/src/assembler_arm.cc
@@ -1563,8 +1563,7 @@
   } else if (dest.IsRegisterPair()) {
     CHECK_EQ(8u, size);
     LoadFromOffset(kLoadWord, dest.AsRegisterPairLow(), SP, src.Int32Value());
-    LoadFromOffset(kLoadWord, dest.AsRegisterPairHigh(),
-                   SP, src.Int32Value() + 4);
+    LoadFromOffset(kLoadWord, dest.AsRegisterPairHigh(), SP, src.Int32Value() + 4);
   } else if (dest.IsSRegister()) {
     LoadSFromOffset(dest.AsSRegister(), SP, src.Int32Value());
   } else {
@@ -1573,6 +1572,25 @@
   }
 }
 
+void ArmAssembler::Load(ManagedRegister mdest, ThreadOffset src, size_t size) {
+  ArmManagedRegister dest = mdest.AsArm();
+  if (dest.IsNoRegister()) {
+    CHECK_EQ(0u, size);
+  } else if (dest.IsCoreRegister()) {
+    CHECK_EQ(4u, size);
+    LoadFromOffset(kLoadWord, dest.AsCoreRegister(), TR, src.Int32Value());
+  } else if (dest.IsRegisterPair()) {
+    CHECK_EQ(8u, size);
+    LoadFromOffset(kLoadWord, dest.AsRegisterPairLow(), TR, src.Int32Value());
+    LoadFromOffset(kLoadWord, dest.AsRegisterPairHigh(), TR, src.Int32Value() + 4);
+  } else if (dest.IsSRegister()) {
+    LoadSFromOffset(dest.AsSRegister(), TR, src.Int32Value());
+  } else {
+    CHECK(dest.IsDRegister());
+    LoadDFromOffset(dest.AsDRegister(), TR, src.Int32Value());
+  }
+}
+
 void ArmAssembler::LoadRawPtrFromThread(ManagedRegister mdest,
                                         ThreadOffset offs) {
   ArmManagedRegister dest = mdest.AsArm();
@@ -1668,19 +1686,31 @@
   StoreToOffset(kStoreWord, scratch, SP, dest.Int32Value());
 }
 
+void ArmAssembler::Copy(ManagedRegister dest_base, Offset dest_offset, FrameOffset src,
+                        ManagedRegister mscratch, size_t size) {
+  Register scratch = mscratch.AsArm().AsCoreRegister();
+  CHECK_EQ(size, 4u);
+  LoadFromOffset(kLoadWord, scratch, SP, src.Int32Value());
+  StoreToOffset(kStoreWord, scratch, dest_base.AsArm().AsCoreRegister(), dest_offset.Int32Value());
+}
+
 void ArmAssembler::Copy(FrameOffset dest, FrameOffset src_base, Offset src_offset,
                         ManagedRegister mscratch, size_t size) {
   UNIMPLEMENTED(FATAL);
 }
 
-void ArmAssembler::Copy(ThreadOffset dest_base, Offset dest_offset, FrameOffset src,
-                        ManagedRegister mscratch, ManagedRegister mscratch2, size_t size) {
-  Register scratch = mscratch.AsArm().AsCoreRegister();
-  Register scratch2 = mscratch2.AsArm().AsCoreRegister();
+void ArmAssembler::Copy(ManagedRegister dest, Offset dest_offset,
+                        ManagedRegister src, Offset src_offset,
+                        ManagedRegister mscratch, size_t size) {
   CHECK_EQ(size, 4u);
-  LoadFromOffset(kLoadWord, scratch, TR, dest_base.Int32Value());
-  LoadFromOffset(kLoadWord, scratch2, SP, src.Int32Value());
-  StoreToOffset(kStoreWord, scratch2, scratch, dest_offset.Int32Value());
+  Register scratch = mscratch.AsArm().AsCoreRegister();
+  LoadFromOffset(kLoadWord, scratch, src.AsArm().AsCoreRegister(), src_offset.Int32Value());
+  StoreToOffset(kStoreWord, scratch, dest.AsArm().AsCoreRegister(), dest_offset.Int32Value());
+}
+
+void ArmAssembler::Copy(FrameOffset dest, Offset dest_offset, FrameOffset src, Offset src_offset,
+                        ManagedRegister scratch, size_t size) {
+  UNIMPLEMENTED(FATAL);
 }
 
 
diff --git a/src/assembler_arm.h b/src/assembler_arm.h
index 69b39c0..3f979c6 100644
--- a/src/assembler_arm.h
+++ b/src/assembler_arm.h
@@ -455,6 +455,8 @@
   // Load routines
   virtual void Load(ManagedRegister dest, FrameOffset src, size_t size);
 
+  virtual void Load(ManagedRegister dest, ThreadOffset src, size_t size);
+
   virtual void LoadRef(ManagedRegister dest, FrameOffset  src);
 
   virtual void LoadRef(ManagedRegister dest, ManagedRegister base,
@@ -484,11 +486,18 @@
   virtual void Copy(FrameOffset dest, ManagedRegister src_base, Offset src_offset,
                     ManagedRegister scratch, size_t size);
 
+  virtual void Copy(ManagedRegister dest_base, Offset dest_offset, FrameOffset src,
+                    ManagedRegister scratch, size_t size);
+
   virtual void Copy(FrameOffset dest, FrameOffset src_base, Offset src_offset,
                     ManagedRegister scratch, size_t size);
 
-  virtual void Copy(ThreadOffset dest_base, Offset dest_offset, FrameOffset src,
-                    ManagedRegister scratch, ManagedRegister scratch2, size_t size);
+  virtual void Copy(ManagedRegister dest, Offset dest_offset,
+                    ManagedRegister src, Offset src_offset,
+                    ManagedRegister scratch, size_t size);
+
+  virtual void Copy(FrameOffset dest, Offset dest_offset, FrameOffset src, Offset src_offset,
+                    ManagedRegister scratch, size_t size);
 
   virtual void MemoryBarrier(ManagedRegister scratch);
 
diff --git a/src/assembler_x86.cc b/src/assembler_x86.cc
index 1fce408..5df8704 100644
--- a/src/assembler_x86.cc
+++ b/src/assembler_x86.cc
@@ -1504,6 +1504,33 @@
   }
 }
 
+void X86Assembler::Load(ManagedRegister mdest, ThreadOffset src, size_t size) {
+  X86ManagedRegister dest = mdest.AsX86();
+  if (dest.IsNoRegister()) {
+    CHECK_EQ(0u, size);
+  } else if (dest.IsCpuRegister()) {
+    CHECK_EQ(4u, size);
+    fs()->movl(dest.AsCpuRegister(), Address::Absolute(src));
+  } else if (dest.IsRegisterPair()) {
+    CHECK_EQ(8u, size);
+    fs()->movl(dest.AsRegisterPairLow(), Address::Absolute(src));
+    fs()->movl(dest.AsRegisterPairHigh(), Address::Absolute(ThreadOffset(src.Int32Value()+4)));
+  } else if (dest.IsX87Register()) {
+    if (size == 4) {
+      fs()->flds(Address::Absolute(src));
+    } else {
+      fs()->fldl(Address::Absolute(src));
+    }
+  } else {
+    CHECK(dest.IsXmmRegister());
+    if (size == 4) {
+      fs()->movss(dest.AsXmmRegister(), Address::Absolute(src));
+    } else {
+      fs()->movsd(dest.AsXmmRegister(), Address::Absolute(src));
+    }
+  }
+}
+
 void X86Assembler::LoadRef(ManagedRegister mdest, FrameOffset  src) {
   X86ManagedRegister dest = mdest.AsX86();
   CHECK(dest.IsCpuRegister());
@@ -1590,6 +1617,14 @@
   UNIMPLEMENTED(FATAL);
 }
 
+void X86Assembler::Copy(ManagedRegister dest_base, Offset dest_offset, FrameOffset src,
+                        ManagedRegister scratch, size_t size) {
+  CHECK(scratch.IsNoRegister());
+  CHECK_EQ(size, 4u);
+  pushl(Address(ESP, src));
+  popl(Address(dest_base.AsX86().AsCpuRegister(), dest_offset));
+}
+
 void X86Assembler::Copy(FrameOffset dest, FrameOffset src_base, Offset src_offset,
                         ManagedRegister mscratch, size_t size) {
   Register scratch = mscratch.AsX86().AsCpuRegister();
@@ -1599,13 +1634,22 @@
   movl(Address(ESP, dest), scratch);
 }
 
-void X86Assembler::Copy(ThreadOffset dest_base, Offset dest_offset, FrameOffset src,
-                        ManagedRegister mscratch, ManagedRegister mscratch2, size_t size) {
-  Register scratch = mscratch.AsX86().AsCpuRegister();
-  CHECK(mscratch2.IsNoRegister());
+void X86Assembler::Copy(ManagedRegister dest, Offset dest_offset,
+                        ManagedRegister src, Offset src_offset,
+                        ManagedRegister scratch, size_t size) {
   CHECK_EQ(size, 4u);
-  fs()->movl(scratch, Address::Absolute(dest_base));
-  pushl(Address(ESP, src));
+  CHECK(scratch.IsNoRegister());
+  pushl(Address(src.AsX86().AsCpuRegister(), src_offset));
+  popl(Address(dest.AsX86().AsCpuRegister(), dest_offset));
+}
+
+void X86Assembler::Copy(FrameOffset dest, Offset dest_offset, FrameOffset src, Offset src_offset,
+                        ManagedRegister mscratch, size_t size) {
+  Register scratch = mscratch.AsX86().AsCpuRegister();
+  CHECK_EQ(size, 4u);
+  CHECK_EQ(dest.Int32Value(), src.Int32Value());
+  movl(scratch, Address(ESP, src));
+  pushl(Address(scratch, src_offset));
   popl(Address(scratch, dest_offset));
 }
 
diff --git a/src/assembler_x86.h b/src/assembler_x86.h
index 5936aac..6781e70 100644
--- a/src/assembler_x86.h
+++ b/src/assembler_x86.h
@@ -479,6 +479,8 @@
   // Load routines
   virtual void Load(ManagedRegister dest, FrameOffset src, size_t size);
 
+  virtual void Load(ManagedRegister dest, ThreadOffset src, size_t size);
+
   virtual void LoadRef(ManagedRegister dest, FrameOffset  src);
 
   virtual void LoadRef(ManagedRegister dest, ManagedRegister base,
@@ -508,11 +510,18 @@
   virtual void Copy(FrameOffset dest, ManagedRegister src_base, Offset src_offset,
                     ManagedRegister scratch, size_t size);
 
+  virtual void Copy(ManagedRegister dest_base, Offset dest_offset, FrameOffset src,
+                    ManagedRegister scratch, size_t size);
+
   virtual void Copy(FrameOffset dest, FrameOffset src_base, Offset src_offset,
                     ManagedRegister scratch, size_t size);
 
-  virtual void Copy(ThreadOffset dest_base, Offset dest_offset, FrameOffset src,
-                    ManagedRegister scratch, ManagedRegister scratch2, size_t size);
+  virtual void Copy(ManagedRegister dest, Offset dest_offset,
+                    ManagedRegister src, Offset src_offset,
+                    ManagedRegister scratch, size_t size);
+
+  virtual void Copy(FrameOffset dest, Offset dest_offset, FrameOffset src, Offset src_offset,
+                    ManagedRegister scratch, size_t size);
 
   virtual void MemoryBarrier(ManagedRegister);
 
diff --git a/src/calling_convention.cc b/src/calling_convention.cc
index c0556c5..17bfcff 100644
--- a/src/calling_convention.cc
+++ b/src/calling_convention.cc
@@ -82,7 +82,7 @@
   return method->NumReferenceArgs() + (method->IsStatic() ? 1 : 0);
 }
 
-FrameOffset JniCallingConvention::LocalReferenceTable_SegmentStatesOffset() const {
+FrameOffset JniCallingConvention::SavedLocalReferenceCookieOffset() const {
   size_t start_of_sirt = SirtLinkOffset().Int32Value() +  kPointerSize;
   size_t references_size = kPointerSize * ReferenceCount();  // size excluding header
   return FrameOffset(start_of_sirt + references_size);
@@ -90,7 +90,7 @@
 
 FrameOffset JniCallingConvention::ReturnValueSaveLocation() const {
   // Segment state is 4 bytes long
-  return FrameOffset(LocalReferenceTable_SegmentStatesOffset().Int32Value() + 4);
+  return FrameOffset(SavedLocalReferenceCookieOffset().Int32Value() + 4);
 }
 
 bool JniCallingConvention::HasNext() {
diff --git a/src/calling_convention.h b/src/calling_convention.h
index cf03bb8..de6cf37 100644
--- a/src/calling_convention.h
+++ b/src/calling_convention.h
@@ -130,7 +130,7 @@
   // Number of references in stack indirect reference table
   size_t ReferenceCount() const;
   // Location where the segment state of the local indirect reference table is saved
-  FrameOffset LocalReferenceTable_SegmentStatesOffset() const;
+  FrameOffset SavedLocalReferenceCookieOffset() const;
   // Location where the return value of a call can be squirreled if another
   // call is made following the native call
   FrameOffset ReturnValueSaveLocation() const;
diff --git a/src/indirect_reference_table.cc b/src/indirect_reference_table.cc
index f1e903c..e78b362 100644
--- a/src/indirect_reference_table.cc
+++ b/src/indirect_reference_table.cc
@@ -18,6 +18,7 @@
 #include "jni_internal.h"
 #include "reference_table.h"
 #include "runtime.h"
+#include "thread.h"
 #include "utils.h"
 
 #include <cstdlib>
@@ -148,6 +149,10 @@
     table_[topIndex++] = obj;
     segment_state_.parts.topIndex = topIndex;
   }
+  if (false) {
+    LOG(INFO) << "+++ added at " << ExtractIndex(result) << " top=" << segment_state_.parts.topIndex
+              << " holes=" << segment_state_.parts.numHoles;
+  }
 
   DCHECK(result != NULL);
   return result;
@@ -173,7 +178,8 @@
   int idx = ExtractIndex(iref);
   if (idx >= topIndex) {
     /* bad -- stale reference? */
-    LOG(ERROR) << "JNI ERROR (app bug): accessed stale " << kind_ << " " << iref << " (index " << idx << " in a table of size " << topIndex << ")";
+    LOG(ERROR) << "JNI ERROR (app bug): accessed stale " << kind_ << " "
+               << iref << " (index " << idx << " in a table of size " << topIndex << ")";
     AbortMaybe();
     return false;
   }
@@ -230,6 +236,11 @@
   int idx = ExtractIndex(iref);
 
   JavaVMExt* vm = Runtime::Current()->GetJavaVM();
+  if (GetIndirectRefKind(iref) == kSirtOrInvalid &&
+      Thread::Current()->SirtContains(reinterpret_cast<jobject>(iref))) {
+    LOG(WARNING) << "Attempt to remove local SIRT entry from IRT, ignoring";
+    return true;
+  }
   if (GetIndirectRefKind(iref) == kSirtOrInvalid || vm->work_around_app_jni_bugs) {
     idx = LinearScan(iref, bottomIndex, topIndex, table_);
     if (idx == -1) {
@@ -240,12 +251,14 @@
 
   if (idx < bottomIndex) {
     /* wrong segment */
-    LOG(INFO) << "Attempt to remove index outside index area (" << idx << " vs " << bottomIndex << "-" << topIndex << ")";
+    LOG(INFO) << "Attempt to remove index outside index area (" << idx
+              << " vs " << bottomIndex << "-" << topIndex << ")";
     return false;
   }
   if (idx >= topIndex) {
     /* bad -- stale reference? */
-    LOG(INFO) << "Attempt to remove invalid index " << idx << " (bottom=" << bottomIndex << " top=" << topIndex << ")";
+    LOG(INFO) << "Attempt to remove invalid index " << idx
+              << " (bottom=" << bottomIndex << " top=" << topIndex << ")";
     return false;
   }
 
@@ -260,18 +273,25 @@
     int numHoles = segment_state_.parts.numHoles - prevState.parts.numHoles;
     if (numHoles != 0) {
       while (--topIndex > bottomIndex && numHoles != 0) {
-        //LOG(INFO) << "+++ checking for hole at " << topIndex-1 << " (cookie=" << cookie << ") val=" << table_[topIndex-1];
+        if (false) {
+          LOG(INFO) << "+++ checking for hole at " << topIndex-1
+                    << " (cookie=" << cookie << ") val=" << table_[topIndex - 1];
+        }
         if (table_[topIndex-1] != NULL) {
           break;
         }
-        //LOG(INFO) << "+++ ate hole at " << (topIndex-1);
+        if (false) {
+          LOG(INFO) << "+++ ate hole at " << (topIndex - 1);
+        }
         numHoles--;
       }
       segment_state_.parts.numHoles = numHoles + prevState.parts.numHoles;
       segment_state_.parts.topIndex = topIndex;
     } else {
       segment_state_.parts.topIndex = topIndex-1;
-      //LOG(INFO) << "+++ ate last entry " << topIndex-1;
+      if (false) {
+        LOG(INFO) << "+++ ate last entry " << topIndex - 1;
+      }
     }
   } else {
     /*
@@ -289,7 +309,9 @@
 
     table_[idx] = NULL;
     segment_state_.parts.numHoles++;
-    //LOG(INFO) << "+++ left hole at " << idx << ", holes=" << segment_state_.parts.numHoles;
+    if (false) {
+      LOG(INFO) << "+++ left hole at " << idx << ", holes=" << segment_state_.parts.numHoles;
+    }
   }
 
   return true;
diff --git a/src/jni_compiler.cc b/src/jni_compiler.cc
index 6404b05..6c09f58 100644
--- a/src/jni_compiler.cc
+++ b/src/jni_compiler.cc
@@ -116,8 +116,10 @@
       // Compute SIRT entry, note null is placed in the SIRT but its boxed value
       // must be NULL
       FrameOffset sirt_offset = jni_conv->CurrentParamSirtEntryOffset();
-      // Check SIRT offset is within frame
+      // Check SIRT offset is within frame and doesn't run into the saved segment state
       CHECK_LT(sirt_offset.Uint32Value(), frame_size);
+      CHECK_NE(sirt_offset.Uint32Value(),
+               jni_conv->SavedLocalReferenceCookieOffset().Uint32Value());
       bool input_in_reg = mr_conv->IsCurrentParamInRegister();
       bool input_on_stack = mr_conv->IsCurrentParamOnStack();
       CHECK(input_in_reg || input_on_stack);
@@ -267,14 +269,26 @@
     ManagedRegister jni_env = jni_conv->CurrentParamRegister();
     DCHECK(!jni_env.Equals(jni_conv->InterproceduralScratchRegister()));
     __ LoadRawPtrFromThread(jni_env, Thread::JniEnvOffset());
-    __ Copy(jni_conv->LocalReferenceTable_SegmentStatesOffset(), jni_env,
-            JNIEnvExt::SegmentStateOffset(), jni_conv->InterproceduralScratchRegister(), 4);
+    // Frame[saved_local_ref_cookie_offset] = env->local_ref_cookie
+    __ Copy(jni_conv->SavedLocalReferenceCookieOffset(),
+            jni_env, JNIEnvExt::LocalRefCookieOffset(),
+            jni_conv->InterproceduralScratchRegister(), 4);
+    // env->local_ref_cookie = env->locals.segment_state
+    __ Copy(jni_env, JNIEnvExt::LocalRefCookieOffset(),
+            jni_env, JNIEnvExt::SegmentStateOffset(),
+            jni_conv->InterproceduralScratchRegister(), 4);
   } else {
     FrameOffset jni_env = jni_conv->CurrentParamStackOffset();
     __ CopyRawPtrFromThread(jni_env, Thread::JniEnvOffset(),
                             jni_conv->InterproceduralScratchRegister());
-    __ Copy(jni_conv->LocalReferenceTable_SegmentStatesOffset(), jni_env,
-            JNIEnvExt::SegmentStateOffset(), jni_conv->InterproceduralScratchRegister(), 4);
+    // Frame[saved_local_ref_cookie_offset] = env->local_ref_cookie
+    __ Copy(jni_conv->SavedLocalReferenceCookieOffset(),
+            jni_env, JNIEnvExt::LocalRefCookieOffset(),
+            jni_conv->InterproceduralScratchRegister(), 4);
+    // env->local_ref_cookie = env->locals.segment_state
+    __ Copy(jni_env, JNIEnvExt::LocalRefCookieOffset(),
+            jni_env, JNIEnvExt::SegmentStateOffset(),
+            jni_conv->InterproceduralScratchRegister(), 4);
   }
 
   // 9. Plant call to native code associated with method
@@ -387,10 +401,18 @@
   __ Move(mr_conv->ReturnRegister(), jni_conv->ReturnRegister());
 
   // 15. Restore segment state and remove SIRT from thread
-  __ Copy(Thread::JniEnvOffset(), JNIEnvExt::SegmentStateOffset(),
-          jni_conv->LocalReferenceTable_SegmentStatesOffset(),
-          jni_conv->InterproceduralScratchRegister(),
-          jni_conv->ReturnScratchRegister(), 4);
+  {
+    ManagedRegister jni_env = jni_conv->InterproceduralScratchRegister();
+    __ LoadRawPtrFromThread(jni_env, Thread::JniEnvOffset());
+    // env->locals.segment_state = env->local_ref_cookie
+    __ Copy(jni_env, JNIEnvExt::SegmentStateOffset(),
+            jni_env, JNIEnvExt::LocalRefCookieOffset(),
+            jni_conv->ReturnScratchRegister(), 4);
+    // env->local_ref_cookie = Frame[saved_local_ref_cookie_offset]
+    __ Copy(jni_env, JNIEnvExt::LocalRefCookieOffset(),
+            jni_conv->SavedLocalReferenceCookieOffset(),
+            jni_conv->ReturnScratchRegister(), 4);
+  }
   __ CopyRawPtrToThread(Thread::TopSirtOffset(), jni_conv->SirtLinkOffset(),
                         jni_conv->InterproceduralScratchRegister());
 
diff --git a/src/jni_compiler_test.cc b/src/jni_compiler_test.cc
index c2949c9..ecc2f88 100644
--- a/src/jni_compiler_test.cc
+++ b/src/jni_compiler_test.cc
@@ -516,7 +516,7 @@
 
 jint local_ref_test(JNIEnv* env, jobject thisObj, jint x) {
   // Add 10 local references
-  for(int i = 0; i < 10; i++) {
+  for (int i = 0; i < 10; i++) {
     AddLocalReference<jobject>(env, Decode<Object*>(env, thisObj));
   }
   return x+1;
@@ -525,7 +525,7 @@
 TEST_F(JniCompilerTest, LocalReferenceTableClearingTest) {
   SetupForTest(false, "fooI", "(I)I", reinterpret_cast<void*>(&local_ref_test));
   // 1000 invocations of a method that adds 10 local references
-  for(int i=0; i < 1000; i++) {
+  for (int i=0; i < 1000; i++) {
     jint result = env_->CallIntMethod(jobj_, jmethod_, i);
     EXPECT_TRUE(result == i + 1);
   }
diff --git a/src/jni_internal.cc b/src/jni_internal.cc
index 0436500..7adbaeb 100644
--- a/src/jni_internal.cc
+++ b/src/jni_internal.cc
@@ -68,7 +68,7 @@
   JNIEnvExt* env = reinterpret_cast<JNIEnvExt*>(public_env);
   IndirectReferenceTable& locals = env->locals;
 
-  uint32_t cookie = IRT_FIRST_SEGMENT; // TODO
+  uint32_t cookie = env->local_ref_cookie;
   IndirectRef ref = locals.Add(cookie, obj);
   if (ref == NULL) {
     // TODO: just change Add's DCHECK to CHECK and lose this?
@@ -873,7 +873,7 @@
 
     IndirectReferenceTable& locals = ts.Env()->locals;
 
-    uint32_t cookie = IRT_FIRST_SEGMENT; // TODO
+    uint32_t cookie = ts.Env()->local_ref_cookie;
     IndirectRef ref = locals.Add(cookie, Decode<Object*>(ts, obj));
     return reinterpret_cast<jobject>(ref);
   }
@@ -886,7 +886,7 @@
 
     IndirectReferenceTable& locals = ts.Env()->locals;
 
-    uint32_t cookie = IRT_FIRST_SEGMENT; // TODO
+    uint32_t cookie = ts.Env()->local_ref_cookie;
     if (!locals.Remove(cookie, obj)) {
       // Attempting to delete a local reference that is not in the
       // topmost local reference frame is a no-op.  DeleteLocalRef returns
@@ -2508,15 +2508,20 @@
 JNIEnvExt::JNIEnvExt(Thread* self, JavaVMExt* vm)
     : self(self),
       vm(vm),
+      local_ref_cookie(IRT_FIRST_SEGMENT),
+      locals(kLocalsInitial, kLocalsMax, kLocal),
       check_jni(vm->check_jni),
       work_around_app_jni_bugs(vm->work_around_app_jni_bugs),
       critical(false),
-      monitors("monitors", kMonitorsInitial, kMonitorsMax),
-      locals(kLocalsInitial, kLocalsMax, kLocal) {
+      monitors("monitors", kMonitorsInitial, kMonitorsMax) {
   functions = unchecked_functions = &gNativeInterface;
   if (check_jni) {
     functions = GetCheckJniNativeInterface();
   }
+  // The JniEnv local reference values must be at a consistent offset or else cross-compilation
+  // errors will ensue.
+  CHECK_EQ(JNIEnvExt::LocalRefCookieOffset().Int32Value(), 12);
+  CHECK_EQ(JNIEnvExt::SegmentStateOffset().Int32Value(), 16);
 }
 
 JNIEnvExt::~JNIEnvExt() {
diff --git a/src/jni_internal.h b/src/jni_internal.h
index e5bcb75..332f3a0 100644
--- a/src/jni_internal.h
+++ b/src/jni_internal.h
@@ -130,9 +130,19 @@
                   IndirectReferenceTable::SegmentStateOffset().Int32Value());
   }
 
+  static Offset LocalRefCookieOffset() {
+    return Offset(OFFSETOF_MEMBER(JNIEnvExt, local_ref_cookie));
+  }
+
   Thread* const self;
   JavaVMExt* vm;
 
+  // Cookie used when using the local indirect reference table
+  uint32_t local_ref_cookie;
+
+  // JNI local references.
+  IndirectReferenceTable locals;
+
   // Frequently-accessed fields cached from JavaVM.
   bool check_jni;
   bool work_around_app_jni_bugs;
@@ -143,9 +153,6 @@
   // Entered JNI monitors, for bulk exit on thread detach.
   ReferenceTable monitors;
 
-  // JNI local references.
-  IndirectReferenceTable locals;
-
   // Used by -Xcheck:jni.
   const JNINativeInterface* unchecked_functions;
 };
diff --git a/src/thread.h b/src/thread.h
index 123a376..ebab184 100644
--- a/src/thread.h
+++ b/src/thread.h
@@ -432,6 +432,7 @@
     native_to_managed_record_ = record;
     top_of_managed_stack_.SetSP(NULL);
   }
+
   void PopNativeToManagedRecord(const NativeToManagedRecord& record) {
     native_to_managed_record_ = record.link_;
     top_of_managed_stack_.SetSP(reinterpret_cast<Method**>(record.last_top_of_managed_stack_));