Rewrite class table construction in ImageWriter.

Make sure that class tables in images are at maximum load
factor (full) and make that maximum load factor independent
of runtime parameters. As we pre-allocate a class table
buffer of the right size in ImageWriter, we also avoid
unnecessary resizing of the temporary class table.

Make sure that app image class tables are deterministic.
We previously just copied the class table from the app class
loader even though some entries may have been inserted there
during multi-threaded phases of the compilation, causing
non-deterministic contents based on insertion order.

Remove obsolete comment related to patchoat relocations.

Test: m test-art-host-gtest
Test: testrunner.py --host --optimizing
Test: boots.
Bug: 175869411
Change-Id: I605048b639f67a5ed4b03eb8888cbaafa9ba4091
diff --git a/dex2oat/linker/image_writer.cc b/dex2oat/linker/image_writer.cc
index f4155bc..7bdeac2 100644
--- a/dex2oat/linker/image_writer.cc
+++ b/dex2oat/linker/image_writer.cc
@@ -90,6 +90,15 @@
 namespace art {
 namespace linker {
 
+// The actual value of `kImageClassTableMinLoadFactor` is irrelevant because image class tables
+// are never resized, but we still need to pass a reasonable value to the constructor.
+constexpr double kImageClassTableMinLoadFactor = 0.5;
+// We use `kImageClassTableMaxLoadFactor` to determine the buffer size for image class tables
+// to make them full. We never insert additional elements to them, so we do not want to waste
+// extra memory. And unlike runtime class tables, we do not want this to depend on runtime
+// properties (see `Runtime::GetHashTableMaxLoadFactor()` checking for low memory mode).
+constexpr double kImageClassTableMaxLoadFactor = 0.7;
+
 static ArrayRef<const uint8_t> MaybeCompressData(ArrayRef<const uint8_t> source,
                                                  ImageHeader::StorageMode image_storage_mode,
                                                  /*out*/ std::vector<uint8_t>* storage) {
@@ -281,27 +290,6 @@
     CalculateNewObjectOffsets();
   }
 
-  // Obtain class count for debugging purposes
-  if (VLOG_IS_ON(compiler) && compiler_options_.IsAppImage()) {
-    ScopedObjectAccess soa(self);
-
-    size_t app_image_class_count  = 0;
-
-    for (ImageInfo& info : image_infos_) {
-      info.class_table_->Visit([&](ObjPtr<mirror::Class> klass)
-                                   REQUIRES_SHARED(Locks::mutator_lock_) {
-        if (!IsInBootImage(klass.Ptr())) {
-          ++app_image_class_count;
-        }
-
-        // Indicate that we would like to continue visiting classes.
-        return true;
-      });
-    }
-
-    VLOG(compiler) << "Dex2Oat:AppImage:classCount = " << app_image_class_count;
-  }
-
   // This needs to happen after CalculateNewObjectOffsets since it relies on intern_table_bytes_ and
   // bin size sums being calculated.
   TimingLogger::ScopedTiming t("AllocMemory", timings);
@@ -1384,11 +1372,6 @@
         as_klass->GetSFieldsPtr(), as_klass->GetIFieldsPtr(),
     };
     ImageInfo& image_info = GetImageInfo(oat_index);
-    if (!compiler_options_.IsAppImage()) {
-      // Note: Avoid locking to prevent lock order violations from root visiting;
-      // image_info.class_table_ is only accessed from the image writer.
-      image_info.class_table_->InsertWithoutLocks(as_klass);
-    }
     for (LengthPrefixedArray<ArtField>* cur_fields : fields) {
       // Total array length including header.
       if (cur_fields != nullptr) {
@@ -1465,20 +1448,6 @@
         }
       }
     }
