[ELF] -r: rewrite SHT_GROUP content if some members are combined or discarded

* If two group members are combined, we should leave just one index in the SHT_GROUP content.
* If a group member is discarded (/DISCARD/ or upcoming -r --gc-sections combination),
  we should drop its index in the SHT_GROUP content. LLD currently crashes (`getOutputSection()` is null).

Reviewed By: psmith

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

lld/ELF/InputSection.cpp

@@ -28,6 +28,7 @@
 #include <algorithm>
 #include <mutex>
 #include <set>
+#include <unordered_set>
 #include <vector>
 
 using namespace llvm;
@@ -389,11 +390,16 @@ template <class ELFT> void InputSection::copyShtGroup(uint8_t *buf) {
   // The first entry is not a section number but a flag.
   *to++ = from[0];
 
-  // Adjust section numbers because section numbers in an input object
-  // files are different in the output.
+  // Adjust section numbers because section numbers in an input object files are
+  // different in the output. We also need to handle combined or discarded
+  // members.
   ArrayRef<InputSectionBase *> sections = file->getSections();
-  for (uint32_t idx : from.slice(1))
-    *to++ = sections[idx]->getOutputSection()->sectionIndex;
+  std::unordered_set<uint32_t> seen;
+  for (uint32_t idx : from.slice(1)) {
+    OutputSection *osec = sections[idx]->getOutputSection();
+    if (osec && seen.insert(osec->sectionIndex).second)
+      *to++ = osec->sectionIndex;
+  }
 }
 
 InputSectionBase *InputSection::getRelocatedSection() const {

lld/ELF/OutputSections.cpp

@@ -21,6 +21,7 @@
 #include "llvm/Support/Parallel.h"
 #include "llvm/Support/SHA1.h"
 #include <regex>
+#include <unordered_set>
 
 using namespace llvm;
 using namespace llvm::dwarf;
@@ -376,6 +377,15 @@ static void finalizeShtGroup(OutputSection *os,
   // provides signature of the section group.
   ArrayRef<Symbol *> symbols = section->file->getSymbols();
   os->info = in.symTab->getSymbolIndex(symbols[section->info]);
+
+  // Some group members may be combined or discarded, so we need to compute the
+  // new size. The content will be rewritten in InputSection::copyShtGroup.
+  std::unordered_set<uint32_t> seen;
+  ArrayRef<InputSectionBase *> sections = section->file->getSections();
+  for (const uint32_t &idx : section->getDataAs<uint32_t>().slice(1))
+    if (OutputSection *osec = sections[read32(&idx)]->getOutputSection())
+      seen.insert(osec->sectionIndex);
+  os->size = (1 + seen.size()) * sizeof(uint32_t);
 }
 
 void OutputSection::finalize() {

lld/test/ELF/relocatable-comdat.s

@@ -1,45 +1,52 @@
 # REQUIRES: x86
+## Test that SHT_GROUP sections are retained in relocatable output. The content
+## may be rewritten because group members may change their indices. Additionally,
+## group member may be combined or discarded (e.g. /DISCARD/ or --gc-sections).
+
 # RUN: llvm-mc -filetype=obj -triple=x86_64-pc-linux %s -o %t.o
-# RUN: ld.lld -r %t.o %t.o -o %t
-# RUN: llvm-readobj --elf-section-groups --sections %t | FileCheck %s
+# RUN: ld.lld -r %t.o %t.o -o %t.ro
+# RUN: llvm-readelf -g -S %t.ro | FileCheck %s
 
-# CHECK:        Name: .text.bar
-# CHECK-NEXT:   Type: SHT_PROGBITS
-# CHECK-NEXT:   Flags [
-# CHECK-NEXT:     SHF_ALLOC
-# CHECK-NEXT:     SHF_EXECINSTR
-# CHECK-NEXT:     SHF_GROUP
-# CHECK-NEXT:   ]
-# CHECK-NEXT:   Address:
-# CHECK-NEXT:   Offset:
-# CHECK-NEXT:   Size: 8
-# CHECK:      Section {
-# CHECK-NEXT:   Index: 4
-# CHECK-NEXT:   Name: .text.foo
-# CHECK-NEXT:   Type: SHT_PROGBITS
-# CHECK-NEXT:   Flags [
-# CHECK-NEXT:     SHF_ALLOC
-# CHECK-NEXT:     SHF_EXECINSTR
-# CHECK-NEXT:     SHF_GROUP
-# CHECK-NEXT:   ]
-# CHECK-NEXT:   Address:
-# CHECK-NEXT:   Offset:
-# CHECK-NEXT:   Size: 4
+# CHECK: Name    Type  Address          Off    Size   ES Flg Lk        Inf Al
+# CHECK: .group  GROUP 0000000000000000 {{.*}} 000014 04     {{[1-9]}}   1  4
 
-# CHECK:       Groups {
-# CHECK-NEXT:    Group {
-# CHECK-NEXT:      Name: .group
-# CHECK-NEXT:      Index: 2
-# CHECK-NEXT:      Link: 6
-# CHECK-NEXT:      Info: 1 
-# CHECK-NEXT:      Type: COMDAT
-# CHECK-NEXT:      Signature: abc
-# CHECK-NEXT:      Section(s) in group [
-# CHECK-NEXT:        .text.bar
-# CHECK-NEXT:        .text.foo
-# CHECK-NEXT:      ]
-# CHECK-NEXT:    }
-# CHECK-NEXT:  }
+# CHECK: COMDAT group section [{{.*}}] `.group' [abc] contains 4 sections:
+# CHECK-NEXT: Name
+# CHECK-NEXT: .rodata.bar
+# CHECK-NEXT: .rodata.foo
+# CHECK-NEXT: .text.bar
+# CHECK-NEXT: .text.foo
+
+## Rewrite SHT_GROUP content if some members are combined.
+# RUN: echo 'SECTIONS { .rodata : {*(.rodata.*)} .text : {*(.text.*)} }' > %t1.lds
+# RUN: ld.lld -r -T %t1.lds %t.o %t.o -o %t1.ro
+# RUN: llvm-readelf -g -S %t1.ro | FileCheck %s --check-prefix=SCRIPT1
+
+# SCRIPT1: Name    Type  Address          Off    Size   ES Flg Lk        Inf Al
+# SCRIPT1: .group  GROUP 0000000000000000 {{.*}} 00000c 04     {{[1-9]}}   1  4
+
+# SCRIPT1: COMDAT group section [{{.*}}] `.group' [abc] contains 2 sections:
+# SCRIPT1-NEXT: Name
+# SCRIPT1-NEXT: .rodata
+# SCRIPT1-NEXT: .text
+
+# RUN: echo 'SECTIONS { /DISCARD/ : {*(.rodata.*)} }' > %t2.lds
+# RUN: ld.lld -r -T %t2.lds %t.o %t.o -o %t2.ro
+# RUN: llvm-readelf -g -S %t2.ro | FileCheck %s --check-prefix=SCRIPT2
+
+## Handle discarded group members.
+# SCRIPT2: [Nr] Name    Type  Address          Off    Size   ES Flg Lk        Inf Al
+# SCRIPT2: [ 2] .group  GROUP 0000000000000000 {{.*}} 00000c 04     {{[1-9]}}   1  4
+
+# SCRIPT2: COMDAT group section [{{.*}}] `.group' [abc] contains 2 sections:
+# SCRIPT2-NEXT: Name
+# SCRIPT2-NEXT: .text.bar
+# SCRIPT2-NEXT: .text.foo
+
+.section .rodata.bar,"aG",@progbits,abc,comdat
+.byte 42
+.section .rodata.foo,"aG",@progbits,abc,comdat
+.byte 42
 
 .section .text.bar,"axG",@progbits,abc,comdat
 .quad 42