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
diff --git a/compiler/dex/inline_method_analyser.cc b/compiler/dex/inline_method_analyser.cc
index 54ddc21..c8e3d5e 100644
--- a/compiler/dex/inline_method_analyser.cc
+++ b/compiler/dex/inline_method_analyser.cc
@@ -509,7 +509,7 @@
 }
 
 bool InlineMethodAnalyser::IsSyntheticAccessor(MethodReference ref) {
-  const DexFile::MethodId& method_id = ref.dex_file->GetMethodId(ref.dex_method_index);
+  const DexFile::MethodId& method_id = ref.dex_file->GetMethodId(ref.index);
   const char* method_name = ref.dex_file->GetMethodName(method_id);
   // javac names synthetic accessors "access$nnn",
   // jack names them "-getN", "-putN", "-wrapN".
diff --git a/compiler/dex/verification_results.cc b/compiler/dex/verification_results.cc
index cfb56e3..03c90d8 100644
--- a/compiler/dex/verification_results.cc
+++ b/compiler/dex/verification_results.cc
@@ -52,18 +52,16 @@
     // We'll punt this later.
     return;
   }
-  AtomicMap::InsertResult result = atomic_verified_methods_.Insert(
-      DexFileReference(ref.dex_file, ref.dex_method_index),
-      /*expected*/ nullptr,
-      verified_method.get());
+  AtomicMap::InsertResult result = atomic_verified_methods_.Insert(ref,
+                                                                   /*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(DexFileReference(ref.dex_file, ref.dex_method_index),
-                                         &existing));
+      CHECK(atomic_verified_methods_.Get(ref, &existing));
       CHECK_NE(verified_method.get(), existing);
     }
   } else {
@@ -100,7 +98,7 @@
 
 const VerifiedMethod* VerificationResults::GetVerifiedMethod(MethodReference ref) {
   const VerifiedMethod* ret = nullptr;
-  if (atomic_verified_methods_.Get(DexFileReference(ref.dex_file, ref.dex_method_index), &ret)) {
+  if (atomic_verified_methods_.Get(ref, &ret)) {
     return ret;
   }
   ReaderMutexLock mu(Thread::Current(), verified_methods_lock_);
@@ -114,7 +112,7 @@
   // 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(DexFileReference(ref.dex_file, ref.dex_method_index),
+  if (atomic_verified_methods_.Insert(ref,
                                       /*expected*/ nullptr,
                                       verified_method.get()) ==
           AtomicMap::InsertResult::kInsertResultSuccess) {
@@ -149,7 +147,7 @@
 }
 
 void VerificationResults::AddDexFile(const DexFile* dex_file) {
-  atomic_verified_methods_.AddDexFile(dex_file, dex_file->NumMethodIds());
+  atomic_verified_methods_.AddDexFile(dex_file);
   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
@@ -157,9 +155,7 @@
   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(DexFileReference(ref.dex_file, ref.dex_method_index),
-                                            nullptr,
-                                            it->second) ==
+      CHECK(atomic_verified_methods_.Insert(ref, 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 5a03599..d19e993 100644
--- a/compiler/dex/verification_results.h
+++ b/compiler/dex/verification_results.h
@@ -64,10 +64,8 @@
 
  private:
   // Verified methods. The method array is fixed to avoid needing a lock to extend it.
-  using AtomicMap = AtomicDexRefMap<const VerifiedMethod*>;
-  using VerifiedMethodMap = SafeMap<MethodReference,
-                                    const VerifiedMethod*,
-                                    MethodReferenceComparator>;
+  using AtomicMap = AtomicDexRefMap<MethodReference, const VerifiedMethod*>;
+  using VerifiedMethodMap = SafeMap<MethodReference, const VerifiedMethod*>;
 
   VerifiedMethodMap verified_methods_ GUARDED_BY(verified_methods_lock_);
   const CompilerOptions* const compiler_options_;
diff --git a/compiler/driver/compiler_driver.cc b/compiler/driver/compiler_driver.cc
index b9c0314..678f090 100644
--- a/compiler/driver/compiler_driver.cc
+++ b/compiler/driver/compiler_driver.cc
@@ -896,7 +896,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, dex_file->NumMethodIds());
+        compiled_methods_.AddDexFile(dex_file);
       }
     }
     // Resolve eagerly to prepare for compilation.
@@ -964,7 +964,7 @@
     return true;
   }
 
-  std::string tmp = method_ref.dex_file->PrettyMethod(method_ref.dex_method_index, true);
+  std::string tmp = method_ref.PrettyMethod();
   return methods_to_compile_->find(tmp.c_str()) != methods_to_compile_->end();
 }
 
@@ -985,8 +985,7 @@
 
   if (kDebugProfileGuidedCompilation) {
     LOG(INFO) << "[ProfileGuidedCompilation] "
-        << (result ? "Compiled" : "Skipped") << " method:"
-        << method_ref.dex_file->PrettyMethod(method_ref.dex_method_index, true);
+        << (result ? "Compiled" : "Skipped") << " method:" << method_ref.PrettyMethod(true);
   }
   return result;
 }
@@ -1359,7 +1358,7 @@
       &dex_to_dex_references_.back().GetDexFile() != method_ref.dex_file) {
     dex_to_dex_references_.emplace_back(*method_ref.dex_file);
   }