-  } else if (obj->IsClassLoader()) {
-    // Register the class loader if it has a class table.
-    // The fake boot class loader should not get registered.
-    ObjPtr<mirror::ClassLoader> class_loader = obj->AsClassLoader();
-    if (class_loader->GetClassTable() != nullptr) {
-      DCHECK(compiler_options_.IsAppImage());
-      if (class_loader == GetAppClassLoader()) {
-        ImageInfo& image_info = GetImageInfo(oat_index);
-        // Note: Avoid locking to prevent lock order violations from root visiting;
-        // image_info.class_table_ table is only accessed from the image writer
-        // and class_loader->GetClassTable() is iterated but not modified.
-        image_info.class_table_->CopyWithoutLocks(*class_loader->GetClassTable());
-      }
-    }
   }
 }
 
@@ -1627,11 +1596,11 @@
     return true;
   }
 
-  WorkQueue SortAndReleaseClasses()
-      REQUIRES_SHARED(Locks::mutator_lock_) {
+  WorkQueue ProcessCollectedClasses(Thread* self) REQUIRES_SHARED(Locks::mutator_lock_) {
     std::sort(klasses_.begin(), klasses_.end());
 
-    WorkQueue result;
+    ImageWriter* image_writer = image_writer_;
+    WorkQueue work_queue;
     size_t last_dex_file_index = static_cast<size_t>(-1);
     size_t last_oat_index = static_cast<size_t>(-1);
     for (const ClassEntry& entry : klasses_) {
@@ -1640,14 +1609,93 @@
           last_oat_index = GetDefaultOatIndex();  // Primitive type.
         } else {
           uint32_t dex_file_index = entry.dex_file_index - 1u;  // 0 is for primitive types.
-          last_oat_index = image_writer_->GetOatIndexForDexFile(dex_files_[dex_file_index]);
+          last_oat_index = image_writer->GetOatIndexForDexFile(dex_files_[dex_file_index]);
         }
         last_dex_file_index = entry.dex_file_index;
       }
-      result.emplace_back(entry.klass, last_oat_index);
+      // Count the number of classes for class tables.
+      image_writer->image_infos_[last_oat_index].class_table_size_ += 1u;
+      work_queue.emplace_back(entry.klass, last_oat_index);
     }
     klasses_.clear();
