[WebAssembly] Add segment RETAIN flag to support private retained data (#81539)

In WebAssembly, we have `WASM_SYMBOL_NO_STRIP` symbol flag to mark the
referenced content as retained. However, the flag is not enough to
express retained data that is not referenced by any symbol. This patch
adds a new segment flag`WASM_SEG_FLAG_RETAIN` to support "private"
linkage data that is retained by llvm.used.

This kind of data that is not referenced but must be retained is usually
used with encapsulation symbols (__start/__stop). Swift runtime uses
this technique and depends on the fact "all metadata sections in live
objects are retained", which was not guaranteed with `--gc-sections`
before this patch.

This is a revised version of https://reviews.llvm.org/D126950 (has been
reverted) based on @MaskRay's comments

Author: kateinoigakukun

lld/test/wasm/no-strip-segment.s

@@ -0,0 +1,62 @@
+# RUN: split-file %s %t
+# RUN: llvm-mc -filetype=obj --triple=wasm32-unknown-unknown -o %t/main.o %t/main.s
+# RUN: llvm-mc -filetype=obj --triple=wasm32-unknown-unknown -o %t/liba_x.o %t/liba_x.s
+# RUN: llvm-mc -filetype=obj --triple=wasm32-unknown-unknown -o %t/liba_y.o %t/liba_y.s
+# RUN: rm -f %t/liba.a
+# RUN: llvm-ar rcs %t/liba.a %t/liba_x.o %t/liba_y.o
+# RUN: wasm-ld %t/main.o %t/liba.a --gc-sections -o %t/main.wasm --print-gc-sections | FileCheck %s --check-prefix=GC
+# RUN: obj2yaml %t/main.wasm | FileCheck %s
+
+# --gc-sections should remove non-retained and unused "weathers" section from live object liba_x.o
+# GC: removing unused section {{.*}}/liba.a(liba_x.o):(weathers)
+# Should not remove retained "greetings" sections from live objects main.o and liba_x.o
+# GC-NOT: removing unused section %t/main.o:(greetings)
+# GC-NOT: removing unused section %t/liba_x.o:(greetings)
+
+# Note: All symbols are private so that they don't join the symbol table.
+
+#--- main.s
+  .functype grab_liba () -> ()
+  .globl  _start
+_start:
+  .functype _start () -> ()
+  call grab_liba
+  end_function
+
+  .section greetings,"R",@
+  .asciz  "hello"
+  .section weathers,"R",@
+  .asciz  "cloudy"
+
+#--- liba_x.s
+  .globl  grab_liba
+grab_liba:
+  .functype grab_liba () -> ()
+  end_function
+
+  .section greetings,"R",@
+  .asciz  "world"
+  .section weathers,"",@
+  .asciz  "rainy"
+
+#--- liba_y.s
+  .section        greetings,"R",@
+  .asciz  "bye"
+
+
+# "greetings" section
+# CHECK: - Type:            DATA
+# CHECK:   Segments:
+# CHECK:     - SectionOffset:   7
+# CHECK:       InitFlags:       0
+# CHECK:       Offset:
+# CHECK:         Opcode:          I32_CONST
+# CHECK:         Value:           1024
+# CHECK:       Content:         68656C6C6F00776F726C6400
+# "weahters" section.
+# CHECK: - SectionOffset:   25
+# CHECK:   InitFlags:       0
+# CHECK:   Offset:
+# CHECK:     Opcode:          I32_CONST
+# CHECK:     Value:           1036
+# CHECK:   Content:         636C6F75647900

lld/wasm/InputChunks.h