-  dex_to_dex_references_.back().GetMethodIndexes().SetBit(method_ref.dex_method_index);
+  dex_to_dex_references_.back().GetMethodIndexes().SetBit(method_ref.index);
 }
 
 bool CompilerDriver::CanAccessTypeWithoutChecks(ObjPtr<mirror::Class> referrer_class,
@@ -1944,7 +1943,7 @@
         if (compiler_only_verifies) {
           // Just update the compiled_classes_ map. The compiler doesn't need to resolve
           // the type.
-          DexFileReference ref(dex_file, i);
+          ClassReference ref(dex_file, i);
           mirror::Class::Status existing = mirror::Class::kStatusNotReady;
           DCHECK(compiled_classes_.Get(ref, &existing)) << ref.dex_file->GetLocation();
           ClassStateTable::InsertResult result =
@@ -2220,7 +2219,7 @@
                                         TimingLogger* timings) {
   TimingLogger::ScopedTiming t("Verify Dex File", timings);
   if (!compiled_classes_.HaveDexFile(&dex_file)) {
-    compiled_classes_.AddDexFile(&dex_file, dex_file.NumClassDefs());
+    compiled_classes_.AddDexFile(&dex_file);
   }
   ClassLinker* class_linker = Runtime::Current()->GetClassLinker();
   ParallelCompilationManager context(class_linker, class_loader, this, &dex_file, dex_files,
@@ -2877,38 +2876,36 @@
 void CompilerDriver::AddCompiledMethod(const MethodReference& method_ref,
                                        CompiledMethod* const compiled_method,
                                        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(
-      DexFileReference(method_ref.dex_file, method_ref.dex_method_index),
-      /*expected*/ nullptr,
-      compiled_method);
+  DCHECK(GetCompiledMethod(method_ref) == nullptr) << method_ref.PrettyMethod();
+  MethodTable::InsertResult result = compiled_methods_.Insert(method_ref,
+                                                              /*expected*/ nullptr,
+                                                              compiled_method);
   CHECK(result == MethodTable::kInsertResultSuccess);
   non_relative_linker_patch_count_.FetchAndAddRelaxed(non_relative_linker_patch_count);
-  DCHECK(GetCompiledMethod(method_ref) != nullptr)
-      << method_ref.dex_file->PrettyMethod(method_ref.dex_method_index);
+  DCHECK(GetCompiledMethod(method_ref) != nullptr) << method_ref.PrettyMethod();
 }
 
-bool CompilerDriver::GetCompiledClass(ClassReference ref, mirror::Class::Status* status) const {
+bool CompilerDriver::GetCompiledClass(const ClassReference& ref,
+                                      mirror::Class::Status* status) const {
   DCHECK(status != nullptr);
   // The table doesn't know if something wasn't inserted. For this case it will return
   // kStatusNotReady. To handle this, just assume anything we didn't try to verify is not compiled.
-  if (!compiled_classes_.Get(DexFileReference(ref.first, ref.second), status) ||
+  if (!compiled_classes_.Get(ref, status) ||
       *status < mirror::Class::kStatusRetryVerificationAtRuntime) {
     return false;
   }
   return true;
 }
 
-mirror::Class::Status CompilerDriver::GetClassStatus(ClassReference ref) const {
+mirror::Class::Status CompilerDriver::GetClassStatus(const ClassReference& ref) const {
   mirror::Class::Status status = ClassStatus::kStatusNotReady;
   if (!GetCompiledClass(ref, &status)) {
-    classpath_classes_.Get(DexFileReference(ref.first, ref.second), &status);
+    classpath_classes_.Get(ref, &status);
   }
   return status;
 }
 
-void CompilerDriver::RecordClassStatus(ClassReference ref, mirror::Class::Status status) {
+void CompilerDriver::RecordClassStatus(const ClassReference& ref, mirror::Class::Status status) {
   switch (status) {
     case mirror::Class::kStatusErrorResolved:
     case mirror::Class::kStatusErrorUnresolved:
@@ -2921,30 +2918,30 @@
       break;  // Expected states.
     default:
       LOG(FATAL) << "Unexpected class status for class "
-          << PrettyDescriptor(ref.first->GetClassDescriptor(ref.first->GetClassDef(ref.second)))
+          << PrettyDescriptor(
+              ref.dex_file->GetClassDescriptor(ref.dex_file->GetClassDef(ref.index)))
           << " of " << status;
   }
 
   ClassStateTable::InsertResult result;
   ClassStateTable* table = &compiled_classes_;
   do {
-    DexFileReference dex_ref(ref.first, ref.second);
     mirror::Class::Status existing = mirror::Class::kStatusNotReady;
-    if (!table->Get(dex_ref, &existing)) {
+    if (!table->Get(ref, &existing)) {
       // A classpath class.
       if (kIsDebugBuild) {
         // Check to make sure it's not a dex file for an oat file we are compiling since these
         // should always succeed. These do not include classes in for used libraries.
         for (const DexFile* dex_file : GetDexFilesForOatFile()) {
-          CHECK_NE(dex_ref.dex_file, dex_file) << dex_ref.dex_file->GetLocation();
+          CHECK_NE(ref.dex_file, dex_file) << ref.dex_file->GetLocation();
         }
       }
-      if (!classpath_classes_.HaveDexFile(ref.first)) {
+      if (!classpath_classes_.HaveDexFile(ref.dex_file)) {
         // Boot classpath dex file.
         return;
       }
       table = &classpath_classes_;
-      table->Get(dex_ref, &existing);
+      table->Get(ref, &existing);
     }
     if (existing >= status) {
       // Existing status is already better than we expect, break.
@@ -2952,14 +2949,14 @@
     }
     // Update the status if we now have a greater one. This happens with vdex,
     // which records a class is verified, but does not resolve it.
-    result = table->Insert(dex_ref, existing, status);
-    CHECK(result != ClassStateTable::kInsertResultInvalidDexFile) << ref.first->GetLocation();
+    result = table->Insert(ref, existing, status);
+    CHECK(result != ClassStateTable::kInsertResultInvalidDexFile) << ref.dex_file->GetLocation();
   } while (result != ClassStateTable::kInsertResultSuccess);
 }
 
 CompiledMethod* CompilerDriver::GetCompiledMethod(MethodReference ref) const {
   CompiledMethod* compiled_method = nullptr;
-  compiled_methods_.Get(DexFileReference(ref.dex_file, ref.dex_method_index), &compiled_method);
+  compiled_methods_.Get(ref, &compiled_method);
   return compiled_method;
 }
 
diff --git a/compiler/driver/compiler_driver.h b/compiler/driver/compiler_driver.h
index 5043884..0d1cce7 100644
--- a/compiler/driver/compiler_driver.h
+++ b/compiler/driver/compiler_driver.h
@@ -154,8 +154,8 @@
   std::unique_ptr<const std::vector<uint8_t>> CreateQuickResolutionTrampoline() const;
   std::unique_ptr<const std::vector<uint8_t>> CreateQuickToInterpreterBridge() const;
 
-  bool GetCompiledClass(ClassReference ref, mirror::Class::Status* status) const;
-  mirror::Class::Status GetClassStatus(ClassReference ref) const;
+  mirror::Class::Status GetClassStatus(const ClassReference& ref) const;
+  bool GetCompiledClass(const ClassReference& ref, mirror::Class::Status* status) const;
 
   CompiledMethod* GetCompiledMethod(MethodReference ref) const;
   size_t GetNonRelativeLinkerPatchCount() const;
@@ -338,7 +338,7 @@
   // according to the profile file.
   bool ShouldVerifyClassBasedOnProfile(const DexFile& dex_file, uint16_t class_idx) const;
 
-  void RecordClassStatus(ClassReference ref, mirror::Class::Status status);
+  void RecordClassStatus(const ClassReference& ref, mirror::Class::Status status);
 
   // Checks if the specified method has been verified without failures. Returns
   // false if the method is not in the verification results (GetVerificationResults).
@@ -489,13 +489,13 @@
   std::map<ClassReference, bool> requires_constructor_barrier_
       GUARDED_BY(requires_constructor_barrier_lock_);
 
-  using ClassStateTable = AtomicDexRefMap<mirror::Class::Status>;
   // All class references that this compiler has compiled. Indexed by class defs.
+  using ClassStateTable = AtomicDexRefMap<ClassReference, mirror::Class::Status>;
   ClassStateTable compiled_classes_;
   // All class references that are in the classpath. Indexed by class defs.
   ClassStateTable classpath_classes_;
 
-  typedef AtomicDexRefMap<CompiledMethod*> MethodTable;
+  typedef AtomicDexRefMap<MethodReference, CompiledMethod*> MethodTable;
 
  private:
   // All method references that this compiler has compiled.
diff --git a/compiler/linker/arm/relative_patcher_arm_base.cc b/compiler/linker/arm/relative_patcher_arm_base.cc
index 6e63bf8..1b061b3 100644
--- a/compiler/linker/arm/relative_patcher_arm_base.cc
+++ b/compiler/linker/arm/relative_patcher_arm_base.cc
@@ -469,8 +469,7 @@
     if (!method_call_thunk_->HasReservedOffset() ||
         patch_offset - method_call_thunk_->LastReservedOffset() > max_negative_displacement) {
       // No previous thunk in range, check if we can reach the target directly.
-      if (target_method.dex_file == method_ref.dex_file &&
-          target_method.dex_method_index == method_ref.dex_method_index) {
+      if (target_method == method_ref) {
         DCHECK_GT(quick_code_offset, patch_offset);
         if (quick_code_offset - patch_offset > max_positive_displacement) {
           break;
diff --git a/compiler/linker/arm64/relative_patcher_arm64_test.cc b/compiler/linker/arm64/relative_patcher_arm64_test.cc
index d6919e9..8a5b4cc 100644
--- a/compiler/linker/arm64/relative_patcher_arm64_test.cc
+++ b/compiler/linker/arm64/relative_patcher_arm64_test.cc
@@ -344,7 +344,7 @@
                                        uint32_t use_insn) {
     uint32_t method1_offset = GetMethodOffset(1u);
     CHECK(!compiled_method_refs_.empty());
-    CHECK_EQ(compiled_method_refs_[0].dex_method_index, 1u);
+    CHECK_EQ(compiled_method_refs_[0].index, 1u);
     CHECK_EQ(compiled_method_refs_.size(), compiled_methods_.size());
     uint32_t method1_size = compiled_methods_[0]->GetQuickCode().size();
     uint32_t thunk_offset = CompiledCode::AlignCode(method1_offset + method1_size, kArm64);
diff --git a/compiler/linker/multi_oat_relative_patcher.h b/compiler/linker/multi_oat_relative_patcher.h
index 1c5c8a0..02cd4b0 100644
--- a/compiler/linker/multi_oat_relative_patcher.h
+++ b/compiler/linker/multi_oat_relative_patcher.h
@@ -37,8 +37,7 @@
 // to the value set by SetAdjustment().
 class MultiOatRelativePatcher FINAL {
  public:
-  using const_iterator =
-      SafeMap<MethodReference, uint32_t, MethodReferenceComparator>::const_iterator;
+  using const_iterator = SafeMap<MethodReference, uint32_t>::const_iterator;
 
   MultiOatRelativePatcher(InstructionSet instruction_set, const InstructionSetFeatures* features);
 
@@ -136,7 +135,7 @@
   class MethodOffsetMap : public linker::RelativePatcherTargetProvider {
    public:
     std::pair<bool, uint32_t> FindMethodOffset(MethodReference ref) OVERRIDE;
-    SafeMap<MethodReference, uint32_t, MethodReferenceComparator> map;
+    SafeMap<MethodReference, uint32_t> map;
   };
 
   MethodOffsetMap method_offset_map_;
diff --git a/compiler/linker/multi_oat_relative_patcher_test.cc b/compiler/linker/multi_oat_relative_patcher_test.cc
index f89fba6..5c359dc 100644
--- a/compiler/linker/multi_oat_relative_patcher_test.cc
+++ b/compiler/linker/multi_oat_relative_patcher_test.cc
@@ -26,10 +26,6 @@
 
 static const MethodReference kNullMethodRef = MethodReference(nullptr, 0u);
 
-static bool EqualRef(MethodReference lhs, MethodReference rhs) {
-  return lhs.dex_file == rhs.dex_file && lhs.dex_method_index == rhs.dex_method_index;
-}
-
 class MultiOatRelativePatcherTest : public testing::Test {
  protected:
   class MockPatcher : public RelativePatcher {
@@ -182,7 +178,7 @@
   uint32_t method1_offset = 0x100;
   uint32_t method1_offset_check = patcher_.ReserveSpace(method1_offset, method, ref1);
   ASSERT_EQ(adjustment1 + method1_offset, mock_->last_reserve_offset_);
-  ASSERT_TRUE(EqualRef(ref1, mock_->last_reserve_method_));
+  ASSERT_TRUE(ref1 == mock_->last_reserve_method_);
   ASSERT_EQ(method1_offset, method1_offset_check);
 
   uint32_t method2_offset = 0x1230;
@@ -190,13 +186,13 @@
   mock_->next_reserve_adjustment_ = method2_reserve_adjustment;
   uint32_t method2_offset_adjusted = patcher_.ReserveSpace(method2_offset, method, ref2);
   ASSERT_EQ(adjustment1 + method2_offset, mock_->last_reserve_offset_);
-  ASSERT_TRUE(EqualRef(ref2, mock_->last_reserve_method_));
+  ASSERT_TRUE(ref2 == mock_->last_reserve_method_);
   ASSERT_EQ(method2_offset + method2_reserve_adjustment, method2_offset_adjusted);
 
   uint32_t end1_offset = 0x4320;
   uint32_t end1_offset_check = patcher_.ReserveSpaceEnd(end1_offset);
   ASSERT_EQ(adjustment1 + end1_offset, mock_->last_reserve_offset_);
-  ASSERT_TRUE(EqualRef(kNullMethodRef, mock_->last_reserve_method_));
+  ASSERT_TRUE(kNullMethodRef == mock_->last_reserve_method_);
   ASSERT_EQ(end1_offset, end1_offset_check);
 
   uint32_t adjustment2 = 0xd000;
@@ -205,7 +201,7 @@
   uint32_t method3_offset = 0xf00;
   uint32_t method3_offset_check = patcher_.ReserveSpace(method3_offset, method, ref3);
   ASSERT_EQ(adjustment2 + method3_offset, mock_->last_reserve_offset_);
-  ASSERT_TRUE(EqualRef(ref3, mock_->last_reserve_method_));
+  ASSERT_TRUE(ref3 == mock_->last_reserve_method_);
   ASSERT_EQ(method3_offset, method3_offset_check);
 
   uint32_t end2_offset = 0x2400;
@@ -213,7 +209,7 @@
   mock_->next_reserve_adjustment_ = end2_reserve_adjustment;
   uint32_t end2_offset_adjusted = patcher_.ReserveSpaceEnd(end2_offset);
   ASSERT_EQ(adjustment2 + end2_offset, mock_->last_reserve_offset_);
-  ASSERT_TRUE(EqualRef(kNullMethodRef, mock_->last_reserve_method_));
+  ASSERT_TRUE(kNullMethodRef == mock_->last_reserve_method_);
   ASSERT_EQ(end2_offset + end2_reserve_adjustment, end2_offset_adjusted);
 }
 
diff --git a/compiler/linker/relative_patcher_test.h b/compiler/linker/relative_patcher_test.h
index bff6808..ca8743a 100644
--- a/compiler/linker/relative_patcher_test.h
+++ b/compiler/linker/relative_patcher_test.h
@@ -194,8 +194,7 @@
     // Sanity check: original code size must match linked_code.size().
     size_t idx = 0u;
     for (auto ref : compiled_method_refs_) {
-      if (ref.dex_file == method_ref.dex_file &&
-          ref.dex_method_index == method_ref.dex_method_index) {
+      if (ref == method_ref) {
         break;
       }
       ++idx;
@@ -264,7 +263,7 @@
         return std::pair<bool, uint32_t>(true, it->second);
       }
     }
-    SafeMap<MethodReference, uint32_t, MethodReferenceComparator> map;
+    SafeMap<MethodReference, uint32_t> map;
   };
 
   static const uint32_t kTrampolineSize = 4u;
diff --git a/compiler/oat_writer.cc b/compiler/oat_writer.cc
index 7ae3866..0bd3587 100644
--- a/compiler/oat_writer.cc
+++ b/compiler/oat_writer.cc
@@ -765,7 +765,7 @@
                           Allocator::GetMallocAllocator()));
             refs_it->second.ClearAllBits();
           }
-          refs_it->second.SetBit(target_method.dex_method_index);
+          refs_it->second.SetBit(target_method.index);
           writer_->bss_method_entries_.Overwrite(target_method, /* placeholder */ 0u);
         } else if (patch.GetType() == LinkerPatch::Type::kTypeBssEntry) {
           TypeReference ref(patch.TargetTypeDexFile(), patch.TargetTypeIndex());
@@ -922,7 +922,7 @@
         if (relative_patcher_->GetOffset(method_ref) != 0u) {
           // TODO: Should this be a hard failure?
           LOG(WARNING) << "Multiple definitions of "
-              << method_ref.dex_file->PrettyMethod(method_ref.dex_method_index)
+              << method_ref.PrettyMethod()
               << " offsets " << relative_patcher_->GetOffset(method_ref)
               << " " << quick_code_offset;
         } else {
@@ -1526,8 +1526,7 @@
     ObjPtr<mirror::DexCache> dex_cache =
         (dex_file_ == ref.dex_file) ? dex_cache_ : class_linker_->FindDexCache(
             Thread::Current(), *ref.dex_file);
-    ArtMethod* method =
-        class_linker_->LookupResolvedMethod(ref.dex_method_index, dex_cache, class_loader_);
+    ArtMethod* method = class_linker_->LookupResolvedMethod(ref.index, dex_cache, class_loader_);
     CHECK(method != nullptr);
     return method;
   }
diff --git a/compiler/optimizing/code_generator_arm64.cc b/compiler/optimizing/code_generator_arm64.cc
index 19e5d06..c61ef0a 100644
--- a/compiler/optimizing/code_generator_arm64.cc
+++ b/compiler/optimizing/code_generator_arm64.cc
@@ -4641,7 +4641,7 @@
     MethodReference target_method,
     vixl::aarch64::Label* adrp_label) {
   return NewPcRelativePatch(*target_method.dex_file,
-                            target_method.dex_method_index,
+                            target_method.index,
                             adrp_label,
                             &pc_relative_method_patches_);
 }
@@ -4650,7 +4650,7 @@
     MethodReference target_method,
     vixl::aarch64::Label* adrp_label) {
   return NewPcRelativePatch(*target_method.dex_file,
-                            target_method.dex_method_index,
+                            target_method.index,
                             adrp_label,
                             &method_bss_entry_patches_);
 }
diff --git a/compiler/optimizing/code_generator_arm_vixl.cc b/compiler/optimizing/code_generator_arm_vixl.cc
index 8b9495d..6147259 100644
--- a/compiler/optimizing/code_generator_arm_vixl.cc
+++ b/compiler/optimizing/code_generator_arm_vixl.cc
@@ -9119,14 +9119,14 @@
 CodeGeneratorARMVIXL::PcRelativePatchInfo* CodeGeneratorARMVIXL::NewPcRelativeMethodPatch(
     MethodReference target_method) {
   return NewPcRelativePatch(*target_method.dex_file,
-                            target_method.dex_method_index,
+                            target_method.index,
                             &pc_relative_method_patches_);
 }
 
 CodeGeneratorARMVIXL::PcRelativePatchInfo* CodeGeneratorARMVIXL::NewMethodBssEntryPatch(
     MethodReference target_method) {
   return NewPcRelativePatch(*target_method.dex_file,
-                            target_method.dex_method_index,
+                            target_method.index,
                             &method_bss_entry_patches_);
 }
 
diff --git a/compiler/optimizing/code_generator_mips.cc b/compiler/optimizing/code_generator_mips.cc
index f0ef007..9db2bd3 100644
--- a/compiler/optimizing/code_generator_mips.cc
+++ b/compiler/optimizing/code_generator_mips.cc
@@ -1683,7 +1683,7 @@
     MethodReference target_method,
     const PcRelativePatchInfo* info_high) {
   return NewPcRelativePatch(*target_method.dex_file,
-                            target_method.dex_method_index,
+                            target_method.index,
                             info_high,
                             &pc_relative_method_patches_);
 }
@@ -1692,7 +1692,7 @@
     MethodReference target_method,
     const PcRelativePatchInfo* info_high) {
   return NewPcRelativePatch(*target_method.dex_file,
-                            target_method.dex_method_index,
+                            target_method.index,
                             info_high,
                             &method_bss_entry_patches_);
 }
diff --git a/compiler/optimizing/code_generator_mips64.cc b/compiler/optimizing/code_generator_mips64.cc
index 201b1b0..a27cbce 100644
--- a/compiler/optimizing/code_generator_mips64.cc
+++ b/compiler/optimizing/code_generator_mips64.cc
@@ -1592,7 +1592,7 @@
     MethodReference target_method,
     const PcRelativePatchInfo* info_high) {
   return NewPcRelativePatch(*target_method.dex_file,
-                            target_method.dex_method_index,
+                            target_method.index,
                             info_high,
                             &pc_relative_method_patches_);
 }
@@ -1601,7 +1601,7 @@
     MethodReference target_method,
     const PcRelativePatchInfo* info_high) {
   return NewPcRelativePatch(*target_method.dex_file,
-                            target_method.dex_method_index,
+                            target_method.index,
                             info_high,
                             &method_bss_entry_patches_);
 }
diff --git a/compiler/optimizing/code_generator_x86.cc b/compiler/optimizing/code_generator_x86.cc
index e45ad0a..c153cf7 100644
--- a/compiler/optimizing/code_generator_x86.cc
+++ b/compiler/optimizing/code_generator_x86.cc
@@ -4623,7 +4623,7 @@
       invoke->InputAt(invoke->GetSpecialInputIndex())->AsX86ComputeBaseMethodAddress();
   boot_image_method_patches_.emplace_back(address,
                                           *invoke->GetTargetMethod().dex_file,
-                                          invoke->GetTargetMethod().dex_method_index);
+                                          invoke->GetTargetMethod().index);
   __ Bind(&boot_image_method_patches_.back().label);
 }
 
@@ -4633,7 +4633,7 @@
   // Add the patch entry and bind its label at the end of the instruction.
   method_bss_entry_patches_.emplace_back(method_address,
                                          *target_method.dex_file,
-                                         target_method.dex_method_index);
+                                         target_method.index);
   return &method_bss_entry_patches_.back().label;
 }
 
diff --git a/compiler/optimizing/code_generator_x86_64.cc b/compiler/optimizing/code_generator_x86_64.cc
index 8c4374d..bbf05a7 100644
--- a/compiler/optimizing/code_generator_x86_64.cc
+++ b/compiler/optimizing/code_generator_x86_64.cc
@@ -1068,13 +1068,13 @@
 
 void CodeGeneratorX86_64::RecordBootMethodPatch(HInvokeStaticOrDirect* invoke) {
   boot_image_method_patches_.emplace_back(*invoke->GetTargetMethod().dex_file,
-                                          invoke->GetTargetMethod().dex_method_index);
+                                          invoke->GetTargetMethod().index);
   __ Bind(&boot_image_method_patches_.back().label);
 }
 
 Label* CodeGeneratorX86_64::NewMethodBssEntryPatch(MethodReference target_method) {
   // Add a patch entry and return the label.
-  method_bss_entry_patches_.emplace_back(*target_method.dex_file, target_method.dex_method_index);
+  method_bss_entry_patches_.emplace_back(*target_method.dex_file, target_method.index);
   return &method_bss_entry_patches_.back().label;
 }
 
