Postpone DexCache preloading in the ImageWriter.

Preload the DexCache with deterministic contents just before
we're about to copy it to the image memory.

Previously the contents of the preloaded dex cache arrays
were further modified by the OatWriter. As long as OatWriter
was processing the same data on a single thread, this was
deterministic but it was not working as originally intended.

Also fix OatWriter::InitImageMethodVisitor::StartClass() to
use the deterministic ClassLinker::LookupResolvedType()
rather than checking only the DexCache.

Clean up OatWriter::InitImageMethodVisitor::VisitMethod(),
call LookupResolvedMethod() instead of ResolveMethod() for
boot image methods as they should be already resolved.

Test: m test-art-host-gtest
Test: testrunner.py --host --optimizing
Test: Manually verify with oatdump that dex cache array
      contents are the same when compiling on host and
      on the device.
Bug: 119436742
Change-Id: Ifa324beabac2482621529655f72b4f734bd62030
diff --git a/dex2oat/linker/image_writer.cc b/dex2oat/linker/image_writer.cc
index 2b2b029..e4fbd17 100644
--- a/dex2oat/linker/image_writer.cc
+++ b/dex2oat/linker/image_writer.cc
@@ -613,8 +613,22 @@
   CHECK(!oat_filenames.empty());
   CHECK_EQ(image_filenames.size(), oat_filenames.size());
 