-    return result;
+
+    // Prepare image class tables.
+    std::vector<mirror::Class*> boot_image_classes;
+    if (image_writer->compiler_options_.IsAppImage()) {
+      DCHECK_EQ(image_writer->image_infos_.size(), 1u);
+      ImageInfo& image_info = image_writer->image_infos_[0];
+      // Log the non-boot image class count for app image for debugging purposes.
+      VLOG(compiler) << "Dex2Oat:AppImage:classCount = " << image_info.class_table_size_;
+      // Collect boot image classes referenced by app class loader's class table.
+      ClassTable* app_class_table = image_writer->GetAppClassLoader()->GetClassTable();
+      ReaderMutexLock lock(self, app_class_table->lock_);
+      DCHECK_EQ(app_class_table->classes_.size(), 1u);
+      const ClassTable::ClassSet& app_class_set = app_class_table->classes_[0];
+      DCHECK_GE(app_class_set.size(), image_info.class_table_size_);
+      boot_image_classes.reserve(app_class_set.size() - image_info.class_table_size_);
+      for (const ClassTable::TableSlot& slot : app_class_set) {
+        mirror::Class* klass = slot.Read<kWithoutReadBarrier>().Ptr();
+        if (image_writer->IsInBootImage(klass)) {
+          boot_image_classes.push_back(klass);
+        }
+      }
+      DCHECK_EQ(app_class_set.size() - image_info.class_table_size_, boot_image_classes.size());
+      // Increase the app class table size to include referenced boot image classes.
+      image_info.class_table_size_ = app_class_set.size();
+    }
+    for (ImageInfo& image_info : image_writer->image_infos_) {
+      if (image_info.class_table_size_ != 0u) {
+        // Make sure the class table shall be full by allocating a buffer of the right size.
+        size_t buffer_size = static_cast<size_t>(
+            ceil(image_info.class_table_size_ / kImageClassTableMaxLoadFactor));
+        image_info.class_table_buffer_.reset(new ClassTable::TableSlot[buffer_size]);
+        DCHECK(image_info.class_table_buffer_ != nullptr);
+        image_info.class_table_.emplace(kImageClassTableMinLoadFactor,
+                                        kImageClassTableMaxLoadFactor,
+                                        image_info.class_table_buffer_.get(),
+                                        buffer_size);
+      }
+    }
+    for (const auto& pair : work_queue) {
+      ObjPtr<mirror::Class> klass = pair.first->AsClass();
+      size_t oat_index = pair.second;
+      DCHECK(image_writer->image_infos_[oat_index].class_table_.has_value());
+      ClassTable::ClassSet& class_table = *image_writer->image_infos_[oat_index].class_table_;
+      bool inserted = class_table.insert(ClassTable::TableSlot(klass)).second;
+      DCHECK(inserted) << "Class " << klass->PrettyDescriptor()
+          << " (" << klass.Ptr() << ") already inserted";
+    }
+    if (image_writer->compiler_options_.IsAppImage()) {
+      DCHECK_EQ(image_writer->image_infos_.size(), 1u);
+      ImageInfo& image_info = image_writer->image_infos_[0];
+      if (image_info.class_table_size_ != 0u) {
+        // Insert boot image class references to the app class table.
+        // The order of insertion into the app class loader's ClassTable is non-deterministic,
+        // so sort the boot image classes by the boot image address to get deterministic table.
+        std::sort(boot_image_classes.begin(), boot_image_classes.end());
+        DCHECK(image_info.class_table_.has_value());
+        ClassTable::ClassSet& table = *image_info.class_table_;
+        for (mirror::Class* klass : boot_image_classes) {
+          bool inserted = table.insert(ClassTable::TableSlot(klass)).second;
+          DCHECK(inserted) << "Boot image class " << klass->PrettyDescriptor()
+              << " (" << klass << ") already inserted";
+        }
+        DCHECK_EQ(table.size(), image_info.class_table_size_);
+      }
+    }
+    for (ImageInfo& image_info : image_writer->image_infos_) {
+      DCHECK_EQ(image_info.class_table_bytes_, 0u);
+      if (image_info.class_table_size_ != 0u) {
+        DCHECK(image_info.class_table_.has_value());
+        DCHECK_EQ(image_info.class_table_->size(), image_info.class_table_size_);
+        image_info.class_table_bytes_ = image_info.class_table_->WriteToMemory(nullptr);
+        DCHECK_NE(image_info.class_table_bytes_, 0u);
+      } else {
+        DCHECK(!image_info.class_table_.has_value());
+      }
+    }
+
+    return work_queue;
   }
 
  private:
@@ -1807,7 +1855,7 @@
   CollectClassesVisitor visitor(image_writer_);
   class_linker->VisitClasses(&visitor);
   DCHECK(work_queue_.empty());
-  work_queue_ = visitor.SortAndReleaseClasses();
+  work_queue_ = visitor.ProcessCollectedClasses(self);
   for (const std::pair<ObjPtr<mirror::Object>, size_t>& entry : work_queue_) {
     DCHECK(entry.first->IsClass());
     bool assigned = TryAssignBinSlot(entry.first, entry.second);
@@ -2204,13 +2252,6 @@
     if (intern_table->StrongSize() != 0u) {
       image_info.intern_table_bytes_ = intern_table->WriteToMemory(nullptr);
     }
-
-    // Calculate the size of the class table.
-    ReaderMutexLock mu(self, *Locks::classlinker_classes_lock_);
-    DCHECK_EQ(image_info.class_table_->NumReferencedZygoteClasses(), 0u);
-    if (image_info.class_table_->NumReferencedNonZygoteClasses() != 0u) {
-      image_info.class_table_bytes_ += image_info.class_table_->WriteToMemory(nullptr);
-    }
   }
 
   // Finalize bin slot offsets. This may add padding for regions.
@@ -2625,27 +2666,27 @@
     const ImageSection& class_table_section = image_header->GetClassTableSection();
     uint8_t* const class_table_memory_ptr =
         image_info.image_.Begin() + class_table_section.Offset();
