4 files changed, 187 insertions, 11 deletions

diff --git a/dex2oat/linker/image_writer.cc b/dex2oat/linker/image_writer.cc
index 2c3c2206bf..577143cdb5 100644
--- a/dex2oat/linker/image_writer.cc
+++ b/dex2oat/linker/image_writer.cc
@@ -21,10 +21,12 @@
 #include <sys/stat.h>
 #include <zlib.h>
 
+#include <charconv>
 #include <memory>
 #include <numeric>
 #include <vector>
 
+#include "android-base/strings.h"
 #include "art_field-inl.h"
 #include "art_method-inl.h"
 #include "base/callee_save_type.h"
@@ -334,6 +336,13 @@ bool ImageWriter::PrepareImageAddressSpace(TimingLogger* timings) {
     TimingLogger::ScopedTiming t("CalculateNewObjectOffsets", timings);
     ScopedObjectAccess soa(self);
     CalculateNewObjectOffsets();
+
+    // If dirty_image_objects_ is present - try optimizing object layout.
+    // It can only be done after the first CalculateNewObjectOffsets,
+    // because calculated offsets are used to match dirty objects between imgdiag and dex2oat.
+    if (compiler_options_.IsBootImage() && dirty_image_objects_ != nullptr) {
+      TryRecalculateOffsetsWithDirtyObjects();
+    }
   }
 
   // This needs to happen after CalculateNewObjectOffsets since it relies on intern_table_bytes_ and
@@ -729,7 +738,13 @@ ImageWriter::Bin ImageWriter::AssignImageBinSlot(mirror::Object* object, size_t
     // so packing them together will not result in a noticeably tighter dirty-to-clean ratio.
     //
     ObjPtr<mirror::Class> klass = object->GetClass<kVerifyNone, kWithoutReadBarrier>();
-    if (klass->IsClassClass()) {
+    if (klass->IsStringClass<kVerifyNone>()) {
+      // Assign strings to their bin before checking dirty objects, because
+      // string intern processing expects strings to be in Bin::kString.
+      bin = Bin::kString;  // Strings are almost always immutable (except for object header).
+    } else if (dirty_objects_.find(object) != dirty_objects_.end()) {
+      bin = Bin::kKnownDirty;
+    } else if (klass->IsClassClass()) {
       bin = Bin::kClassVerified;
       ObjPtr<mirror::Class> as_klass = object->AsClass<kVerifyNone>();
 
@@ -766,8 +781,6 @@ ImageWriter::Bin ImageWriter::AssignImageBinSlot(mirror::Object* object, size_t
           }
         }
       }
-    } else if (klass->IsStringClass<kVerifyNone>()) {
-      bin = Bin::kString;  // Strings are almost always immutable (except for object header).
     } else if (!klass->HasSuperClass()) {
       // Only `j.l.Object` and primitive classes lack the superclass and
       // there are no instances of primitive classes.
@@ -2533,6 +2546,142 @@ void ImageWriter::CalculateNewObjectOffsets() {
   }
 }
 