diff --git a/compiler/utils/atomic_dex_ref_map-inl.h b/compiler/utils/atomic_dex_ref_map-inl.h
index 14f1f0b..33d59f9 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 @@
       : 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 @@
   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 @@
   }
 }
 
-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 b02c9b6..fad056c 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 @@
 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 @@
     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 @@
 
  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 ae19a9c..8fce36f 100644
--- a/compiler/utils/atomic_dex_ref_map_test.cc
+++ b/compiler/utils/atomic_dex_ref_map_test.cc
@@ -31,40 +31,40 @@
   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);
 }
 
diff --git a/compiler/verifier_deps_test.cc b/compiler/verifier_deps_test.cc
index 5c097da..c706370 100644
--- a/compiler/verifier_deps_test.cc
+++ b/compiler/verifier_deps_test.cc
@@ -97,7 +97,7 @@
     callbacks_->SetVerifierDeps(nullptr);
     // Clear entries in the verification results to avoid hitting a DCHECK that
     // we always succeed inserting a new entry after verifying.
-    AtomicDexRefMap<const VerifiedMethod*>* map =
+    AtomicDexRefMap<MethodReference, const VerifiedMethod*>* map =
         &compiler_driver_->GetVerificationResults()->atomic_verified_methods_;
     map->Visit([](const DexFileReference& ref ATTRIBUTE_UNUSED, const VerifiedMethod* method) {
       delete method;
diff --git a/dex2oat/dex2oat_image_test.cc b/dex2oat/dex2oat_image_test.cc
index 04c2fcd..7ddf1c1 100644
--- a/dex2oat/dex2oat_image_test.cc
+++ b/dex2oat/dex2oat_image_test.cc
@@ -87,16 +87,14 @@
   void GenerateClasses(File* out_file, size_t frequency = 1) {
     VisitLibcoreDexes(VoidFunctor(),
                       [out_file](TypeReference ref) {
-      WriteLine(out_file,
-                ref.dex_file->PrettyType(ref.type_index));
+      WriteLine(out_file, ref.dex_file->PrettyType(ref.TypeIndex()));
     }, frequency, frequency);
     EXPECT_EQ(out_file->Flush(), 0);
   }
 
   void GenerateMethods(File* out_file, size_t frequency = 1) {
     VisitLibcoreDexes([out_file](MethodReference ref) {
-      WriteLine(out_file,
-                ref.dex_file->PrettyMethod(ref.dex_method_index));
+      WriteLine(out_file, ref.PrettyMethod());
     }, VoidFunctor(), frequency, frequency);
     EXPECT_EQ(out_file->Flush(), 0);
   }
@@ -315,7 +313,7 @@
     VisitLibcoreDexes([&profile](MethodReference ref) {
       EXPECT_TRUE(profile.AddMethodIndex(ProfileCompilationInfo::MethodHotness::kFlagHot, ref));
     }, [&profile](TypeReference ref) {
-      EXPECT_TRUE(profile.AddClassesForDex(ref.dex_file, &ref.type_index, &ref.type_index + 1));
+      EXPECT_TRUE(profile.AddClassForDex(ref));
     }, kMethodFrequency, kTypeFrequency);
     ScratchFile profile_file;
     profile.Save(profile_file.GetFile()->Fd());
diff --git a/oatdump/oatdump.cc b/oatdump/oatdump.cc
index 13f7211..7b11258 100644
--- a/oatdump/oatdump.cc
+++ b/oatdump/oatdump.cc
@@ -936,7 +936,7 @@
     }
 
     // Unique string ids loaded from dex code.
-    std::set<StringReference, StringReferenceComparator> unique_string_ids_from_code_;
+    std::set<StringReference> unique_string_ids_from_code_;
 
     // Total string ids loaded from dex code.
     size_t num_string_ids_from_code_ = 0;
diff --git a/profman/boot_image_profile.cc b/profman/boot_image_profile.cc
index 21de083..4092f6e 100644
--- a/profman/boot_image_profile.cc
+++ b/profman/boot_image_profile.cc
@@ -60,8 +60,7 @@
         if (hotness.IsInProfile()) {
           ++counter;
           out_profile->AddMethodHotness(ref, hotness);
-          inferred_classes.emplace(profile.get(),
-                                   dex_file->GetMethodId(ref.dex_method_index).class_idx_);
+          inferred_classes.emplace(profile.get(), ref.GetMethodId().class_idx_);
         }
       }
       // If the counter is greater or equal to the compile threshold, mark the method as hot.
