Add atomic method reference map

Less RAM usage and faster than using a normal map with
MethodReference. Speed is faster by avoiding locking and tree
traversal. RAM usage is lower since the map usually had a value
for most method references.

Plan on using for marking methods for dex2dex, storing compiled
methods. Also use the new map for VerifiedMethods (refactoring).

Added test.

Bug: 32641252

Test: test-art-host-run-test

Change-Id: I46268031b8e0daf9be3597145cf6ecf579a039e2
diff --git a/build/Android.gtest.mk b/build/Android.gtest.mk
index 9902628..1691dbb 100644
--- a/build/Android.gtest.mk
+++ b/build/Android.gtest.mk
@@ -84,6 +84,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_class_linker_test_DEX_DEPS := Interfaces MethodTypes MultiDex MyClass Nested Statics StaticsFromCode
 ART_GTEST_compiler_driver_test_DEX_DEPS := AbstractMethod StaticLeafMethods ProfileTestMultiDex
 ART_GTEST_dex_cache_test_DEX_DEPS := Main Packages MethodTypes
diff --git a/compiler/Android.bp b/compiler/Android.bp
index 1737376..721986f 100644
--- a/compiler/Android.bp
+++ b/compiler/Android.bp
@@ -348,6 +348,7 @@
         "optimizing/ssa_test.cc",
         "optimizing/stack_map_test.cc",
         "optimizing/suspend_check_test.cc",
+        "utils/atomic_method_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 3fb10d8..669d8cd 100644
--- a/compiler/dex/verification_results.cc
+++ b/compiler/dex/verification_results.cc
@@ -23,6 +23,7 @@
 #include "driver/compiler_options.h"
 #include "thread.h"
 #include "thread-inl.h"
+#include "utils/atomic_method_ref_map-inl.h"
 #include "verified_method.h"
 #include "verifier/method_verifier-inl.h"
 
@@ -35,8 +36,11 @@
 
 VerificationResults::~VerificationResults() {
   WriterMutexLock mu(Thread::Current(), verified_methods_lock_);
-  DeleteResults(preregistered_dex_files_);
   STLDeleteValues(&verified_methods_);
+  atomic_verified_methods_.Visit([](const MethodReference& ref ATTRIBUTE_UNUSED,
+                                    const VerifiedMethod* method) {
+    delete method;
+  });
 }
 
 void VerificationResults::ProcessVerifiedMethod(verifier::MethodVerifier* method_verifier) {
@@ -49,16 +53,17 @@
     // We'll punt this later.
     return;
   }
-  bool inserted;
-  DexFileMethodArray* const array = GetMethodArray(ref.dex_file);
+  AtomicMap::InsertResult result = atomic_verified_methods_.Insert(ref,
+                                                                   /*expected*/ nullptr,
+                                                                   verified_method.get());
   const VerifiedMethod* existing = nullptr;
