4 files changed, 28 insertions, 10 deletions

diff --git a/lld/ELF/Config.h b/lld/ELF/Config.h
index 25ec18b64d2..6eec7e1430f 100644
--- a/lld/ELF/Config.h
+++ b/lld/ELF/Config.h
@@ -33,6 +33,9 @@ enum ELFKind {
 // For --build-id.
 enum class BuildIdKind { None, Fnv1, Md5, Sha1, Hexstring, Uuid };
 
+// For --discard-{all,locals,none}.
+enum class DiscardPolicy { Default, All, Locals, None };
+
 // For --strip-{all,debug}.
 enum class StripPolicy { None, All, Debug };
 
@@ -84,9 +87,6 @@ struct Configuration {
   bool BsymbolicFunctions;
   bool Demangle = true;
   bool DisableVerify;
-  bool DiscardAll;
-  bool DiscardLocals;
-  bool DiscardNone;
   bool EhFrameHdr;
   bool EnableNewDtags;
   bool ExportDynamic;
@@ -118,6 +118,7 @@ struct Configuration {
   bool ZNow;
   bool ZOrigin;
   bool ZRelro;
+  DiscardPolicy Discard;
   StripPolicy Strip = StripPolicy::None;
   UnresolvedPolicy UnresolvedSymbols;
   BuildIdKind BuildId = BuildIdKind::None;
diff --git a/lld/ELF/Driver.cpp b/lld/ELF/Driver.cpp
index 3f2071b1fdf..24166cc405a 100644
--- a/lld/ELF/Driver.cpp
+++ b/lld/ELF/Driver.cpp
@@ -344,6 +344,25 @@ static bool isOutputFormatBinary(opt::InputArgList &Args) {
   return false;
 }
 
+static DiscardPolicy getDiscardOption(opt::InputArgList &Args) {
+  auto *Arg =
+      Args.getLastArg(OPT_discard_all, OPT_discard_locals, OPT_discard_none);
+  if (!Arg)
+    return DiscardPolicy::Default;
+
+  switch (Arg->getOption().getID()) {
+  case OPT_discard_all:
+    return DiscardPolicy::All;
+  case OPT_discard_locals:
+    return DiscardPolicy::Locals;
+  case OPT_discard_none:
+    return DiscardPolicy::None;
+  default:
+    llvm_unreachable("unknown discard option");
+  }
+}
+
+
 static StripPolicy getStripOption(opt::InputArgList &Args) {
   if (auto *Arg = Args.getLastArg(OPT_strip_all, OPT_strip_debug)) {
     if (Arg->getOption().getID() == OPT_strip_all)
@@ -376,9 +395,7 @@ void LinkerDriver::readConfigs(opt::InputArgList &Args) {
   Config->BsymbolicFunctions = Args.hasArg(OPT_Bsymbolic_functions);
   Config->Demangle = !Args.hasArg(OPT_no_demangle);
   Config->DisableVerify = Args.hasArg(OPT_disable_verify);
-  Config->DiscardAll = Args.hasArg(OPT_discard_all);
-  Config->DiscardLocals = Args.hasArg(OPT_discard_locals);
-  Config->DiscardNone = Args.hasArg(OPT_discard_none);
+  Config->Discard = getDiscardOption(Args);
   Config->EhFrameHdr = Args.hasArg(OPT_eh_frame_hdr);
   Config->EnableNewDtags = !Args.hasArg(OPT_disable_new_dtags);
   Config->ExportDynamic = Args.hasArg(OPT_export_dynamic);
diff --git a/lld/ELF/OutputSections.cpp b/lld/ELF/OutputSections.cpp
index 6fdaf2d24ce..05c4fa903a7 100644
--- a/lld/ELF/OutputSections.cpp
+++ b/lld/ELF/OutputSections.cpp
@@ -1373,7 +1373,7 @@ template <class ELFT> void SymbolTableSection<ELFT>::writeTo(uint8_t *Buf) {
 
   // All symbols with STB_LOCAL binding precede the weak and global symbols.
   // .dynsym only contains global symbols.
-  if (!Config->DiscardAll && !StrTabSec.isDynamic())
+  if (Config->Discard != DiscardPolicy::All && !StrTabSec.isDynamic())
     writeLocalSymbols(Buf);
 
   writeGlobalSymbols(Buf);
diff --git a/lld/ELF/Writer.cpp b/lld/ELF/Writer.cpp
index 1c1f02c22c0..11c15a3990f 100644
--- a/lld/ELF/Writer.cpp
+++ b/lld/ELF/Writer.cpp
@@ -240,7 +240,7 @@ static std::vector<DefinedCommon<ELFT> *> getCommonSymbols() {
 
 // The main function of the writer.
 template <class ELFT> void Writer<ELFT>::run() {
-  if (!Config->DiscardAll)
+  if (Config->Discard != DiscardPolicy::All)
     copyLocalSymbols();
   addReservedSymbols();
 
@@ -329,7 +329,7 @@ static bool shouldKeepInSymtab(InputSectionBase<ELFT> *Sec, StringRef SymName,
   if (Sec == &InputSection<ELFT>::Discarded)
     return false;
 
-  if (Config->DiscardNone)
+  if (Config->Discard == DiscardPolicy::None)
     return true;
 
   // In ELF assembly .L symbols are normally discarded by the assembler.
@@ -340,7 +340,7 @@ static bool shouldKeepInSymtab(InputSectionBase<ELFT> *Sec, StringRef SymName,
   if (!SymName.startswith(".L") && !SymName.empty())
     return true;
 
-  if (Config->DiscardLocals)
+  if (Config->Discard == DiscardPolicy::Locals)
     return false;
 
   return !(Sec->getSectionHdr()->sh_flags & SHF_MERGE);