Clean up AtomicDexRefMap

Make ClassReference, TypeReference, and MethodReference extend
DexFileReference. This enables using all of these types as the key
for AtomicDexRefMap.

Test: test-art-host
Bug: 63851220
Bug: 63756964

Change-Id: Ida3c94cadb53272cb5057e5cebc5971c1ab4d366

3 files changed, 73 insertions, 42 deletions

diff --git a/compiler/utils/atomic_dex_ref_map-inl.h b/compiler/utils/atomic_dex_ref_map-inl.h
index 14f1f0b981..33d59f9d42 100644
--- a/compiler/utils/atomic_dex_ref_map-inl.h
+++ b/compiler/utils/atomic_dex_ref_map-inl.h
@@ -19,15 +19,40 @@
 
 #include "atomic_dex_ref_map.h"
 
+#include <type_traits>
+
 #include "dex_file-inl.h"
+#include "class_reference.h"
+#include "method_reference.h"
+#include "type_reference.h"
 
 namespace art {
 
-template <typename T>
-inline typename AtomicDexRefMap<T>::InsertResult AtomicDexRefMap<T>::Insert(
-    DexFileReference ref,
-    const T& expected,
-    const T& desired) {
+template <typename DexFileReferenceType, typename Value>
+inline size_t AtomicDexRefMap<DexFileReferenceType, Value>::NumberOfDexIndices(
+    const DexFile* dex_file) {
+  // TODO: Use specialization for this? Not sure if worth it.
+  static_assert(std::is_same<DexFileReferenceType, MethodReference>::value ||
+                std::is_same<DexFileReferenceType, ClassReference>::value ||
+                std::is_same<DexFileReferenceType, TypeReference>::value,
+                "invalid index type");
+  if (std::is_same<DexFileReferenceType, MethodReference>::value) {
+    return dex_file->NumMethodIds();
+  }
+  if (std::is_same<DexFileReferenceType, ClassReference>::value) {
+    return dex_file->NumClassDefs();
+  }
+  if (std::is_same<DexFileReferenceType, TypeReference>::value) {
+    return dex_file->NumTypeIds();
+  }
+  UNREACHABLE();
+}
+
+template <typename DexFileReferenceType, typename Value>
+inline typename AtomicDexRefMap<DexFileReferenceType, Value>::InsertResult
+    AtomicDexRefMap<DexFileReferenceType, Value>::Insert(const DexFileReferenceType& ref,
+                                                         const Value& expected,
+                                                         const Value& desired) {
   ElementArray* const array = GetArray(ref.dex_file);
   if (array == nullptr) {
     return kInsertResultInvalidDexFile;
@@ -38,8 +63,9 @@ inline typename AtomicDexRefMap<T>::InsertResult AtomicDexRefMap<T>::Insert(
       : kInsertResultCASFailure;
 }
 
-template <typename T>
-inline bool AtomicDexRefMap<T>::Get(DexFileReference ref, T* out) const {
+template <typename DexFileReferenceType, typename Value>
+inline bool AtomicDexRefMap<DexFileReferenceType, Value>::Get(const DexFileReferenceType& ref,
+                                                              Value* out) const {
   const ElementArray* const array = GetArray(ref.dex_file);
   if (array == nullptr) {
     return false;
@@ -48,36 +74,37 @@ inline bool AtomicDexRefMap<T>::Get(DexFileReference ref, T* out) const {
   return true;
 }
 
-template <typename T>
-inline void AtomicDexRefMap<T>::AddDexFile(const DexFile* dex_file, size_t max_index) {
-  arrays_.Put(dex_file, std::move(ElementArray(max_index)));
+template <typename DexFileReferenceType, typename Value>
+inline void AtomicDexRefMap<DexFileReferenceType, Value>::AddDexFile(const DexFile* dex_file) {
+  arrays_.Put(dex_file, std::move(ElementArray(NumberOfDexIndices(dex_file))));
 }
 
-template <typename T>
-inline void AtomicDexRefMap<T>::AddDexFiles(const std::vector<const DexFile*>& dex_files) {
+template <typename DexFileReferenceType, typename Value>
+inline void AtomicDexRefMap<DexFileReferenceType, Value>::AddDexFiles(
+    const std::vector<const DexFile*>& dex_files) {
   for (const DexFile* dex_file : dex_files) {
     if (!HaveDexFile(dex_file)) {
-      AddDexFile(dex_file, dex_file->NumClassDefs());
+      AddDexFile(dex_file);
     }
   }
 }
 
-template <typename T>
-inline typename AtomicDexRefMap<T>::ElementArray* AtomicDexRefMap<T>::GetArray(
-    const DexFile* dex_file) {
+template <typename DexFileReferenceType, typename Value>
+inline typename AtomicDexRefMap<DexFileReferenceType, Value>::ElementArray*
+    AtomicDexRefMap<DexFileReferenceType, Value>::GetArray(const DexFile* dex_file) {
   auto it = arrays_.find(dex_file);
   return (it != arrays_.end()) ? &it->second : nullptr;
 }
 
-template <typename T>
-inline const typename AtomicDexRefMap<T>::ElementArray* AtomicDexRefMap<T>::GetArray(
-    const DexFile* dex_file) const {
+template <typename DexFileReferenceType, typename Value>
+inline const typename AtomicDexRefMap<DexFileReferenceType, Value>::ElementArray*
+    AtomicDexRefMap<DexFileReferenceType, Value>::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 AtomicDexRefMap<T>::Visit(const Visitor& visitor) {
+template <typename DexFileReferenceType, typename Value> template <typename Visitor>
+inline void AtomicDexRefMap<DexFileReferenceType, Value>::Visit(const Visitor& visitor) {
   for (auto& pair : arrays_) {
     const DexFile* dex_file = pair.first;
     const ElementArray& elements = pair.second;
@@ -87,8 +114,8 @@ inline void AtomicDexRefMap<T>::Visit(const Visitor& visitor) {
   }
 }
 
-template <typename T>
-inline void AtomicDexRefMap<T>::ClearEntries() {
+template <typename DexFileReferenceType, typename Value>
+inline void AtomicDexRefMap<DexFileReferenceType, Value>::ClearEntries() {
   for (auto& it : arrays_) {
     for (auto& element : it.second) {
       element.StoreRelaxed(nullptr);
diff --git a/compiler/utils/atomic_dex_ref_map.h b/compiler/utils/atomic_dex_ref_map.h
index b02c9b634e..fad056c191 100644
--- a/compiler/utils/atomic_dex_ref_map.h
+++ b/compiler/utils/atomic_dex_ref_map.h
@@ -18,7 +18,7 @@
 #define ART_COMPILER_UTILS_ATOMIC_DEX_REF_MAP_H_
 
 #include "base/dchecked_vector.h"
-#include "dex_file.h"
+#include "dex_file_reference.h"
 #include "safe_map.h"
 
 namespace art {
@@ -26,7 +26,7 @@ namespace art {
 class DexFile;
 
 // Used by CompilerCallbacks to track verification information from the Runtime.
-template <typename T>
+template <typename DexFileReferenceType, typename Value>
 class AtomicDexRefMap {
  public:
   explicit AtomicDexRefMap() {}
@@ -38,14 +38,16 @@ class AtomicDexRefMap {
     kInsertResultCASFailure,
     kInsertResultSuccess,
   };
-  InsertResult Insert(DexFileReference ref, const T& expected, const T& desired);
+  InsertResult Insert(const DexFileReferenceType& ref,
+                      const Value& expected,
+                      const Value& desired);
 
   // Retreive an item, returns false if the dex file is not added.
-  bool Get(DexFileReference ref, T* out) const;
+  bool Get(const DexFileReferenceType& ref, Value* 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, size_t max_index);
+  void AddDexFile(const DexFile* dex_file);
   void AddDexFiles(const std::vector<const DexFile*>& dex_files);
 
   bool HaveDexFile(const DexFile* dex_file) const {
@@ -60,12 +62,14 @@ class AtomicDexRefMap {
 
  private:
   // Verified methods. The method array is fixed to avoid needing a lock to extend it.
-  using ElementArray = dchecked_vector<Atomic<T>>;
+  using ElementArray = dchecked_vector<Atomic<Value>>;
   using DexFileArrays = SafeMap<const DexFile*, ElementArray>;
 
   const ElementArray* GetArray(const DexFile* dex_file) const;
   ElementArray* GetArray(const DexFile* dex_file);
 
+  static size_t NumberOfDexIndices(const DexFile* dex_file);
+
   DexFileArrays arrays_;
 };
 
diff --git a/compiler/utils/atomic_dex_ref_map_test.cc b/compiler/utils/atomic_dex_ref_map_test.cc
index ae19a9c6da..8fce36f021 100644
--- a/compiler/utils/atomic_dex_ref_map_test.cc
+++ b/compiler/utils/atomic_dex_ref_map_test.cc
@@ -31,40 +31,40 @@ TEST_F(AtomicDexRefMapTest, RunTests) {
   ScopedObjectAccess soa(Thread::Current());
   std::unique_ptr<const DexFile> dex(OpenTestDexFile("Interfaces"));
   ASSERT_TRUE(dex != nullptr);
-  using Map = AtomicDexRefMap<int>;
+  using Map = AtomicDexRefMap<MethodReference, int>;
   Map map;
   int value = 123;
   // Error case: Not already inserted.
-  EXPECT_FALSE(map.Get(DexFileReference(dex.get(), 1), &value));
+  EXPECT_FALSE(map.Get(MethodReference(dex.get(), 1), &value));
   EXPECT_FALSE(map.HaveDexFile(dex.get()));
   // Error case: Dex file not registered.
-  EXPECT_TRUE(map.Insert(DexFileReference(dex.get(), 1), 0, 1) == Map::kInsertResultInvalidDexFile);
-  map.AddDexFile(dex.get(), dex->NumMethodIds());
+  EXPECT_TRUE(map.Insert(MethodReference(dex.get(), 1), 0, 1) == Map::kInsertResultInvalidDexFile);
+  map.AddDexFile(dex.get());
   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(DexFileReference(dex.get(), 1), &value));
+  EXPECT_TRUE(map.Get(MethodReference(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(DexFileReference(dex.get(), 1), 0, kInsertValue) ==
+  EXPECT_TRUE(map.Insert(MethodReference(dex.get(), 1), 0, kInsertValue) ==
               Map::kInsertResultSuccess);
-  EXPECT_TRUE(map.Get(DexFileReference(dex.get(), 1), &value));
+  EXPECT_TRUE(map.Get(MethodReference(dex.get(), 1), &value));
   EXPECT_EQ(value, kInsertValue);
   static const int kInsertValue2 = 123;
-  EXPECT_TRUE(map.Insert(DexFileReference(dex.get(), 2), 0, kInsertValue2) ==
+  EXPECT_TRUE(map.Insert(MethodReference(dex.get(), 2), 0, kInsertValue2) ==
               Map::kInsertResultSuccess);
-  EXPECT_TRUE(map.Get(DexFileReference(dex.get(), 1), &value));
+  EXPECT_TRUE(map.Get(MethodReference(dex.get(), 1), &value));
   EXPECT_EQ(value, kInsertValue);
-  EXPECT_TRUE(map.Get(DexFileReference(dex.get(), 2), &value));
+  EXPECT_TRUE(map.Get(MethodReference(dex.get(), 2), &value));
   EXPECT_EQ(value, kInsertValue2);
   // Error case: Incorrect expected value for CAS.
-  EXPECT_TRUE(map.Insert(DexFileReference(dex.get(), 1), 0, kInsertValue + 1) ==
+  EXPECT_TRUE(map.Insert(MethodReference(dex.get(), 1), 0, kInsertValue + 1) ==
       Map::kInsertResultCASFailure);
   // Correctly overwrite the value and verify.
-  EXPECT_TRUE(map.Insert(DexFileReference(dex.get(), 1), kInsertValue, kInsertValue + 1) ==
+  EXPECT_TRUE(map.Insert(MethodReference(dex.get(), 1), kInsertValue, kInsertValue + 1) ==
       Map::kInsertResultSuccess);
-  EXPECT_TRUE(map.Get(DexFileReference(dex.get(), 1), &value));
+  EXPECT_TRUE(map.Get(MethodReference(dex.get(), 1), &value));
   EXPECT_EQ(value, kInsertValue + 1);
 }