Add image compressed blocks

Add support for splitting the image into a set of solid blocks.

Added dex2oat option --max-image-block-size and correspodning image
unit test.

Motivation: Enable parallel image decompression in the future.

Bug: 116052292
Test: test-art-host

Change-Id: I37c6c6a43ef94c4a62bf38a0cf51f26ce06347ac
diff --git a/compiler/driver/compiler_options.cc b/compiler/driver/compiler_options.cc
index bef5be1..b28c7e0 100644
--- a/compiler/driver/compiler_options.cc
+++ b/compiler/driver/compiler_options.cc
@@ -71,6 +71,7 @@
       count_hotness_in_compiled_code_(false),
       resolve_startup_const_strings_(false),
       check_profiled_methods_(ProfileMethodsCheck::kNone),
+      max_image_block_size_(std::numeric_limits<uint32_t>::max()),
       register_allocation_strategy_(RegisterAllocator::kRegisterAllocatorDefault),
       passes_to_run_(nullptr) {
 }
diff --git a/compiler/driver/compiler_options.h b/compiler/driver/compiler_options.h
index f0970a9..17a779c 100644
--- a/compiler/driver/compiler_options.h
+++ b/compiler/driver/compiler_options.h
@@ -335,6 +335,14 @@
     return check_profiled_methods_;
   }
 
+  uint32_t MaxImageBlockSize() const {
+    return max_image_block_size_;
+  }
+
+  void SetMaxImageBlockSize(uint32_t size) {
+    max_image_block_size_ = size;
+  }
+
  private:
   bool ParseDumpInitFailures(const std::string& option, std::string* error_msg);
   void ParseDumpCfgPasses(const StringPiece& option, UsageFn Usage);
@@ -424,6 +432,9 @@
   // up compiled and are not punted.
   ProfileMethodsCheck check_profiled_methods_;
 
+  // Maximum solid block size in the generated image.
+  uint32_t max_image_block_size_;
+
   RegisterAllocator::Strategy register_allocation_strategy_;
 
   // If not null, specifies optimization passes which will be run instead of defaults.
diff --git a/compiler/driver/compiler_options_map-inl.h b/compiler/driver/compiler_options_map-inl.h
index c7334a7..7e2a64b 100644
--- a/compiler/driver/compiler_options_map-inl.h
+++ b/compiler/driver/compiler_options_map-inl.h
@@ -84,6 +84,7 @@
   if (map.Exists(Base::CheckProfiledMethods)) {
     options->check_profiled_methods_ = *map.Get(Base::CheckProfiledMethods);
   }
+  map.AssignIfExists(Base::MaxImageBlockSize, &options->max_image_block_size_);
 
   if (map.Exists(Base::DumpTimings)) {
     options->dump_timings_ = true;
@@ -201,7 +202,11 @@
 
       .Define("--verbose-methods=_")
           .template WithType<ParseStringList<','>>()
-          .IntoKey(Map::VerboseMethods);
+          .IntoKey(Map::VerboseMethods)
+
+      .Define("--max-image-block-size=_")
+          .template WithType<unsigned int>()
+          .IntoKey(Map::MaxImageBlockSize);
 }
 
 #pragma GCC diagnostic pop
diff --git a/compiler/driver/compiler_options_map.def b/compiler/driver/compiler_options_map.def
index c2fac5e..0a9c873 100644
--- a/compiler/driver/compiler_options_map.def
+++ b/compiler/driver/compiler_options_map.def
@@ -65,5 +65,6 @@
 COMPILER_OPTIONS_KEY (Unit,                        DumpTimings)
 COMPILER_OPTIONS_KEY (Unit,                        DumpPassTimings)
 COMPILER_OPTIONS_KEY (Unit,                        DumpStats)
+COMPILER_OPTIONS_KEY (unsigned int,                MaxImageBlockSize)
 
 #undef COMPILER_OPTIONS_KEY
diff --git a/dex2oat/dex2oat.cc b/dex2oat/dex2oat.cc
index 3a24542..31832b4 100644
--- a/dex2oat/dex2oat.cc
+++ b/dex2oat/dex2oat.cc
@@ -482,6 +482,8 @@
   UsageError("  --resolve-startup-const-strings=true|false: If true, the compiler eagerly");
   UsageError("      resolves strings referenced from const-string of startup methods.");
   UsageError("");
