[lld][WebAssembly] Add --keep-section flag

This flag causes wasm-ld preserve a section even in the face of
`--strip-all`.  This is useful, for example, to preserve the
target_features section in the ase of clang (which can run wasm-opt
after linking), and emcc (which performs a bunch of post-link work).

Fixes: https://github.com/llvm/llvm-project/issues/60613
Fixes: https://github.com/llvm/llvm-project/issues/55781

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

clang/lib/Driver/ToolChains/WebAssembly.cpp

@@ -141,14 +141,25 @@ void wasm::Linker::ConstructJob(Compilation &C, const JobAction &JA,
   CmdArgs.push_back("-o");
   CmdArgs.push_back(Output.getFilename());
 
+  // When optimizing, if wasm-opt is available, run it.
+  std::string WasmOptPath;
+  if (Arg *A = Args.getLastArg(options::OPT_O_Group)) {
+    WasmOptPath = ToolChain.GetProgramPath("wasm-opt");
+    if (WasmOptPath == "wasm-opt") {
+      WasmOptPath = {};
+    }
+  }
+
+  if (!WasmOptPath.empty()) {
+    CmdArgs.push_back("--keep-section=target_features");
+  }
+
   C.addCommand(std::make_unique<Command>(JA, *this,
                                          ResponseFileSupport::AtFileCurCP(),
                                          Linker, CmdArgs, Inputs, Output));
 
-  // When optimizing, if wasm-opt is available, run it.
   if (Arg *A = Args.getLastArg(options::OPT_O_Group)) {
-    auto WasmOptPath = ToolChain.GetProgramPath("wasm-opt");
-    if (WasmOptPath != "wasm-opt") {
+    if (!WasmOptPath.empty()) {
       StringRef OOpt = "s";
       if (A->getOption().matches(options::OPT_O4) ||
           A->getOption().matches(options::OPT_Ofast))
@@ -160,13 +171,13 @@ void wasm::Linker::ConstructJob(Compilation &C, const JobAction &JA,
 
       if (OOpt != "0") {
         const char *WasmOpt = Args.MakeArgString(WasmOptPath);
-        ArgStringList CmdArgs;
-        CmdArgs.push_back(Output.getFilename());
-        CmdArgs.push_back(Args.MakeArgString(llvm::Twine("-O") + OOpt));
-        CmdArgs.push_back("-o");
-        CmdArgs.push_back(Output.getFilename());
+        ArgStringList OptArgs;
+        OptArgs.push_back(Output.getFilename());
+        OptArgs.push_back(Args.MakeArgString(llvm::Twine("-O") + OOpt));
+        OptArgs.push_back("-o");
+        OptArgs.push_back(Output.getFilename());
         C.addCommand(std::make_unique<Command>(
-            JA, *this, ResponseFileSupport::AtFileCurCP(), WasmOpt, CmdArgs,
+            JA, *this, ResponseFileSupport::AtFileCurCP(), WasmOpt, OptArgs,
             Inputs, Output));
       }
     }

lld/test/wasm/strip-all.s

@@ -0,0 +1,36 @@
+# RUN: llvm-mc -filetype=obj -triple=wasm32-unknown-unknown %s -o %t.start.o
+# RUN: wasm-ld --strip-all -o %t.wasm %t.start.o
+# RUN: obj2yaml %t.wasm | FileCheck %s
+#
+## Test alias -s
+# RUN: wasm-ld -s -o %t2.wasm %t.start.o
+# RUN: obj2yaml %t2.wasm | FileCheck %s
+#
+## Check that there is no name section
+# CHECK-NOT:   Name:    name
+# CHECK-NOT:   Name:    target_features
+#
+## Test --keep-section=name preserver the name section
+# RUN: wasm-ld --strip-all --keep-section=name -o %t3.wasm %t.start.o
+# RUN: obj2yaml %t3.wasm | FileCheck --check-prefix=CHECK-NAME %s
+#
+# CHECK-NAME:   Name:    name
+# CHECK-NAME-NOT:   Name:    target_features
+#
+## Test --keep-section can be specified more than once
+# RUN: wasm-ld --strip-all --keep-section=name --keep-section=target_features -o %t4.wasm %t.start.o
+# RUN: obj2yaml %t4.wasm | FileCheck --check-prefix=CHECK-FEATURES %s
+#
+# CHECK-FEATURES:   Name:    name
+# CHECK-FEATURES:   Name:    target_features
+
+  .globl  _start
+_start:
+  .functype _start () -> ()
+  end_function
+
+.section .custom_section.target_features,"",@
+.int8 1
+.int8 43
+.int8 15
+.ascii "mutable-globals"

lld/test/wasm/strip-all.test

@@ -1,10 +0,0 @@
-RUN: llvm-mc -filetype=obj -triple=wasm32-unknown-unknown %p/Inputs/start.s -o %t.start.o
-RUN: wasm-ld --strip-all -o %t.wasm %t.start.o
-RUN: obj2yaml %t.wasm | FileCheck %s
-
-# Test alias -s
-RUN: wasm-ld -s -o %t.wasm %t.start.o
-RUN: obj2yaml %t.wasm | FileCheck %s
-
-# Check that there is no name section
-CHECK-NOT:   Name:    name

lld/wasm/Config.h

@@ -49,6 +49,7 @@ struct Configuration {
   bool extendedConst;
   bool growableTable;
   bool gcSections;
+  llvm::StringSet<> keepSections;
   std::optional<std::pair<llvm::StringRef, llvm::StringRef>> memoryImport;
   std::optional<llvm::StringRef> memoryExport;
   bool sharedMemory;

lld/wasm/Driver.cpp

@@ -478,6 +478,8 @@ static void readConfigs(opt::InputArgList &args) {
   config->relocatable = args.hasArg(OPT_relocatable);
   config->gcSections =
       args.hasFlag(OPT_gc_sections, OPT_no_gc_sections, !config->relocatable);
+  for (auto *arg : args.filtered(OPT_keep_section))
+    config->keepSections.insert(arg->getValue());
   config->mergeDataSegments =
       args.hasFlag(OPT_merge_data_segments, OPT_no_merge_data_segments,
                    !config->relocatable);

lld/wasm/Options.td

@@ -193,6 +193,9 @@ def growable_table: FF<"growable-table">,
 def global_base: JJ<"global-base=">,
   HelpText<"Memory offset at which to place global data (Defaults to 1024)">;
 
+defm keep_section: Eq<"keep-section",
+     "Preserve a section even when --strip-all is given. This is useful for compiler drivers such as clang or emcc that, for example, depend on the features section for post-link processing. Can be specified multiple times to keep multiple sections">;
+
 def import_memory: FF<"import-memory">,
   HelpText<"Import the module's memory from the default module of \"env\" with the name \"memory\".">;
 def import_memory_with_name: JJ<"import-memory=">,

lld/wasm/SyntheticSections.h

@@ -373,7 +373,11 @@ public:
   NameSection(ArrayRef<OutputSegment *> segments)
       : SyntheticSection(llvm::wasm::WASM_SEC_CUSTOM, "name"),
         segments(segments) {}
-  bool isNeeded() const override { return !config->stripAll && numNames() > 0; }
+  bool isNeeded() const override {
+    if (config->stripAll && !config->keepSections.count(name))
+      return false;
+    return numNames() > 0;
+  }
   void writeBody() override;
   unsigned numNames() const {
     // We always write at least one name which is the name of the
@@ -393,7 +397,9 @@ public:
   ProducersSection()
       : SyntheticSection(llvm::wasm::WASM_SEC_CUSTOM, "producers") {}
   bool isNeeded() const override {
-    return !config->stripAll && fieldCount() > 0;
+    if (config->stripAll && !config->keepSections.count(name))
+      return false;
+    return fieldCount() > 0;
   }
   void writeBody() override;
   void addInfo(const llvm::wasm::WasmProducerInfo &info);
@@ -412,7 +418,9 @@ public:
   TargetFeaturesSection()
       : SyntheticSection(llvm::wasm::WASM_SEC_CUSTOM, "target_features") {}
   bool isNeeded() const override {
-    return !config->stripAll && features.size() > 0;
+    if (config->stripAll && !config->keepSections.count(name))
+      return false;
+    return features.size() > 0;
   }
   void writeBody() override;