Implement EnsureCapacity, PushLocalFrame, and PopLocalFrame.

These are as good as the old implementations, except that unbalanced usages
won't be cleaned up completely (you'll slowly grow the vector in your JNIEnv).

This patch also renames IndirectReferenceTable::Contains to the less misleading
ContainsDirectPointer, and fixes JNI::GetObjectRefType to not claim that
invalid local references are locals indefinitely.

We also now include detail messages in OOMEs where possible. (Test 061 still
passes.) We still log regardless, since OOME should be a rare thing.

Change-Id: I77b2f44ea066e92c517e5c96700ec533727b9c78
diff --git a/src/indirect_reference_table.cc b/src/indirect_reference_table.cc
index 716c214..5f5541e 100644
--- a/src/indirect_reference_table.cc
+++ b/src/indirect_reference_table.cc
@@ -204,17 +204,17 @@
   return true;
 }
 
-static int LinearScan(IndirectRef iref, int bottomIndex, int topIndex, const Object** table) {
+static int Find(Object* direct_pointer, int bottomIndex, int topIndex, const Object** table) {
   for (int i = bottomIndex; i < topIndex; ++i) {
-    if (table[i] == reinterpret_cast<const Object*>(iref)) {
+    if (table[i] == direct_pointer) {
       return i;
     }
   }
   return -1;
 }
 
-bool IndirectReferenceTable::Contains(IndirectRef iref) const {
-  return LinearScan(iref, 0, segment_state_.parts.topIndex, table_) != -1;
+bool IndirectReferenceTable::ContainsDirectPointer(Object* direct_pointer) const {
+  return Find(direct_pointer, 0, segment_state_.parts.topIndex, table_) != -1;
 }
 
 /*
@@ -249,7 +249,8 @@
     return true;
   }
   if (GetIndirectRefKind(iref) == kSirtOrInvalid || vm->work_around_app_jni_bugs) {
-    idx = LinearScan(iref, bottomIndex, topIndex, table_);
+    Object* direct_pointer = reinterpret_cast<Object*>(iref);
+    idx = Find(direct_pointer, bottomIndex, topIndex, table_);
     if (idx == -1) {
       LOG(WARNING) << "trying to work around app JNI bugs, but didn't find " << iref << " in table!";
       return false;
diff --git a/src/indirect_reference_table.h b/src/indirect_reference_table.h
index ba20d4d..f6cab95 100644
--- a/src/indirect_reference_table.h
+++ b/src/indirect_reference_table.h
@@ -281,7 +281,7 @@
   }
 
   // TODO: remove when we remove work_around_app_jni_bugs support.
-  bool Contains(IndirectRef iref) const;
+  bool ContainsDirectPointer(Object* direct_pointer) const;
 
   /*
    * Remove an existing entry.
diff --git a/src/jni_internal.cc b/src/jni_internal.cc
index 1d46526..9480644 100644
--- a/src/jni_internal.cc
+++ b/src/jni_internal.cc
@@ -24,6 +24,21 @@
 #include "stringpiece.h"
 #include "thread.h"
 
+static const size_t kMonitorsInitial = 32; // Arbitrary.
+static const size_t kMonitorsMax = 4096; // Arbitrary sanity check.
+
+static const size_t kLocalsInitial = 64; // Arbitrary.
+static const size_t kLocalsMax = 512; // Arbitrary sanity check.
+
+static const size_t kPinTableInitial = 16; // Arbitrary.
+static const size_t kPinTableMax = 1024; // Arbitrary sanity check.
+
+static const size_t kGlobalsInitial = 512; // Arbitrary.
+static const size_t kGlobalsMax = 51200; // Arbitrary sanity check.
+
+static const size_t kWeakGlobalsInitial = 16; // Arbitrary.
+static const size_t kWeakGlobalsMax = 51200; // Arbitrary sanity check.
+
 namespace art {
 
 /*
@@ -780,22 +795,40 @@
     LOG(FATAL) << "JNI FatalError called: " << msg;
   }
 
-  static jint PushLocalFrame(JNIEnv* env, jint cap) {
+  static jint PushLocalFrame(JNIEnv* env, jint capacity) {
     ScopedJniThreadState ts(env);
-    UNIMPLEMENTED(WARNING) << "ignoring PushLocalFrame(" << cap << ")";
+    if (EnsureLocalCapacity(ts, capacity, "PushLocalFrame") != JNI_OK) {
+      return JNI_ERR;
+    }
+    ts.Env()->PushFrame(capacity);
     return JNI_OK;
   }
 
-  static jobject PopLocalFrame(JNIEnv* env, jobject res) {
+  static jobject PopLocalFrame(JNIEnv* env, jobject java_survivor) {
     ScopedJniThreadState ts(env);
-    UNIMPLEMENTED(WARNING) << "ignoring PopLocalFrame " << res;
-    return res;
+    Object* survivor = Decode<Object*>(ts, java_survivor);
+    ts.Env()->PopFrame();
+    return AddLocalReference<jobject>(env, survivor);
   }
 
-  static jint EnsureLocalCapacity(JNIEnv* env, jint cap) {
+  static jint EnsureLocalCapacity(JNIEnv* env, jint desired_capacity) {
     ScopedJniThreadState ts(env);
-    UNIMPLEMENTED(WARNING) << "ignoring EnsureLocalCapacity(" << cap << ")";
-    return 0;
+    return EnsureLocalCapacity(ts, desired_capacity, "EnsureLocalCapacity");
+  }
+
+  static jint EnsureLocalCapacity(ScopedJniThreadState& ts, jint desired_capacity, const char* caller) {
+    // TODO: we should try to expand the table if necessary.
+    if (desired_capacity < 1 || desired_capacity > static_cast<jint>(kLocalsMax)) {
+      LOG(ERROR) << "Invalid capacity given to " << caller << ": " << desired_capacity;
+      return JNI_ERR;
+    }
+    // TODO: this isn't quite right, since "capacity" includes holes.
+    size_t capacity = ts.Env()->locals.Capacity();
+    bool okay = (static_cast<jint>(kLocalsMax - capacity) >= desired_capacity);
+    if (!okay) {
+      ts.Self()->ThrowOutOfMemoryError(caller);
+    }
+    return okay ? JNI_OK : JNI_ERR;
   }
 
   static jobject NewGlobalRef(JNIEnv* env, jobject obj) {
@@ -2213,7 +2246,10 @@
     IndirectRefKind kind = GetIndirectRefKind(ref);
     switch (kind) {
     case kLocal:
-      return JNILocalRefType;
+      if (ts.Env()->locals.Get(ref) != kInvalidIndirectRefObject) {
+        return JNILocalRefType;
+      }
+      return JNIInvalidRefType;
     case kGlobal:
       return JNIGlobalRefType;
     case kWeakGlobal:
@@ -2231,7 +2267,7 @@
       // If we're handing out direct pointers, check whether it's a direct pointer
       // to a local reference.
       if (Decode<Object*>(ts, java_object) == reinterpret_cast<Object*>(java_object)) {
-        if (ts.Env()->locals.Contains(java_object)) {
+        if (ts.Env()->locals.ContainsDirectPointer(reinterpret_cast<Object*>(java_object))) {
           return JNILocalRefType;
         }
       }
@@ -2477,12 +2513,6 @@
   JNI::GetObjectRefType,
 };
 
-static const size_t kMonitorsInitial = 32; // Arbitrary.
-static const size_t kMonitorsMax = 4096; // Arbitrary sanity check.
-
-static const size_t kLocalsInitial = 64; // Arbitrary.
-static const size_t kLocalsMax = 512; // Arbitrary sanity check.
-
 JNIEnvExt::JNIEnvExt(Thread* self, JavaVMExt* vm)
     : self(self),
       vm(vm),
@@ -2510,6 +2540,17 @@
   monitors.Dump();
 }
 
+void JNIEnvExt::PushFrame(int capacity) {
+  stacked_local_ref_cookies.push_back(local_ref_cookie);
+  local_ref_cookie = locals.GetSegmentState();
+}
+
+void JNIEnvExt::PopFrame() {
+  locals.SetSegmentState(local_ref_cookie);
+  local_ref_cookie = stacked_local_ref_cookies.back();
+  stacked_local_ref_cookies.pop_back();
+}
+
 // JNI Invocation interface.
 
 extern "C" jint JNI_CreateJavaVM(JavaVM** p_vm, void** p_env, void* vm_args) {
@@ -2609,15 +2650,6 @@
   JII::AttachCurrentThreadAsDaemon
 };
 
-static const size_t kPinTableInitialSize = 16;
-static const size_t kPinTableMaxSize = 1024;
-
-static const size_t kGlobalsInitial = 512; // Arbitrary.
-static const size_t kGlobalsMax = 51200; // Arbitrary sanity check.
-
-static const size_t kWeakGlobalsInitial = 16; // Arbitrary.
-static const size_t kWeakGlobalsMax = 51200; // Arbitrary sanity check.
-
 JavaVMExt::JavaVMExt(Runtime* runtime, Runtime::ParsedOptions* options)
     : runtime(runtime),
       check_jni_abort_hook(NULL),
@@ -2628,7 +2660,7 @@
       trace(options->jni_trace_),
       work_around_app_jni_bugs(false), // TODO: add a way to enable this
       pins_lock("JNI pin table lock"),
-      pin_table("pin table", kPinTableInitialSize, kPinTableMaxSize),
+      pin_table("pin table", kPinTableInitial, kPinTableMax),
       globals_lock("JNI global reference table lock"),
       globals(kGlobalsInitial, kGlobalsMax, kGlobal),
       weak_globals_lock("JNI weak global reference table lock"),
diff --git a/src/jni_internal.h b/src/jni_internal.h
index e715bdf..6f9b755 100644
--- a/src/jni_internal.h
+++ b/src/jni_internal.h
@@ -125,6 +125,9 @@
 
   void DumpReferenceTables();
 
+  void PushFrame(int capacity);
+  void PopFrame();
+
   static Offset SegmentStateOffset() {
     return Offset(OFFSETOF_MEMBER(JNIEnvExt, locals) +
                   IndirectReferenceTable::SegmentStateOffset().Int32Value());
@@ -137,12 +140,17 @@
   Thread* const self;
   JavaVMExt* vm;
 
-  // Cookie used when using the local indirect reference table
+  // Cookie used when using the local indirect reference table.
   uint32_t local_ref_cookie;
 
   // JNI local references.
   IndirectReferenceTable locals;
 
+  // Stack of cookies corresponding to PushLocalFrame/PopLocalFrame calls.
+  // TODO: to avoid leaks (and bugs), we need to clear this vector on entry (or return)
+  // to a native method.
+  std::vector<uint32_t> stacked_local_ref_cookies;
+
   // Frequently-accessed fields cached from JavaVM.
   bool check_jni;
   bool work_around_app_jni_bugs;
diff --git a/src/jni_internal_test.cc b/src/jni_internal_test.cc
index ddd39b5..59dca3a 100644
--- a/src/jni_internal_test.cc
+++ b/src/jni_internal_test.cc
@@ -720,7 +720,7 @@
   EXPECT_TRUE(o != NULL);
   EXPECT_TRUE(o != s);
 
-  // TODO: check that o is a local reference.
+  EXPECT_EQ(JNILocalRefType, env_->GetObjectRefType(o));
 }
 
 TEST_F(JniInternalTest, DeleteLocalRef_NULL) {
@@ -746,6 +746,43 @@
   env_->DeleteLocalRef(o);
 }
 
+TEST_F(JniInternalTest, PushLocalFrame_PopLocalFrame) {
+  jobject original = env_->NewStringUTF("");
+  ASSERT_TRUE(original != NULL);
+
+  jobject outer;
+  jobject inner1, inner2;
+  Object* inner2_direct_pointer;
+  {
+    env_->PushLocalFrame(4);
+    outer = env_->NewLocalRef(original);
+
+    {
+      env_->PushLocalFrame(4);
+      inner1 = env_->NewLocalRef(outer);
+      inner2 = env_->NewStringUTF("survivor");
+      inner2_direct_pointer = Decode<Object*>(env_, inner2);
+      env_->PopLocalFrame(inner2);
+    }
+
+    EXPECT_EQ(JNILocalRefType, env_->GetObjectRefType(original));
+    EXPECT_EQ(JNILocalRefType, env_->GetObjectRefType(outer));
+    EXPECT_EQ(JNIInvalidRefType, env_->GetObjectRefType(inner1));
+
+    // Our local reference for the survivor is invalid because the survivor
+    // gets a new local reference...
+    EXPECT_EQ(JNIInvalidRefType, env_->GetObjectRefType(inner2));
+    // ...but the survivor should be in the local reference table.
+    EXPECT_TRUE(env_->locals.ContainsDirectPointer(inner2_direct_pointer));
+
+    env_->PopLocalFrame(NULL);
+  }
+  EXPECT_EQ(JNILocalRefType, env_->GetObjectRefType(original));
+  EXPECT_EQ(JNIInvalidRefType, env_->GetObjectRefType(outer));
+  EXPECT_EQ(JNIInvalidRefType, env_->GetObjectRefType(inner1));
+  EXPECT_EQ(JNIInvalidRefType, env_->GetObjectRefType(inner2));
+}
+
 TEST_F(JniInternalTest, NewGlobalRef_NULL) {
   EXPECT_TRUE(env_->NewGlobalRef(NULL) == NULL);
 }
diff --git a/src/thread.cc b/src/thread.cc
index e1a8841..d68d1db 100644
--- a/src/thread.cc
+++ b/src/thread.cc
@@ -1176,12 +1176,17 @@
 }
 
 void Thread::ThrowOutOfMemoryError(Class* c, size_t byte_count) {
-  LOG(ERROR) << "Failed to allocate a " << byte_count << "-byte "
-             << PrettyDescriptor(c->GetDescriptor())
-             << (throwing_OutOfMemoryError_ ? " (recursive case)" : "");
+  std::string msg(StringPrintf("Failed to allocate a %zd-byte %s", byte_count,
+      PrettyDescriptor(c->GetDescriptor()).c_str()));
+  ThrowOutOfMemoryError(msg.c_str());
+}
+
+void Thread::ThrowOutOfMemoryError(const char* msg) {
+  LOG(ERROR) << StringPrintf("Throwing OutOfMemoryError \"%s\"%s",
+      msg, (throwing_OutOfMemoryError_ ? " (recursive case)" : ""));
   if (!throwing_OutOfMemoryError_) {
     throwing_OutOfMemoryError_ = true;
-    ThrowNewException("Ljava/lang/OutOfMemoryError;", NULL);
+    ThrowNewException("Ljava/lang/OutOfMemoryError;", msg);
   } else {
     SetException(pre_allocated_OutOfMemoryError_);
   }
diff --git a/src/thread.h b/src/thread.h
index 3a8b2a8..a7097ce 100644
--- a/src/thread.h
+++ b/src/thread.h
@@ -283,7 +283,8 @@
 
   void ThrowNewExceptionV(const char* exception_class_descriptor, const char* fmt, va_list ap);
 
-  // This exception is special, because we need to pre-allocate an instance.
+  // OutOfMemoryError is special, because we need to pre-allocate an instance.
+  void ThrowOutOfMemoryError(const char* msg);
   void ThrowOutOfMemoryError(Class* c, size_t byte_count);
 
   Frame FindExceptionHandler(void* throw_pc, void** handler_pc);