+  UsageError("  --max-image-block-size=<size>: Maximum solid block size for compressed images.");
+  UsageError("");
   UsageError("      Example: --compilation-reason=install");
   UsageError("");
   std::cerr << "See log for usage error information\n";
diff --git a/dex2oat/linker/image_test.cc b/dex2oat/linker/image_test.cc
index 64b98cd..ebd829b 100644
--- a/dex2oat/linker/image_test.cc
+++ b/dex2oat/linker/image_test.cc
@@ -32,7 +32,11 @@
   // Compile multi-image with ImageLayoutA being the last image.
   {
     CompilationHelper helper;
-    Compile(ImageHeader::kStorageModeUncompressed, helper, "ImageLayoutA", {"LMyClass;"});
+    Compile(ImageHeader::kStorageModeUncompressed,
+            /*max_image_block_size=*/std::numeric_limits<uint32_t>::max(),
+            helper,
+            "ImageLayoutA",
+            {"LMyClass;"});
     image_sizes = helper.GetImageObjectSectionSizes();
   }
   TearDown();
@@ -41,7 +45,11 @@
   // Compile multi-image with ImageLayoutB being the last image.
   {
     CompilationHelper helper;
-    Compile(ImageHeader::kStorageModeUncompressed, helper, "ImageLayoutB", {"LMyClass;"});
+    Compile(ImageHeader::kStorageModeUncompressed,
+            /*max_image_block_size=*/std::numeric_limits<uint32_t>::max(),
+            helper,
+            "ImageLayoutB",
+            {"LMyClass;"});
     image_sizes_extra = helper.GetImageObjectSectionSizes();
   }
   // Make sure that the new stuff in the clinit in ImageLayoutB is in the last image and not in the
@@ -76,9 +84,8 @@
                              oat_file_end,
                              /*boot_image_begin=*/ 0u,
                              /*boot_image_size=*/ 0u,
-                             sizeof(void*),
-                             ImageHeader::kDefaultStorageMode,
-                             /*data_size=*/ 0u);
+                             sizeof(void*));
+
     ASSERT_TRUE(image_header.IsValid());
     ASSERT_TRUE(!image_header.IsAppImage());
 