@@ -110,9 +109,9 @@
       // This counter is how many profiles contain the class.
       size_t counter = 0;
       for (const std::unique_ptr<const ProfileCompilationInfo>& profile : profiles) {
-        auto it = inferred_classes.find(std::make_pair(profile.get(), ref.type_index));
+        auto it = inferred_classes.find(std::make_pair(profile.get(), ref.TypeIndex()));
         if (it != inferred_classes.end() ||
-            profile->ContainsClass(*ref.dex_file, ref.type_index)) {
+            profile->ContainsClass(*ref.dex_file, ref.TypeIndex())) {
           ++counter;
         }
       }
@@ -121,10 +120,10 @@
       }
       if (counter >= options.image_class_theshold) {
         ++class_count;
-        out_profile->AddClassesForDex(ref.dex_file, &ref.type_index, &ref.type_index + 1);
+        out_profile->AddClassForDex(ref);
       } else if (is_clean && counter >= options.image_class_clean_theshold) {
         ++clean_class_count;
-        out_profile->AddClassesForDex(ref.dex_file, &ref.type_index, &ref.type_index + 1);
+        out_profile->AddClassForDex(ref);
       }
     }
   }
diff --git a/profman/profman.cc b/profman/profman.cc
index d0c99e0..5e840e6 100644
--- a/profman/profman.cc
+++ b/profman/profman.cc
@@ -634,8 +634,7 @@
         if (kInvalidTypeIndex >= dex_file->NumTypeIds()) {
           // The dex file does not contain all possible type ids which leaves us room
           // to add an "invalid" type id.
-          class_ref->dex_file = dex_file;
-          class_ref->type_index = dex::TypeIndex(kInvalidTypeIndex);
+          *class_ref = TypeReference(dex_file, dex::TypeIndex(kInvalidTypeIndex));
           return true;
         } else {
           // The dex file contains all possible type ids. We don't have any free type id
@@ -653,8 +652,7 @@
         // Class is only referenced in the current dex file but not defined in it.
         continue;
       }