-  if (array != nullptr) {
-    DCHECK(array != nullptr);
-    Atomic<const VerifiedMethod*>* slot = &(*array)[ref.dex_method_index];
-    inserted = slot->CompareExchangeStrongSequentiallyConsistent(nullptr, verified_method.get());
+  bool inserted;
+  if (result != AtomicMap::kInsertResultInvalidDexFile) {
+    inserted = (result == AtomicMap::kInsertResultSuccess);
     if (!inserted) {
-      existing = slot->LoadSequentiallyConsistent();
-      DCHECK_NE(verified_method.get(), existing);
+      // Rare case.
+      CHECK(atomic_verified_methods_.Get(ref, &existing));
+      CHECK_NE(verified_method.get(), existing);
     }
   } else {
     WriterMutexLock mu(Thread::Current(), verified_methods_lock_);
@@ -89,9 +94,9 @@
 }
 
 const VerifiedMethod* VerificationResults::GetVerifiedMethod(MethodReference ref) {
-  DexFileMethodArray* array = GetMethodArray(ref.dex_file);
-  if (array != nullptr) {
-    return (*array)[ref.dex_method_index].LoadRelaxed();
+  const VerifiedMethod* ret = nullptr;
+  if (atomic_verified_methods_.Get(ref, &ret)) {
+    return ret;
   }
   ReaderMutexLock mu(Thread::Current(), verified_methods_lock_);
   auto it = verified_methods_.find(ref);
@@ -124,10 +129,8 @@
   return true;
 }
 
-void VerificationResults::PreRegisterDexFile(const DexFile* dex_file) {
-  CHECK(preregistered_dex_files_.find(dex_file) == preregistered_dex_files_.end())
-      << dex_file->GetLocation();
-  DexFileMethodArray array(dex_file->NumMethodIds());
+void VerificationResults::AddDexFile(const DexFile* dex_file) {
+  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
@@ -135,31 +138,13 @@
   for (auto it = verified_methods_.begin(); it != verified_methods_.end(); ) {
     MethodReference ref = it->first;
     if (ref.dex_file == dex_file) {
-      array[ref.dex_method_index].StoreSequentiallyConsistent(it->second);
+      CHECK(atomic_verified_methods_.Insert(ref, nullptr, it->second) ==
+          AtomicMap::kInsertResultSuccess);
       it = verified_methods_.erase(it);
     } else {
       ++it;
     }
   }
-  preregistered_dex_files_.emplace(dex_file, std::move(array));
-}
-
-void VerificationResults::DeleteResults(DexFileResults& array) {
-  for (auto& pair : array) {
-    for (Atomic<const VerifiedMethod*>& method : pair.second) {
-      delete method.LoadSequentiallyConsistent();
-    }
-  }
-  array.clear();
-}
-
-VerificationResults::DexFileMethodArray* VerificationResults::GetMethodArray(
-    const DexFile* dex_file) {
-  auto it = preregistered_dex_files_.find(dex_file);
-  if (it != preregistered_dex_files_.end()) {
-    return &it->second;
-  }
-  return nullptr;
 }
 
 }  // namespace art
diff --git a/compiler/dex/verification_results.h b/compiler/dex/verification_results.h
index b3356e0..ea38f4d 100644
--- a/compiler/dex/verification_results.h
+++ b/compiler/dex/verification_results.h
@@ -26,6 +26,7 @@
 #include "class_reference.h"
 #include "method_reference.h"
 #include "safe_map.h"
+#include "utils/atomic_method_ref_map.h"
 
 namespace art {
 
@@ -54,26 +55,22 @@
 
   bool IsCandidateForCompilation(MethodReference& method_ref, const uint32_t access_flags);
 
-  // Add a dex file array to the preregistered_dex_files_ array. These dex files require no locks to
-  // access. It is not safe to call if other callers are calling GetVerifiedMethod concurrently.
-  void PreRegisterDexFile(const DexFile* dex_file) REQUIRES(!verified_methods_lock_);
+  // Add a dex file to enable using the atomic map.
+  void AddDexFile(const DexFile* dex_file) REQUIRES(!verified_methods_lock_);
 
  private:
   // Verified methods. The method array is fixed to avoid needing a lock to extend it.
-  using DexFileMethodArray = dchecked_vector<Atomic<const VerifiedMethod*>>;
-  using DexFileResults = std::map<const DexFile*, DexFileMethodArray>;
+  using AtomicMap = AtomicMethodRefMap<const VerifiedMethod*>;
   using VerifiedMethodMap = SafeMap<MethodReference,
                                     const VerifiedMethod*,
                                     MethodReferenceComparator>;
 
-  static void DeleteResults(DexFileResults& array);
-
-  DexFileMethodArray* GetMethodArray(const DexFile* dex_file) REQUIRES(!verified_methods_lock_);
   VerifiedMethodMap verified_methods_ GUARDED_BY(verified_methods_lock_);
   const CompilerOptions* const compiler_options_;
 
-  // Dex2oat can preregister dex files to avoid locking when calling GetVerifiedMethod.
-  DexFileResults preregistered_dex_files_;
+  // Dex2oat can add dex files to atomic_verified_methods_ to avoid locking when calling
+  // GetVerifiedMethod.
+  AtomicMap atomic_verified_methods_;
 
   ReaderWriterMutex verified_methods_lock_ DEFAULT_MUTEX_ACQUIRED_AFTER;
 
diff --git a/compiler/utils/atomic_method_ref_map-inl.h b/compiler/utils/atomic_method_ref_map-inl.h
new file mode 100644
index 0000000..70ea028
--- /dev/null
+++ b/compiler/utils/atomic_method_ref_map-inl.h
@@ -0,0 +1,83 @@
+/*
+ * Copyright (C) 2016 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_COMPILER_UTILS_ATOMIC_METHOD_REF_MAP_INL_H_
+#define ART_COMPILER_UTILS_ATOMIC_METHOD_REF_MAP_INL_H_
+
+#include "atomic_method_ref_map.h"
+
+#include "dex_file-inl.h"
+
+namespace art {
+
+template <typename T>
+inline typename AtomicMethodRefMap<T>::InsertResult AtomicMethodRefMap<T>::Insert(
+    MethodReference 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)
+      ? kInsertResultSuccess
+      : kInsertResultCASFailure;
+}
+
+template <typename T>
+inline bool AtomicMethodRefMap<T>::Get(MethodReference ref, T* out) const {
+  const ElementArray* const array = GetArray(ref.dex_file);
+  if (array == nullptr) {
+    return kInsertResultInvalidDexFile;
+  }
+  *out = (*array)[ref.dex_method_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())));
+}
+
+template <typename T>
+inline typename AtomicMethodRefMap<T>::ElementArray* AtomicMethodRefMap<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(
+    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) {
+  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());
+    }
+  }
+}
+
+}  // namespace art
+
+#endif  // ART_COMPILER_UTILS_ATOMIC_METHOD_REF_MAP_INL_H_
diff --git a/compiler/utils/atomic_method_ref_map.h b/compiler/utils/atomic_method_ref_map.h
new file mode 100644
index 0000000..f0db231
--- /dev/null
+++ b/compiler/utils/atomic_method_ref_map.h
@@ -0,0 +1,67 @@
+/*
+ * Copyright (C) 2016 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_COMPILER_UTILS_ATOMIC_METHOD_REF_MAP_H_
+#define ART_COMPILER_UTILS_ATOMIC_METHOD_REF_MAP_H_
+
+#include "base/dchecked_vector.h"
+#include "method_reference.h"
+#include "safe_map.h"
+
+namespace art {
+
+class DexFile;
+
+// Used by CompilerCallbacks to track verification information from the Runtime.
+template <typename T>
+class AtomicMethodRefMap {
+ public:
+  explicit AtomicMethodRefMap() {}
+  ~AtomicMethodRefMap() {}
+
+  // Atomically swap the element in if the existing value matches expected.
+  enum InsertResult {
+    kInsertResultInvalidDexFile,
+    kInsertResultCASFailure,
+    kInsertResultSuccess,
+  };
+  InsertResult Insert(MethodReference 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;
+
+  // 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);
+
+  // Visit all of the dex files and elements.
+  template <typename Visitor>
+  void Visit(const Visitor& visitor);
+
+ private:
+  // Verified methods. The method array is fixed to avoid needing a lock to extend it.
+  using ElementArray = dchecked_vector<Atomic<T>>;
+  using DexFileArrays = SafeMap<const DexFile*, ElementArray>;
+
+  const ElementArray* GetArray(const DexFile* dex_file) const;
+  ElementArray* GetArray(const DexFile* dex_file);
+
+  DexFileArrays arrays_;
+};
+
+}  // namespace art
+
+#endif  // ART_COMPILER_UTILS_ATOMIC_METHOD_REF_MAP_H_
diff --git a/compiler/utils/atomic_method_ref_map_test.cc b/compiler/utils/atomic_method_ref_map_test.cc
new file mode 100644
index 0000000..c3e48ff
--- /dev/null
+++ b/compiler/utils/atomic_method_ref_map_test.cc
@@ -0,0 +1,69 @@
+/*
+ * Copyright (C) 2016 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.
+ */
+
+#include "atomic_method_ref_map-inl.h"
+
+#include <memory>
+
+#include "common_runtime_test.h"
+#include "dex_file-inl.h"
+#include "method_reference.h"
+#include "scoped_thread_state_change-inl.h"
+
+namespace art {
+
+class AtomicMethodRefMapTest : public CommonRuntimeTest {};
+
+TEST_F(AtomicMethodRefMapTest, RunTests) {
+  ScopedObjectAccess soa(Thread::Current());
+  std::unique_ptr<const DexFile> dex(OpenTestDexFile("Interfaces"));
+  ASSERT_TRUE(dex != nullptr);
+  using Map = AtomicMethodRefMap<int>;
+  Map map;
+  int value = 123;
+  // Error case: Not already inserted.
+  EXPECT_FALSE(map.Get(MethodReference(dex.get(), 1), &value));
+  // Error case: Dex file not registered.
+  EXPECT_TRUE(map.Insert(MethodReference(dex.get(), 1), 0, 1) == Map::kInsertResultInvalidDexFile);
+  map.AddDexFile(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_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) ==
+              Map::kInsertResultSuccess);
+  EXPECT_TRUE(map.Get(MethodReference(dex.get(), 1), &value));
+  EXPECT_EQ(value, kInsertValue);
+  static const int kInsertValue2 = 123;
+  EXPECT_TRUE(map.Insert(MethodReference(dex.get(), 2), 0, kInsertValue2) ==
+              Map::kInsertResultSuccess);
+  EXPECT_TRUE(map.Get(MethodReference(dex.get(), 1), &value));
+  EXPECT_EQ(value, kInsertValue);
+  EXPECT_TRUE(map.Get(MethodReference(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) ==
+      Map::kInsertResultCASFailure);
+  // Correctly overwrite the value and verify.
+  EXPECT_TRUE(map.Insert(MethodReference(dex.get(), 1), kInsertValue, kInsertValue + 1) ==
+      Map::kInsertResultSuccess);
+  EXPECT_TRUE(map.Get(MethodReference(dex.get(), 1), &value));
+  EXPECT_EQ(value, kInsertValue + 1);
+}
+
+}  // namespace art
diff --git a/dex2oat/dex2oat.cc b/dex2oat/dex2oat.cc
index 9e6032f..91a32f9 100644
--- a/dex2oat/dex2oat.cc
+++ b/dex2oat/dex2oat.cc
@@ -1636,7 +1636,7 @@
                                         soa.Decode<mirror::ClassLoader>(class_loader_).Ptr())));
       // Pre-register dex files so that we can access verification results without locks during
       // compilation and verification.
-      verification_results_->PreRegisterDexFile(dex_file);
+      verification_results_->AddDexFile(dex_file);
     }
 
     return true;