@@ -81,6 +81,7 @@ class InputChunk {
   void generateRelocationCode(raw_ostream &os) const;
 
   bool isTLS() const { return flags & llvm::wasm::WASM_SEG_FLAG_TLS; }
+  bool isRetained() const { return flags & llvm::wasm::WASM_SEG_FLAG_RETAIN; }
 
   ObjFile *file;
   OutputSection *outputSec = nullptr;

lld/wasm/MarkLive.cpp

@@ -40,7 +40,9 @@ class MarkLive {
 
 private:
   void enqueue(Symbol *sym);
+  void enqueue(InputChunk *chunk);
   void enqueueInitFunctions(const ObjFile *sym);
+  void enqueueRetainedSegments(const ObjFile *file);
   void mark();
   bool isCallCtorsLive();
 
@@ -56,21 +58,30 @@ void MarkLive::enqueue(Symbol *sym) {
   LLVM_DEBUG(dbgs() << "markLive: " << sym->getName() << "\n");
 
   InputFile *file = sym->getFile();
-  bool needInitFunctions = file && !file->isLive() && sym->isDefined();
+  bool markImplicitDeps = file && !file->isLive() && sym->isDefined();
 
   sym->markLive();
 
-  // Mark ctor functions in the object that defines this symbol live.
-  // The ctor functions are all referenced by the synthetic callCtors
-  // function. However, this function does not contain relocations so we
-  // have to manually mark the ctors as live.
-  if (needInitFunctions)
+  if (markImplicitDeps) {
+    // Mark ctor functions in the object that defines this symbol live.
+    // The ctor functions are all referenced by the synthetic callCtors
+    // function. However, this function does not contain relocations so we
+    // have to manually mark the ctors as live.
     enqueueInitFunctions(cast<ObjFile>(file));
+    // Mark retained segments in the object that defines this symbol live.
+    enqueueRetainedSegments(cast<ObjFile>(file));
+  }
 
   if (InputChunk *chunk = sym->getChunk())
     queue.push_back(chunk);
 }
 
+void MarkLive::enqueue(InputChunk *chunk) {
+  LLVM_DEBUG(dbgs() << "markLive: " << toString(chunk) << "\n");
+  chunk->live = true;
+  queue.push_back(chunk);
+}
+
 // The ctor functions are all referenced by the synthetic callCtors
 // function.  However, this function does not contain relocations so we
 // have to manually mark the ctors as live.
@@ -83,6 +94,14 @@ void MarkLive::enqueueInitFunctions(const ObjFile *obj) {
   }
 }
 
+// Mark segments flagged by segment-level no-strip. Segment-level no-strip is
+// usually used to retain segments without having symbol table entry.
+void MarkLive::enqueueRetainedSegments(const ObjFile *file) {
+  for (InputChunk *chunk : file->segments)
+    if (chunk->isRetained())
+      enqueue(chunk);
+}
+
 void MarkLive::run() {
   // Add GC root symbols.
   if (!config->entry.empty())
@@ -96,10 +115,14 @@ void MarkLive::run() {
   if (WasmSym::callDtors)
     enqueue(WasmSym::callDtors);
 
-  // Enqueue constructors in objects explicitly live from the command-line.
   for (const ObjFile *obj : ctx.objectFiles)
-    if (obj->isLive())
+    if (obj->isLive()) {
+      // Enqueue constructors in objects explicitly live from the command-line.
       enqueueInitFunctions(obj);
+      // Enqueue retained segments in objects explicitly live from the
+      // command-line.
+      enqueueRetainedSegments(obj);
+    }
 
   mark();
 

llvm/include/llvm/BinaryFormat/Wasm.h

@@ -216,6 +216,7 @@ enum WasmSymbolType : unsigned {
 enum WasmSegmentFlag : unsigned {
   WASM_SEG_FLAG_STRINGS = 0x1,
   WASM_SEG_FLAG_TLS = 0x2,
+  WASM_SEG_FLAG_RETAIN = 0x4,
 };
 
 // Kinds of tag attributes.

llvm/include/llvm/CodeGen/TargetLoweringObjectFileImpl.h

@@ -207,11 +207,14 @@ class TargetLoweringObjectFileCOFF : public TargetLoweringObjectFile {
 
 class TargetLoweringObjectFileWasm : public TargetLoweringObjectFile {
   mutable unsigned NextUniqueID = 0;
+  SmallPtrSet<GlobalObject *, 2> Used;
 
 public:
   TargetLoweringObjectFileWasm() = default;
   ~TargetLoweringObjectFileWasm() override = default;
 
+  void getModuleMetadata(Module &M) override;
+
   MCSection *getExplicitSectionGlobal(const GlobalObject *GO, SectionKind Kind,
                                       const TargetMachine &TM) const override;
 

llvm/lib/CodeGen/TargetLoweringObjectFileImpl.cpp

@@ -2142,7 +2142,7 @@ static const Comdat *getWasmComdat(const GlobalValue *GV) {
   return C;
 }
 
-static unsigned getWasmSectionFlags(SectionKind K) {
+static unsigned getWasmSectionFlags(SectionKind K, bool Retain) {
   unsigned Flags = 0;
 
   if (K.isThreadLocal())
@@ -2151,11 +2151,22 @@ static unsigned getWasmSectionFlags(SectionKind K) {
   if (K.isMergeableCString())
     Flags |= wasm::WASM_SEG_FLAG_STRINGS;
 
+  if (Retain)
+    Flags |= wasm::WASM_SEG_FLAG_RETAIN;
+
   // TODO(sbc): Add suport for K.isMergeableConst()
 
   return Flags;
 }
 
+void TargetLoweringObjectFileWasm::getModuleMetadata(Module &M) {
+  SmallVector<GlobalValue *, 4> Vec;
+  collectUsedGlobalVariables(M, Vec, false);
+  for (GlobalValue *GV : Vec)
+    if (auto *GO = dyn_cast<GlobalObject>(GV))
+      Used.insert(GO);
+}
+
 MCSection *TargetLoweringObjectFileWasm::getExplicitSectionGlobal(
     const GlobalObject *GO, SectionKind Kind, const TargetMachine &TM) const {
   // We don't support explict section names for functions in the wasm object
@@ -2179,16 +2190,18 @@ MCSection *TargetLoweringObjectFileWasm::getExplicitSectionGlobal(
     Group = C->getName();
   }
 
-  unsigned Flags = getWasmSectionFlags(Kind);
+  unsigned Flags = getWasmSectionFlags(Kind, Used.count(GO));
   MCSectionWasm *Section = getContext().getWasmSection(
       Name, Kind, Flags, Group, MCContext::GenericSectionID);
 
   return Section;
 }
 
-static MCSectionWasm *selectWasmSectionForGlobal(
-    MCContext &Ctx, const GlobalObject *GO, SectionKind Kind, Mangler &Mang,
-    const TargetMachine &TM, bool EmitUniqueSection, unsigned *NextUniqueID) {
+static MCSectionWasm *
+selectWasmSectionForGlobal(MCContext &Ctx, const GlobalObject *GO,
+                           SectionKind Kind, Mangler &Mang,
+                           const TargetMachine &TM, bool EmitUniqueSection,
+                           unsigned *NextUniqueID, bool Retain) {
   StringRef Group = "";
   if (const Comdat *C = getWasmComdat(GO)) {
     Group = C->getName();
@@ -2213,7 +2226,7 @@ static MCSectionWasm *selectWasmSectionForGlobal(
     (*NextUniqueID)++;
   }
 
-  unsigned Flags = getWasmSectionFlags(Kind);
+  unsigned Flags = getWasmSectionFlags(Kind, Retain);
   return Ctx.getWasmSection(Name, Kind, Flags, Group, UniqueID);
 }
 
@@ -2231,9 +2244,11 @@ MCSection *TargetLoweringObjectFileWasm::SelectSectionForGlobal(
   else
     EmitUniqueSection = TM.getDataSections();
   EmitUniqueSection |= GO->hasComdat();
+  bool Retain = Used.count(GO);
+  EmitUniqueSection |= Retain;
 
   return selectWasmSectionForGlobal(getContext(), GO, Kind, getMangler(), TM,
-                                    EmitUniqueSection, &NextUniqueID);
+                                    EmitUniqueSection, &NextUniqueID, Retain);
 }
 
 bool TargetLoweringObjectFileWasm::shouldPutJumpTableInFunctionSection(

llvm/lib/MC/MCParser/WasmAsmParser.cpp

@@ -115,6 +115,9 @@ class WasmAsmParser : public MCAsmParserExtension {
       case 'S':
         flags |= wasm::WASM_SEG_FLAG_STRINGS;
         break;
+      case 'R':
+        flags |= wasm::WASM_SEG_FLAG_RETAIN;
+        break;
       default:
         return -1U;
       }

llvm/lib/MC/MCSectionWasm.cpp

@@ -70,6 +70,8 @@ void MCSectionWasm::printSwitchToSection(const MCAsmInfo &MAI, const Triple &T,
     OS << 'S';
   if (SegmentFlags & wasm::WASM_SEG_FLAG_TLS)
     OS << 'T';
+  if (SegmentFlags & wasm::WASM_SEG_FLAG_RETAIN)
+    OS << 'R';
 
   OS << '"';
 

llvm/lib/ObjectYAML/WasmYAML.cpp

@@ -561,6 +561,7 @@ void ScalarBitSetTraits<WasmYAML::SegmentFlags>::bitset(
 #define BCase(X) IO.bitSetCase(Value, #X, wasm::WASM_SEG_FLAG_##X)
   BCase(STRINGS);
   BCase(TLS);
+  BCase(RETAIN);
 #undef BCase
 }
 

llvm/test/CodeGen/WebAssembly/no-strip.ll

@@ -0,0 +1,22 @@
+; RUN: llc < %s --mtriple=wasm32-unknown-unknown | FileCheck %s
+
+@llvm.used = appending global [
+  5 x ptr
+] [
+  ptr @ga, ptr @gb, ptr @gc, ptr @gd, ptr @ge
+], section "llvm.metadata"
+
+; CHECK: .section        .data.ga,"R",@
+@ga = global i32 42
+; CHECK: .section        .data.gb,"R",@
+@gb = internal global i32 41
+; CHECK: .section        .data..Lgc,"R",@
+@gc = private global i32 40
+; CHECK: .section        .rodata.gd,"R",@
+@gd = constant i32 39
+
+; All sections with the same explicit name are flagged as retained if a part of them is retained.
+; CHECK: .section        dddd,"R",@
+@ge = global i32 38, section "dddd"
+; CHECK: .section        dddd,"R",@
+@gg = global i32 37, section "dddd"