-      class_ref->dex_file = dex_file;
-      class_ref->type_index = type_index;
+      *class_ref = TypeReference(dex_file, type_index);
       return true;
     }
     return false;
@@ -698,7 +696,7 @@
       return DexFile::kDexNoIndex;
     }
     const DexFile::MethodId* method_id = dex_file->FindMethodId(
-        dex_file->GetTypeId(class_ref.type_index), *name_id, *proto_id);
+        dex_file->GetTypeId(class_ref.TypeIndex()), *name_id, *proto_id);
     if (method_id == nullptr) {
       LOG(WARNING) << "Could not find method_id: " << name;
       return DexFile::kDexNoIndex;
@@ -718,7 +716,7 @@
                        /*out*/uint32_t* dex_pc) {
     const DexFile* dex_file = class_ref.dex_file;
     uint32_t offset = dex_file->FindCodeItemOffset(
-        *dex_file->FindClassDef(class_ref.type_index),
+        *dex_file->FindClassDef(class_ref.TypeIndex()),
         method_index);
     const DexFile::CodeItem* code_item = dex_file->GetCodeItem(offset);
 
@@ -799,11 +797,11 @@
             dex_file->GetBaseLocation(),
             dex_file->GetLocationChecksum(),
             dex_file->NumMethodIds());