+  Thread* self = Thread::Current();
   {
-    ScopedObjectAccess soa(Thread::Current());
+    // Preload deterministic contents to the dex cache arrays we're going to write.
+    ScopedObjectAccess soa(self);
+    ObjPtr<mirror::ClassLoader> class_loader = GetClassLoader();
+    std::vector<ObjPtr<mirror::DexCache>> dex_caches = FindDexCaches(self);
+    for (ObjPtr<mirror::DexCache> dex_cache : dex_caches) {
+      if (compile_app_image_ && IsInBootImage(dex_cache.Ptr())) {
+        continue;  // Boot image DexCache is not written to the app image.
+      }
+      PreloadDexCache(dex_cache, class_loader);
+    }
+  }
+
+  {
+    ScopedObjectAccess soa(self);
     for (size_t i = 0; i < oat_filenames.size(); ++i) {
       CreateHeader(i);
       CopyAndFixupNativeData(i);
@@ -623,7 +637,7 @@
 
   {
     // TODO: heap validation can't handle these fix up passes.
-    ScopedObjectAccess soa(Thread::Current());
+    ScopedObjectAccess soa(self);
     Runtime::Current()->GetHeap()->DisableObjectValidation();
     CopyAndFixupObjects();
   }
@@ -1409,8 +1423,79 @@
   Runtime::Current()->GetClassLinker()->VisitClassLoaders(visitor);
 }
 
-void ImageWriter::PruneAndPreloadDexCache(ObjPtr<mirror::DexCache> dex_cache,
-                                          ObjPtr<mirror::ClassLoader> class_loader) {
+void ImageWriter::PruneDexCache(ObjPtr<mirror::DexCache> dex_cache,
+                                ObjPtr<mirror::ClassLoader> class_loader) {
+  Runtime* runtime = Runtime::Current();
+  ClassLinker* class_linker = runtime->GetClassLinker();
+  const DexFile& dex_file = *dex_cache->GetDexFile();
+  // Prune methods.
+  dex::TypeIndex last_class_idx;  // Initialized to invalid index.
+  ObjPtr<mirror::Class> last_class = nullptr;
+  mirror::MethodDexCacheType* resolved_methods = dex_cache->GetResolvedMethods();
+  for (size_t slot_idx = 0, num = dex_cache->NumResolvedMethods(); slot_idx != num; ++slot_idx) {
+    auto pair =
+        mirror::DexCache::GetNativePairPtrSize(resolved_methods, slot_idx, target_ptr_size_);
+    uint32_t stored_index = pair.index;
+    ArtMethod* method = pair.object;
+    if (method == nullptr) {
+      continue;  // Empty entry.
+    }
+    // Check if the referenced class is in the image. Note that we want to check the referenced
+    // class rather than the declaring class to preserve the semantics, i.e. using a MethodId
+    // results in resolving the referenced class and that can for example throw OOME.
+    const DexFile::MethodId& method_id = dex_file.GetMethodId(stored_index);
+    if (method_id.class_idx_ != last_class_idx) {
+      last_class_idx = method_id.class_idx_;
+      last_class = class_linker->LookupResolvedType(last_class_idx, dex_cache, class_loader);
+      if (last_class != nullptr && !KeepClass(last_class)) {
+        last_class = nullptr;
+      }
+    }
+    if (last_class == nullptr) {
+      dex_cache->ClearResolvedMethod(stored_index, target_ptr_size_);
+    }
+  }
+  // Prune fields.
+  mirror::FieldDexCacheType* resolved_fields = dex_cache->GetResolvedFields();
+  last_class_idx = dex::TypeIndex();  // Initialized to invalid index.
+  last_class = nullptr;
+  for (size_t slot_idx = 0, num = dex_cache->NumResolvedFields(); slot_idx != num; ++slot_idx) {
+    auto pair = mirror::DexCache::GetNativePairPtrSize(resolved_fields, slot_idx, target_ptr_size_);
+    uint32_t stored_index = pair.index;
+    ArtField* field = pair.object;
+    if (field == nullptr) {
+      continue;  // Empty entry.
+    }
+    // Check if the referenced class is in the image. Note that we want to check the referenced
+    // class rather than the declaring class to preserve the semantics, i.e. using a FieldId
+    // results in resolving the referenced class and that can for example throw OOME.
+    const DexFile::FieldId& field_id = dex_file.GetFieldId(stored_index);
+    if (field_id.class_idx_ != last_class_idx) {
+      last_class_idx = field_id.class_idx_;
+      last_class = class_linker->LookupResolvedType(last_class_idx, dex_cache, class_loader);
+      if (last_class != nullptr && !KeepClass(last_class)) {
+        last_class = nullptr;
+      }
+    }
+    if (last_class == nullptr) {
+      dex_cache->ClearResolvedField(stored_index, target_ptr_size_);
+    }
+  }
+  // Prune types.
+  for (size_t slot_idx = 0, num = dex_cache->NumResolvedTypes(); slot_idx != num; ++slot_idx) {
+    mirror::TypeDexCachePair pair =
+        dex_cache->GetResolvedTypes()[slot_idx].load(std::memory_order_relaxed);
+    uint32_t stored_index = pair.index;
+    ObjPtr<mirror::Class> klass = pair.object.Read();
+    if (klass != nullptr && !KeepClass(klass)) {
+      dex_cache->ClearResolvedType(dex::TypeIndex(stored_index));
+    }
+  }
+  // Strings do not need pruning.
+}
+
+void ImageWriter::PreloadDexCache(ObjPtr<mirror::DexCache> dex_cache,
+                                  ObjPtr<mirror::ClassLoader> class_loader) {
   // To ensure deterministic contents of the hash-based arrays, each slot shall contain
   // the candidate with the lowest index. As we're processing entries in increasing index
   // order, this means trying to look up the entry for the current index if the slot is
@@ -1419,7 +1504,7 @@
   Runtime* runtime = Runtime::Current();
   ClassLinker* class_linker = runtime->GetClassLinker();
   const DexFile& dex_file = *dex_cache->GetDexFile();
-  // Prune methods.
+  // Preload the methods array and make the contents deterministic.
   mirror::MethodDexCacheType* resolved_methods = dex_cache->GetResolvedMethods();
   dex::TypeIndex last_class_idx;  // Initialized to invalid index.
   ObjPtr<mirror::Class> last_class = nullptr;
@@ -1439,27 +1524,20 @@
     if (method_id.class_idx_ != last_class_idx) {
       last_class_idx = method_id.class_idx_;
       last_class = class_linker->LookupResolvedType(last_class_idx, dex_cache, class_loader);
-      if (last_class != nullptr && !KeepClass(last_class)) {
-        last_class = nullptr;
-      }
     }
     if (method == nullptr || i < stored_index) {
       if (last_class != nullptr) {
         // Try to resolve the method with the class linker, which will insert
         // it into the dex cache if successful.
         method = class_linker->FindResolvedMethod(last_class, dex_cache, class_loader, i);
-        // If the referenced class is in the image, the defining class must also be there.
-        DCHECK(method == nullptr || KeepClass(method->GetDeclaringClass()));
         DCHECK(method == nullptr || dex_cache->GetResolvedMethod(i, target_ptr_size_) == method);
       }
     } else {
       DCHECK_EQ(i, stored_index);
-      if (last_class == nullptr) {
-        dex_cache->ClearResolvedMethod(stored_index, target_ptr_size_);
-      }
+      DCHECK(last_class != nullptr);
     }
   }
-  // Prune fields and make the contents of the field array deterministic.
+  // Preload the fields array and make the contents deterministic.
   mirror::FieldDexCacheType* resolved_fields = dex_cache->GetResolvedFields();
   last_class_idx = dex::TypeIndex();  // Initialized to invalid index.
   last_class = nullptr;
@@ -1484,19 +1562,17 @@
     }
     if (field == nullptr || i < stored_index) {
       if (last_class != nullptr) {
+        // Try to resolve the field with the class linker, which will insert
+        // it into the dex cache if successful.
         field = class_linker->FindResolvedFieldJLS(last_class, dex_cache, class_loader, i);
-        // If the referenced class is in the image, the defining class must also be there.
-        DCHECK(field == nullptr || KeepClass(field->GetDeclaringClass()));
         DCHECK(field == nullptr || dex_cache->GetResolvedField(i, target_ptr_size_) == field);
       }
     } else {
       DCHECK_EQ(i, stored_index);
-      if (last_class == nullptr) {
-        dex_cache->ClearResolvedField(stored_index, target_ptr_size_);
-      }
+      DCHECK(last_class != nullptr);
     }
   }
