Store additional information in image header.

Store the reservation needed for loading the boot image and
the number of boot class path components. The former helps
avoid one extra file open+read+close when loading the boot
image and the latter shall be used in subsequent CLs to load
boot image for a partial boot class path.

Test: m test-art-host-gtest
Test: testrunner.py --host --optimizing
Test: Pixel 2 XL boots.
Bug: 119868597
Change-Id: I01dba923cfa3f8faf0e41a4139b8913c78404d54
diff --git a/dex2oat/linker/image_test.cc b/dex2oat/linker/image_test.cc
index ebd829b..1a5701d 100644
--- a/dex2oat/linker/image_test.cc
+++ b/dex2oat/linker/image_test.cc
@@ -73,7 +73,10 @@
     uint32_t oat_data_end = ART_BASE_ADDRESS + (9 * KB);
     uint32_t oat_file_end = ART_BASE_ADDRESS + (10 * KB);
     ImageSection sections[ImageHeader::kSectionCount];
-    ImageHeader image_header(image_begin,
+    uint32_t image_reservation_size = RoundUp(oat_file_end - image_begin, kPageSize);
+    ImageHeader image_header(image_reservation_size,
+                             /*component_count=*/ 1u,
+                             image_begin,
                              image_size_,
                              sections,
                              image_roots,
diff --git a/dex2oat/linker/image_writer.cc b/dex2oat/linker/image_writer.cc
index 231ac04..61d105f 100644
--- a/dex2oat/linker/image_writer.cc
+++ b/dex2oat/linker/image_writer.cc
@@ -2583,6 +2583,23 @@
   const uint8_t* oat_file_end = oat_file_begin + image_info.oat_loaded_size_;
   const uint8_t* oat_data_end = image_info.oat_data_begin_ + image_info.oat_size_;
 
+  uint32_t image_reservation_size = image_info.image_size_;
+  DCHECK_ALIGNED(image_reservation_size, kPageSize);
+  uint32_t component_count = 1u;
+  if (!compiler_options_.IsAppImage()) {
+    if (oat_index == 0u) {
+      const ImageInfo& last_info = image_infos_.back();
+      const uint8_t* end = last_info.oat_file_begin_ + last_info.oat_loaded_size_;
+      DCHECK_ALIGNED(image_info.image_begin_, kPageSize);
+      image_reservation_size =
+          dchecked_integral_cast<uint32_t>(RoundUp(end - image_info.image_begin_, kPageSize));
+      component_count = image_infos_.size();
+    } else {
+      image_reservation_size = 0u;
+      component_count = 0u;
+    }
+  }
+
   // Create the image sections.
   auto section_info_pair = image_info.CreateImageSections();
   const size_t image_end = section_info_pair.first;
@@ -2619,6 +2636,8 @@
   // Create the header, leave 0 for data size since we will fill this in as we are writing the
   // image.
   new (image_info.image_.Begin()) ImageHeader(
+      image_reservation_size,
+      component_count,
       PointerToLowMemUInt32(image_info.image_begin_),
       image_end,
       sections.data(),
diff --git a/runtime/gc/collector/immune_spaces_test.cc b/runtime/gc/collector/immune_spaces_test.cc
index c29b79c..b0d09ba 100644
--- a/runtime/gc/collector/immune_spaces_test.cc
+++ b/runtime/gc/collector/immune_spaces_test.cc
@@ -112,6 +112,8 @@
     // Create image header.
     ImageSection sections[ImageHeader::kSectionCount];
     new (image_map.Begin()) ImageHeader(
+        /*image_reservation_size=*/ image_size,
+        /*component_count=*/ 1u,
         /*image_begin=*/ PointerToLowMemUInt32(image_map.Begin()),
         /*image_size=*/ image_size,
         sections,
diff --git a/runtime/gc/space/image_space.cc b/runtime/gc/space/image_space.cc
index 02ab50b..b621599 100644
--- a/runtime/gc/space/image_space.cc
+++ b/runtime/gc/space/image_space.cc
@@ -396,6 +396,13 @@
                                              image_reservation,
                                              error_msg);
     if (space != nullptr) {
+      uint32_t expected_reservation_size =
+          RoundUp(space->GetImageHeader().GetImageSize(), kPageSize);
+      if (!CheckImageReservationSize(*space, expected_reservation_size, error_msg) ||
+          !CheckImageComponentCount(*space, /*expected_component_count=*/ 1u, error_msg)) {
+        return nullptr;
+      }
+
       TimingLogger::ScopedTiming timing("RelocateImage", &logger);
       ImageHeader* image_header = reinterpret_cast<ImageHeader*>(space->GetMemMap()->Begin());
       if (!RelocateInPlace(*image_header,
@@ -581,6 +588,34 @@
     return space;
   }
 
+  static bool CheckImageComponentCount(const ImageSpace& space,
+                                       uint32_t expected_component_count,
+                                       /*out*/std::string* error_msg) {
+    const ImageHeader& header = space.GetImageHeader();
+    if (header.GetComponentCount() != expected_component_count) {
+      *error_msg = StringPrintf("Unexpected component count in %s, received %u, expected %u",
+                                space.GetImageFilename().c_str(),
+                                header.GetComponentCount(),
+                                expected_component_count);
+      return false;
+    }
+    return true;
+  }
+
+  static bool CheckImageReservationSize(const ImageSpace& space,
+                                        uint32_t expected_reservation_size,
+                                        /*out*/std::string* error_msg) {
+    const ImageHeader& header = space.GetImageHeader();
+    if (header.GetImageReservationSize() != expected_reservation_size) {
+      *error_msg = StringPrintf("Unexpected reservation size in %s, received %u, expected %u",
+                                space.GetImageFilename().c_str(),
+                                header.GetImageReservationSize(),
+                                expected_reservation_size);
+      return false;
+    }
+    return true;
+  }
+
  private:
   static MemMap LoadImageFile(const char* image_filename,
                               const char* image_location,
@@ -1261,62 +1296,17 @@
                       /*out*/std::string* error_msg) REQUIRES_SHARED(Locks::mutator_lock_) {
     TimingLogger logger(__PRETTY_FUNCTION__, /*precise=*/ true, VLOG_IS_ON(image));
     std::string filename = GetSystemImageFilename(image_location_.c_str(), image_isa_);
-    std::vector<std::string> locations =
-        ExpandMultiImageLocations(boot_class_path_locations_, image_location_);
-    uint32_t image_start;
-    uint32_t image_end;
-    if (!GetBootImageAddressRange(filename, &image_start, &image_end, error_msg)) {
-      return false;
-    }
-    if (locations.size() > 1u) {
-      std::string last_filename = GetSystemImageFilename(locations.back().c_str(), image_isa_);
-      uint32_t dummy;
-      if (!GetBootImageAddressRange(last_filename, &dummy, &image_end, error_msg)) {
-        return false;
-      }
-    }
-    MemMap image_reservation;
-    MemMap local_extra_reservation;
-    if (!ReserveBootImageMemory(/*reservation_size=*/ image_end - image_start,
-                                image_start,
-                                extra_reservation_size,
-                                &image_reservation,
-                                &local_extra_reservation,
-                                error_msg)) {
-      return false;
-    }
 
-    std::vector<std::unique_ptr<ImageSpace>> spaces;
-    spaces.reserve(locations.size());
-    for (const std::string& location : locations) {
-      filename = GetSystemImageFilename(location.c_str(), image_isa_);
-      spaces.push_back(Load(location, filename, &logger, &image_reservation, error_msg));
-      if (spaces.back() == nullptr) {
-        return false;
-      }
-    }
-    for (size_t i = 0u, size = spaces.size(); i != size; ++i) {
-      std::string expected_boot_class_path =
-          (i == 0u) ? android::base::Join(boot_class_path_locations_, ':') : std::string();
-      if (!OpenOatFile(spaces[i].get(),
-                       boot_class_path_[i],
-                       expected_boot_class_path,
-                       /*validate_oat_file=*/ false,
-                       &logger,
-                       &image_reservation,
-                       error_msg)) {
-        return false;
-      }
-    }
-    if (!CheckReservationExhausted(image_reservation, error_msg)) {
+    if (!LoadFromFile(filename,
+                      /*validate_oat_file=*/ false,
+                      extra_reservation_size,
+                      &logger,
+                      boot_image_spaces,
+                      extra_reservation,
+                      error_msg)) {
       return false;
     }
 
-    MaybeRelocateSpaces(spaces, &logger);
-    InitRuntimeMethods(spaces);
-    boot_image_spaces->swap(spaces);
-    *extra_reservation = std::move(local_extra_reservation);
-
     if (VLOG_IS_ON(image)) {
       LOG(INFO) << "ImageSpace::BootImageLoader::LoadFromSystem exiting "
           << boot_image_spaces->front();
@@ -1333,71 +1323,17 @@
       /*out*/std::string* error_msg) REQUIRES_SHARED(Locks::mutator_lock_) {
     TimingLogger logger(__PRETTY_FUNCTION__, /*precise=*/ true, VLOG_IS_ON(image));
     DCHECK(DalvikCacheExists());
-    std::vector<std::string> locations =
-        ExpandMultiImageLocations(boot_class_path_locations_, image_location_);
-    uint32_t image_start;
-    uint32_t image_end;
-    if (!GetBootImageAddressRange(cache_filename_, &image_start, &image_end, error_msg)) {
-      return false;
-    }
-    if (locations.size() > 1u) {
-      std::string last_filename;
-      if (!GetDalvikCacheFilename(locations.back().c_str(),
-                                  dalvik_cache_.c_str(),
-                                  &last_filename,
-                                  error_msg)) {
-        return false;
-      }
-      uint32_t dummy;
-      if (!GetBootImageAddressRange(last_filename, &dummy, &image_end, error_msg)) {
-        return false;
-      }
-    }
-    MemMap image_reservation;
-    MemMap local_extra_reservation;
-    if (!ReserveBootImageMemory(/*reservation_size=*/ image_end - image_start,
-                                image_start,
-                                extra_reservation_size,
-                                &image_reservation,
-                                &local_extra_reservation,
-                                error_msg)) {
-      return false;
-    }
 
-    std::vector<std::unique_ptr<ImageSpace>> spaces;
-    spaces.reserve(locations.size());
-    for (const std::string& location : locations) {
-      std::string filename;
-      if (!GetDalvikCacheFilename(location.c_str(), dalvik_cache_.c_str(), &filename, error_msg)) {
-        return false;
-      }
-      spaces.push_back(Load(location, filename, &logger, &image_reservation, error_msg));
-      if (spaces.back() == nullptr) {
-        return false;
-      }
-    }
-    for (size_t i = 0u, size = spaces.size(); i != size; ++i) {
-      std::string expected_boot_class_path =
-          (i == 0u) ? android::base::Join(boot_class_path_locations_, ':') : std::string();
-      if (!OpenOatFile(spaces[i].get(),
-                       boot_class_path_[i],
-                       expected_boot_class_path,
-                       validate_oat_file,
-                       &logger,
-                       &image_reservation,
-                       error_msg)) {
-        return false;
-      }
-    }
-    if (!CheckReservationExhausted(image_reservation, error_msg)) {
+    if (!LoadFromFile(cache_filename_,
+                      validate_oat_file,
+                      extra_reservation_size,
+                      &logger,
+                      boot_image_spaces,
+                      extra_reservation,
+                      error_msg)) {
       return false;
     }
 
-    MaybeRelocateSpaces(spaces, &logger);
-    InitRuntimeMethods(spaces);
-    boot_image_spaces->swap(spaces);
-    *extra_reservation = std::move(local_extra_reservation);
-
     if (VLOG_IS_ON(image)) {
       LOG(INFO) << "ImageSpace::BootImageLoader::LoadFromDalvikCache exiting "
           << boot_image_spaces->front();
@@ -1407,6 +1343,81 @@
   }
 
  private:
+  bool LoadFromFile(
+      const std::string& filename,
+      bool validate_oat_file,
+      size_t extra_reservation_size,
+      TimingLogger* logger,
+      /*out*/std::vector<std::unique_ptr<space::ImageSpace>>* boot_image_spaces,
+      /*out*/MemMap* extra_reservation,
+      /*out*/std::string* error_msg) REQUIRES_SHARED(Locks::mutator_lock_) {
+    ImageHeader system_hdr;
+    if (!ReadSpecificImageHeader(filename.c_str(), &system_hdr)) {
+      *error_msg = StringPrintf("Cannot read header of %s", filename.c_str());
+      return false;
+    }
+    if (system_hdr.GetComponentCount() != boot_class_path_.size()) {
+      *error_msg = StringPrintf("Unexpected component count in %s, received %u, expected %zu",
+                                filename.c_str(),
+                                system_hdr.GetComponentCount(),
+                                boot_class_path_.size());
+      return false;
+    }
+    MemMap image_reservation;
+    MemMap local_extra_reservation;
+    if (!ReserveBootImageMemory(system_hdr.GetImageReservationSize(),
+                                reinterpret_cast32<uint32_t>(system_hdr.GetImageBegin()),
+                                extra_reservation_size,
+                                &image_reservation,
+                                &local_extra_reservation,
+                                error_msg)) {
+      return false;
+    }
+
+    std::vector<std::string> locations =
+        ExpandMultiImageLocations(boot_class_path_locations_, image_location_);
+    std::vector<std::string> filenames =
+        ExpandMultiImageLocations(boot_class_path_locations_, filename);
+    DCHECK_EQ(locations.size(), filenames.size());
+    std::vector<std::unique_ptr<ImageSpace>> spaces;
+    spaces.reserve(locations.size());
+    for (std::size_t i = 0u, size = locations.size(); i != size; ++i) {
+      spaces.push_back(Load(locations[i], filenames[i], logger, &image_reservation, error_msg));
+      const ImageSpace* space = spaces.back().get();
+      if (space == nullptr) {
+        return false;
+      }
+      uint32_t expected_component_count = (i == 0u) ? system_hdr.GetComponentCount() : 0u;
+      uint32_t expected_reservation_size = (i == 0u) ? system_hdr.GetImageReservationSize() : 0u;
+      if (!Loader::CheckImageReservationSize(*space, expected_reservation_size, error_msg) ||
+          !Loader::CheckImageComponentCount(*space, expected_component_count, error_msg)) {
+        return false;
+      }
+    }
+    for (size_t i = 0u, size = spaces.size(); i != size; ++i) {
+      std::string expected_boot_class_path =
+          (i == 0u) ? android::base::Join(boot_class_path_locations_, ':') : std::string();
+      if (!OpenOatFile(spaces[i].get(),
+                       boot_class_path_[i],
+                       expected_boot_class_path,
+                       validate_oat_file,
+                       logger,
+                       &image_reservation,
+                       error_msg)) {
+        return false;
+      }
+    }
+    if (!CheckReservationExhausted(image_reservation, error_msg)) {
+      return false;
+    }
+
+    MaybeRelocateSpaces(spaces, logger);
+    InitRuntimeMethods(spaces);
+    boot_image_spaces->swap(spaces);
+    *extra_reservation = std::move(local_extra_reservation);
+    return true;
+  }
+
   template <typename T>
   ALWAYS_INLINE static T* RelocatedAddress(T* src, uint32_t diff) {
     DCHECK(src != nullptr);
@@ -2050,27 +2061,14 @@
     return true;
   }
 
-  bool GetBootImageAddressRange(const std::string& filename,
-                                /*out*/uint32_t* start,
-                                /*out*/uint32_t* end,
-                                /*out*/std::string* error_msg) {
-    ImageHeader system_hdr;
-    if (!ReadSpecificImageHeader(filename.c_str(), &system_hdr)) {
-      *error_msg = StringPrintf("Cannot read header of %s", filename.c_str());
-      return false;
-    }
-    *start = reinterpret_cast32<uint32_t>(system_hdr.GetImageBegin());
-    CHECK_ALIGNED(*start, kPageSize);
-    *end = RoundUp(reinterpret_cast32<uint32_t>(system_hdr.GetOatFileEnd()), kPageSize);
-    return true;
-  }
-
   bool ReserveBootImageMemory(uint32_t reservation_size,
                               uint32_t image_start,
                               size_t extra_reservation_size,
                               /*out*/MemMap* image_reservation,
                               /*out*/MemMap* extra_reservation,
                               /*out*/std::string* error_msg) {
+    DCHECK_ALIGNED(reservation_size, kPageSize);
+    DCHECK_ALIGNED(image_start, kPageSize);
     DCHECK(!image_reservation->IsValid());
     DCHECK_LT(extra_reservation_size, std::numeric_limits<uint32_t>::max() - reservation_size);
     size_t total_size = reservation_size + extra_reservation_size;
diff --git a/runtime/image.cc b/runtime/image.cc
index fb581f9..fe1d88f 100644
--- a/runtime/image.cc
+++ b/runtime/image.cc
@@ -29,9 +29,11 @@
 namespace art {
 
 const uint8_t ImageHeader::kImageMagic[] = { 'a', 'r', 't', '\n' };
-const uint8_t ImageHeader::kImageVersion[] = { '0', '7', '2', '\0' };  // CRC32UpdateBytes intrinsic
+const uint8_t ImageHeader::kImageVersion[] = { '0', '7', '3', '\0' };  // Image reservation.
 
-ImageHeader::ImageHeader(uint32_t image_begin,
+ImageHeader::ImageHeader(uint32_t image_reservation_size,
+                         uint32_t component_count,
+                         uint32_t image_begin,
                          uint32_t image_size,
                          ImageSection* sections,
                          uint32_t image_roots,
@@ -43,7 +45,9 @@
                          uint32_t boot_image_begin,
                          uint32_t boot_image_size,
                          uint32_t pointer_size)
-  : image_begin_(image_begin),
+  : image_reservation_size_(image_reservation_size),
+    component_count_(component_count),
+    image_begin_(image_begin),
     image_size_(image_size),
     image_checksum_(0u),
     oat_checksum_(oat_checksum),
@@ -96,6 +100,9 @@
   if (memcmp(version_, kImageVersion, sizeof(kImageVersion)) != 0) {
     return false;
   }
+  if (!IsAligned<kPageSize>(image_reservation_size_)) {
+    return false;
+  }
   // Unsigned so wraparound is well defined.
   if (image_begin_ >= image_begin_ + image_size_) {
     return false;
diff --git a/runtime/image.h b/runtime/image.h
index 9d98431..140f504 100644
--- a/runtime/image.h
+++ b/runtime/image.h
@@ -138,7 +138,9 @@
   };
 
   ImageHeader() {}
-  ImageHeader(uint32_t image_begin,
+  ImageHeader(uint32_t image_reservation_size,
+              uint32_t component_count,
+              uint32_t image_begin,
               uint32_t image_size,
               ImageSection* sections,
               uint32_t image_roots,
@@ -154,6 +156,14 @@
   bool IsValid() const;
   const char* GetMagic() const;
 
+  uint32_t GetImageReservationSize() const {
+    return image_reservation_size_;
+  }
+
+  uint32_t GetComponentCount() const {
+    return component_count_;
+  }
+
   uint8_t* GetImageBegin() const {
     return reinterpret_cast<uint8_t*>(image_begin_);
   }
@@ -423,6 +433,19 @@
   uint8_t magic_[4];
   uint8_t version_[4];
 
+  // The total memory reservation size for the image.
+  // For boot image or boot image extension, the primary image includes the reservation
+  // for all image files and oat files, secondary images have the reservation set to 0.
+  // App images have reservation equal to `image_size_` rounded up to page size because
+  // their oat files are mmapped independently.
+  uint32_t image_reservation_size_ = 0u;
+
+  // The number of components.
+  // For boot image or boot image extension, the primary image stores the total number
+  // of images, secondary images have this set to 0.
+  // App images have 1 component.
+  uint32_t component_count_ = 0u;
+
   // Required base address for mapping the image.
   uint32_t image_begin_ = 0u;