1 files changed, 49 insertions, 58 deletions

diff --git a/llvm/lib/Support/CommandLine.cpp b/llvm/lib/Support/CommandLine.cpp
index 25510fa58ff..20f129a4f3d 100644
--- a/llvm/lib/Support/CommandLine.cpp
+++ b/llvm/lib/Support/CommandLine.cpp
@@ -142,7 +142,7 @@ public:
   // This collects Options added with the cl::DefaultOption flag. Since they can
   // be overridden, they are not added to the appropriate SubCommands until
   // ParseCommandLineOptions actually runs.
-  SmallVector<Option*, 4> DefaultOptions;
+  SmallVector<std::pair<Option*, SubCommand*>, 4> DefaultOptions;
 
   // This collects the different option categories that have been registered.
   SmallPtrSet<OptionCategory *, 16> RegisteredOptionCategories;
@@ -182,15 +182,16 @@ public:
   }
 
   void addLiteralOption(Option &Opt, StringRef Name) {
-    if (Opt.Subs.empty())
-      addLiteralOption(Opt, &*TopLevelSubCommand, Name);
-    else {
-      for (auto SC : Opt.Subs)
-        addLiteralOption(Opt, SC, Name);
-    }
+    for(SubCommand *SC: Opt.getSubCommands())
+      addLiteralOption(Opt, SC, Name);
   }
 