+std::optional<HashSet<mirror::Object*>> ImageWriter::MatchDirtyObjectOffsets(
+    const HashMap<uint32_t, DirtyEntry>& dirty_entries) REQUIRES_SHARED(Locks::mutator_lock_) {
+  HashSet<mirror::Object*> dirty_objects;
+  bool mismatch_found = false;
+
+  auto visitor = [&](Object* obj) REQUIRES_SHARED(Locks::mutator_lock_) {
+    DCHECK(obj != nullptr);
+    if (mismatch_found) {
+      return;
+    }
+    if (!IsImageBinSlotAssigned(obj)) {
+      return;
+    }
+
+    uint8_t* image_address = reinterpret_cast<uint8_t*>(GetImageAddress(obj));
+    uint32_t offset = static_cast<uint32_t>(image_address - global_image_begin_);
+
+    auto entry_it = dirty_entries.find(offset);
+    if (entry_it == dirty_entries.end()) {
+      return;
+    }
+
+    const DirtyEntry& entry = entry_it->second;
+
+    const bool is_class = obj->IsClass();
+    const uint32_t descriptor_hash =
+        is_class ? obj->AsClass()->DescriptorHash() : obj->GetClass()->DescriptorHash();
+
+    if (is_class != entry.is_class || descriptor_hash != entry.descriptor_hash) {
+      LOG(WARNING) << "Dirty image objects offset mismatch (outdated file?)";
+      mismatch_found = true;
+      return;
+    }
+
+    dirty_objects.insert(obj);
+  };
+  Runtime::Current()->GetHeap()->VisitObjects(visitor);
+
+  // A single mismatch indicates that dirty-image-objects layout differs from
+  // current ImageWriter layout. In this case any "valid" matches are likely to be accidental,
+  // so there's no point in optimizing the layout with such data.
+  if (mismatch_found) {
+    return {};
+  }
+  if (dirty_objects.size() != dirty_entries.size()) {
+    LOG(WARNING) << "Dirty image objects missing offsets (outdated file?)";
+    return {};
+  }
+  return dirty_objects;
+}
+
+void ImageWriter::ResetObjectOffsets() {
+  const size_t image_infos_size = image_infos_.size();
+  image_infos_.clear();
+  image_infos_.resize(image_infos_size);
+
+  native_object_relocations_.clear();
+
+  // CalculateNewObjectOffsets stores image offsets of the objects in lock words,
+  // while original lock words are preserved in saved_hashcode_map.
+  // Restore original lock words.
+  auto visitor = [&](Object* obj) REQUIRES_SHARED(Locks::mutator_lock_) {
+    DCHECK(obj != nullptr);
+    const auto it = saved_hashcode_map_.find(obj);
+    obj->SetLockWord(it != saved_hashcode_map_.end() ? LockWord::FromHashCode(it->second, 0u) :
+                                                       LockWord::Default(),
+                     false);
+  };
+  Runtime::Current()->GetHeap()->VisitObjects(visitor);
+
+  saved_hashcode_map_.clear();
+}
+
+void ImageWriter::TryRecalculateOffsetsWithDirtyObjects() {
+  const std::optional<HashMap<uint32_t, ImageWriter::DirtyEntry>> dirty_entries =
+      ParseDirtyObjectOffsets(*dirty_image_objects_);
+  if (!dirty_entries || dirty_entries->empty()) {
+    return;
+  }
+
+  std::optional<HashSet<mirror::Object*>> dirty_objects = MatchDirtyObjectOffsets(*dirty_entries);
+  if (!dirty_objects || dirty_objects->empty()) {
+    return;
+  }
+  // Calculate offsets again, now with dirty object offsets.
+  LOG(INFO) << "Recalculating object offsets using dirty-image-objects";
+  dirty_objects_ = std::move(*dirty_objects);
+  ResetObjectOffsets();
+  CalculateNewObjectOffsets();
+}
+
+std::optional<HashMap<uint32_t, ImageWriter::DirtyEntry>> ImageWriter::ParseDirtyObjectOffsets(
+    const HashSet<std::string>& dirty_image_objects) REQUIRES_SHARED(Locks::mutator_lock_) {
+  HashMap<uint32_t, DirtyEntry> dirty_entries;
+
+  // Go through each dirty-image-object line, parse only lines of the format:
+  // "dirty_obj: <offset> <type> <descriptor_hash>"
+  // <offset> -- decimal uint32.
+  // <type> -- "class" or "instance" (defines if descriptor is referring to a class or an instance).
+  // <descriptor_hash> -- decimal uint32 (from DescriptorHash() method).
+  const std::string prefix = "dirty_obj:";
+  for (const std::string& entry_str : dirty_image_objects) {
+    // Skip the lines of old dirty-image-object format.
+    if (std::strncmp(entry_str.data(), prefix.data(), prefix.size()) != 0) {
+      continue;
+    }
+
+    const std::vector<std::string> tokens = android::base::Split(entry_str, " ");
+    if (tokens.size() != 4) {
+      LOG(WARNING) << "Invalid dirty image objects format: \"" << entry_str << "\"";
+      return {};
+    }
+
+    uint32_t offset = 0;
+    std::from_chars_result res =
+        std::from_chars(tokens[1].data(), tokens[1].data() + tokens[1].size(), offset);
+    if (res.ec != std::errc()) {
+      LOG(WARNING) << "Couldn't parse dirty object offset: \"" << entry_str << "\"";
+      return {};
+    }
+
+    DirtyEntry entry;
+    entry.is_class = (tokens[2] == "class");
+    res = std::from_chars(
+        tokens[3].data(), tokens[3].data() + tokens[3].size(), entry.descriptor_hash);
+    if (res.ec != std::errc()) {
+      LOG(WARNING) << "Couldn't parse dirty object descriptor hash: \"" << entry_str << "\"";
+      return {};
+    }
+
+    dirty_entries.insert(std::make_pair(offset, entry));
+  }
+
+  return dirty_entries;
+}
+
 std::pair<size_t, dchecked_vector<ImageSection>>
 ImageWriter::ImageInfo::CreateImageSections() const {
   dchecked_vector<ImageSection> sections(ImageHeader::kSectionCount);
diff --git a/dex2oat/linker/image_writer.h b/dex2oat/linker/image_writer.h
index e5eeacc02d..b6535c363d 100644
--- a/dex2oat/linker/image_writer.h
+++ b/dex2oat/linker/image_writer.h
@@ -462,6 +462,23 @@ class ImageWriter final {
       REQUIRES_SHARED(Locks::mutator_lock_);
   void CalculateObjectBinSlots(mirror::Object* obj)
       REQUIRES_SHARED(Locks::mutator_lock_);
+  // Undo the changes of CalculateNewObjectOffsets.
+  void ResetObjectOffsets() REQUIRES_SHARED(Locks::mutator_lock_);
+  // Reset and calculate new offsets with dirty objects optimization.
+  // Does nothing if dirty object offsets don't match with current offsets.
+  void TryRecalculateOffsetsWithDirtyObjects() REQUIRES_SHARED(Locks::mutator_lock_);
+
+  // Dirty object data from dirty-image-objects.
+  struct DirtyEntry {
+    uint32_t descriptor_hash = 0;
+    bool is_class = false;
+  };
+  // Parse dirty-image-objects into (offset->entry) map. Returns nullopt on parse error.
+  static std::optional<HashMap<uint32_t, DirtyEntry>> ParseDirtyObjectOffsets(
+      const HashSet<std::string>& dirty_image_objects) REQUIRES_SHARED(Locks::mutator_lock_);
+  // Get all objects that match dirty_entries by offset. Returns nullopt if there is a mismatch.
+  std::optional<HashSet<mirror::Object*>> MatchDirtyObjectOffsets(
+      const HashMap<uint32_t, DirtyEntry>& dirty_entries) REQUIRES_SHARED(Locks::mutator_lock_);
 
   // Creates the contiguous image in memory and adjusts pointers.
   void CopyAndFixupNativeData(size_t oat_index) REQUIRES_SHARED(Locks::mutator_lock_);
@@ -685,9 +702,14 @@ class ImageWriter final {
   // Map of dex files to the indexes of oat files that they were compiled into.
   const HashMap<const DexFile*, size_t>& dex_file_oat_index_map_;
 
-  // Set of objects known to be dirty in the image. Can be nullptr if there are none.
+  // Set of classes/objects known to be dirty in the image. Can be nullptr if there are none.
+  // For old dirty-image-objects format this set contains descriptors of dirty classes.
+  // For new format -- a set of dirty object offsets and descriptor hashes.
   const HashSet<std::string>* dirty_image_objects_;
 
+  // Dirty object instances parsed from dirty_image_object_
+  HashSet<mirror::Object*> dirty_objects_;
+
   // Objects are guaranteed to not cross the region size boundary.
   size_t region_size_ = 0u;
 
diff --git a/imgdiag/Android.bp b/imgdiag/Android.bp
index afd86a0f32..a0354ab6c5 100644
--- a/imgdiag/Android.bp
+++ b/imgdiag/Android.bp
@@ -84,6 +84,7 @@ art_cc_binary {
         "libartd",
         "libartbased",
         "libartd-compiler",
+        "libdexfiled",
     ],
     apex_available: [
         "com.android.art.debug",
diff --git a/imgdiag/imgdiag.cc b/imgdiag/imgdiag.cc
index fe70b10fb6..ed2ff727fe 100644
--- a/imgdiag/imgdiag.cc
+++ b/imgdiag/imgdiag.cc
@@ -450,17 +450,22 @@ class RegionSpecializedBase<mirror::Object> : public RegionCommon<mirror::Object
   void DiffEntryContents(mirror::Object* entry,
                          uint8_t* remote_bytes,
                          const uint8_t* base_ptr,
-                         bool log_dirty_objects)
-      REQUIRES_SHARED(Locks::mutator_lock_) {
+                         bool log_dirty_objects,
+                         size_t entry_offset) REQUIRES_SHARED(Locks::mutator_lock_) {
     const char* tabs = "    ";
     // Attempt to find fields for all dirty bytes.
     mirror::Class* klass = entry->GetClass();
+    std::string temp;
     if (entry->IsClass()) {
       os_ << tabs
           << "Class " << mirror::Class::PrettyClass(entry->AsClass()) << " " << entry << "\n";
+      os_ << tabs << "dirty_obj: " << entry_offset << " class "
+          << entry->AsClass()->DescriptorHash() << "\n";
     } else {
       os_ << tabs
           << "Instance of " << mirror::Class::PrettyClass(klass) << " " << entry << "\n";
+      os_ << tabs << "dirty_obj: " << entry_offset << " instance " << klass->DescriptorHash()
+          << "\n";
     }
     PrintEntryPages(reinterpret_cast<uintptr_t>(entry), EntrySize(entry), os_);
 
@@ -776,7 +781,8 @@ class RegionSpecializedBase<ArtMethod> : public RegionCommon<ArtMethod> {
   void DiffEntryContents(ArtMethod* method,
                          uint8_t* remote_bytes,
                          const uint8_t* base_ptr,
-                         bool log_dirty_objects ATTRIBUTE_UNUSED)
+                         bool log_dirty_objects ATTRIBUTE_UNUSED,
+                         size_t entry_offset ATTRIBUTE_UNUSED)
       REQUIRES_SHARED(Locks::mutator_lock_) {
     const char* tabs = "    ";
     os_ << tabs << "ArtMethod " << ArtMethod::PrettyMethod(method) << "\n";
@@ -1047,10 +1053,8 @@ class RegionData : public RegionSpecializedBase<T> {
       uint8_t* entry_bytes = reinterpret_cast<uint8_t*>(entry);
       ptrdiff_t offset = entry_bytes - begin_image_ptr;
       uint8_t* remote_bytes = &contents[offset];
-      RegionSpecializedBase<T>::DiffEntryContents(entry,
-                                                  remote_bytes,
-                                                  &base_ptr[offset],
-                                                  log_dirty_objects);
+      RegionSpecializedBase<T>::DiffEntryContents(
+          entry, remote_bytes, &base_ptr[offset], log_dirty_objects, static_cast<size_t>(offset));
     }
   }