Add owned section for CompactDex

The owned section is the part of the shared data section owned by
a given dex file. This enables efficiently attributing an offset to
a dex file.

Bug: 74443371
Bug: 63756964
Test: test-art-host

(cherry-picked from commit 2a66809d28c9ee1cbabe52991dbf526095c96c04)

Merged-In: I2de9a281e18b02a20c3dcf5f484eacb591220cdc
Change-Id: I2de9a281e18b02a20c3dcf5f484eacb591220cdc

7 files changed, 80 insertions, 16 deletions

diff --git a/dex2oat/dex2oat_test.cc b/dex2oat/dex2oat_test.cc
index 094dfee3a6..09ff14e4ba 100644
--- a/dex2oat/dex2oat_test.cc
+++ b/dex2oat/dex2oat_test.cc
@@ -1465,14 +1465,14 @@ TEST_F(Dex2oatTest, LayoutSections) {
 
 // Test that generating compact dex works.
 TEST_F(Dex2oatTest, GenerateCompactDex) {
-  std::unique_ptr<const DexFile> dex(OpenTestDexFile("ManyMethods"));
   // Generate a compact dex based odex.
   const std::string dir = GetScratchDir();
   const std::string oat_filename = dir + "/base.oat";
   const std::string vdex_filename = dir + "/base.vdex";
+  const std::string dex_location = GetTestDexFileName("MultiDex");
   std::string error_msg;
   const int res = GenerateOdexForTestWithStatus(
-      {dex->GetLocation()},
+      { dex_location },
       oat_filename,
       CompilerFilter::Filter::kQuicken,
       &error_msg,
@@ -1485,16 +1485,43 @@ TEST_F(Dex2oatTest, GenerateCompactDex) {
                                                    nullptr,
                                                    false,
                                                    /*low_4gb*/false,
-                                                   dex->GetLocation().c_str(),
+                                                   dex_location.c_str(),
                                                    &error_msg));
   ASSERT_TRUE(odex_file != nullptr);
   std::vector<const OatDexFile*> oat_dex_files = odex_file->GetOatDexFiles();
-  ASSERT_EQ(oat_dex_files.size(), 1u);
-  // Check that each dex is a compact dex.
+  ASSERT_GT(oat_dex_files.size(), 1u);
+  // Check that each dex is a compact dex file.
+  std::vector<std::unique_ptr<const CompactDexFile>> compact_dex_files;
   for (const OatDexFile* oat_dex : oat_dex_files) {
     std::unique_ptr<const DexFile> dex_file(oat_dex->OpenDexFile(&error_msg));
     ASSERT_TRUE(dex_file != nullptr) << error_msg;
     ASSERT_TRUE(dex_file->IsCompactDexFile());
+    compact_dex_files.push_back(
+        std::unique_ptr<const CompactDexFile>(dex_file.release()->AsCompactDexFile()));
+  }
+  for (const std::unique_ptr<const CompactDexFile>& dex_file : compact_dex_files) {
+    // Test that every code item is in the owned section.
+    const CompactDexFile::Header& header = dex_file->GetHeader();
+    EXPECT_LE(header.OwnedDataBegin(), header.OwnedDataEnd());
+    EXPECT_LE(header.OwnedDataBegin(), header.data_size_);
+    EXPECT_LE(header.OwnedDataEnd(), header.data_size_);
+    for (uint32_t i = 0; i < dex_file->NumClassDefs(); ++i) {
+      const DexFile::ClassDef& class_def = dex_file->GetClassDef(i);
+      class_def.VisitMethods(dex_file.get(), [&](const ClassDataItemIterator& it) {
+        if (it.GetMethodCodeItemOffset() != 0u) {
+          ASSERT_GE(it.GetMethodCodeItemOffset(), header.OwnedDataBegin());
+          ASSERT_LT(it.GetMethodCodeItemOffset(), header.OwnedDataEnd());
+        }
+      });
+    }
+    // Test that the owned sections don't overlap.
+    for (const std::unique_ptr<const CompactDexFile>& other_dex : compact_dex_files) {
+      if (dex_file != other_dex) {
+        ASSERT_TRUE(
+            (dex_file->GetHeader().OwnedDataBegin() >= other_dex->GetHeader().OwnedDataEnd()) ||
+            (dex_file->GetHeader().OwnedDataEnd() <= other_dex->GetHeader().OwnedDataBegin()));
+      }
+    }
   }
 }
 
diff --git a/dexlayout/compact_dex_writer.cc b/dexlayout/compact_dex_writer.cc
index bd76bf11d3..2b4144c611 100644
--- a/dexlayout/compact_dex_writer.cc
+++ b/dexlayout/compact_dex_writer.cc
@@ -298,6 +298,8 @@ void CompactDexWriter::WriteHeader(Stream* stream) {
   header.class_defs_off_ = collections.ClassDefsOffset();
   header.data_size_ = header_->DataSize();
   header.data_off_ = header_->DataOffset();
+  header.owned_data_begin_ = owned_data_begin_;
+  header.owned_data_end_ = owned_data_end_;
 
   // Compact dex specific flags.
   header.debug_info_offsets_pos_ = debug_info_offsets_pos_;
@@ -426,6 +428,7 @@ bool CompactDexWriter::Write(DexContainer* output, std::string* error_msg)  {
     // Data section.
     data_stream->AlignTo(kDataSectionAlignment);
   }
+  owned_data_begin_ = data_stream->Tell();
 
   // Write code item first to minimize the space required for encoded methods.
   // For cdex, the code items don't depend on the debug info.
@@ -490,6 +493,7 @@ bool CompactDexWriter::Write(DexContainer* output, std::string* error_msg)  {
   WriteDebugInfoOffsetTable(data_stream);
 
   data_stream->AlignTo(kDataSectionAlignment);
+  owned_data_end_ = data_stream->Tell();
   if (compute_offsets_) {
     header_->SetDataSize(data_stream->Tell());
     if (header_->DataSize() != 0) {
@@ -497,7 +501,6 @@ bool CompactDexWriter::Write(DexContainer* output, std::string* error_msg)  {
       main_stream->AlignTo(kDataSectionAlignment);
       // For now, default to saying the data is right after the main stream.
       header_->SetDataOffset(main_stream->Tell());
-      header_->SetDataOffset(0u);
     } else {
       header_->SetDataOffset(0u);
     }
diff --git a/dexlayout/compact_dex_writer.h b/dexlayout/compact_dex_writer.h
index eaf85185f1..4b142a85bb 100644
--- a/dexlayout/compact_dex_writer.h
+++ b/dexlayout/compact_dex_writer.h
@@ -169,6 +169,10 @@ class CompactDexWriter : public DexWriter {
   // Base offset of where debug info starts in the dex file.
   uint32_t debug_info_base_ = 0u;
 
+  // Part of the shared data section owned by this file.
+  uint32_t owned_data_begin_ = 0u;
+  uint32_t owned_data_end_ = 0u;
+
   // State for where we are deduping.
   Deduper* code_item_dedupe_ = nullptr;
   Deduper* data_item_dedupe_ = nullptr;
diff --git a/libdexfile/dex/compact_dex_file.h b/libdexfile/dex/compact_dex_file.h
index 78cd76818a..affc9a20b0 100644
--- a/libdexfile/dex/compact_dex_file.h
+++ b/libdexfile/dex/compact_dex_file.h
@@ -51,6 +51,16 @@ class CompactDexFile : public DexFile {
       return data_size_;
     }
 
+    // Range of the shared data section owned by the dex file. Owned in this context refers to data
+    // for this DEX that was not deduplicated to another DEX.
+    uint32_t OwnedDataBegin() const {
+      return owned_data_begin_;
+    }
+
+    uint32_t OwnedDataEnd() const {
+      return owned_data_end_;
+    }
+
    private:
     uint32_t feature_flags_ = 0u;
 
@@ -63,6 +73,10 @@ class CompactDexFile : public DexFile {
     // Base offset of where debug info starts in the dex file.
     uint32_t debug_info_base_ = 0u;
 
+    // Range of the shared data section owned by the dex file.
+    uint32_t owned_data_begin_ = 0u;
+    uint32_t owned_data_end_ = 0u;
+
     friend class CompactDexFile;
     friend class CompactDexWriter;
   };
diff --git a/libdexfile/dex/dex_file-inl.h b/libdexfile/dex/dex_file-inl.h
index ae0c2f415b..d1b32007c3 100644
--- a/libdexfile/dex/dex_file-inl.h
+++ b/libdexfile/dex/dex_file-inl.h
@@ -515,6 +515,18 @@ inline const uint8_t* DexFile::GetCatchHandlerData(const DexInstructionIterator&
   return handler_data + offset;
 }
 
+template <typename Visitor>
+inline void DexFile::ClassDef::VisitMethods(const DexFile* dex_file, const Visitor& visitor) const {
+  const uint8_t* class_data = dex_file->GetClassData(*this);
+  if (class_data != nullptr) {
+    ClassDataItemIterator it(*dex_file, class_data);
+    it.SkipAllFields();
+    for (; it.HasNext(); it.Next()) {
+      visitor(it);
+    }
+  }
+}
+
 }  // namespace art
 
 #endif  // ART_LIBDEXFILE_DEX_DEX_FILE_INL_H_
diff --git a/libdexfile/dex/dex_file.h b/libdexfile/dex/dex_file.h
index 5560cf19c0..aeb49d2c25 100644
--- a/libdexfile/dex/dex_file.h
+++ b/libdexfile/dex/dex_file.h
@@ -196,6 +196,15 @@ class DexFile {
     DISALLOW_COPY_AND_ASSIGN(MethodId);
   };
 
+  // Base code_item, compact dex and standard dex have different code item layouts.
+  struct CodeItem {
+   protected:
+    CodeItem() = default;
+
+   private:
+    DISALLOW_COPY_AND_ASSIGN(CodeItem);
+  };
+
   // Raw class_def_item.
   struct ClassDef {
     dex::TypeIndex class_idx_;  // index into type_ids_ array for this class
@@ -227,6 +236,9 @@ class DexFile {
       }
     }
 
+    template <typename Visitor>
+    void VisitMethods(const DexFile* dex_file, const Visitor& visitor) const;
+
    private:
     DISALLOW_COPY_AND_ASSIGN(ClassDef);
   };
@@ -300,15 +312,6 @@ class DexFile {
     DISALLOW_COPY_AND_ASSIGN(CallSiteIdItem);
   };
 
-  // Base code_item, compact dex and standard dex have different code item layouts.
-  struct CodeItem {
-   protected:
-    CodeItem() = default;
-
-   private:
-    DISALLOW_COPY_AND_ASSIGN(CodeItem);
-  };
-
   // Raw try_item.
   struct TryItem {
     static constexpr size_t kAlignment = sizeof(uint32_t);
diff --git a/runtime/vdex_file.h b/runtime/vdex_file.h
index 77e1f2ccfe..326fcbc1fe 100644
--- a/runtime/vdex_file.h
+++ b/runtime/vdex_file.h
@@ -97,7 +97,8 @@ class VdexFile {
 
     // The format version of the dex section header and the dex section, containing
     // both the dex code and the quickening data.
-    static constexpr uint8_t kDexSectionVersion[] = { '0', '0', '1', '\0' };
+    // Last update: Add owned section for CompactDex.
+    static constexpr uint8_t kDexSectionVersion[] = { '0', '0', '2', '\0' };
 
     // If the .vdex file has no dex section (hence no dex code nor quickening data),
     // we encode this magic version.