Faster annotations processing in `ClassLinker::LoadClass()`.

Iterate annotations alongside other method data, using the
fact that both are ordered by method index. We use separate
iterators for loading direct and virtual methods.

For a test class with 500 unannotated static methods and
200 @NeverCompile static methods, this change reduces the
time needed by `ClassLinker::LoadClass` approximately 2x.
For a test class with twice that many methods, the time
without this change increases by ~3x but with this change
the time increases by ~2x, showing nice linear scaling.

Test: m test-art-host-gtest
Test: testrunner.py --host --optimizing
Bug: 329196666
Change-Id: I800b8e8e9f9e5a4fbfbda1361f99feb8c43ff1d5

5 files changed, 95 insertions, 27 deletions

diff --git a/libdexfile/dex/dex_file.h b/libdexfile/dex/dex_file.h
index 8e399af812..4186a80b43 100644
--- a/libdexfile/dex/dex_file.h
+++ b/libdexfile/dex/dex_file.h
@@ -704,12 +704,14 @@ class DexFile {
 
   const dex::AnnotationSetItem* GetFieldAnnotationSetItem(
       const dex::FieldAnnotationsItem& anno_item) const {
-    return DataPointer<dex::AnnotationSetItem>(anno_item.annotations_off_);
+    // `DexFileVerifier` checks that the offset is not zero.
+    return NonNullDataPointer<dex::AnnotationSetItem>(anno_item.annotations_off_);
   }
 
   const dex::AnnotationSetItem* GetMethodAnnotationSetItem(
       const dex::MethodAnnotationsItem& anno_item) const {
-    return DataPointer<dex::AnnotationSetItem>(anno_item.annotations_off_);
+    // `DexFileVerifier` checks that the offset is not zero.
+    return NonNullDataPointer<dex::AnnotationSetItem>(anno_item.annotations_off_);
   }
 
   const dex::AnnotationSetRefList* GetParameterAnnotationSetRefList(
@@ -821,8 +823,14 @@ class DexFile {
 
   template <typename T>
   const T* DataPointer(size_t offset) const {
+    return (offset != 0u) ? NonNullDataPointer<T>(offset) : nullptr;
+  }
+
+  template <typename T>
+  const T* NonNullDataPointer(size_t offset) const {
+    DCHECK_NE(offset, 0u);
     DCHECK_LT(offset, DataSize()) << "Offset past end of data section";
-    return (offset != 0u) ? reinterpret_cast<const T*>(DataBegin() + offset) : nullptr;
+    return reinterpret_cast<const T*>(DataBegin() + offset);
   }
 
   const OatDexFile* GetOatDexFile() const {
diff --git a/runtime/class_linker.cc b/runtime/class_linker.cc
index 9d405207a5..9f463ada75 100644
--- a/runtime/class_linker.cc
+++ b/runtime/class_linker.cc
@@ -3868,6 +3868,29 @@ LinearAlloc* ClassLinker::GetOrCreateAllocatorForClassLoader(ObjPtr<mirror::Clas
   return allocator;
 }
 
+// Helper class for iterating over method annotations, using their ordering in the dex file.
+// Since direct and virtual methods are separated (but each section is ordered), we shall use
+// separate iterators for loading direct and virtual methods.
+class ClassLinker::MethodAnnotationsIterator {
+ public:
+  MethodAnnotationsIterator(const DexFile& dex_file,
+                            const dex::AnnotationsDirectoryItem* annotations_dir)
+      : current_((annotations_dir != nullptr) ? dex_file.GetMethodAnnotations(annotations_dir)
+                                              : nullptr),
+        end_((annotations_dir != nullptr) ? current_ + annotations_dir->methods_size_ : nullptr) {}
+
+  const dex::MethodAnnotationsItem* AdvanceTo(uint32_t method_idx) {
+    while (current_ != end_ && current_->method_idx_ < method_idx) {
+      ++current_;
+    }
+    return (current_ != end_ && current_->method_idx_ == method_idx) ? current_ : nullptr;
+  }
+
+ private:
+  const dex::MethodAnnotationsItem* current_;
+  const dex::MethodAnnotationsItem* const end_;
+};
+
 void ClassLinker::LoadClass(Thread* self,
                             const DexFile& dex_file,
                             const dex::ClassDef& dex_class_def,
@@ -3912,6 +3935,10 @@ void ClassLinker::LoadClass(Thread* self,
     uint32_t last_dex_method_index = dex::kDexNoIndex;
     size_t last_class_def_method_index = 0;
 
+    // Initialize separate `MethodAnnotationsIterator`s for direct and virtual methods.
+    MethodAnnotationsIterator mai_direct(dex_file, dex_file.GetAnnotationsDirectory(dex_class_def));
+    MethodAnnotationsIterator mai_virtual = mai_direct;
+
     uint16_t hotness_threshold = runtime->GetJITOptions()->GetWarmupThreshold();
     // Use the visitor since the ranged based loops are bit slower from seeking. Seeking to the
     // methods needs to decode all of the fields.
@@ -3935,7 +3962,7 @@ void ClassLinker::LoadClass(Thread* self,
         }, [&](const ClassAccessor::Method& method) REQUIRES_SHARED(Locks::mutator_lock_) {
           ArtMethod* art_method = klass->GetDirectMethodUnchecked(class_def_method_index,
               image_pointer_size_);
-          LoadMethod(dex_file, method, klass.Get(), art_method);
+          LoadMethod(dex_file, method, klass.Get(), &mai_direct, art_method);
           LinkCode(this, art_method, oat_class_ptr, class_def_method_index);
           uint32_t it_method_index = method.GetIndex();
           if (last_dex_method_index == it_method_index) {
@@ -3953,7 +3980,7 @@ void ClassLinker::LoadClass(Thread* self,
               class_def_method_index - accessor.NumDirectMethods(),
               image_pointer_size_);
           art_method->ResetCounter(hotness_threshold);
-          LoadMethod(dex_file, method, klass.Get(), art_method);
+          LoadMethod(dex_file, method, klass.Get(), &mai_virtual, art_method);
           LinkCode(this, art_method, oat_class_ptr, class_def_method_index);
           ++class_def_method_index;
         });
@@ -3996,7 +4023,8 @@ void ClassLinker::LoadField(const ClassAccessor::Field& field,
 void ClassLinker::LoadMethod(const DexFile& dex_file,
                              const ClassAccessor::Method& method,
                              ObjPtr<mirror::Class> klass,
-                             ArtMethod* dst) {
+                             /*inout*/ MethodAnnotationsIterator* mai,
+                             /*out*/ ArtMethod* dst) {
   ScopedAssertNoThreadSuspension sants(__FUNCTION__);
 
   const uint32_t dex_method_idx = method.GetIndex();
@@ -4045,9 +4073,11 @@ void ClassLinker::LoadMethod(const DexFile& dex_file,
 
   if (UNLIKELY((access_flags & kAccNative) != 0u)) {
     // Check if the native method is annotated with @FastNative or @CriticalNative.
-    const dex::ClassDef& class_def = dex_file.GetClassDef(klass->GetDexClassDefIndex());
-    access_flags |=
-        annotations::GetNativeMethodAnnotationAccessFlags(dex_file, class_def, dex_method_idx);
+    const dex::MethodAnnotationsItem* method_annotations = mai->AdvanceTo(dex_method_idx);
+    if (method_annotations != nullptr) {
+      access_flags |=
+          annotations::GetNativeMethodAnnotationAccessFlags(dex_file, *method_annotations);
+    }
     dst->SetAccessFlags(access_flags);
     DCHECK(!dst->IsAbstract());
     DCHECK(!dst->HasCodeItem());
@@ -4064,8 +4094,9 @@ void ClassLinker::LoadMethod(const DexFile& dex_file,
     DCHECK_EQ(method.GetCodeItemOffset(), 0u);
     dst->SetDataPtrSize(nullptr, image_pointer_size_);  // Single implementation not set yet.
   } else {
-    const dex::ClassDef& class_def = dex_file.GetClassDef(klass->GetDexClassDefIndex());
-    if (annotations::MethodIsNeverCompile(dex_file, class_def, dex_method_idx)) {
+    const dex::MethodAnnotationsItem* method_annotations = mai->AdvanceTo(dex_method_idx);
+    if (method_annotations != nullptr &&
+        annotations::MethodIsNeverCompile(dex_file, *method_annotations)) {
       access_flags |= kAccCompileDontBother;
     }
     dst->SetAccessFlags(access_flags);
diff --git a/runtime/class_linker.h b/runtime/class_linker.h
index f964a5573c..edb866adc7 100644
--- a/runtime/class_linker.h
+++ b/runtime/class_linker.h
@@ -924,6 +924,7 @@ class EXPORT ClassLinker {
   class LinkFieldsHelper;
   template <PointerSize kPointerSize>
   class LinkMethodsHelper;
+  class MethodAnnotationsIterator;
   class VisiblyInitializedCallback;
 
   struct ClassLoaderData {
@@ -1043,7 +1044,8 @@ class EXPORT ClassLinker {
   void LoadMethod(const DexFile& dex_file,
                   const ClassAccessor::Method& method,
                   ObjPtr<mirror::Class> klass,
-                  ArtMethod* dst)
+                  /*inout*/ MethodAnnotationsIterator* mai,
+                  /*out*/ ArtMethod* dst)
       REQUIRES_SHARED(Locks::mutator_lock_);
 
   void FixupStaticTrampolines(Thread* self, ObjPtr<mirror::Class> klass)
diff --git a/runtime/dex/dex_file_annotations.cc b/runtime/dex/dex_file_annotations.cc
index ab2384ee95..8488993ae4 100644
--- a/runtime/dex/dex_file_annotations.cc
+++ b/runtime/dex/dex_file_annotations.cc
@@ -1295,25 +1295,19 @@ static bool IsMethodBuildAnnotationPresent(const DexFile& dex_file,
   return false;
 }
 
-uint32_t GetNativeMethodAnnotationAccessFlags(const DexFile& dex_file,
-                                              const dex::ClassDef& class_def,
-                                              uint32_t method_index) {
-  const dex::AnnotationSetItem* annotation_set =
-      FindAnnotationSetForMethod(dex_file, class_def, method_index);
-  if (annotation_set == nullptr) {
-    return 0u;
-  }
+static uint32_t GetNativeMethodAnnotationAccessFlags(const DexFile& dex_file,
+                                                     const dex::AnnotationSetItem& annotation_set) {
   uint32_t access_flags = 0u;
   if (IsMethodBuildAnnotationPresent(
           dex_file,
-          *annotation_set,
+          annotation_set,
           "Ldalvik/annotation/optimization/FastNative;",
           WellKnownClasses::dalvik_annotation_optimization_FastNative)) {
     access_flags |= kAccFastNative;
   }
   if (IsMethodBuildAnnotationPresent(
           dex_file,
-          *annotation_set,
+          annotation_set,
           "Ldalvik/annotation/optimization/CriticalNative;",
           WellKnownClasses::dalvik_annotation_optimization_CriticalNative)) {
     access_flags |= kAccCriticalNative;
@@ -1322,21 +1316,48 @@ uint32_t GetNativeMethodAnnotationAccessFlags(const DexFile& dex_file,
   return access_flags;
 }
 
-bool MethodIsNeverCompile(const DexFile& dex_file,
-                          const dex::ClassDef& class_def,
-                          uint32_t method_index) {
+uint32_t GetNativeMethodAnnotationAccessFlags(const DexFile& dex_file,
+                                              const dex::ClassDef& class_def,
+                                              uint32_t method_index) {
   const dex::AnnotationSetItem* annotation_set =
       FindAnnotationSetForMethod(dex_file, class_def, method_index);
   if (annotation_set == nullptr) {
-    return false;
+    return 0u;
   }
+  return GetNativeMethodAnnotationAccessFlags(dex_file, *annotation_set);
+}
+
+uint32_t GetNativeMethodAnnotationAccessFlags(const DexFile& dex_file,
+                                              const dex::MethodAnnotationsItem& method_annotations) {
+  return GetNativeMethodAnnotationAccessFlags(
+      dex_file, *dex_file.GetMethodAnnotationSetItem(method_annotations));
+}
+
+static bool MethodIsNeverCompile(const DexFile& dex_file,
+                                 const dex::AnnotationSetItem& annotation_set) {
   return IsMethodBuildAnnotationPresent(
       dex_file,
-      *annotation_set,
+      annotation_set,
       "Ldalvik/annotation/optimization/NeverCompile;",
       WellKnownClasses::dalvik_annotation_optimization_NeverCompile);
 }
 
+bool MethodIsNeverCompile(const DexFile& dex_file,
+                          const dex::ClassDef& class_def,
+                          uint32_t method_index) {
+  const dex::AnnotationSetItem* annotation_set =
+      FindAnnotationSetForMethod(dex_file, class_def, method_index);
+  if (annotation_set == nullptr) {
+    return false;
+  }
+  return MethodIsNeverCompile(dex_file, *annotation_set);
+}
+
+bool MethodIsNeverCompile(const DexFile& dex_file,
+                          const dex::MethodAnnotationsItem& method_annotations) {
+  return MethodIsNeverCompile(dex_file, *dex_file.GetMethodAnnotationSetItem(method_annotations));
+}
+
 bool MethodIsNeverInline(const DexFile& dex_file,
                          const dex::ClassDef& class_def,
                          uint32_t method_index) {
diff --git a/runtime/dex/dex_file_annotations.h b/runtime/dex/dex_file_annotations.h
index 3a79e11256..6f6c6ddc1f 100644
--- a/runtime/dex/dex_file_annotations.h
+++ b/runtime/dex/dex_file_annotations.h
@@ -86,11 +86,17 @@ bool IsMethodAnnotationPresent(ArtMethod* method,
 EXPORT uint32_t GetNativeMethodAnnotationAccessFlags(const DexFile& dex_file,
                                                      const dex::ClassDef& class_def,
                                                      uint32_t method_index);
+// An overload of `GetNativeMethodAnnotationAccessFlags()` that takes a `MethodAnnotationsItem`.
+uint32_t GetNativeMethodAnnotationAccessFlags(const DexFile& dex_file,
+                                              const dex::MethodAnnotationsItem& method_annotations);
 // Is the method from the `dex_file` with the given `field_index`
 // annotated with @dalvik.annotation.optimization.NeverCompile?
 EXPORT bool MethodIsNeverCompile(const DexFile& dex_file,
                                  const dex::ClassDef& class_def,
                                  uint32_t method_index);
+// An overload of `MethodIsNeverCompile()` that takes a `MethodAnnotationsItem`.
+bool MethodIsNeverCompile(const DexFile& dex_file,
+                          const dex::MethodAnnotationsItem& method_annotations);
 // Is the method from the `dex_file` with the given `field_index`
 // annotated with @dalvik.annotation.optimization.NeverInline?
 bool MethodIsNeverInline(const DexFile& dex_file,