-      dex_resolved_classes.first->AddClass(class_ref.type_index);
+      dex_resolved_classes.first->AddClass(class_ref.TypeIndex());
       std::vector<ProfileMethodInfo> methods;
       if (method_str == kClassAllMethods) {
         // Add all of the methods.
-        const DexFile::ClassDef* class_def = dex_file->FindClassDef(class_ref.type_index);
+        const DexFile::ClassDef* class_def = dex_file->FindClassDef(class_ref.TypeIndex());
         const uint8_t* class_data = dex_file->GetClassData(*class_def);
         if (class_data != nullptr) {
           ClassDataItemIterator it(*dex_file, class_data);
diff --git a/runtime/class_reference.h b/runtime/class_reference.h
index 7b20612..2ef9ab8 100644
--- a/runtime/class_reference.h
+++ b/runtime/class_reference.h
@@ -20,22 +20,22 @@
 #include <stdint.h>
 #include <utility>
 
+#include "dex_file_reference.h"
+
 namespace art {
 
 class DexFile;
 
 // A class is uniquely located by its DexFile and the class_defs_ table index into that DexFile
-typedef std::pair<const DexFile*, uint32_t> ClassReference;
+class ClassReference : public DexFileReference {
+ public:
+  ClassReference(const DexFile* file, uint32_t class_def_idx)
+     : DexFileReference(file, class_def_idx) {}
 
-inline bool operator<(const ClassReference& lhs, const ClassReference& rhs) {
-  if (lhs.second < rhs.second) {
-    return true;
-  } else if (lhs.second > rhs.second) {
-    return false;
-  } else {
-    return (lhs.first < rhs.first);
+  uint32_t ClassDefIdx() const {
+    return index;
   }
-}
+};
 
 }  // namespace art
 
diff --git a/runtime/dex_file.h b/runtime/dex_file.h
index 5f81b98..1061fd3 100644
--- a/runtime/dex_file.h
+++ b/runtime/dex_file.h
@@ -1238,12 +1238,6 @@
   ART_FRIEND_TEST(ClassLinkerTest, RegisterDexFileName);  // for constructor
 };
 
-struct DexFileReference {
-  DexFileReference(const DexFile* file, uint32_t idx) : dex_file(file), index(idx) { }
-  const DexFile* dex_file;
-  uint32_t index;
-};
-
 std::ostream& operator<<(std::ostream& os, const DexFile& dex_file);
 
 // Iterate over a dex file's ProtoId's paramters
diff --git a/runtime/dex_file_reference.h b/runtime/dex_file_reference.h
new file mode 100644
index 0000000..01a6425
--- /dev/null
+++ b/runtime/dex_file_reference.h
@@ -0,0 +1,52 @@
+/*
+ * Copyright (C) 2017 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#ifndef ART_RUNTIME_DEX_FILE_REFERENCE_H_
+#define ART_RUNTIME_DEX_FILE_REFERENCE_H_
+
+#include <cstdint>
+
+namespace art {
+
+class DexFile;
+
+class DexFileReference {
+ public:
+  DexFileReference(const DexFile* file, uint32_t idx) : dex_file(file), index(idx) {}
+  const DexFile* dex_file;
+  uint32_t index;
+
+  struct Comparator {
+    bool operator()(const DexFileReference& a, const DexFileReference& b) const {
+      if (a.dex_file != b.dex_file) {
+        return a.dex_file < b.dex_file;
+      }
+      return a.index < b.index;
+    }
+  };
+};
+
+// Default comparators, compares the indicies, not the backing data.
+inline bool operator<(const DexFileReference& a, const DexFileReference& b) {
+  return DexFileReference::Comparator()(a, b);
+}
+inline bool operator==(const DexFileReference& a, const DexFileReference& b) {
+  return a.dex_file == b.dex_file && a.index == b.index;
+}
+
+}  // namespace art
+
+#endif  // ART_RUNTIME_DEX_FILE_REFERENCE_H_
diff --git a/runtime/entrypoints/quick/quick_trampoline_entrypoints.cc b/runtime/entrypoints/quick/quick_trampoline_entrypoints.cc
index ad65304..7d7c5a5 100644
--- a/runtime/entrypoints/quick/quick_trampoline_entrypoints.cc
+++ b/runtime/entrypoints/quick/quick_trampoline_entrypoints.cc
@@ -1156,34 +1156,33 @@
           LOG(FATAL) << "Unexpected call into trampoline: " << instr->DumpString(nullptr);
           UNREACHABLE();
       }
-      called_method.dex_method_index = (is_range) ? instr->VRegB_3rc() : instr->VRegB_35c();
+      called_method.index = (is_range) ? instr->VRegB_3rc() : instr->VRegB_35c();
       // Check that the invoke matches what we expected, note that this path only happens for debug
       // builds.
       if (found_stack_map) {
         DCHECK_EQ(stack_map_invoke_type, invoke_type);
         if (invoke_type != kSuper) {
           // Super may be sharpened.
-          DCHECK_EQ(stack_map_dex_method_idx, called_method.dex_method_index)
+          DCHECK_EQ(stack_map_dex_method_idx, called_method.index)
               << called_method.dex_file->PrettyMethod(stack_map_dex_method_idx) << " "
-              << called_method.dex_file->PrettyMethod(called_method.dex_method_index);
+              << called_method.PrettyMethod();
         }
       } else {
         VLOG(dex) << "Accessed dex file for invoke " << invoke_type << " "
-                  << called_method.dex_method_index;
+                  << called_method.index;
       }
     } else {
       invoke_type = stack_map_invoke_type;
-      called_method.dex_method_index = stack_map_dex_method_idx;
+      called_method.index = stack_map_dex_method_idx;
     }
   } else {
     invoke_type = kStatic;
     called_method.dex_file = called->GetDexFile();
-    called_method.dex_method_index = called->GetDexMethodIndex();
+    called_method.index = called->GetDexMethodIndex();
   }
   uint32_t shorty_len;
   const char* shorty =
-      called_method.dex_file->GetMethodShorty(
-          called_method.dex_file->GetMethodId(called_method.dex_method_index), &shorty_len);
+      called_method.dex_file->GetMethodShorty(called_method.GetMethodId(), &shorty_len);
   RememberForGcArgumentVisitor visitor(sp, invoke_type == kStatic, shorty, shorty_len, &soa);
   visitor.VisitArguments();
   self->EndAssertNoThreadSuspension(old_cause);
@@ -1196,7 +1195,7 @@
         hs.NewHandleWrapper(virtual_or_interface ? &receiver : &dummy));
     DCHECK_EQ(caller->GetDexFile(), called_method.dex_file);
     called = linker->ResolveMethod<ClassLinker::ResolveMode::kCheckICCEAndIAE>(
-        self, called_method.dex_method_index, caller, invoke_type);
+        self, called_method.index, caller, invoke_type);
 
     // Update .bss entry in oat file if any.
     if (called != nullptr && called_method.dex_file->GetOatDexFile() != nullptr) {
@@ -1207,10 +1206,10 @@
             mapping->begin(),
             mapping->end(),
             [called_method](const MethodBssMappingEntry& entry) {
-              return entry.method_index < called_method.dex_method_index;
+              return entry.method_index < called_method.index;
             });
-        if (pp != mapping->end() && pp->CoversIndex(called_method.dex_method_index)) {
-          size_t bss_offset = pp->GetBssOffset(called_method.dex_method_index,
+        if (pp != mapping->end() && pp->CoversIndex(called_method.index)) {
+          size_t bss_offset = pp->GetBssOffset(called_method.index,
                                                static_cast<size_t>(kRuntimePointerSize));
           DCHECK_ALIGNED(bss_offset, static_cast<size_t>(kRuntimePointerSize));
           const OatFile* oat_file = called_method.dex_file->GetOatDexFile()->GetOatFile();
@@ -1250,7 +1249,7 @@
         // TODO Maybe put this into a mirror::Class function.
         ObjPtr<mirror::Class> ref_class = linker->LookupResolvedType(
             *dex_cache->GetDexFile(),
-            dex_cache->GetDexFile()->GetMethodId(called_method.dex_method_index).class_idx_,
+            dex_cache->GetDexFile()->GetMethodId(called_method.index).class_idx_,
             dex_cache.Get(),
             class_loader.Get());
         if (ref_class->IsInterface()) {
diff --git a/runtime/jit/profile_compilation_info.cc b/runtime/jit/profile_compilation_info.cc
index 1f2163f..57fc4976 100644
--- a/runtime/jit/profile_compilation_info.cc
+++ b/runtime/jit/profile_compilation_info.cc
@@ -143,7 +143,7 @@
   if (data == nullptr) {
     return false;
   }
-  data->AddMethod(flags, ref.dex_method_index);
+  data->AddMethod(flags, ref.index);
   return true;
 }
 
@@ -683,7 +683,7 @@
   if (data == nullptr) {  // checksum mismatch
     return false;
   }
-  InlineCacheMap* inline_cache = data->FindOrAddMethod(pmi.ref.dex_method_index);
+  InlineCacheMap* inline_cache = data->FindOrAddMethod(pmi.ref.index);
 
   for (const ProfileMethodInfo::ProfileInlineCache& cache : pmi.inline_caches) {
     if (cache.is_missing_types) {
@@ -700,7 +700,7 @@
         // Don't bother adding classes if we are missing types.
         break;
       }
-      dex_pc_data->AddClass(class_dex_data->profile_index, class_ref.type_index);
+      dex_pc_data->AddClass(class_dex_data->profile_index, class_ref.TypeIndex());
     }
   }
   return true;
@@ -1333,7 +1333,7 @@
     const MethodReference& method_ref) const {
   const DexFileData* dex_data = FindDexData(method_ref.dex_file);
   return dex_data != nullptr
-      ? dex_data->GetHotnessInfo(method_ref.dex_method_index)
+      ? dex_data->GetHotnessInfo(method_ref.index)
       : MethodHotness();
 }
 
@@ -1342,8 +1342,7 @@
   DexFileData* dex_data = GetOrAddDexFileData(method_ref.dex_file);
   if (dex_data != nullptr) {
     // TODO: Add inline caches.
-    dex_data->AddMethod(static_cast<MethodHotness::Flag>(hotness.GetFlags()),
-                        method_ref.dex_method_index);
+    dex_data->AddMethod(static_cast<MethodHotness::Flag>(hotness.GetFlags()), method_ref.index);
     return true;
   }
   return false;
diff --git a/runtime/jit/profile_compilation_info.h b/runtime/jit/profile_compilation_info.h
index 7fd7a2d..5c7448f 100644
--- a/runtime/jit/profile_compilation_info.h
+++ b/runtime/jit/profile_compilation_info.h
@@ -254,6 +254,16 @@
     data->class_set.insert(index_begin, index_end);
     return true;
   }
+  // Add a single type id for a dex file.
+  bool AddClassForDex(const TypeReference& ref) {
+    DexFileData* data = GetOrAddDexFileData(ref.dex_file);
+    if (data == nullptr) {
+      return false;
+    }
+    data->class_set.insert(ref.TypeIndex());
+    return true;
+  }
+
 
   // Add a method index to the profile (without inline caches). The method flags determine if it is
   // hot, startup, or post startup, or a combination of the previous.
diff --git a/runtime/jit/profile_compilation_info_test.cc b/runtime/jit/profile_compilation_info_test.cc
index 40d303f..2cb8294 100644
--- a/runtime/jit/profile_compilation_info_test.cc
+++ b/runtime/jit/profile_compilation_info_test.cc
@@ -195,7 +195,7 @@
       for (const auto& class_ref : inline_cache.classes) {
         uint8_t dex_profile_index = dex_map.FindOrAdd(const_cast<DexFile*>(class_ref.dex_file),
                                                       static_cast<uint8_t>(dex_map.size()))->second;
-        dex_pc_data.AddClass(dex_profile_index, class_ref.type_index);
+        dex_pc_data.AddClass(dex_profile_index, class_ref.TypeIndex());
         if (dex_profile_index >= offline_pmi.dex_references.size()) {
           // This is a new dex.
           const std::string& dex_key = ProfileCompilationInfo::GetProfileDexFileKey(
diff --git a/runtime/method_reference.h b/runtime/method_reference.h
index 3948ed5..31f3b8e 100644
--- a/runtime/method_reference.h
+++ b/runtime/method_reference.h
@@ -20,27 +20,19 @@
 #include <stdint.h>
 #include <string>
 #include "dex_file.h"
+#include "dex_file_reference.h"
 
 namespace art {
 
 // A method is uniquely located by its DexFile and the method_ids_ table index into that DexFile
-struct MethodReference {
-  MethodReference(const DexFile* file, uint32_t index) : dex_file(file), dex_method_index(index) {
+class MethodReference : public DexFileReference {
+ public:
+  MethodReference(const DexFile* file, uint32_t index) : DexFileReference(file, index) {}
+  std::string PrettyMethod(bool with_signature = true) const {
+    return dex_file->PrettyMethod(index, with_signature);
   }
-  std::string PrettyMethod(bool with_signature = true) {
-    return dex_file->PrettyMethod(dex_method_index, with_signature);
-  }
-  const DexFile* dex_file;
-  uint32_t dex_method_index;
-};
-
-struct MethodReferenceComparator {
-  bool operator()(MethodReference mr1, MethodReference mr2) const {
-    if (mr1.dex_file == mr2.dex_file) {
-      return mr1.dex_method_index < mr2.dex_method_index;
-    } else {
-      return mr1.dex_file < mr2.dex_file;
-    }
+  const DexFile::MethodId& GetMethodId() const {
+    return dex_file->GetMethodId(index);
   }
 };
 
@@ -48,8 +40,8 @@
 struct MethodReferenceValueComparator {
   bool operator()(MethodReference mr1, MethodReference mr2) const {
     if (mr1.dex_file == mr2.dex_file) {
-      DCHECK_EQ(mr1.dex_method_index < mr2.dex_method_index, SlowCompare(mr1, mr2));
-      return mr1.dex_method_index < mr2.dex_method_index;
+      DCHECK_EQ(mr1.index < mr2.index, SlowCompare(mr1, mr2));
+      return mr1.index < mr2.index;
     } else {
       return SlowCompare(mr1, mr2);
     }
@@ -58,8 +50,8 @@
   bool SlowCompare(MethodReference mr1, MethodReference mr2) const {
     // The order is the same as for method ids in a single dex file.
     // Compare the class descriptors first.
-    const DexFile::MethodId& mid1 = mr1.dex_file->GetMethodId(mr1.dex_method_index);
-    const DexFile::MethodId& mid2 = mr2.dex_file->GetMethodId(mr2.dex_method_index);
+    const DexFile::MethodId& mid1 = mr1.GetMethodId();
+    const DexFile::MethodId& mid2 = mr2.GetMethodId();
     int descriptor_diff = strcmp(mr1.dex_file->StringByTypeIdx(mid1.class_idx_),
                                  mr2.dex_file->StringByTypeIdx(mid2.class_idx_));
     if (descriptor_diff != 0) {
diff --git a/runtime/string_reference.h b/runtime/string_reference.h
index 6ba4773..d0ab4e4 100644
--- a/runtime/string_reference.h
+++ b/runtime/string_reference.h
@@ -21,37 +21,30 @@
 
 #include "base/logging.h"
 #include "dex_file-inl.h"
+#include "dex_file_reference.h"
 #include "dex_file_types.h"
 #include "utf-inl.h"
 
 namespace art {
 
 // A string is located by its DexFile and the string_ids_ table index into that DexFile.
-struct StringReference {
+class StringReference : public DexFileReference {
+ public:
   StringReference(const DexFile* file, dex::StringIndex index)
-      : dex_file(file), string_index(index) { }
+     : DexFileReference(file, index.index_) {}
 
-  const char* GetStringData() const {
-    return dex_file->GetStringData(dex_file->GetStringId(string_index));
+  dex::StringIndex StringIndex() const {
+    return dex::StringIndex(index);
   }
 
-  const DexFile* dex_file;
-  dex::StringIndex string_index;
-};
-
-// Compare only the reference and not the string contents.
-struct StringReferenceComparator {
-  bool operator()(const StringReference& a, const StringReference& b) const {
-    if (a.dex_file != b.dex_file) {
-      return a.dex_file < b.dex_file;
-    }
-    return a.string_index < b.string_index;
+  const char* GetStringData() const {
+    return dex_file->GetStringData(dex_file->GetStringId(StringIndex()));
   }
 };
 
 // Compare the actual referenced string values. Used for string reference deduplication.
 struct StringReferenceValueComparator {
-  bool operator()(StringReference sr1, StringReference sr2) const {
+  bool operator()(const StringReference& sr1, const StringReference& sr2) const {
     // Note that we want to deduplicate identical strings even if they are referenced
     // by different dex files, so we need some (any) total ordering of strings, rather
     // than references. However, the references should usually be from the same dex file,
@@ -60,10 +53,10 @@
     if (sr1.dex_file == sr2.dex_file) {
       // Use the string order enforced by the dex file verifier.
       DCHECK_EQ(
-          sr1.string_index < sr2.string_index,
+          sr1.index < sr2.index,
           CompareModifiedUtf8ToModifiedUtf8AsUtf16CodePointValues(sr1.GetStringData(),
                                                                   sr2.GetStringData()) < 0);
-      return sr1.string_index < sr2.string_index;
+      return sr1.index < sr2.index;
     } else {
       // Cannot compare indexes, so do the string comparison.
       return CompareModifiedUtf8ToModifiedUtf8AsUtf16CodePointValues(sr1.GetStringData(),
diff --git a/runtime/type_reference.h b/runtime/type_reference.h
index c44019d..5ddc9d0 100644
--- a/runtime/type_reference.h
+++ b/runtime/type_reference.h
@@ -28,31 +28,23 @@
 class DexFile;
 
 // A type is located by its DexFile and the string_ids_ table index into that DexFile.
-struct TypeReference {
+class TypeReference : public DexFileReference {
+ public:
   TypeReference(const DexFile* file = nullptr, dex::TypeIndex index = dex::TypeIndex())
-      : dex_file(file),
-        type_index(index) {}
+      : DexFileReference(file, index.index_) {}
 
-  const DexFile* dex_file;
-  dex::TypeIndex type_index;
-};
-
-struct TypeReferenceComparator {
-  bool operator()(TypeReference mr1, TypeReference mr2) const {
-    if (mr1.dex_file != mr2.dex_file) {
-      return mr1.dex_file < mr2.dex_file;
-    }
-    return mr1.type_index < mr2.type_index;
+  dex::TypeIndex TypeIndex() const {
+    return dex::TypeIndex(index);
   }
 };
 
 // Compare the actual referenced type names. Used for type reference deduplication.
 struct TypeReferenceValueComparator {
-  bool operator()(TypeReference tr1, TypeReference tr2) const {
+  bool operator()(const TypeReference& tr1, const TypeReference& tr2) const {
     // Note that we want to deduplicate identical boot image types even if they are
     // referenced by different dex files, so we simply compare the descriptors.
-    StringReference sr1(tr1.dex_file, tr1.dex_file->GetTypeId(tr1.type_index).descriptor_idx_);
-    StringReference sr2(tr2.dex_file, tr2.dex_file->GetTypeId(tr2.type_index).descriptor_idx_);
+    StringReference sr1(tr1.dex_file, tr1.dex_file->GetTypeId(tr1.TypeIndex()).descriptor_idx_);
+    StringReference sr2(tr2.dex_file, tr2.dex_file->GetTypeId(tr2.TypeIndex()).descriptor_idx_);
     return StringReferenceValueComparator()(sr1, sr2);
   }
 };