-    Thread* self = Thread::Current();
-    ReaderMutexLock mu(self, *Locks::classlinker_classes_lock_);
 
-    ClassTable* table = image_info.class_table_.get();
-    CHECK(table != nullptr);
-    const size_t class_table_bytes = table->WriteToMemory(class_table_memory_ptr);
+    DCHECK(image_info.class_table_.has_value());
+    const ClassTable::ClassSet& table = *image_info.class_table_;
+    CHECK_EQ(table.size(), image_info.class_table_size_);
+    const size_t class_table_bytes = table.WriteToMemory(class_table_memory_ptr);
     CHECK_EQ(class_table_bytes, image_info.class_table_bytes_);
+
     // Fixup the pointers in the newly written class table to contain image addresses. See
     // above comment for intern tables.
     ClassTable temp_class_table;
     temp_class_table.ReadFromMemory(class_table_memory_ptr);
-    CHECK_EQ(temp_class_table.NumReferencedZygoteClasses(),
-             table->NumReferencedNonZygoteClasses() + table->NumReferencedZygoteClasses());
+    CHECK_EQ(temp_class_table.NumReferencedZygoteClasses(), table.size());
     UnbufferedRootVisitor visitor(&root_visitor, RootInfo(kRootUnknown));
     temp_class_table.VisitRoots(visitor);
-    // Record relocations. (The root visitor does not get to see the slot addresses.)
-    // Note that the low bits in the slots contain bits of the descriptors' hash codes
-    // but the relocation works fine for these "adjusted" references.
-    ReaderMutexLock lock(self, temp_class_table.lock_);
-    DCHECK(!temp_class_table.classes_.empty());
-    DCHECK(!temp_class_table.classes_[0].empty());  // The ClassSet was inserted at the beginning.
+
+    if (kIsDebugBuild) {
+      ReaderMutexLock lock(Thread::Current(), temp_class_table.lock_);
+      CHECK(!temp_class_table.classes_.empty());
+      // The ClassSet was inserted at the beginning.
+      CHECK_EQ(temp_class_table.classes_[0].size(), table.size());
+    }
   }
 }
 
@@ -3258,7 +3299,7 @@
 
 ImageWriter::ImageInfo::ImageInfo()
     : intern_table_(new InternTable),
-      class_table_(new ClassTable) {}
+      class_table_() {}
 
 template <typename DestType>
 void ImageWriter::CopyAndFixupReference(DestType* dest, ObjPtr<mirror::Object> src) {
diff --git a/dex2oat/linker/image_writer.h b/dex2oat/linker/image_writer.h
index 6f847a2..7a11fe6 100644
--- a/dex2oat/linker/image_writer.h
+++ b/dex2oat/linker/image_writer.h
@@ -390,7 +390,9 @@
     std::unique_ptr<InternTable> intern_table_;
 
     // Class table associated with this image for serialization.
-    std::unique_ptr<ClassTable> class_table_;
+    size_t class_table_size_ = 0;
+    std::unique_ptr<ClassTable::ClassSet::value_type[]> class_table_buffer_;
+    std::optional<ClassTable::ClassSet> class_table_;
 
     // Padding offsets to ensure region alignment (if required).
     // Objects need to be added from the recorded offset until the end of the region.
diff --git a/runtime/class_table.cc b/runtime/class_table.cc
index cb87ee2..65d9768 100644
--- a/runtime/class_table.cc
+++ b/runtime/class_table.cc
@@ -146,24 +146,6 @@
   classes_.back().InsertWithHash(TableSlot(klass, hash), hash);
 }
 