-  // Prune types and make the contents of the type array deterministic.
+  // Preload the types array and make the contents deterministic.
   // This is done after fields and methods as their lookup can touch the types array.
   for (size_t i = 0, end = dex_cache->GetDexFile()->NumTypeIds(); i < end; ++i) {
     dex::TypeIndex type_idx(i);
@@ -1507,15 +1583,10 @@
     ObjPtr<mirror::Class> klass = pair.object.Read();
     if (klass == nullptr || i < stored_index) {
       klass = class_linker->LookupResolvedType(type_idx, dex_cache, class_loader);
-      if (klass != nullptr) {
-        DCHECK_EQ(dex_cache->GetResolvedType(type_idx), klass);
-        stored_index = i;  // For correct clearing below if not keeping the `klass`.
-      }
-    } else if (i == stored_index && !KeepClass(klass)) {
-      dex_cache->ClearResolvedType(dex::TypeIndex(stored_index));
+      DCHECK(klass == nullptr || dex_cache->GetResolvedType(type_idx) == klass);
     }
   }
-  // Strings do not need pruning, but the contents of the string array must be deterministic.
+  // Preload the strings array and make the contents deterministic.
   for (size_t i = 0, end = dex_cache->GetDexFile()->NumStringIds(); i < end; ++i) {
     dex::StringIndex string_idx(i);
     uint32_t slot_idx = dex_cache->StringSlotIndex(string_idx);
@@ -1564,21 +1635,11 @@
   }
 
   // Clear references to removed classes from the DexCaches.
-  std::vector<ObjPtr<mirror::DexCache>> dex_caches;
-  {
-    ReaderMutexLock mu2(self, *Locks::dex_lock_);
-    dex_caches.reserve(class_linker->GetDexCachesData().size());
-    for (const ClassLinker::DexCacheData& data : class_linker->GetDexCachesData()) {
-      if (self->IsJWeakCleared(data.weak_root)) {
-        continue;
-      }
-      dex_caches.push_back(self->DecodeJObject(data.weak_root)->AsDexCache());
-    }
-  }
+  std::vector<ObjPtr<mirror::DexCache>> dex_caches = FindDexCaches(self);
   for (ObjPtr<mirror::DexCache> dex_cache : dex_caches) {
     // Pass the class loader associated with the DexCache. This can either be
     // the app's `class_loader` or `nullptr` if boot class loader.
-    PruneAndPreloadDexCache(dex_cache, IsInBootImage(dex_cache.Ptr()) ? nullptr : class_loader);
+    PruneDexCache(dex_cache, IsInBootImage(dex_cache.Ptr()) ? nullptr : class_loader);
   }
 
   // Drop the array class cache in the ClassLinker, as these are roots holding those classes live.
@@ -1588,6 +1649,20 @@
   prune_class_memo_.clear();
 }
 
