1 files changed, 43 insertions, 20 deletions

diff --git a/llvm/lib/CGData/CodeGenDataReader.cpp b/llvm/lib/CGData/CodeGenDataReader.cpp
index 2f2481ea60f8..ebeb4ae36f99 100644
--- a/llvm/lib/CGData/CodeGenDataReader.cpp
+++ b/llvm/lib/CGData/CodeGenDataReader.cpp
@@ -32,10 +32,40 @@ setupMemoryBuffer(const Twine &Filename, vfs::FileSystem &FS) {
 
 Error CodeGenDataReader::mergeFromObjectFile(
     const object::ObjectFile *Obj, OutlinedHashTreeRecord &GlobalOutlineRecord,
+    StableFunctionMapRecord &GlobalFunctionMapRecord,
     stable_hash *CombinedHash) {
   Triple TT = Obj->makeTriple();
   auto CGOutLineName =
       getCodeGenDataSectionName(CG_outline, TT.getObjectFormat(), false);
+  auto CGMergeName =
+      getCodeGenDataSectionName(CG_merge, TT.getObjectFormat(), false);
+
+  auto processSectionContents = [&](const StringRef &Name,
+                                    const StringRef &Contents) {
+    if (Name != CGOutLineName && Name != CGMergeName)
+      return;
+    if (CombinedHash)
+      *CombinedHash = stable_hash_combine(*CombinedHash, xxh3_64bits(Contents));
+    auto *Data = reinterpret_cast<const unsigned char *>(Contents.data());
+    auto *EndData = Data + Contents.size();
+    // In case dealing with an executable that has concatenated cgdata,
+    // we want to merge them into a single cgdata.
+    // Although it's not a typical workflow, we support this scenario
+    // by looping over all data in the sections.
+    if (Name == CGOutLineName) {
+      while (Data != EndData) {
+        OutlinedHashTreeRecord LocalOutlineRecord;
+        LocalOutlineRecord.deserialize(Data);
+        GlobalOutlineRecord.merge(LocalOutlineRecord);
+      }
+    } else if (Name == CGMergeName) {
+      while (Data != EndData) {
+        StableFunctionMapRecord LocalFunctionMapRecord;
+        LocalFunctionMapRecord.deserialize(Data);
+        GlobalFunctionMapRecord.merge(LocalFunctionMapRecord);
+      }
+    }
+  };
 
   for (auto &Section : Obj->sections()) {
     Expected<StringRef> NameOrErr = Section.getName();
@@ -44,23 +74,7 @@ Error CodeGenDataReader::mergeFromObjectFile(
     Expected<StringRef> ContentsOrErr = Section.getContents();
     if (!ContentsOrErr)
       return ContentsOrErr.takeError();
-    auto *Data = reinterpret_cast<const unsigned char *>(ContentsOrErr->data());
-    auto *EndData = Data + ContentsOrErr->size();
-
-    if (*NameOrErr == CGOutLineName) {
-      if (CombinedHash)
-        *CombinedHash =
-            stable_hash_combine(*CombinedHash, xxh3_64bits(*ContentsOrErr));
-      // In case dealing with an executable that has concatenated cgdata,
-      // we want to merge them into a single cgdata.
-      // Although it's not a typical workflow, we support this scenario.
-      while (Data != EndData) {
-        OutlinedHashTreeRecord LocalOutlineRecord;
-        LocalOutlineRecord.deserialize(Data);
-        GlobalOutlineRecord.merge(LocalOutlineRecord);
-      }
-    }
-    // TODO: Add support for other cgdata sections.
+    processSectionContents(*NameOrErr, *ContentsOrErr);
   }
 
   return Error::success();
@@ -69,7 +83,8 @@ Error CodeGenDataReader::mergeFromObjectFile(
 Error IndexedCodeGenDataReader::read() {
   using namespace support;
 
-  // The smallest header with the version 1 is 24 bytes
+  // The smallest header with the version 1 is 24 bytes.
+  // Do not update this value even with the new version of the header.
   const unsigned MinHeaderSize = 24;
   if (DataBuffer->getBufferSize() < MinHeaderSize)
     return error(cgdata_error::bad_header);
@@ -87,6 +102,12 @@ Error IndexedCodeGenDataReader::read() {
       return error(cgdata_error::eof);
     HashTreeRecord.deserialize(Ptr);
   }
+  if (hasStableFunctionMap()) {
+    const unsigned char *Ptr = Start + Header.StableFunctionMapOffset;
+    if (Ptr >= End)
+      return error(cgdata_error::eof);
+    FunctionMapRecord.deserialize(Ptr);
+  }
 
   return success();
 }
@@ -152,6 +173,8 @@ Error TextCodeGenDataReader::read() {
     StringRef Str = Line->drop_front().rtrim();
     if (Str.equals_insensitive("outlined_hash_tree"))
       DataKind |= CGDataKind::FunctionOutlinedHashTree;
+    else if (Str.equals_insensitive("stable_function_map"))
+      DataKind |= CGDataKind::StableFunctionMergingMap;
     else
       return error(cgdata_error::bad_header);
   }
@@ -170,8 +193,8 @@ Error TextCodeGenDataReader::read() {
   yaml::Input YOS(StringRef(Pos, Size));
   if (hasOutlinedHashTree())
     HashTreeRecord.deserializeYAML(YOS);
-
-  // TODO: Add more yaml cgdata in order
+  if (hasStableFunctionMap())
+    FunctionMapRecord.deserializeYAML(YOS);
 
   return Error::success();
 }