@@ -97,6 +104,7 @@
 TEST_F(ImageTest, TestDefaultMethods) {
   CompilationHelper helper;
   Compile(ImageHeader::kStorageModeUncompressed,
+          /*max_image_block_size=*/std::numeric_limits<uint32_t>::max(),
           helper,
           "DefaultMethods",
           {"LIface;", "LImpl;", "LIterableBase;"});
@@ -156,6 +164,7 @@
 TEST_F(ImageTest, TestSoftVerificationFailureDuringClassInitialization) {
   CompilationHelper helper;
   Compile(ImageHeader::kStorageModeUncompressed,
+          /*max_image_block_size=*/std::numeric_limits<uint32_t>::max(),
           helper,
           "VerifySoftFailDuringClinit",
           /*image_classes=*/ {"LClassToInitialize;"},
diff --git a/dex2oat/linker/image_test.h b/dex2oat/linker/image_test.h
index c90eadd..1400779 100644
--- a/dex2oat/linker/image_test.h
+++ b/dex2oat/linker/image_test.h
@@ -77,9 +77,10 @@
     CommonCompilerTest::SetUp();
   }
 
-  void TestWriteRead(ImageHeader::StorageMode storage_mode);
+  void TestWriteRead(ImageHeader::StorageMode storage_mode, uint32_t max_image_block_size);
 
   void Compile(ImageHeader::StorageMode storage_mode,
+               uint32_t max_image_block_size,
                /*out*/ CompilationHelper& out_helper,
                const std::string& extra_dex = "",
                const std::initializer_list<std::string>& image_classes = {},
@@ -374,6 +375,7 @@
 
 inline void ImageTest::Compile(
     ImageHeader::StorageMode storage_mode,
+    uint32_t max_image_block_size,
     CompilationHelper& helper,
     const std::string& extra_dex,
     const std::initializer_list<std::string>& image_classes,
@@ -388,6 +390,7 @@
   CreateCompilerDriver();
   // Set inline filter values.
   compiler_options_->SetInlineMaxCodeUnits(CompilerOptions::kDefaultInlineMaxCodeUnits);
+  compiler_options_->SetMaxImageBlockSize(max_image_block_size);
   image_classes_.clear();
   if (!extra_dex.empty()) {
     helper.extra_dex_files = OpenTestDexFiles(extra_dex.c_str());
@@ -411,9 +414,10 @@
   }
 }
 
-inline void ImageTest::TestWriteRead(ImageHeader::StorageMode storage_mode) {
+inline void ImageTest::TestWriteRead(ImageHeader::StorageMode storage_mode,
+                                     uint32_t max_image_block_size) {
   CompilationHelper helper;
-  Compile(storage_mode, /*out*/ helper);
+  Compile(storage_mode, max_image_block_size, /*out*/ helper);
   std::vector<uint64_t> image_file_sizes;
   for (ScratchFile& image_file : helper.image_files) {
     std::unique_ptr<File> file(OS::OpenFileForReading(image_file.GetFilename().c_str()));
@@ -488,6 +492,10 @@
     } else if (image_file_size > 16 * KB) {
       // Compressed, file should be smaller than image. Not really valid for small images.
       ASSERT_LE(image_file_size, image_space->GetImageHeader().GetImageSize());
+      // TODO: Actually validate the blocks, this is hard since the blocks are not copied over for
+      // compressed images. Add kPageSize since image_size is rounded up to this.
+      ASSERT_GT(image_space->GetImageHeader().GetBlockCount() * max_image_block_size,
+                image_space->GetImageHeader().GetImageSize() - kPageSize);
     }
 
     image_space->VerifyImageAllocations();
diff --git a/dex2oat/linker/image_write_read_test.cc b/dex2oat/linker/image_write_read_test.cc
index 30996b5..5ddbd09 100644
--- a/dex2oat/linker/image_write_read_test.cc
+++ b/dex2oat/linker/image_write_read_test.cc
@@ -20,15 +20,23 @@
 namespace linker {
 
 TEST_F(ImageTest, WriteReadUncompressed) {
-  TestWriteRead(ImageHeader::kStorageModeUncompressed);
+  TestWriteRead(ImageHeader::kStorageModeUncompressed,
+                /*max_image_block_size=*/std::numeric_limits<uint32_t>::max());
 }
 
 TEST_F(ImageTest, WriteReadLZ4) {
-  TestWriteRead(ImageHeader::kStorageModeLZ4);
+  TestWriteRead(ImageHeader::kStorageModeLZ4,
+                /*max_image_block_size=*/std::numeric_limits<uint32_t>::max());
 }
 
 TEST_F(ImageTest, WriteReadLZ4HC) {
-  TestWriteRead(ImageHeader::kStorageModeLZ4HC);
+  TestWriteRead(ImageHeader::kStorageModeLZ4HC,
+                /*max_image_block_size=*/std::numeric_limits<uint32_t>::max());
+}
+
+
+TEST_F(ImageTest, WriteReadLZ4HCKBBlock) {
+  TestWriteRead(ImageHeader::kStorageModeLZ4HC, /*max_image_block_size=*/KB);
 }
 
 }  // namespace linker
diff --git a/dex2oat/linker/image_writer.cc b/dex2oat/linker/image_writer.cc
index 75b3555..3ab1cf9 100644
--- a/dex2oat/linker/image_writer.cc
+++ b/dex2oat/linker/image_writer.cc
@@ -745,32 +745,93 @@
 
     // Image data size excludes the bitmap and the header.
     ImageHeader* const image_header = reinterpret_cast<ImageHeader*>(image_info.image_.Begin());
-    ArrayRef<const uint8_t> raw_image_data(image_info.image_.Begin() + sizeof(ImageHeader),
-                                           image_header->GetImageSize() - sizeof(ImageHeader));
 
-    CHECK_EQ(image_header->storage_mode_, image_storage_mode_);
-    std::vector<uint8_t> compressed_data;
-    ArrayRef<const uint8_t> image_data =
-        MaybeCompressData(raw_image_data, image_storage_mode_, &compressed_data);
-    image_header->data_size_ = image_data.size();  // Fill in the data size.
+    // Block sources (from the image).
+    const bool is_compressed = image_storage_mode_ != ImageHeader::kStorageModeUncompressed;
+    std::vector<std::pair<uint32_t, uint32_t>> block_sources;
+    std::vector<ImageHeader::Block> blocks;
 
-    // Write out the image + fields + methods.
-    if (!image_file->PwriteFully(image_data.data(), image_data.size(), sizeof(ImageHeader))) {
-      PLOG(ERROR) << "Failed to write image file data " << image_filename;
-      return false;
+    // Add a set of solid blocks such that no block is larger than the maximum size. A solid block
+    // is a block that must be decompressed all at once.
+    auto add_blocks = [&](uint32_t offset, uint32_t size) {
+      while (size != 0u) {
+        const uint32_t cur_size = std::min(size, compiler_options_.MaxImageBlockSize());
+        block_sources.emplace_back(offset, cur_size);
+        offset += cur_size;
+        size -= cur_size;
+      }
+    };
+
+    add_blocks(sizeof(ImageHeader), image_header->GetImageSize() - sizeof(ImageHeader));
+
+    // Checksum of compressed image data and header.
+    uint32_t image_checksum = adler32(0L, Z_NULL, 0);
+    image_checksum = adler32(image_checksum,
+                             reinterpret_cast<const uint8_t*>(image_header),
+                             sizeof(ImageHeader));
+    // Copy and compress blocks.
+    size_t out_offset = sizeof(ImageHeader);
+    for (const std::pair<uint32_t, uint32_t> block : block_sources) {
+      ArrayRef<const uint8_t> raw_image_data(image_info.image_.Begin() + block.first,
+                                             block.second);
+      std::vector<uint8_t> compressed_data;
+      ArrayRef<const uint8_t> image_data =
+          MaybeCompressData(raw_image_data, image_storage_mode_, &compressed_data);
+
+      if (!is_compressed) {
+        // For uncompressed, preserve alignment since the image will be directly mapped.
+        out_offset = block.first;
+      }
+
+      // Fill in the compressed location of the block.
+      blocks.emplace_back(ImageHeader::Block(
+          image_storage_mode_,
+          /*data_offset=*/ out_offset,
+          /*data_size=*/ image_data.size(),
+          /*image_offset=*/ block.first,
+          /*image_size=*/ block.second));
+
+      // Write out the image + fields + methods.
+      if (!image_file->PwriteFully(image_data.data(), image_data.size(), out_offset)) {
+        PLOG(ERROR) << "Failed to write image file data " << image_filename;
+        image_file->Erase();
+        return false;
+      }
+      out_offset += image_data.size();
+      image_checksum = adler32(image_checksum, image_data.data(), image_data.size());
     }
 
-    // Write out the image bitmap at the page aligned start of the image end, also uncompressed for
-    // convenience.
-    const ImageSection& bitmap_section = image_header->GetImageBitmapSection();
+    // Write the block metadata directly after the image sections.
+    // Note: This is not part of the mapped image and is not preserved after decompressing, it's
+    // only used for image loading. For this reason, only write it out for compressed images.
+    if (is_compressed) {
+      // Align up since the compressed data is not necessarily aligned.
+      out_offset = RoundUp(out_offset, alignof(ImageHeader::Block));
+      CHECK(!blocks.empty());
+      const size_t blocks_bytes = blocks.size() * sizeof(blocks[0]);
+      if (!image_file->PwriteFully(&blocks[0], blocks_bytes, out_offset)) {
+        PLOG(ERROR) << "Failed to write image blocks " << image_filename;
+        image_file->Erase();
+        return false;
+      }
+      image_header->blocks_offset_ = out_offset;
+      image_header->blocks_count_ = blocks.size();
+      out_offset += blocks_bytes;
+    }
+
+    // Data size includes everything except the bitmap.
+    image_header->data_size_ = out_offset - sizeof(ImageHeader);
+
+    // Update and write the bitmap section. Note that the bitmap section is relative to the
+    // possibly compressed image.
+    ImageSection& bitmap_section = image_header->GetImageSection(ImageHeader::kSectionImageBitmap);
     // Align up since data size may be unaligned if the image is compressed.
-    size_t bitmap_position_in_file = RoundUp(sizeof(ImageHeader) + image_data.size(), kPageSize);
-    if (image_storage_mode_ == ImageHeader::kDefaultStorageMode) {
-      CHECK_EQ(bitmap_position_in_file, bitmap_section.Offset());
-    }
+    out_offset = RoundUp(out_offset, kPageSize);
+    bitmap_section = ImageSection(out_offset, bitmap_section.Size());
+
     if (!image_file->PwriteFully(image_info.image_bitmap_->Begin(),
                                  bitmap_section.Size(),
-                                 bitmap_position_in_file)) {
+                                 bitmap_section.Offset())) {
       PLOG(ERROR) << "Failed to write image file bitmap " << image_filename;
       return false;
     }
@@ -781,22 +842,17 @@
       return false;
     }
 
-    // Calculate the image checksum.
-    uint32_t image_checksum = adler32(0L, Z_NULL, 0);
-    image_checksum = adler32(image_checksum,
-                             reinterpret_cast<const uint8_t*>(image_header),
-                             sizeof(ImageHeader));
-    image_checksum = adler32(image_checksum, image_data.data(), image_data.size());
+    // Calculate the image checksum of the remaining data.
     image_checksum = adler32(image_checksum,
                              reinterpret_cast<const uint8_t*>(image_info.image_bitmap_->Begin()),
                              bitmap_section.Size());
     image_header->SetImageChecksum(image_checksum);
 
     if (VLOG_IS_ON(compiler)) {
-      size_t separately_written_section_size = bitmap_section.Size() + sizeof(ImageHeader);
-
-      size_t total_uncompressed_size = raw_image_data.size() + separately_written_section_size,
-             total_compressed_size   = image_data.size() + separately_written_section_size;
+      const size_t separately_written_section_size = bitmap_section.Size();
+      const size_t total_uncompressed_size = image_info.image_size_ +
+          separately_written_section_size;
+      const size_t total_compressed_size = out_offset + separately_written_section_size;
 
       VLOG(compiler) << "Dex2Oat:uncompressedImageSize = " << total_uncompressed_size;
       if (total_uncompressed_size != total_compressed_size) {
@@ -804,8 +860,8 @@
       }
     }
 
-    CHECK_EQ(bitmap_position_in_file + bitmap_section.Size(),
-             static_cast<size_t>(image_file->GetLength()));
+    CHECK_EQ(bitmap_section.End(), static_cast<size_t>(image_file->GetLength()))
+        << "Bitmap should be at the end of the file";
 
     // Write header last in case the compiler gets killed in the middle of image writing.
     // We do not want to have a corrupted image with a valid header.
@@ -2574,9 +2630,7 @@
       PointerToLowMemUInt32(oat_file_end),
       boot_image_begin,
       boot_oat_end - boot_image_begin,
-      static_cast<uint32_t>(target_ptr_size_),
-      image_storage_mode_,
-      /*data_size*/0u);
+      static_cast<uint32_t>(target_ptr_size_));
 }
 
 ArtMethod* ImageWriter::GetImageMethodAddress(ArtMethod* method) {
diff --git a/oatdump/oatdump.cc b/oatdump/oatdump.cc
index 51f6008..1c74a92 100644
--- a/oatdump/oatdump.cc
+++ b/oatdump/oatdump.cc
@@ -1936,7 +1936,7 @@
       stats_.file_bytes = file->GetLength();
       // If the image is compressed, adjust to decompressed size.
       size_t uncompressed_size = image_header_.GetImageSize() - sizeof(ImageHeader);
-      if (image_header_.GetStorageMode() == ImageHeader::kStorageModeUncompressed) {
+      if (image_header_.HasCompressedBlock()) {
         DCHECK_EQ(uncompressed_size, data_size) << "Sizes should match for uncompressed image";
       }
       stats_.file_bytes += uncompressed_size - data_size;
diff --git a/runtime/gc/collector/immune_spaces_test.cc b/runtime/gc/collector/immune_spaces_test.cc
index 9f98f6c..c29b79c 100644
--- a/runtime/gc/collector/immune_spaces_test.cc
+++ b/runtime/gc/collector/immune_spaces_test.cc
@@ -124,9 +124,7 @@
         /*oat_file_end=*/ PointerToLowMemUInt32(oat_map.Begin() + oat_size),
         /*boot_image_begin=*/ 0u,
         /*boot_image_size=*/ 0u,
-        /*pointer_size=*/ sizeof(void*),
-        ImageHeader::kStorageModeUncompressed,
-        /*data_size=*/ 0u);
+        /*pointer_size=*/ sizeof(void*));
     return new DummyImageSpace(std::move(image_map),
                                std::move(live_bitmap),
                                std::move(oat_file),
diff --git a/runtime/gc/space/image_space.cc b/runtime/gc/space/image_space.cc
index bfb3746..5f131d3 100644
--- a/runtime/gc/space/image_space.cc
+++ b/runtime/gc/space/image_space.cc
@@ -16,7 +16,6 @@
 
 #include "image_space.h"
 
-#include <lz4.h>
 #include <sys/statvfs.h>
 #include <sys/types.h>
 #include <unistd.h>
@@ -466,9 +465,10 @@
     // Check that the file is larger or equal to the header size + data size.
     const uint64_t image_file_size = static_cast<uint64_t>(file->GetLength());
     if (image_file_size < sizeof(ImageHeader) + image_header->GetDataSize()) {
-      *error_msg = StringPrintf("Image file truncated: %" PRIu64 " vs. %" PRIu64 ".",
-                                image_file_size,
-                                sizeof(ImageHeader) + image_header->GetDataSize());
+      *error_msg = StringPrintf(
+          "Image file truncated: %" PRIu64 " vs. %" PRIu64 ".",
+           image_file_size,
+           static_cast<uint64_t>(sizeof(ImageHeader) + image_header->GetDataSize()));
       return nullptr;
     }
 
@@ -588,8 +588,9 @@
                               /*inout*/MemMap* image_reservation,
                               /*out*/std::string* error_msg) {
     TimingLogger::ScopedTiming timing("MapImageFile", logger);
-    const ImageHeader::StorageMode storage_mode = image_header.GetStorageMode();
-    if (storage_mode == ImageHeader::kStorageModeUncompressed) {
+    std::string temp_error_msg;
+    const bool is_compressed = image_header.HasCompressedBlock();
+    if (!is_compressed) {
       uint8_t* address = (image_reservation != nullptr) ? image_reservation->Begin() : nullptr;
       return MemMap::MapFileAtAddress(address,
                                       image_header.GetImageSize(),
@@ -604,15 +605,6 @@
                                       error_msg);
     }
 
-    if (storage_mode != ImageHeader::kStorageModeLZ4 &&
-        storage_mode != ImageHeader::kStorageModeLZ4HC) {
-      if (error_msg != nullptr) {
-        *error_msg = StringPrintf("Invalid storage mode in image header %d",
-                                  static_cast<int>(storage_mode));
-      }
-      return MemMap::Invalid();
-    }
-
     // Reserve output and decompress into it.
     MemMap map = MemMap::MapAnonymous(image_location,
                                       image_header.GetImageSize(),
@@ -622,7 +614,6 @@
                                       error_msg);
     if (map.IsValid()) {
       const size_t stored_size = image_header.GetDataSize();
-      const size_t decompress_offset = sizeof(ImageHeader);  // Skip the header.
       MemMap temp_map = MemMap::MapFile(sizeof(ImageHeader) + stored_size,
                                         PROT_READ,
                                         MAP_PRIVATE,
@@ -637,27 +628,20 @@
       }
       memcpy(map.Begin(), &image_header, sizeof(ImageHeader));
       const uint64_t start = NanoTime();
-      // LZ4HC and LZ4 have same internal format, both use LZ4_decompress.
-      TimingLogger::ScopedTiming timing2("LZ4 decompress image", logger);
-      const size_t decompressed_size = LZ4_decompress_safe(
-          reinterpret_cast<char*>(temp_map.Begin()) + sizeof(ImageHeader),
-          reinterpret_cast<char*>(map.Begin()) + decompress_offset,
-          stored_size,
-          map.Size() - decompress_offset);
+      for (const ImageHeader::Block& block : image_header.GetBlocks(temp_map.Begin())) {
+        TimingLogger::ScopedTiming timing2("LZ4 decompress image", logger);
+        if (!block.Decompress(/*out_ptr=*/map.Begin(), /*in_ptr=*/temp_map.Begin(), error_msg)) {
+          if (error_msg != nullptr) {
+            *error_msg = "Failed to decompress image block " + *error_msg;
+          }
+          return MemMap::Invalid();
+        }
+      }
       const uint64_t time = NanoTime() - start;
       // Add one 1 ns to prevent possible divide by 0.
       VLOG(image) << "Decompressing image took " << PrettyDuration(time) << " ("
                   << PrettySize(static_cast<uint64_t>(map.Size()) * MsToNs(1000) / (time + 1))
                   << "/s)";
-      if (decompressed_size + sizeof(ImageHeader) != image_header.GetImageSize()) {
-        if (error_msg != nullptr) {
-          *error_msg = StringPrintf(
-              "Decompressed size does not match expected image size %zu vs %zu",
-              decompressed_size + sizeof(ImageHeader),
-              image_header.GetImageSize());
-        }
-        return MemMap::Invalid();
-      }
     }
 
     return map;
@@ -766,6 +750,7 @@
 
     ALWAYS_INLINE void operator()(ObjPtr<mirror::Object> obj,
                                   MemberOffset offset,
+
                                   bool is_static ATTRIBUTE_UNUSED) const
         NO_THREAD_SAFETY_ANALYSIS {
       // There could be overlap between ranges, we must avoid visiting the same reference twice.
diff --git a/runtime/image.cc b/runtime/image.cc
index f50c39c..ae3d8e3 100644
--- a/runtime/image.cc
+++ b/runtime/image.cc
@@ -16,6 +16,9 @@
 
 #include "image.h"
 
+#include <lz4.h>
+#include <sstream>
+
 #include "base/bit_utils.h"
 #include "base/length_prefixed_array.h"
 #include "base/utils.h"
@@ -26,7 +29,7 @@
 namespace art {
 
 const uint8_t ImageHeader::kImageMagic[] = { 'a', 'r', 't', '\n' };
-const uint8_t ImageHeader::kImageVersion[] = { '0', '7', '0', '\0' };  // Store ImtIndex.
+const uint8_t ImageHeader::kImageVersion[] = { '0', '7', '1', '\0' };  // Add image blocks.
 
 ImageHeader::ImageHeader(uint32_t image_begin,
                          uint32_t image_size,
@@ -39,9 +42,7 @@
                          uint32_t oat_file_end,
                          uint32_t boot_image_begin,
                          uint32_t boot_image_size,
-                         uint32_t pointer_size,
-                         StorageMode storage_mode,
-                         size_t data_size)
+                         uint32_t pointer_size)
   : image_begin_(image_begin),
     image_size_(image_size),
     image_checksum_(0u),
@@ -53,9 +54,7 @@
     boot_image_begin_(boot_image_begin),
     boot_image_size_(boot_image_size),
     image_roots_(image_roots),
-    pointer_size_(pointer_size),
-    storage_mode_(storage_mode),
-    data_size_(data_size) {
+    pointer_size_(pointer_size) {
   CHECK_EQ(image_begin, RoundUp(image_begin, kPageSize));
   CHECK_EQ(oat_file_begin, RoundUp(oat_file_begin, kPageSize));
   CHECK_EQ(oat_data_begin, RoundUp(oat_data_begin, kPageSize));
@@ -144,4 +143,34 @@
   return ConvertToPointerSize(pointer_size_);
 }
 
+bool ImageHeader::Block::Decompress(uint8_t* out_ptr,
+                                    const uint8_t* in_ptr,
+                                    std::string* error_msg) const {
+  switch (storage_mode_) {
+    case kStorageModeUncompressed: {
+      CHECK_EQ(image_size_, data_size_);
+      memcpy(out_ptr + image_offset_, in_ptr + data_offset_, data_size_);
+      break;
+    }
+    case kStorageModeLZ4:
+    case kStorageModeLZ4HC: {
+      // LZ4HC and LZ4 have same internal format, both use LZ4_decompress.
+      const size_t decompressed_size = LZ4_decompress_safe(
+          reinterpret_cast<const char*>(in_ptr) + data_offset_,
+          reinterpret_cast<char*>(out_ptr) + image_offset_,
+          data_size_,
+          image_size_);
+      CHECK_EQ(decompressed_size, image_size_);
+      break;
+    }
+    default: {
+      if (error_msg != nullptr) {
+        *error_msg = (std::ostringstream() << "Invalid image format " << storage_mode_).str();
+      }
+      return false;
+    }
+  }
+  return true;
+}
+
 }  // namespace art
diff --git a/runtime/image.h b/runtime/image.h
index f33b9b2..76fb3b7 100644
--- a/runtime/image.h
+++ b/runtime/image.h
@@ -21,6 +21,7 @@
 
 #include "base/enums.h"
 #include "base/globals.h"
+#include "base/iteration_range.h"
 #include "mirror/object.h"
 
 namespace art {
@@ -82,8 +83,10 @@
   uint32_t size_;
 };
 
-// header of image files written by ImageWriter, read and validated by Space.
-class PACKED(4) ImageHeader {
+// Header of image files written by ImageWriter, read and validated by Space.
+// Packed to object alignment since the first object follows directly after the header.
+static_assert(kObjectAlignment == 8, "Alignment check");
+class PACKED(8) ImageHeader {
  public:
   enum StorageMode : uint32_t {
     kStorageModeUncompressed,
@@ -93,8 +96,40 @@
   };
   static constexpr StorageMode kDefaultStorageMode = kStorageModeUncompressed;
 
-  ImageHeader() {}
+  // Solid block of the image. May be compressed or uncompressed.
+  class PACKED(4) Block final {
+   public:
+    Block(StorageMode storage_mode,
+          uint32_t data_offset,
+          uint32_t data_size,
+          uint32_t image_offset,
+          uint32_t image_size)
+        : storage_mode_(storage_mode),
+          data_offset_(data_offset),
+          data_size_(data_size),
+          image_offset_(image_offset),
+          image_size_(image_size) {}
 
+    bool Decompress(uint8_t* out_ptr, const uint8_t* in_ptr, std::string* error_msg) const;
+
+    StorageMode GetStorageMode() const {
+      return storage_mode_;
+    }
+
+   private:
+    // Storage method for the image, the image may be compressed.
+    StorageMode storage_mode_ = kDefaultStorageMode;
+
+    // Compressed offset and size.
+    uint32_t data_offset_ = 0u;
+    uint32_t data_size_ = 0u;
+
+    // Image offset and size (decompressed or mapped location).
+    uint32_t image_offset_ = 0u;
+    uint32_t image_size_ = 0u;
+  };
+
+  ImageHeader() {}
   ImageHeader(uint32_t image_begin,
               uint32_t image_size,
               ImageSection* sections,
@@ -106,9 +141,7 @@
               uint32_t oat_file_end,
               uint32_t boot_image_begin,
               uint32_t boot_image_size,
-              uint32_t pointer_size,
-              StorageMode storage_mode,
-              size_t data_size);
+              uint32_t pointer_size);
 
   bool IsValid() const;
   const char* GetMagic() const;
@@ -231,6 +264,11 @@
 
   ArtMethod* GetImageMethod(ImageMethod index) const;
 
+  ImageSection& GetImageSection(ImageSections index) {
+    DCHECK_LT(static_cast<size_t>(index), kSectionCount);
+    return sections_[index];
+  }
+
   const ImageSection& GetImageSection(ImageSections index) const {
     DCHECK_LT(static_cast<size_t>(index), kSectionCount);
     return sections_[index];
@@ -304,10 +342,6 @@
     return boot_image_size_;
   }
 
-  StorageMode GetStorageMode() const {
-    return storage_mode_;
-  }
-
   uint64_t GetDataSize() const {
     return data_size_;
   }
@@ -345,6 +379,24 @@
                                     uint8_t* base,
                                     PointerSize pointer_size) const;
 
+  IterationRange<const Block*> GetBlocks() const {
+    return GetBlocks(GetImageBegin());
+  }
+
+  IterationRange<const Block*> GetBlocks(const uint8_t* image_begin) const {
+    const Block* begin = reinterpret_cast<const Block*>(image_begin + blocks_offset_);
+    return {begin, begin + blocks_count_};
+  }
+
+  // Return true if the image has any compressed blocks.
+  bool HasCompressedBlock() const {
+    return blocks_count_ != 0u;
+  }
+
+  uint32_t GetBlockCount() const {
+    return blocks_count_;
+  }
+
  private:
   static const uint8_t kImageMagic[4];
   static const uint8_t kImageVersion[4];
@@ -404,13 +456,14 @@
   // Image methods, may be inside of the boot image for app images.
   uint64_t image_methods_[kImageMethodsCount];
 
-  // Storage method for the image, the image may be compressed.
-  StorageMode storage_mode_ = kDefaultStorageMode;
-
   // Data size for the image data excluding the bitmap and the header. For compressed images, this
   // is the compressed size in the file.
   uint32_t data_size_ = 0u;
 
+  // Image blocks, only used for compressed images.
+  uint32_t blocks_offset_ = 0u;
+  uint32_t blocks_count_ = 0u;
+
   friend class linker::ImageWriter;
 };