+std::vector<ObjPtr<mirror::DexCache>> ImageWriter::FindDexCaches(Thread* self) {
+  std::vector<ObjPtr<mirror::DexCache>> dex_caches;
+  ClassLinker* class_linker = Runtime::Current()->GetClassLinker();
+  ReaderMutexLock mu2(self, *Locks::dex_lock_);
+  dex_caches.reserve(class_linker->GetDexCachesData().size());
+  for (const ClassLinker::DexCacheData& data : class_linker->GetDexCachesData()) {
+    if (self->IsJWeakCleared(data.weak_root)) {
+      continue;
+    }
+    dex_caches.push_back(self->DecodeJObject(data.weak_root)->AsDexCache());
+  }
+  return dex_caches;
+}
+
 void ImageWriter::CheckNonImageClassesRemoved() {
   auto visitor = [&](Object* obj) REQUIRES_SHARED(Locks::mutator_lock_) {
     if (obj->IsClass() && !IsInBootImage(obj)) {
diff --git a/dex2oat/linker/image_writer.h b/dex2oat/linker/image_writer.h
index ccd529a..bf89665 100644
--- a/dex2oat/linker/image_writer.h
+++ b/dex2oat/linker/image_writer.h
@@ -460,9 +460,18 @@
   // Remove unwanted classes from various roots.
   void PruneNonImageClasses() REQUIRES_SHARED(Locks::mutator_lock_);
 
-  // Remove unwanted classes from the DexCache roots and preload deterministic DexCache contents.
-  void PruneAndPreloadDexCache(ObjPtr<mirror::DexCache> dex_cache,
-                               ObjPtr<mirror::ClassLoader> class_loader)
+  // Remove unwanted classes from the DexCache roots.
+  void PruneDexCache(ObjPtr<mirror::DexCache> dex_cache, ObjPtr<mirror::ClassLoader> class_loader)
+      REQUIRES_SHARED(Locks::mutator_lock_)
+      REQUIRES(!Locks::classlinker_classes_lock_);
+
+  // Preload deterministic DexCache contents.
+  void PreloadDexCache(ObjPtr<mirror::DexCache> dex_cache, ObjPtr<mirror::ClassLoader> class_loader)
+      REQUIRES_SHARED(Locks::mutator_lock_)
+      REQUIRES(!Locks::classlinker_classes_lock_);
+
+  // Find dex caches for pruning or preloading.
+  std::vector<ObjPtr<mirror::DexCache>> FindDexCaches(Thread* self)
       REQUIRES_SHARED(Locks::mutator_lock_)
       REQUIRES(!Locks::classlinker_classes_lock_);
 
diff --git a/dex2oat/linker/oat_writer.cc b/dex2oat/linker/oat_writer.cc
index 28287bd..7f2877f 100644
--- a/dex2oat/linker/oat_writer.cc
+++ b/dex2oat/linker/oat_writer.cc
@@ -1500,7 +1500,8 @@
     }
     ObjPtr<mirror::DexCache> dex_cache = class_linker_->FindDexCache(Thread::Current(), *dex_file);
     const DexFile::ClassDef& class_def = dex_file->GetClassDef(class_def_index);
-    mirror::Class* klass = dex_cache->GetResolvedType(class_def.class_idx_);
+    ObjPtr<mirror::Class> klass =
+        class_linker_->LookupResolvedType(class_def.class_idx_, dex_cache, class_loader_);
     if (klass != nullptr) {
       for (ArtMethod& method : klass->GetCopiedMethods(pointer_size_)) {
         // Find origin method. Declaring class and dex_method_idx
@@ -1550,24 +1551,11 @@
     ObjPtr<mirror::DexCache> dex_cache = class_linker_->FindDexCache(self, *dex_file_);
     ArtMethod* resolved_method;
     if (writer_->GetCompilerOptions().IsBootImage()) {
-      const InvokeType invoke_type = method.GetInvokeType(
-          dex_file_->GetClassDef(class_def_index_).access_flags_);
-      // Unchecked as we hold mutator_lock_ on entry.
-      ScopedObjectAccessUnchecked soa(self);
-      StackHandleScope<1> hs(self);
-      resolved_method = class_linker_->ResolveMethod<ClassLinker::ResolveMode::kNoChecks>(
-          method.GetIndex(),
-          hs.NewHandle(dex_cache),
-          ScopedNullHandle<mirror::ClassLoader>(),
-          /* referrer */ nullptr,
-          invoke_type);
+      resolved_method = class_linker_->LookupResolvedMethod(
+          method.GetIndex(), dex_cache, /*class_loader=*/ nullptr);
       if (resolved_method == nullptr) {
-        LOG(FATAL_WITHOUT_ABORT) << "Unexpected failure to resolve a method: "
+        LOG(FATAL) << "Unexpected failure to look up a method: "
             << dex_file_->PrettyMethod(method.GetIndex(), true);
-        self->AssertPendingException();
-        mirror::Throwable* exc = self->GetException();
-        std::string dump = exc->Dump();
-        LOG(FATAL) << dump;
         UNREACHABLE();
       }
     } else {