ART: Refactor known-classloader visits

Refactor visiting dex Elements and DexFiles in known classloaders,
unifying the walking code.

Test: m test-art-host
Change-Id: I4203ac4fbb0ee68660aadc0dfbf8affacbc03b8b
diff --git a/openjdkjvmti/ti_class.cc b/openjdkjvmti/ti_class.cc
index 7a94326..e289eb0 100644
--- a/openjdkjvmti/ti_class.cc
+++ b/openjdkjvmti/ti_class.cc
@@ -40,6 +40,7 @@
 #include "base/array_ref.h"
 #include "base/macros.h"
 #include "class_linker.h"
+#include "class_loader_utils.h"
 #include "class_table-inl.h"
 #include "common_throws.h"
 #include "dex/art_dex_file_loader.h"
@@ -942,23 +943,12 @@
   art::Handle<art::mirror::ClassLoader> class_loader(
       hs.NewHandle(soa.Decode<art::mirror::ClassLoader>(loader)));
   std::vector<const art::DexFile*> dex_files;
-  ClassLoaderHelper::VisitDexFileObjects(
-      self,
+  art::VisitClassLoaderDexFiles(
+      soa,
       class_loader,
-      [&] (art::ObjPtr<art::mirror::Object> dex_file) REQUIRES_SHARED(art::Locks::mutator_lock_) {
-        art::StackHandleScope<2> hs(self);
-        art::Handle<art::mirror::Object> h_dex_file(hs.NewHandle(dex_file));
-        art::Handle<art::mirror::LongArray> cookie(
-            hs.NewHandle(ClassLoaderHelper::GetDexFileCookie(h_dex_file)));
-        size_t num_elements = cookie->GetLength();
-        // We need to skip over the oat_file that's the first element. The other elements are all
-        // dex files.
-        for (size_t i = 1; i < num_elements; i++) {
-          dex_files.push_back(
-              reinterpret_cast<const art::DexFile*>(static_cast<uintptr_t>(cookie->Get(i))));
-        }
-        // Iterate over all dex files.
-        return true;
+      [&](const art::DexFile* dex_file) {
+        dex_files.push_back(dex_file);
+        return true;  // Continue with other dex files.
       });
   // We hold the loader so the dex files won't go away until after this call at worst.
   return CopyClassDescriptors(env, dex_files, count_ptr, classes);
diff --git a/runtime/class_linker.cc b/runtime/class_linker.cc
index c667fe2..ebac5c1 100644
--- a/runtime/class_linker.cc
+++ b/runtime/class_linker.cc
@@ -1122,12 +1122,8 @@
   DCHECK(out_dex_file_names != nullptr);
   DCHECK(error_msg != nullptr);
   ScopedObjectAccessUnchecked soa(Thread::Current());
-  ArtField* const dex_path_list_field =
-      jni::DecodeArtField(WellKnownClasses::dalvik_system_BaseDexClassLoader_pathList);
-  ArtField* const dex_elements_field =
-      jni::DecodeArtField(WellKnownClasses::dalvik_system_DexPathList_dexElements);
-  CHECK(dex_path_list_field != nullptr);
-  CHECK(dex_elements_field != nullptr);
+  StackHandleScope<1> hs(soa.Self());
+  Handle<mirror::ClassLoader> handle(hs.NewHandle(class_loader));
   while (!ClassLinker::IsBootClassLoader(soa, class_loader)) {
     if (soa.Decode<mirror::Class>(WellKnownClasses::dalvik_system_PathClassLoader) !=
         class_loader->GetClass()) {
@@ -1136,32 +1132,29 @@
       // Unsupported class loader.
       return false;
     }
-    ObjPtr<mirror::Object> dex_path_list = dex_path_list_field->GetObject(class_loader);
-    if (dex_path_list != nullptr) {
-      // DexPathList has an array dexElements of Elements[] which each contain a dex file.
-      ObjPtr<mirror::Object> dex_elements_obj = dex_elements_field->GetObject(dex_path_list);
-      // Loop through each dalvik.system.DexPathList$Element's dalvik.system.DexFile and look
-      // at the mCookie which is a DexFile vector.
-      if (dex_elements_obj != nullptr) {
-        ObjPtr<mirror::ObjectArray<mirror::Object>> dex_elements =
-            dex_elements_obj->AsObjectArray<mirror::Object>();
-        // Reverse order since we insert the parent at the front.
-        for (int32_t i = dex_elements->GetLength() - 1; i >= 0; --i) {
-          ObjPtr<mirror::Object> element = dex_elements->GetWithoutChecks(i);
-          if (element == nullptr) {
-            *error_msg = StringPrintf("Null dex element at index %d", i);
-            return false;
-          }
-          ObjPtr<mirror::String> name;
-          if (!GetDexPathListElementName(element, &name)) {
-            *error_msg = StringPrintf("Invalid dex path list element at index %d", i);
-            return false;
-          }
-          if (name != nullptr) {
-            out_dex_file_names->push_front(name.Ptr());
-          }
-        }
+    // Get element names. Sets error to true on failure.
+    auto add_element_names = [&](ObjPtr<mirror::Object> element, bool* error)
+        REQUIRES_SHARED(Locks::mutator_lock_) {
+      if (element == nullptr) {
+        *error_msg = "Null dex element";
+        *error = true;  // Null element is a critical error.
+        return false;   // Had an error, stop the visit.
       }
+      ObjPtr<mirror::String> name;
+      if (!GetDexPathListElementName(element, &name)) {
+        *error_msg = "Invalid dex path list element";
+        *error = false;  // Invalid element is not a critical error.
+        return false;    // Stop the visit.
+      }
+      if (name != nullptr) {
+        out_dex_file_names->push_front(name.Ptr());
+      }
+      return true;  // Continue with the next Element.
+    };
+    bool error = VisitClassLoaderDexElements(soa, handle, add_element_names, /* error */ false);
+    if (error) {
+      // An error occurred during DexPathList Element visiting.
+      return false;
     }
     class_loader = class_loader->GetParent();
   }
@@ -2444,71 +2437,33 @@
     const char* descriptor,
     size_t hash,
     Handle<mirror::ClassLoader> class_loader) {
-  CHECK(IsPathOrDexClassLoader(soa, class_loader) || IsDelegateLastClassLoader(soa, class_loader))
+  DCHECK(IsPathOrDexClassLoader(soa, class_loader) || IsDelegateLastClassLoader(soa, class_loader))
       << "Unexpected class loader for descriptor " << descriptor;
 
-  Thread* self = soa.Self();
-  ArtField* const cookie_field =
-      jni::DecodeArtField(WellKnownClasses::dalvik_system_DexFile_cookie);
-  ArtField* const dex_file_field =
-      jni::DecodeArtField(WellKnownClasses::dalvik_system_DexPathList__Element_dexFile);
-  ObjPtr<mirror::Object> dex_path_list =
-      jni::DecodeArtField(WellKnownClasses::dalvik_system_BaseDexClassLoader_pathList)->
-          GetObject(class_loader.Get());
-  if (dex_path_list != nullptr && dex_file_field != nullptr && cookie_field != nullptr) {
-    // DexPathList has an array dexElements of Elements[] which each contain a dex file.
-    ObjPtr<mirror::Object> dex_elements_obj =
-        jni::DecodeArtField(WellKnownClasses::dalvik_system_DexPathList_dexElements)->
-            GetObject(dex_path_list);
-    // Loop through each dalvik.system.DexPathList$Element's dalvik.system.DexFile and look
-    // at the mCookie which is a DexFile vector.
-    if (dex_elements_obj != nullptr) {
-      StackHandleScope<1> hs(self);
-      Handle<mirror::ObjectArray<mirror::Object>> dex_elements =
-          hs.NewHandle(dex_elements_obj->AsObjectArray<mirror::Object>());
-      for (int32_t i = 0; i < dex_elements->GetLength(); ++i) {
-        ObjPtr<mirror::Object> element = dex_elements->GetWithoutChecks(i);
-        if (element == nullptr) {
-          // Should never happen, fall back to java code to throw a NPE.
-          break;
-        }
-        ObjPtr<mirror::Object> dex_file = dex_file_field->GetObject(element);
-        if (dex_file != nullptr) {
-          ObjPtr<mirror::LongArray> long_array = cookie_field->GetObject(dex_file)->AsLongArray();
-          if (long_array == nullptr) {
-            // This should never happen so log a warning.
-            LOG(WARNING) << "Null DexFile::mCookie for " << descriptor;
-            break;
-          }
-          int32_t long_array_size = long_array->GetLength();
-          // First element is the oat file.
-          for (int32_t j = kDexFileIndexStart; j < long_array_size; ++j) {
-            const DexFile* cp_dex_file = reinterpret_cast<const DexFile*>(static_cast<uintptr_t>(
-                long_array->GetWithoutChecks(j)));
-            const DexFile::ClassDef* dex_class_def =
-                OatDexFile::FindClassDef(*cp_dex_file, descriptor, hash);
-            if (dex_class_def != nullptr) {
-              ObjPtr<mirror::Class> klass = DefineClass(self,
-                                                        descriptor,
-                                                        hash,
-                                                        class_loader,
-                                                        *cp_dex_file,
-                                                        *dex_class_def);
-              if (klass == nullptr) {
-                CHECK(self->IsExceptionPending()) << descriptor;
-                self->ClearException();
-                // TODO: Is it really right to break here, and not check the other dex files?
-                return nullptr;
-              }
-              return klass;
-            }
-          }
-        }
+  ObjPtr<mirror::Class> ret;
+  auto define_class = [&](const DexFile* cp_dex_file) REQUIRES_SHARED(Locks::mutator_lock_) {
+    const DexFile::ClassDef* dex_class_def =
+        OatDexFile::FindClassDef(*cp_dex_file, descriptor, hash);
+    if (dex_class_def != nullptr) {
+      ObjPtr<mirror::Class> klass = DefineClass(soa.Self(),
+                                                descriptor,
+                                                hash,
+                                                class_loader,
+                                                *cp_dex_file,
+                                                *dex_class_def);
+      if (klass == nullptr) {
+        CHECK(soa.Self()->IsExceptionPending()) << descriptor;
+        soa.Self()->ClearException();
+        // TODO: Is it really right to break here, and not check the other dex files?
       }
+      ret = klass;
+      return false;  // Found a Class (or error == nullptr), stop visit.
     }
-    self->AssertNoPendingException();
-  }
-  return nullptr;
+    return true;  // Continue with the next DexFile.
+  };
+
+  VisitClassLoaderDexFiles(soa, class_loader, define_class);
+  return ret;
 }
 
 mirror::Class* ClassLinker::FindClass(Thread* self,
@@ -7901,42 +7856,19 @@
     if (loader->GetClass() == path_class_loader ||
         loader->GetClass() == dex_class_loader ||
         loader->GetClass() == delegate_last_class_loader) {
-      ArtField* const cookie_field =
-          jni::DecodeArtField(WellKnownClasses::dalvik_system_DexFile_cookie);
-      ArtField* const dex_file_field =
-          jni::DecodeArtField(WellKnownClasses::dalvik_system_DexPathList__Element_dexFile);
-      ObjPtr<mirror::Object> dex_path_list =
-          jni::DecodeArtField(WellKnownClasses::dalvik_system_BaseDexClassLoader_pathList)->
-              GetObject(loader);
-      if (dex_path_list != nullptr && dex_file_field != nullptr && cookie_field != nullptr) {
-        ObjPtr<mirror::Object> dex_elements_obj =
-            jni::DecodeArtField(WellKnownClasses::dalvik_system_DexPathList_dexElements)->
-            GetObject(dex_path_list);
-        if (dex_elements_obj != nullptr) {
-          ObjPtr<mirror::ObjectArray<mirror::Object>> dex_elements =
-              dex_elements_obj->AsObjectArray<mirror::Object>();
-          oss << "(";
-          const char* path_separator = "";
-          for (int32_t i = 0; i != dex_elements->GetLength(); ++i) {
-            ObjPtr<mirror::Object> element = dex_elements->GetWithoutChecks(i);
-            ObjPtr<mirror::Object> dex_file =
-                (element != nullptr) ? dex_file_field->GetObject(element) : nullptr;
-            ObjPtr<mirror::LongArray> long_array =
-                (dex_file != nullptr) ? cookie_field->GetObject(dex_file)->AsLongArray() : nullptr;
-            if (long_array != nullptr) {
-              int32_t long_array_size = long_array->GetLength();
-              // First element is the oat file.
-              for (int32_t j = kDexFileIndexStart; j < long_array_size; ++j) {
-                const DexFile* cp_dex_file = reinterpret_cast<const DexFile*>(
-                    static_cast<uintptr_t>(long_array->GetWithoutChecks(j)));
-                oss << path_separator << cp_dex_file->GetLocation();
-                path_separator = ":";
-              }
-            }
-          }
-          oss << ")";
-        }
-      }
+      oss << "(";
+      ScopedObjectAccessUnchecked soa(Thread::Current());
+      StackHandleScope<1> hs(soa.Self());
+      Handle<mirror::ClassLoader> handle(hs.NewHandle(loader));
+      const char* path_separator = "";
+      VisitClassLoaderDexFiles(soa,
+                               handle,
+                               [&](const DexFile* dex_file) {
+                                 oss << path_separator << dex_file->GetLocation();
+                                 path_separator = ":";
+                                 return true;  // Continue with the next DexFile.
+                               });
+      oss << ")";
     }
   }
 
diff --git a/runtime/class_loader_utils.h b/runtime/class_loader_utils.h
index d160a51..1439f11 100644
--- a/runtime/class_loader_utils.h
+++ b/runtime/class_loader_utils.h
@@ -17,8 +17,12 @@
 #ifndef ART_RUNTIME_CLASS_LOADER_UTILS_H_
 #define ART_RUNTIME_CLASS_LOADER_UTILS_H_
 
+#include "art_field-inl.h"
+#include "base/mutex.h"
 #include "handle_scope.h"
+#include "jni_internal.h"
 #include "mirror/class_loader.h"
+#include "native/dalvik_system_DexFile.h"
 #include "scoped_thread_state_change-inl.h"
 #include "well_known_classes.h"
 
@@ -26,7 +30,7 @@
 
 // Returns true if the given class loader is either a PathClassLoader or a DexClassLoader.
 // (they both have the same behaviour with respect to class lockup order)
-static bool IsPathOrDexClassLoader(ScopedObjectAccessAlreadyRunnable& soa,
+inline bool IsPathOrDexClassLoader(ScopedObjectAccessAlreadyRunnable& soa,
                                    Handle<mirror::ClassLoader> class_loader)
     REQUIRES_SHARED(Locks::mutator_lock_) {
   mirror::Class* class_loader_class = class_loader->GetClass();
@@ -37,7 +41,7 @@
           soa.Decode<mirror::Class>(WellKnownClasses::dalvik_system_DexClassLoader));
 }
 
-static bool IsDelegateLastClassLoader(ScopedObjectAccessAlreadyRunnable& soa,
+inline bool IsDelegateLastClassLoader(ScopedObjectAccessAlreadyRunnable& soa,
                                       Handle<mirror::ClassLoader> class_loader)
     REQUIRES_SHARED(Locks::mutator_lock_) {
   mirror::Class* class_loader_class = class_loader->GetClass();
@@ -45,6 +49,114 @@
       soa.Decode<mirror::Class>(WellKnownClasses::dalvik_system_DelegateLastClassLoader);
 }
 
+// Visit the DexPathList$Element instances in the given classloader with the given visitor.
+// Constraints on the visitor:
+//   * The visitor should return true to continue visiting more Elements.
+//   * The last argument of the visitor is an out argument of RetType. It will be returned
+//     when the visitor ends the visit (by returning false).
+// This function assumes that the given classloader is a subclass of BaseDexClassLoader!
+template <typename Visitor, typename RetType>
+inline RetType VisitClassLoaderDexElements(ScopedObjectAccessAlreadyRunnable& soa,
+                                           Handle<mirror::ClassLoader> class_loader,
+                                           Visitor fn,
+                                           RetType defaultReturn)
+    REQUIRES_SHARED(Locks::mutator_lock_) {
+  Thread* self = soa.Self();
+  ObjPtr<mirror::Object> dex_path_list =
+      jni::DecodeArtField(WellKnownClasses::dalvik_system_BaseDexClassLoader_pathList)->
+          GetObject(class_loader.Get());
+  if (dex_path_list != nullptr) {
+    // DexPathList has an array dexElements of Elements[] which each contain a dex file.
+    ObjPtr<mirror::Object> dex_elements_obj =
+        jni::DecodeArtField(WellKnownClasses::dalvik_system_DexPathList_dexElements)->
+            GetObject(dex_path_list);
+    // Loop through each dalvik.system.DexPathList$Element's dalvik.system.DexFile and look
+    // at the mCookie which is a DexFile vector.
+    if (dex_elements_obj != nullptr) {
+      StackHandleScope<1> hs(self);
+      Handle<mirror::ObjectArray<mirror::Object>> dex_elements =
+          hs.NewHandle(dex_elements_obj->AsObjectArray<mirror::Object>());
+      for (int32_t i = 0; i < dex_elements->GetLength(); ++i) {
+        ObjPtr<mirror::Object> element = dex_elements->GetWithoutChecks(i);
+        if (element == nullptr) {
+          // Should never happen, fail.
+          break;
+        }
+        RetType ret_value;
+        if (!fn(element, &ret_value)) {
+          return ret_value;
+        }
+      }
+    }
+    self->AssertNoPendingException();
+  }
+  return defaultReturn;
+}
+
+// Visit the DexFiles in the given classloader with the given visitor.
+// Constraints on the visitor:
+//   * The visitor should return true to continue visiting more DexFiles.
+//   * The last argument of the visitor is an out argument of RetType. It will be returned
+//     when the visitor ends the visit (by returning false).
+// This function assumes that the given classloader is a subclass of BaseDexClassLoader!
+template <typename Visitor, typename RetType>
+inline RetType VisitClassLoaderDexFiles(ScopedObjectAccessAlreadyRunnable& soa,
+                                        Handle<mirror::ClassLoader> class_loader,
+                                        Visitor fn,
+                                        RetType defaultReturn)
+    REQUIRES_SHARED(Locks::mutator_lock_) {
+  ArtField* const cookie_field =
+      jni::DecodeArtField(WellKnownClasses::dalvik_system_DexFile_cookie);
+  ArtField* const dex_file_field =
+      jni::DecodeArtField(WellKnownClasses::dalvik_system_DexPathList__Element_dexFile);
+  if (dex_file_field == nullptr || cookie_field == nullptr) {
+    return defaultReturn;
+  }
+  auto visit_dex_files = [&](ObjPtr<mirror::Object> element, RetType* ret)
+      REQUIRES_SHARED(Locks::mutator_lock_) {
+    ObjPtr<mirror::Object> dex_file = dex_file_field->GetObject(element);
+    if (dex_file != nullptr) {
+      ObjPtr<mirror::LongArray> long_array = cookie_field->GetObject(dex_file)->AsLongArray();
+      if (long_array == nullptr) {
+        // This should never happen so log a warning.
+        LOG(WARNING) << "Null DexFile::mCookie";
+        *ret = defaultReturn;
+        return true;
+      }
+      int32_t long_array_size = long_array->GetLength();
+      // First element is the oat file.
+      for (int32_t j = kDexFileIndexStart; j < long_array_size; ++j) {
+        const DexFile* cp_dex_file = reinterpret_cast<const DexFile*>(static_cast<uintptr_t>(
+            long_array->GetWithoutChecks(j)));
+        RetType ret_value;
+        if (!fn(cp_dex_file, /* out */ &ret_value)) {
+          *ret = ret_value;
+          return false;
+        }
+      }
+    }
+    return true;
+  };
+
+  return VisitClassLoaderDexElements(soa, class_loader, visit_dex_files, defaultReturn);
+}
+
+// Simplified version of the above, w/o out argument.
+template <typename Visitor>
+inline void VisitClassLoaderDexFiles(ScopedObjectAccessAlreadyRunnable& soa,
+                                     Handle<mirror::ClassLoader> class_loader,
+                                     Visitor fn)
+    REQUIRES_SHARED(Locks::mutator_lock_) {
+  auto helper = [&fn](const art::DexFile* dex_file, void** ATTRIBUTE_UNUSED)
+      REQUIRES_SHARED(Locks::mutator_lock_) {
+    return fn(dex_file);
+  };
+  VisitClassLoaderDexFiles<decltype(helper), void*>(soa,
+                                                    class_loader,
+                                                    helper,
+                                                    /* default */ nullptr);
+}
+
 }  // namespace art
 
 #endif  // ART_RUNTIME_CLASS_LOADER_UTILS_H_
diff --git a/runtime/common_runtime_test.cc b/runtime/common_runtime_test.cc
index e4fbc86..e7a1374 100644
--- a/runtime/common_runtime_test.cc
+++ b/runtime/common_runtime_test.cc
@@ -34,6 +34,7 @@
 #include "base/stl_util.h"
 #include "base/unix_file/fd_file.h"
 #include "class_linker.h"
+#include "class_loader_utils.h"
 #include "compiler_callbacks.h"
 #include "dex/art_dex_file_loader.h"
 #include "dex/dex_file-inl.h"
@@ -541,58 +542,23 @@
 std::vector<const DexFile*> CommonRuntimeTestImpl::GetDexFiles(
     ScopedObjectAccess& soa,
     Handle<mirror::ClassLoader> class_loader) {
-  std::vector<const DexFile*> ret;
-
   DCHECK(
       (class_loader->GetClass() ==
           soa.Decode<mirror::Class>(WellKnownClasses::dalvik_system_PathClassLoader)) ||
       (class_loader->GetClass() ==
           soa.Decode<mirror::Class>(WellKnownClasses::dalvik_system_DelegateLastClassLoader)));
 
-  // The class loader is a PathClassLoader which inherits from BaseDexClassLoader.
-  // We need to get the DexPathList and loop through it.
-  ArtField* cookie_field = jni::DecodeArtField(WellKnownClasses::dalvik_system_DexFile_cookie);
-  ArtField* dex_file_field =
-      jni::DecodeArtField(WellKnownClasses::dalvik_system_DexPathList__Element_dexFile);
-  ObjPtr<mirror::Object> dex_path_list =
-      jni::DecodeArtField(WellKnownClasses::dalvik_system_BaseDexClassLoader_pathList)->
-          GetObject(class_loader.Get());
-  if (dex_path_list != nullptr && dex_file_field!= nullptr && cookie_field != nullptr) {
-    // DexPathList has an array dexElements of Elements[] which each contain a dex file.
-    ObjPtr<mirror::Object> dex_elements_obj =
-        jni::DecodeArtField(WellKnownClasses::dalvik_system_DexPathList_dexElements)->
-            GetObject(dex_path_list);
-    // Loop through each dalvik.system.DexPathList$Element's dalvik.system.DexFile and look
-    // at the mCookie which is a DexFile vector.
-    if (dex_elements_obj != nullptr) {
-      StackHandleScope<1> hs(soa.Self());
-      Handle<mirror::ObjectArray<mirror::Object>> dex_elements =
-          hs.NewHandle(dex_elements_obj->AsObjectArray<mirror::Object>());
-      for (int32_t i = 0; i < dex_elements->GetLength(); ++i) {
-        ObjPtr<mirror::Object> element = dex_elements->GetWithoutChecks(i);
-        if (element == nullptr) {
-          // Should never happen, fall back to java code to throw a NPE.
-          break;
-        }
-        ObjPtr<mirror::Object> dex_file = dex_file_field->GetObject(element);
-        if (dex_file != nullptr) {
-          ObjPtr<mirror::LongArray> long_array = cookie_field->GetObject(dex_file)->AsLongArray();
-          DCHECK(long_array != nullptr);
-          int32_t long_array_size = long_array->GetLength();
-          for (int32_t j = kDexFileIndexStart; j < long_array_size; ++j) {
-            const DexFile* cp_dex_file = reinterpret_cast<const DexFile*>(static_cast<uintptr_t>(
-                long_array->GetWithoutChecks(j)));
-            if (cp_dex_file == nullptr) {
-              LOG(WARNING) << "Null DexFile";
-              continue;
-            }
-            ret.push_back(cp_dex_file);
-          }
-        }
-      }
-    }
-  }
-
+  std::vector<const DexFile*> ret;
+  VisitClassLoaderDexFiles(soa,
+                           class_loader,
+                           [&](const DexFile* cp_dex_file) {
+                             if (cp_dex_file == nullptr) {
+                               LOG(WARNING) << "Null DexFile";
+                             } else {
+                               ret.push_back(cp_dex_file);
+                             }
+                             return true;
+                           });
   return ret;
 }