-void ClassTable::CopyWithoutLocks(const ClassTable& source_table) {
-  if (kIsDebugBuild) {
-    for (ClassSet& class_set : classes_) {
-      CHECK(class_set.empty());
-    }
-  }
-  for (const ClassSet& class_set : source_table.classes_) {
-    for (const TableSlot& slot : class_set) {
-      classes_.back().insert(slot);
-    }
-  }
-}
-
-void ClassTable::InsertWithoutLocks(ObjPtr<mirror::Class> klass) {
-  const uint32_t hash = TableSlot::HashDescriptor(klass);
-  classes_.back().InsertWithHash(TableSlot(klass, hash), hash);
-}
-
 void ClassTable::InsertWithHash(ObjPtr<mirror::Class> klass, size_t hash) {
   WriterMutexLock mu(Thread::Current(), lock_);
   classes_.back().InsertWithHash(TableSlot(klass, hash), hash);
@@ -255,26 +237,6 @@
   return true;
 }
 
-size_t ClassTable::WriteToMemory(uint8_t* ptr) const {
-  ReaderMutexLock mu(Thread::Current(), lock_);
-  ClassSet combined;
-  // Combine all the class sets in case there are multiple, also adjusts load factor back to
-  // default in case classes were pruned.
-  for (const ClassSet& class_set : classes_) {
-    for (const TableSlot& root : class_set) {
-      combined.insert(root);
-    }
-  }
-  const size_t ret = combined.WriteToMemory(ptr);
-  // Validity check
-  if (kIsDebugBuild && ptr != nullptr) {
-    size_t read_count;
-    ClassSet class_set(ptr, /*make copy*/false, &read_count);
-    class_set.Verify();
-  }
-  return ret;
-}
-
 size_t ClassTable::ReadFromMemory(uint8_t* ptr) {
   size_t read_count = 0;
   AddClassSet(ClassSet(ptr, /*make copy*/false, &read_count));
diff --git a/runtime/class_table.h b/runtime/class_table.h
index 810c09c..544a26b 100644
--- a/runtime/class_table.h
+++ b/runtime/class_table.h
@@ -238,11 +238,6 @@
       REQUIRES(!lock_)
       REQUIRES_SHARED(Locks::mutator_lock_);
 
-  // Combines all of the tables into one class set.
-  size_t WriteToMemory(uint8_t* ptr) const
-      REQUIRES(!lock_)
-      REQUIRES_SHARED(Locks::mutator_lock_);
-
   // Read a table from ptr and put it at the front of the class set.
   size_t ReadFromMemory(uint8_t* ptr)
       REQUIRES(!lock_)
@@ -269,10 +264,6 @@
   }
 
  private:
-  // Only copies classes.
-  void CopyWithoutLocks(const ClassTable& source_table) NO_THREAD_SAFETY_ANALYSIS;
-  void InsertWithoutLocks(ObjPtr<mirror::Class> klass) NO_THREAD_SAFETY_ANALYSIS;
-
   size_t CountDefiningLoaderClasses(ObjPtr<mirror::ClassLoader> defining_loader,
                                     const ClassSet& set) const
       REQUIRES(lock_)
diff --git a/runtime/class_table_test.cc b/runtime/class_table_test.cc
index 5275c7e..642e10a 100644
--- a/runtime/class_table_test.cc
+++ b/runtime/class_table_test.cc
@@ -144,11 +144,16 @@
   table.Remove(descriptor_x);
   EXPECT_FALSE(table.Contains(h_X.Get()));
 
-  // Test that WriteToMemory and ReadFromMemory work.
+  // Test that reading a class set from memory works.
   table.Insert(h_X.Get());
-  const size_t count = table.WriteToMemory(nullptr);
+  ClassTable::ClassSet temp_set;
+  table.Visit([&temp_set](ObjPtr<mirror::Class> klass) REQUIRES_SHARED(Locks::mutator_lock_) {
+    temp_set.insert(ClassTable::TableSlot(klass));
+    return true;
+  });
+  const size_t count = temp_set.WriteToMemory(nullptr);
   std::unique_ptr<uint8_t[]> buffer(new uint8_t[count]());
-  ASSERT_EQ(table.WriteToMemory(&buffer[0]), count);
+  ASSERT_EQ(temp_set.WriteToMemory(&buffer[0]), count);
   ClassTable table2;
   size_t count2 = table2.ReadFromMemory(&buffer[0]);
   EXPECT_EQ(count, count2);