[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

Author: rlavaee

llvm/include/llvm/ObjectYAML/ELFYAML.h

@@ -161,6 +161,7 @@ struct BBAddrMapEntry {
     llvm::yaml::Hex64 Metadata;
   };
   llvm::yaml::Hex64 Address;
+  Optional<uint32_t> NumBlocks;
   Optional<std::vector<BBEntry>> BBEntries;
 };
 

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) +

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(

llvm/test/tools/obj2yaml/ELF/bb-addr-map.yaml

@@ -39,8 +39,10 @@ FileHeader:
 Sections:
   - Name:    .llvm_bb_addr_map
     Type:    SHT_LLVM_BB_ADDR_MAP
+    ShSize:    [[SIZE=<none>]]
     Entries:
       - Address:   0x0
+        NumBlocks: [[NUMBLOCKS=<none>]]
         BBEntries:
           - AddressOffset:    0x1
             Size:             0x2
@@ -57,37 +59,10 @@ Sections:
             Size:             0xB
             Metadata:         0xC
 
-## Check that obj2yaml uses the "Content" tag to describe an .llvm_bb_addr_map section
-## when it can't extract the entries. For instance, when truncated data is given as
-## 'Content'.
-
-# RUN: yaml2obj --docnum=2 %s -o %t2
-# RUN: obj2yaml %t2 | FileCheck %s --check-prefix=INVALID
-
-# INVALID:      --- !ELF
-# INVALID-NEXT: FileHeader:
-# INVALID-NEXT:   Class:   ELFCLASS64
-# INVALID-NEXT:   Data:    ELFDATA2LSB
-# INVALID-NEXT:   Type:    ET_EXEC
-# INVALID-NEXT: Sections:
-# INVALID-NEXT:   - Name: .llvm_bb_addr_map
-# INVALID-NEXT:     Type: SHT_LLVM_BB_ADDR_MAP
-# INVALID-NEXT:     Content: '10000000000000'
-
---- !ELF
-FileHeader:
-  Class: ELFCLASS64
-  Data:  ELFDATA2LSB
-  Type:  ET_EXEC
-Sections:
-  - Name:    .llvm_bb_addr_map
-    Type:    SHT_LLVM_BB_ADDR_MAP
-    Content: '10000000000000'
-
 ## Check obj2yaml can dump empty .llvm_bb_addr_map sections.
 
-# RUN: yaml2obj --docnum=3 %s -o %t3
-# RUN: obj2yaml %t3 | FileCheck %s --check-prefix=EMPTY
+# RUN: yaml2obj --docnum=2 %s -o %t2
+# RUN: obj2yaml %t2 | FileCheck %s --check-prefix=EMPTY
 
 # EMPTY:      --- !ELF
 # EMPTY-NEXT: FileHeader:
@@ -111,8 +86,8 @@ Sections:
 
 ## Check obj2yaml can dump multiple .llvm_bb_addr_map sections.
 
-# RUN: yaml2obj --docnum=4 %s -o %t4
-# RUN: obj2yaml %t4 | FileCheck %s --check-prefix=MULTI
+# RUN: yaml2obj --docnum=3 %s -o %t3
+# RUN: obj2yaml %t3 | FileCheck %s --check-prefix=MULTI
 
 # MULTI:      --- !ELF
 # MULTI-NEXT: FileHeader:
@@ -153,3 +128,24 @@ Sections:
     Type:    SHT_LLVM_BB_ADDR_MAP
     Entries:
       - Address:   0x20
+
+## Check that obj2yaml uses the "Content" tag to describe an .llvm_bb_addr_map section
+## when it can't extract the entries, for example, when section is truncated, or when
+## an invalid NumBlocks field is specified.
+
+# RUN: yaml2obj --docnum=1 -DSIZE=0x8 %s -o %t4
+# RUN: obj2yaml %t4 | FileCheck %s --check-prefixes=TRUNCATED,INVALID
+
+# RUN: yaml2obj --docnum=1 -DNUMBLOCKS=2 %s -o %t5
+# RUN: obj2yaml %t5 | FileCheck %s --check-prefixes=BADNUMBLOCKS,INVALID
+
+# INVALID:           --- !ELF
+# INVALID-NEXT:      FileHeader:
+# INVALID-NEXT:        Class:   ELFCLASS64
+# INVALID-NEXT:        Data:    ELFDATA2LSB
+# INVALID-NEXT:        Type:    ET_EXEC
+# INVALID-NEXT:      Sections:
+# INVALID-NEXT:        - Name: .llvm_bb_addr_map
+# INVALID-NEXT:          Type: SHT_LLVM_BB_ADDR_MAP
+# BADNUMBLOCKS-NEXT:     Content: {{([[:xdigit:]]+)}}{{$}}
+# TRUNCATED-NEXT:        Content: '{{([[:xdigit:]]{16})}}'{{$}}

llvm/test/tools/yaml2obj/ELF/bb-addr-map.yaml

@@ -47,6 +47,12 @@
 # CHECK-NEXT:     0000: 00000000 00000000 01010203
 # CHECK-NEXT:   )
 
+# Case 6: Override the NumBlocks field.
+# CHECK:        Name: .llvm_bb_addr_map (1)
+# CHECK:        SectionData (
+# CHECK-NEXT:     0000: 20000000 00000000 02010203
+# CHECK-NEXT:   )
+
 --- !ELF
 FileHeader:
   Class: ELFCLASS64
@@ -95,6 +101,17 @@ Sections:
             Size:             0x00000002
             Metadata:         0x00000003
 
+## 6) We can override the NumBlocks field with a value different from the
+##    actual number of BB Entries.
+  - Name:    '.llvm_bb_addr_map (6)'
+    Type:    SHT_LLVM_BB_ADDR_MAP
+    Entries:
+      - Address:   0x0000000000000020
+        NumBlocks: 2
+        BBEntries:
+          - AddressOffset:    0x00000001
+            Size:             0x00000002
+            Metadata:         0x00000003
 
 ## Check we can't use Entries at the same time as either Content or Size.
 # RUN: not yaml2obj --docnum=2 -DCONTENT="00" %s 2>&1 | FileCheck %s --check-prefix=INVALID

llvm/tools/obj2yaml/elf2yaml.cpp

@@ -861,7 +861,7 @@ ELFDumper<ELFT>::dumpBBAddrMapSection(const Elf_Shdr *Shdr) {
       uint64_t Metadata = Data.getULEB128(Cur);
       BBEntries.push_back({Offset, Size, Metadata});
     }
-    Entries.push_back({Address, BBEntries});
+    Entries.push_back({Address, /* NumBlocks */ {}, BBEntries});
   }
 
   if (!Cur) {