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

[sbc100]

Summary

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

[sbc100]

Still needs a test..

[dschuff]

Do we want to make this any more general? In the future we might want to preserve other sections, e.g. passing optimization or profiling info from LLVM to Binaryen. Or maybe JSPI info?

[sbc100]

In D149917#4322367, @dschuff wrote:

Do we want to make this any more general? In the future we might want to preserve other sections, e.g. passing optimization or profiling info from LLVM to Binaryen. Or maybe JSPI info?

Do you have something you suggest? Something like --preserve-section=custom_section_name?

[dschuff]

Yeah, I think that would work. Or maybe --preserve-sections=sec1,sec2 since we might want to preserve multiple sections.

[dschuff]

or use --keep-section to match objcopy/strip? https://github.com/llvm/llvm-project/blob/main/llvm/test/tools/llvm-objcopy/wasm/basic-keep.test

[sbc100]

• rename flag

[sbc100]

Updated to use --keep-section. PTAL.

[kripken]

clang/lib/Driver/ToolChains/WebAssembly.cpp

154

Shouldn't secton => section?

[dschuff]

lld/wasm/Options.td

196

Can this flag be used more than once? From the code it looks like maybe it can, but the description should maybe say that and there should be a test (probably I should have written one for objcopy/strip too)

[sbc100]

• feedback

[sbc100]

• feedback

[sbc100]

• feedback

[sbc100]

PTAL

[dschuff]

otherwise LGTM

lld/wasm/Options.td

196

Should this description also say that keep-section can be used more than once? or is there some other way to know that?

clang/lib/Driver/ToolChains/WebAssembly.cpp

if (!Args.hasArg(options::OPT_nostdlib, options::OPT_nodefaultlibs)) {
 
     CmdArgs.push_back("-lc");
     AddRunTimeLibs(ToolChain, ToolChain.getDriver(), CmdArgs, Args);
   }
 
   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))
         OOpt = "4";
       else if (A->getOption().matches(options::OPT_O0))
         OOpt = "0";
       else if (A->getOption().matches(options::OPT_O))
         OOpt = A->getValue();
 
       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));
       }
     }
   }
 }
 
 /// Given a base library directory, append path components to form the
 /// LTO directory.

lld/test/wasm/strip-all.s

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

-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

struct Configuration {
   bool experimentalPic;
   bool emitRelocs;
   bool exportAll;
   bool exportDynamic;
   bool exportTable;
   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;
   bool importTable;
   bool importUndefined;
   std::optional<bool> is64;
   bool mergeDataSegments;
   bool pie;

lld/wasm/Driver.cpp

static void readConfigs(opt::InputArgList &args) {
   config->ltoPartitions = args::getInteger(args, OPT_lto_partitions, 1);
   config->ltoDebugPassManager = args.hasArg(OPT_lto_debug_pass_manager);
   config->mapFile = args.getLastArgValue(OPT_Map);
   config->optimize = args::getInteger(args, OPT_O, 1);
   config->outputFile = args.getLastArgValue(OPT_o);
   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);
   config->pie = args.hasFlag(OPT_pie, OPT_no_pie, false);
   config->printGcSections =
       args.hasFlag(OPT_print_gc_sections, OPT_no_print_gc_sections, false);
   config->saveTemps = args.hasArg(OPT_save_temps);
   config->searchPaths = args::getStrings(args, OPT_library_path);

lld/wasm/Options.td

def export_table: FF<"export-table">,
   HelpText<"Export function table to the environment">;
 
 def growable_table: FF<"growable-table">,
   HelpText<"Remove maximum size from function table, allowing table to grow">;
 
 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=">,
   HelpText<"Import the module's memory from the passed module with the passed name.">,
   MetaVarName<"<module>,<name>">;
 
 def export_memory: FF<"export-memory">,
   HelpText<"Export the module's memory with the default name of \"memory\"">;

lld/wasm/SyntheticSections.h

 };
 
 // Create the custom "name" section containing debug symbol names.
 class NameSection : public SyntheticSection {
 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
     // module itself.
     return 1 + numNamedGlobals() + numNamedFunctions();
   }
   unsigned numNamedGlobals() const;
   unsigned numNamedFunctions() const;
   unsigned numNamedDataSegments() const;
 
 protected:
   ArrayRef<OutputSegment *> segments;
 };
 
 class ProducersSection : public SyntheticSection {
 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);
 
 protected:
   int fieldCount() const {
     return int(!languages.empty()) + int(!tools.empty()) + int(!sDKs.empty());
   }
   SmallVector<std::pair<std::string, std::string>, 8> languages;
   SmallVector<std::pair<std::string, std::string>, 8> tools;
   SmallVector<std::pair<std::string, std::string>, 8> sDKs;
 };
 
 class TargetFeaturesSection : public SyntheticSection {
 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;
 
   llvm::SmallSet<std::string, 8> features;
 };
 
 class RelocSection : public SyntheticSection {
 public: