Generalize atomic_method_ref_map to support dex references

Generalize atomic method ref map to support dex references instead
of only method references.

The goal is to use this in a future CL to replace compiled_classes_.

Test: test-art-host

Change-Id: Ic6d1e619584f790eea68f5160fa0fcd664524cd7
diff --git a/build/Android.gtest.mk b/build/Android.gtest.mk
index ff4f901..cf6d1ec 100644
--- a/build/Android.gtest.mk
+++ b/build/Android.gtest.mk
@@ -102,7 +102,7 @@
 # Dex file dependencies for each gtest.
 ART_GTEST_dex2oat_environment_tests_DEX_DEPS := Main MainStripped MultiDex MultiDexModifiedSecondary Nested
 
-ART_GTEST_atomic_method_ref_map_test_DEX_DEPS := Interfaces
+ART_GTEST_atomic_dex_ref_map_test_DEX_DEPS := Interfaces
 ART_GTEST_class_linker_test_DEX_DEPS := AllFields ErroneousA ErroneousB ErroneousInit ForClassLoaderA ForClassLoaderB ForClassLoaderC ForClassLoaderD Interfaces MethodTypes MultiDex MyClass Nested Statics StaticsFromCode
 ART_GTEST_class_loader_context_test_DEX_DEPS := Main MultiDex MyClass
 ART_GTEST_class_table_test_DEX_DEPS := XandY
diff --git a/compiler/Android.bp b/compiler/Android.bp
index cacb8cb..e2246b1 100644
--- a/compiler/Android.bp
+++ b/compiler/Android.bp
@@ -361,7 +361,7 @@
         "optimizing/ssa_test.cc",
         "optimizing/stack_map_test.cc",
         "optimizing/suspend_check_test.cc",
-        "utils/atomic_method_ref_map_test.cc",
+        "utils/atomic_dex_ref_map_test.cc",
         "utils/dedupe_set_test.cc",
         "utils/intrusive_forward_list_test.cc",
         "utils/string_reference_test.cc",
diff --git a/compiler/dex/verification_results.cc b/compiler/dex/verification_results.cc
index 183c955..beb3439 100644
--- a/compiler/dex/verification_results.cc
+++ b/compiler/dex/verification_results.cc
@@ -24,7 +24,7 @@
 #include "runtime.h"
 #include "thread.h"
 #include "thread-current-inl.h"
-#include "utils/atomic_method_ref_map-inl.h"
+#include "utils/atomic_dex_ref_map-inl.h"
 #include "verified_method.h"
 #include "verifier/method_verifier-inl.h"
 
@@ -38,7 +38,7 @@
 VerificationResults::~VerificationResults() {
   WriterMutexLock mu(Thread::Current(), verified_methods_lock_);
   STLDeleteValues(&verified_methods_);
-  atomic_verified_methods_.Visit([](const MethodReference& ref ATTRIBUTE_UNUSED,
+  atomic_verified_methods_.Visit([](const DexFileReference& ref ATTRIBUTE_UNUSED,
                                     const VerifiedMethod* method) {
     delete method;
   });
@@ -56,16 +56,18 @@
     // We'll punt this later.
     return;
   }
-  AtomicMap::InsertResult result = atomic_verified_methods_.Insert(ref,
-                                                                   /*expected*/ nullptr,
-                                                                   verified_method.get());
+  AtomicMap::InsertResult result = atomic_verified_methods_.Insert(
+      DexFileReference(ref.dex_file, ref.dex_method_index),
+      /*expected*/ nullptr,
+      verified_method.get());
   const VerifiedMethod* existing = nullptr;
   bool inserted;
   if (result != AtomicMap::kInsertResultInvalidDexFile) {
     inserted = (result == AtomicMap::kInsertResultSuccess);
     if (!inserted) {
       // Rare case.
-      CHECK(atomic_verified_methods_.Get(ref, &existing));
+      CHECK(atomic_verified_methods_.Get(DexFileReference(ref.dex_file, ref.dex_method_index),
+                                         &existing));
       CHECK_NE(verified_method.get(), existing);
     }
   } else {
@@ -103,7 +105,7 @@
 const VerifiedMethod* VerificationResults::GetVerifiedMethod(MethodReference ref) {
   const VerifiedMethod* ret = nullptr;
   DCHECK(compiler_options_->IsAnyCompilationEnabled());
-  if (atomic_verified_methods_.Get(ref, &ret)) {
+  if (atomic_verified_methods_.Get(DexFileReference(ref.dex_file, ref.dex_method_index), &ret)) {
     return ret;
   }
   ReaderMutexLock mu(Thread::Current(), verified_methods_lock_);
@@ -117,7 +119,9 @@
   // at runtime.
   std::unique_ptr<VerifiedMethod> verified_method = std::make_unique<VerifiedMethod>(
       /* encountered_error_types */ 0, /* has_runtime_throw */ false);
-  if (atomic_verified_methods_.Insert(ref, /*expected*/ nullptr, verified_method.get()) ==
+  if (atomic_verified_methods_.Insert(DexFileReference(ref.dex_file, ref.dex_method_index),
+                                      /*expected*/ nullptr,
+                                      verified_method.get()) ==
           AtomicMap::InsertResult::kInsertResultSuccess) {
     verified_method.release();
   }
@@ -150,7 +154,7 @@
 }
 
 void VerificationResults::AddDexFile(const DexFile* dex_file) {
-  atomic_verified_methods_.AddDexFile(dex_file);
+  atomic_verified_methods_.AddDexFile(dex_file, dex_file->NumMethodIds());
   WriterMutexLock mu(Thread::Current(), verified_methods_lock_);
   // There can be some verified methods that are already registered for the dex_file since we set
   // up well known classes earlier. Remove these and put them in the array so that we don't
@@ -158,7 +162,9 @@
   for (auto it = verified_methods_.begin(); it != verified_methods_.end(); ) {
     MethodReference ref = it->first;
     if (ref.dex_file == dex_file) {
-      CHECK(atomic_verified_methods_.Insert(ref, nullptr, it->second) ==
+      CHECK(atomic_verified_methods_.Insert(DexFileReference(ref.dex_file, ref.dex_method_index),
+                                            nullptr,
+                                            it->second) ==
           AtomicMap::kInsertResultSuccess);
       it = verified_methods_.erase(it);
     } else {
diff --git a/compiler/dex/verification_results.h b/compiler/dex/verification_results.h
index 22749fa..5a03599 100644
--- a/compiler/dex/verification_results.h
+++ b/compiler/dex/verification_results.h
@@ -26,7 +26,7 @@
 #include "class_reference.h"
 #include "method_reference.h"
 #include "safe_map.h"
-#include "utils/atomic_method_ref_map.h"
+#include "utils/atomic_dex_ref_map.h"
 
 namespace art {
 
@@ -64,7 +64,7 @@
 
  private:
   // Verified methods. The method array is fixed to avoid needing a lock to extend it.
-  using AtomicMap = AtomicMethodRefMap<const VerifiedMethod*>;
+  using AtomicMap = AtomicDexRefMap<const VerifiedMethod*>;
   using VerifiedMethodMap = SafeMap<MethodReference,
                                     const VerifiedMethod*,
                                     MethodReferenceComparator>;
diff --git a/compiler/driver/compiler_driver.cc b/compiler/driver/compiler_driver.cc
index ff92ac4..49a42de 100644
--- a/compiler/driver/compiler_driver.cc
+++ b/compiler/driver/compiler_driver.cc
@@ -71,7 +71,7 @@
 #include "thread_pool.h"
 #include "trampolines/trampoline_compiler.h"
 #include "transaction.h"
-#include "utils/atomic_method_ref_map-inl.h"
+#include "utils/atomic_dex_ref_map-inl.h"
 #include "utils/dex_cache_arrays_layout-inl.h"
 #include "utils/swap_space.h"
 #include "vdex_file.h"
@@ -321,7 +321,7 @@
 }
 
 CompilerDriver::~CompilerDriver() {
-  compiled_methods_.Visit([this](const MethodReference& ref ATTRIBUTE_UNUSED,
+  compiled_methods_.Visit([this](const DexFileReference& ref ATTRIBUTE_UNUSED,
                                  CompiledMethod* method) {
     if (method != nullptr) {
       CompiledMethod::ReleaseSwapAllocatedCompiledMethod(this, method);
@@ -899,7 +899,7 @@
     for (const DexFile* dex_file : dex_files) {
       // Can be already inserted if the caller is CompileOne. This happens for gtests.
       if (!compiled_methods_.HaveDexFile(dex_file)) {
-        compiled_methods_.AddDexFile(dex_file);
+        compiled_methods_.AddDexFile(dex_file, dex_file->NumMethodIds());
       }
     }
     // Resolve eagerly to prepare for compilation.
@@ -2847,9 +2847,10 @@
                                        size_t non_relative_linker_patch_count) {
   DCHECK(GetCompiledMethod(method_ref) == nullptr)
       << method_ref.dex_file->PrettyMethod(method_ref.dex_method_index);
-  MethodTable::InsertResult result = compiled_methods_.Insert(method_ref,
-                                                              /*expected*/ nullptr,
-                                                              compiled_method);
+  MethodTable::InsertResult result = compiled_methods_.Insert(
+      DexFileReference(method_ref.dex_file, method_ref.dex_method_index),
+      /*expected*/ nullptr,
+      compiled_method);
   CHECK(result == MethodTable::kInsertResultSuccess);
   non_relative_linker_patch_count_.FetchAndAddRelaxed(non_relative_linker_patch_count);
   DCHECK(GetCompiledMethod(method_ref) != nullptr)
@@ -2897,7 +2898,7 @@
 
 CompiledMethod* CompilerDriver::GetCompiledMethod(MethodReference ref) const {
   CompiledMethod* compiled_method = nullptr;
-  compiled_methods_.Get(ref, &compiled_method);
+  compiled_methods_.Get(DexFileReference(ref.dex_file, ref.dex_method_index), &compiled_method);
   return compiled_method;
 }
 
diff --git a/compiler/driver/compiler_driver.h b/compiler/driver/compiler_driver.h
index e9e7378..a3272d3 100644
--- a/compiler/driver/compiler_driver.h
+++ b/compiler/driver/compiler_driver.h
@@ -39,7 +39,7 @@
 #include "os.h"
 #include "safe_map.h"
 #include "thread_pool.h"
-#include "utils/atomic_method_ref_map.h"
+#include "utils/atomic_dex_ref_map.h"
 #include "utils/dex_cache_arrays_layout.h"
 
 namespace art {
@@ -489,7 +489,7 @@
   mutable Mutex compiled_classes_lock_ DEFAULT_MUTEX_ACQUIRED_AFTER;
   ClassStateTable compiled_classes_ GUARDED_BY(compiled_classes_lock_);
 
-  typedef AtomicMethodRefMap<CompiledMethod*> MethodTable;
+  typedef AtomicDexRefMap<CompiledMethod*> MethodTable;
 
  private:
   // All method references that this compiler has compiled.
diff --git a/compiler/utils/atomic_method_ref_map-inl.h b/compiler/utils/atomic_dex_ref_map-inl.h
similarity index 62%
rename from compiler/utils/atomic_method_ref_map-inl.h
rename to compiler/utils/atomic_dex_ref_map-inl.h
index ad3a099..c41d8fc 100644
--- a/compiler/utils/atomic_method_ref_map-inl.h
+++ b/compiler/utils/atomic_dex_ref_map-inl.h
@@ -14,72 +14,72 @@
  * limitations under the License.
  */
 
-#ifndef ART_COMPILER_UTILS_ATOMIC_METHOD_REF_MAP_INL_H_
-#define ART_COMPILER_UTILS_ATOMIC_METHOD_REF_MAP_INL_H_
+#ifndef ART_COMPILER_UTILS_ATOMIC_DEX_REF_MAP_INL_H_
+#define ART_COMPILER_UTILS_ATOMIC_DEX_REF_MAP_INL_H_
 
-#include "atomic_method_ref_map.h"
+#include "atomic_dex_ref_map.h"
 
 #include "dex_file-inl.h"
 
 namespace art {
 
 template <typename T>
-inline typename AtomicMethodRefMap<T>::InsertResult AtomicMethodRefMap<T>::Insert(
-    MethodReference ref,
+inline typename AtomicDexRefMap<T>::InsertResult AtomicDexRefMap<T>::Insert(
+    DexFileReference ref,
     const T& expected,
     const T& desired) {
   ElementArray* const array = GetArray(ref.dex_file);
   if (array == nullptr) {
     return kInsertResultInvalidDexFile;
   }
-  return (*array)[ref.dex_method_index].CompareExchangeStrongSequentiallyConsistent(
-      expected, desired)
+  DCHECK_LT(ref.index, array->size());
+  return (*array)[ref.index].CompareExchangeStrongSequentiallyConsistent(expected, desired)
       ? kInsertResultSuccess
       : kInsertResultCASFailure;
 }
 
 template <typename T>
-inline bool AtomicMethodRefMap<T>::Get(MethodReference ref, T* out) const {
+inline bool AtomicDexRefMap<T>::Get(DexFileReference ref, T* out) const {
   const ElementArray* const array = GetArray(ref.dex_file);
   if (array == nullptr) {
     return false;
   }
-  *out = (*array)[ref.dex_method_index].LoadRelaxed();
+  *out = (*array)[ref.index].LoadRelaxed();
   return true;
 }
 
 template <typename T>
-inline void AtomicMethodRefMap<T>::AddDexFile(const DexFile* dex_file) {
-  arrays_.Put(dex_file, std::move(ElementArray(dex_file->NumMethodIds())));
+inline void AtomicDexRefMap<T>::AddDexFile(const DexFile* dex_file, size_t max_index) {
+  arrays_.Put(dex_file, std::move(ElementArray(max_index)));
 }
 
 template <typename T>
-inline typename AtomicMethodRefMap<T>::ElementArray* AtomicMethodRefMap<T>::GetArray(
+inline typename AtomicDexRefMap<T>::ElementArray* AtomicDexRefMap<T>::GetArray(
     const DexFile* dex_file) {
   auto it = arrays_.find(dex_file);
   return (it != arrays_.end()) ? &it->second : nullptr;
 }
 
 template <typename T>
-inline const typename AtomicMethodRefMap<T>::ElementArray* AtomicMethodRefMap<T>::GetArray(
+inline const typename AtomicDexRefMap<T>::ElementArray* AtomicDexRefMap<T>::GetArray(
     const DexFile* dex_file) const {
   auto it = arrays_.find(dex_file);
   return (it != arrays_.end()) ? &it->second : nullptr;
 }
 
 template <typename T> template <typename Visitor>
-inline void AtomicMethodRefMap<T>::Visit(const Visitor& visitor) {
+inline void AtomicDexRefMap<T>::Visit(const Visitor& visitor) {
   for (auto& pair : arrays_) {
     const DexFile* dex_file = pair.first;
     const ElementArray& elements = pair.second;
     for (size_t i = 0; i < elements.size(); ++i) {
-      visitor(MethodReference(dex_file, i), elements[i].LoadRelaxed());
+      visitor(DexFileReference(dex_file, i), elements[i].LoadRelaxed());
     }
   }
 }
 
 template <typename T>
-inline void AtomicMethodRefMap<T>::ClearEntries() {
+inline void AtomicDexRefMap<T>::ClearEntries() {
   for (auto& it : arrays_) {
     for (auto& element : it.second) {
       element.StoreRelaxed(nullptr);
@@ -89,4 +89,4 @@
 
 }  // namespace art
 
-#endif  // ART_COMPILER_UTILS_ATOMIC_METHOD_REF_MAP_INL_H_
+#endif  // ART_COMPILER_UTILS_ATOMIC_DEX_REF_MAP_INL_H_
diff --git a/compiler/utils/atomic_method_ref_map.h b/compiler/utils/atomic_dex_ref_map.h
similarity index 80%
rename from compiler/utils/atomic_method_ref_map.h
rename to compiler/utils/atomic_dex_ref_map.h
index fed848f..2da4ffa 100644
--- a/compiler/utils/atomic_method_ref_map.h
+++ b/compiler/utils/atomic_dex_ref_map.h
@@ -14,11 +14,11 @@
  * limitations under the License.
  */
 
-#ifndef ART_COMPILER_UTILS_ATOMIC_METHOD_REF_MAP_H_
-#define ART_COMPILER_UTILS_ATOMIC_METHOD_REF_MAP_H_
+#ifndef ART_COMPILER_UTILS_ATOMIC_DEX_REF_MAP_H_
+#define ART_COMPILER_UTILS_ATOMIC_DEX_REF_MAP_H_
 
 #include "base/dchecked_vector.h"
-#include "method_reference.h"
+#include "dex_file.h"
 #include "safe_map.h"
 
 namespace art {
@@ -27,10 +27,10 @@
 
 // Used by CompilerCallbacks to track verification information from the Runtime.
 template <typename T>
-class AtomicMethodRefMap {
+class AtomicDexRefMap {
  public:
-  explicit AtomicMethodRefMap() {}
-  ~AtomicMethodRefMap() {}
+  explicit AtomicDexRefMap() {}
+  ~AtomicDexRefMap() {}
 
   // Atomically swap the element in if the existing value matches expected.
   enum InsertResult {
@@ -38,14 +38,14 @@
     kInsertResultCASFailure,
     kInsertResultSuccess,
   };
-  InsertResult Insert(MethodReference ref, const T& expected, const T& desired);
+  InsertResult Insert(DexFileReference ref, const T& expected, const T& desired);
 
   // Retreive an item, returns false if the dex file is not added.
-  bool Get(MethodReference ref, T* out) const;
+  bool Get(DexFileReference ref, T* out) const;
 
   // Dex files must be added before method references belonging to them can be used as keys. Not
   // thread safe.
-  void AddDexFile(const DexFile* dex_file);
+  void AddDexFile(const DexFile* dex_file, size_t max_index);
 
   bool HaveDexFile(const DexFile* dex_file) const {
     return arrays_.find(dex_file) != arrays_.end();
@@ -70,4 +70,4 @@
 
 }  // namespace art
 
-#endif  // ART_COMPILER_UTILS_ATOMIC_METHOD_REF_MAP_H_
+#endif  // ART_COMPILER_UTILS_ATOMIC_DEX_REF_MAP_H_
diff --git a/compiler/utils/atomic_method_ref_map_test.cc b/compiler/utils/atomic_dex_ref_map_test.cc
similarity index 64%
rename from compiler/utils/atomic_method_ref_map_test.cc
rename to compiler/utils/atomic_dex_ref_map_test.cc
index 9e5bf4b..ae19a9c 100644
--- a/compiler/utils/atomic_method_ref_map_test.cc
+++ b/compiler/utils/atomic_dex_ref_map_test.cc
@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-#include "atomic_method_ref_map-inl.h"
+#include "atomic_dex_ref_map-inl.h"
 
 #include <memory>
 
@@ -25,46 +25,46 @@
 
 namespace art {
 
-class AtomicMethodRefMapTest : public CommonRuntimeTest {};
+class AtomicDexRefMapTest : public CommonRuntimeTest {};
 
-TEST_F(AtomicMethodRefMapTest, RunTests) {
+TEST_F(AtomicDexRefMapTest, RunTests) {
   ScopedObjectAccess soa(Thread::Current());
   std::unique_ptr<const DexFile> dex(OpenTestDexFile("Interfaces"));
   ASSERT_TRUE(dex != nullptr);
-  using Map = AtomicMethodRefMap<int>;
+  using Map = AtomicDexRefMap<int>;
   Map map;
   int value = 123;
   // Error case: Not already inserted.
-  EXPECT_FALSE(map.Get(MethodReference(dex.get(), 1), &value));
+  EXPECT_FALSE(map.Get(DexFileReference(dex.get(), 1), &value));
   EXPECT_FALSE(map.HaveDexFile(dex.get()));
   // Error case: Dex file not registered.
-  EXPECT_TRUE(map.Insert(MethodReference(dex.get(), 1), 0, 1) == Map::kInsertResultInvalidDexFile);
-  map.AddDexFile(dex.get());
+  EXPECT_TRUE(map.Insert(DexFileReference(dex.get(), 1), 0, 1) == Map::kInsertResultInvalidDexFile);
+  map.AddDexFile(dex.get(), dex->NumMethodIds());
   EXPECT_TRUE(map.HaveDexFile(dex.get()));
   EXPECT_GT(dex->NumMethodIds(), 10u);
   // After we have added the get should succeed but return the default value.
-  EXPECT_TRUE(map.Get(MethodReference(dex.get(), 1), &value));
+  EXPECT_TRUE(map.Get(DexFileReference(dex.get(), 1), &value));
   EXPECT_EQ(value, 0);
   // Actually insert an item and make sure we can retreive it.
   static const int kInsertValue = 44;
-  EXPECT_TRUE(map.Insert(MethodReference(dex.get(), 1), 0, kInsertValue) ==
+  EXPECT_TRUE(map.Insert(DexFileReference(dex.get(), 1), 0, kInsertValue) ==
               Map::kInsertResultSuccess);
-  EXPECT_TRUE(map.Get(MethodReference(dex.get(), 1), &value));
+  EXPECT_TRUE(map.Get(DexFileReference(dex.get(), 1), &value));
   EXPECT_EQ(value, kInsertValue);
   static const int kInsertValue2 = 123;
-  EXPECT_TRUE(map.Insert(MethodReference(dex.get(), 2), 0, kInsertValue2) ==
+  EXPECT_TRUE(map.Insert(DexFileReference(dex.get(), 2), 0, kInsertValue2) ==
               Map::kInsertResultSuccess);
-  EXPECT_TRUE(map.Get(MethodReference(dex.get(), 1), &value));
+  EXPECT_TRUE(map.Get(DexFileReference(dex.get(), 1), &value));
   EXPECT_EQ(value, kInsertValue);
-  EXPECT_TRUE(map.Get(MethodReference(dex.get(), 2), &value));
+  EXPECT_TRUE(map.Get(DexFileReference(dex.get(), 2), &value));
   EXPECT_EQ(value, kInsertValue2);
   // Error case: Incorrect expected value for CAS.
-  EXPECT_TRUE(map.Insert(MethodReference(dex.get(), 1), 0, kInsertValue + 1) ==
+  EXPECT_TRUE(map.Insert(DexFileReference(dex.get(), 1), 0, kInsertValue + 1) ==
       Map::kInsertResultCASFailure);
   // Correctly overwrite the value and verify.
-  EXPECT_TRUE(map.Insert(MethodReference(dex.get(), 1), kInsertValue, kInsertValue + 1) ==
+  EXPECT_TRUE(map.Insert(DexFileReference(dex.get(), 1), kInsertValue, kInsertValue + 1) ==
       Map::kInsertResultSuccess);
-  EXPECT_TRUE(map.Get(MethodReference(dex.get(), 1), &value));
+  EXPECT_TRUE(map.Get(DexFileReference(dex.get(), 1), &value));
   EXPECT_EQ(value, kInsertValue + 1);
 }
 
diff --git a/compiler/verifier_deps_test.cc b/compiler/verifier_deps_test.cc
index 7e616a7..686da21 100644
--- a/compiler/verifier_deps_test.cc
+++ b/compiler/verifier_deps_test.cc
@@ -33,7 +33,7 @@
 #include "runtime.h"
 #include "scoped_thread_state_change-inl.h"
 #include "thread.h"
-#include "utils/atomic_method_ref_map-inl.h"
+#include "utils/atomic_dex_ref_map-inl.h"
 #include "verifier/method_verifier-inl.h"
 
 namespace art {
@@ -97,9 +97,9 @@
     callbacks_->SetVerifierDeps(nullptr);
     // Clear entries in the verification results to avoid hitting a DCHECK that
     // we always succeed inserting a new entry after verifying.
-    AtomicMethodRefMap<const VerifiedMethod*>* map =
+    AtomicDexRefMap<const VerifiedMethod*>* map =
         &compiler_driver_->GetVerificationResults()->atomic_verified_methods_;
-    map->Visit([](const MethodReference& ref ATTRIBUTE_UNUSED, const VerifiedMethod* method) {
+    map->Visit([](const DexFileReference& ref ATTRIBUTE_UNUSED, const VerifiedMethod* method) {
       delete method;
     });
     map->ClearEntries();