[obj2yaml,yaml2obj] Add NumBlocks to the BBAddrMapEntry yaml field.

As discussed in D95511, this allows us to encode invalid BBAddrMap
sections to be used in more rigorous testing.

Reviewed By: jhenderson

Differential Revision: https://reviews.llvm.org/D96831

2 files changed, 7 insertions, 4 deletions

diff --git a/llvm/lib/ObjectYAML/ELFEmitter.cpp b/llvm/lib/ObjectYAML/ELFEmitter.cpp
index 089ef8e8..a9ff25c 100644
--- a/llvm/lib/ObjectYAML/ELFEmitter.cpp
+++ b/llvm/lib/ObjectYAML/ELFEmitter.cpp
@@ -1358,12 +1358,14 @@ void ELFState<ELFT>::writeSectionContent(
   for (const ELFYAML::BBAddrMapEntry &E : *Section.Entries) {
     // Write the address of the function.
     CBA.write<uintX_t>(E.Address, ELFT::TargetEndianness);
-    // Write number of BBEntries (number of basic blocks in the function).
-    size_t NumBlocks = E.BBEntries ? E.BBEntries->size() : 0;
+    // Write number of BBEntries (number of basic blocks in the function). This
+    // is overriden by the 'NumBlocks' YAML field if specified.
+    uint32_t NumBlocks =
+        E.NumBlocks.getValueOr(E.BBEntries ? E.BBEntries->size() : 0);
     SHeader.sh_size += sizeof(uintX_t) + CBA.writeULEB128(NumBlocks);
-    if (!NumBlocks)
-      continue;
     // Write all BBEntries.
+    if (!E.BBEntries)
+      continue;
     for (const ELFYAML::BBAddrMapEntry::BBEntry &BBE : *E.BBEntries)
       SHeader.sh_size += CBA.writeULEB128(BBE.AddressOffset) +
                          CBA.writeULEB128(BBE.Size) +
diff --git a/llvm/lib/ObjectYAML/ELFYAML.cpp b/llvm/lib/ObjectYAML/ELFYAML.cpp
index bde91c5..1eedc64 100644
--- a/llvm/lib/ObjectYAML/ELFYAML.cpp
+++ b/llvm/lib/ObjectYAML/ELFYAML.cpp
@@ -1665,6 +1665,7 @@ void MappingTraits<ELFYAML::BBAddrMapEntry>::mapping(
   assert(IO.getContext() && "The IO context is not initialized");
   IO.mapOptional("Address", E.Address, Hex64(0));
   IO.mapOptional("BBEntries", E.BBEntries);
+  IO.mapOptional("NumBlocks", E.NumBlocks);
 }
 
 void MappingTraits<ELFYAML::BBAddrMapEntry::BBEntry>::mapping(