-  void addOption(Option *O, SubCommand *SC) {
+  void addOption(Option *O, SubCommand *SC, bool ProcessDefaultOptions = false) {
+    if (!ProcessDefaultOptions && O->isDefaultOption()) {
+      DefaultOptions.push_back(std::make_pair(O, SC));
+      return;
+    }
+
     bool HadErrors = false;
     if (O->hasArgStr()) {
       // If it's a DefaultOption, check to make sure it isn't already there.
@@ -232,22 +233,14 @@ public:
       for (const auto &Sub : RegisteredSubCommands) {
         if (SC == Sub)
           continue;
-        addOption(O, Sub);
+        addOption(O, Sub, ProcessDefaultOptions);
       }
     }
   }
 
-  void addOption(Option *O, bool ProcessDefaultOption = false) {
-    if (!ProcessDefaultOption && O->isDefaultOption()) {
-      DefaultOptions.push_back(O);
-      return;
-    }
-
-    if (O->Subs.empty()) {
-      addOption(O, &*TopLevelSubCommand);
-    } else {
-      for (auto SC : O->Subs)
-        addOption(O, SC);
+  void addDefaultOptions() {
+    for (std::pair<Option *, SubCommand *> &DO : DefaultOptions) {
+      addOption(DO.first, DO.second, true);
     }
   }
 
@@ -285,17 +278,8 @@ public:
   }
 
   void removeOption(Option *O) {
-    if (O->Subs.empty())
-      removeOption(O, &*TopLevelSubCommand);
-    else {
-      if (O->isInAllSubCommands()) {
-        for (auto SC : RegisteredSubCommands)
-          removeOption(O, SC);
-      } else {
-        for (auto SC : O->Subs)
-          removeOption(O, SC);
-      }
-    }
+    for (auto SC : RegisteredSubCommands)
+      removeOption(O, SC);
   }
 
   bool hasOptions(const SubCommand &Sub) const {
@@ -324,17 +308,8 @@ public:
   }
 
   void updateArgStr(Option *O, StringRef NewName) {
-    if (O->Subs.empty())
-      updateArgStr(O, NewName, &*TopLevelSubCommand);
-    else {
-      if (O->isInAllSubCommands()) {
-        for (auto SC : RegisteredSubCommands)
-          updateArgStr(O, NewName, SC);
-      } else {
-        for (auto SC : O->Subs)
-          updateArgStr(O, NewName, SC);
-      }
-    }
+    for (auto SC : RegisteredSubCommands)
+      updateArgStr(O, NewName, SC);
   }
 
   void printOptionValues();
@@ -427,13 +402,38 @@ extrahelp::extrahelp(StringRef Help) : morehelp(Help) {
   GlobalParser->MoreHelp.push_back(Help);
 }
 
-void Option::addArgument() {
-  GlobalParser->addOption(this);
+void Option::addArgument(SubCommand &SC) {
+  GlobalParser->addOption(this, &SC);
   FullyInitialized = true;
 }
 
 void Option::removeArgument() { GlobalParser->removeOption(this); }
 
+SmallPtrSet<OptionCategory *, 1> Option::getCategories() const {
+  SmallPtrSet<OptionCategory *, 1> Cats;
+  for (OptionCategory *C: GlobalParser->RegisteredOptionCategories) {
+    if (C->MemberOptions.find(this) != C->MemberOptions.end())
+      Cats.insert(C);
+  }
+  if (Cats.empty())
+    Cats.insert(&GeneralCategory);
+  return Cats;
+}
+
+SmallPtrSet<SubCommand *, 1> Option::getSubCommands() const {
+  // This can happen for enums and literal options.
+  if (ArgStr.empty())
+    return SmallPtrSet<SubCommand *, 1>{&*TopLevelSubCommand};
+
+  SmallPtrSet<SubCommand *, 1> Subs;
+  for (SubCommand *SC : GlobalParser->getRegisteredSubcommands()) {
+    auto I = SC->OptionsMap.find(ArgStr);
+    if (I != SC->OptionsMap.end() && I->getValue() == this)
+      Subs.insert(SC);
+  }
+  return Subs;
+}
+
 void Option::setArgStr(StringRef S) {
   if (FullyInitialized)
     GlobalParser->updateArgStr(this, S);
@@ -444,14 +444,7 @@ void Option::setArgStr(StringRef S) {
 }
 
 void Option::addCategory(OptionCategory &C) {
-  assert(!Categories.empty() && "Categories cannot be empty.");
-  // Maintain backward compatibility by replacing the default GeneralCategory
-  // if it's still set.  Otherwise, just add the new one.  The GeneralCategory
-  // must be explicitly added if you want multiple categories that include it.
-  if (&C != &GeneralCategory && Categories[0] == &GeneralCategory)
-    Categories[0] = &C;
-  else if (find(Categories, &C) == Categories.end())
-    Categories.push_back(&C);
+  C.MemberOptions.insert(this);
 }
 
 void Option::reset() {
@@ -1302,9 +1295,7 @@ bool CommandLineParser::ParseCommandLineOptions(int argc,
   auto &SinkOpts = ChosenSubCommand->SinkOpts;
   auto &OptionsMap = ChosenSubCommand->OptionsMap;
 
-  for (auto O: DefaultOptions) {
-    addOption(O, true);
-  }
+  addDefaultOptions();
 
   if (ConsumeAfterOpt) {
     assert(PositionalOpts.size() > 0 &&
@@ -2204,7 +2195,7 @@ protected:
     // options within categories will also be alphabetically sorted.
     for (size_t I = 0, E = Opts.size(); I != E; ++I) {
       Option *Opt = Opts[I].second;
-      for (auto &Cat : Opt->Categories) {
+      for (auto *Cat : Opt->getCategories()) {
         assert(CategorizedOptions.count(Cat) > 0 &&
                "Option has an unregistered category");
         CategorizedOptions[Cat].push_back(Opt);
@@ -2465,7 +2456,7 @@ cl::getRegisteredSubcommands() {
 
 void cl::HideUnrelatedOptions(cl::OptionCategory &Category, SubCommand &Sub) {
   for (auto &I : Sub.OptionsMap) {
-    for (auto &Cat : I.second->Categories) {
+    for (OptionCategory *Cat : I.second->getCategories()) {
       if (Cat != &Category &&
           Cat != &GenericCategory)
         I.second->setHiddenFlag(cl::ReallyHidden);
@@ -2476,7 +2467,7 @@ void cl::HideUnrelatedOptions(cl::OptionCategory &Category, SubCommand &Sub) {
 void cl::HideUnrelatedOptions(ArrayRef<const cl::OptionCategory *> Categories,
                               SubCommand &Sub) {
   for (auto &I : Sub.OptionsMap) {
-    for (auto &Cat : I.second->Categories) {
+    for (OptionCategory *Cat : I.second->getCategories()) {
       if (find(Categories, Cat) == Categories.end() && Cat != &GenericCategory)
         I.second->setHiddenFlag(cl::ReallyHidden);
     }