diff --git a/runtime/jit/jit_code_cache.h b/runtime/jit/jit_code_cache.h
index 46a4085..fc011dd 100644
--- a/runtime/jit/jit_code_cache.h
+++ b/runtime/jit/jit_code_cache.h
@@ -35,9 +35,23 @@
 class LinearAlloc;
 class InlineCache;
 class IsMarkedVisitor;
+class JitJniStubTestHelper;
 class OatQuickMethodHeader;
 struct ProfileMethodInfo;
 class ProfilingInfo;
+class Thread;
+
+namespace gc {
+namespace accounting {
+template<size_t kAlignment> class MemoryRangeBitmap;
+}  // namespace accounting
+}  // namespace gc
+
+namespace mirror {
+class Class;
+class Object;
+template<class T> class ObjectArray;
+}  // namespace mirror
 
 namespace gc {
 namespace accounting {
@@ -137,6 +151,9 @@
   // Return true if the code cache contains this method.
   bool ContainsMethod(ArtMethod* method) REQUIRES(!lock_);
 
+  // Return the code pointer for a JNI-compiled stub if the method is in the cache, null otherwise.
+  const void* GetJniStubCode(ArtMethod* method) REQUIRES(!lock_);
+
   // Allocate a region of data that contain `size` bytes, and potentially space
   // for storing `number_of_roots` roots. Returns null if there is no more room.
   // Return the number of bytes allocated.
@@ -160,11 +177,6 @@
     return live_bitmap_.get();
   }
 
-  // Return whether we should do a full collection given the current state of the cache.
-  bool ShouldDoFullCollection()
-      REQUIRES(lock_)
-      REQUIRES_SHARED(Locks::mutator_lock_);
-
   // Perform a collection on the code cache.
   void GarbageCollectCache(Thread* self)
       REQUIRES(!lock_)
@@ -296,6 +308,12 @@
       REQUIRES(!lock_)
       REQUIRES(!Locks::cha_lock_);
 
+  // Removes method from the cache. The caller must ensure that all threads
+  // are suspended and the method should not be in any thread's stack.
+  bool RemoveMethodLocked(ArtMethod* method, bool release_memory)
+      REQUIRES(lock_)
+      REQUIRES(Locks::mutator_lock_);
+
   // Free in the mspace allocations for `code_ptr`.
   void FreeCode(const void* code_ptr) REQUIRES(lock_);
 
@@ -315,6 +333,11 @@
   // Set the footprint limit of the code cache.
   void SetFootprintLimit(size_t new_footprint) REQUIRES(lock_);
 
+  // Return whether we should do a full collection given the current state of the cache.
+  bool ShouldDoFullCollection()
+      REQUIRES(lock_)
+      REQUIRES_SHARED(Locks::mutator_lock_);
+
   void DoCollection(Thread* self, bool collect_profiling_info)
       REQUIRES(!lock_)
       REQUIRES_SHARED(Locks::mutator_lock_);
@@ -341,6 +364,9 @@
       REQUIRES(!lock_)
       REQUIRES_SHARED(Locks::mutator_lock_);
 
+  class JniStubKey;
+  class JniStubData;
+
   // Lock for guarding allocations, collections, and the method_code_map_.
   Mutex lock_;
   // Condition to wait on during collection.
@@ -357,6 +383,8 @@
   void* data_mspace_ GUARDED_BY(lock_);
   // Bitmap for collecting code and data.
   std::unique_ptr<CodeCacheBitmap> live_bitmap_;
+  // Holds compiled code associated with the shorty for a JNI stub.
+  SafeMap<JniStubKey, JniStubData> jni_stubs_map_ GUARDED_BY(lock_);
   // Holds compiled code associated to the ArtMethod.
   SafeMap<const void*, ArtMethod*> method_code_map_ GUARDED_BY(lock_);
   // Holds osr compiled code associated to the ArtMethod.
@@ -418,6 +446,7 @@
   // Condition to wait on for accessing inline caches.
   ConditionVariable inline_cache_cond_ GUARDED_BY(lock_);
 
+  friend class art::JitJniStubTestHelper;
   DISALLOW_IMPLICIT_CONSTRUCTORS(JitCodeCache);
 };