[llvm-install-name-tool] Merge install-name options

This diff merges all options for llvm-install-name-tool under a single
function processLoadCommands. Also adds another test case for -add_rpath
option.

Test plan: make check-all

Reviewed by: jhenderson, alexshap, smeenai, Ktwu

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

4 files changed, 124 insertions, 116 deletions

diff --git a/llvm/test/tools/llvm-objcopy/MachO/install-name-tool-add-rpath.test b/llvm/test/tools/llvm-objcopy/MachO/install-name-tool-add-rpath.test
index 1435c6b..7b21fdc 100644
--- a/llvm/test/tools/llvm-objcopy/MachO/install-name-tool-add-rpath.test
+++ b/llvm/test/tools/llvm-objcopy/MachO/install-name-tool-add-rpath.test
@@ -22,6 +22,13 @@
 
 # NO-INPUT: no input file specified
 
+## Add same RPATH twice:
+# RUN: not llvm-install-name-tool -add_rpath @executable_X \
+# RUN:                            -add_rpath @executable_X %t.i386 2>&1 \
+# RUN: | FileCheck --check-prefix=DOUBLE %s
+
+# DOUBLE: duplicate load command
+
 ## Check that cmdsize accounts for NULL terminator.
 # RUN: yaml2obj %p/Inputs/x86_64.yaml -o %t.x86_64
 # RUN: llvm-install-name-tool -add_rpath abcd %t.x86_64
diff --git a/llvm/tools/llvm-objcopy/CopyConfig.cpp b/llvm/tools/llvm-objcopy/CopyConfig.cpp
index f93406f..1fde54dd 100644
--- a/llvm/tools/llvm-objcopy/CopyConfig.cpp
+++ b/llvm/tools/llvm-objcopy/CopyConfig.cpp
@@ -874,42 +874,39 @@ parseInstallNameToolOptions(ArrayRef<const char *> ArgsArr) {
     auto Match = [=](StringRef RPath) { return RPath == Old || RPath == New; };
 
     // Cannot specify duplicate -rpath entries
-    auto It1 = find_if(Config.RPathsToUpdate,
-                       [&Match](const std::pair<StringRef, StringRef> &OldNew) {
-                         return Match(OldNew.first) || Match(OldNew.second);
-                       });
+    auto It1 = find_if(
+        Config.RPathsToUpdate,
+        [&Match](const DenseMap<StringRef, StringRef>::value_type &OldNew) {
+          return Match(OldNew.getFirst()) || Match(OldNew.getSecond());
+        });
     if (It1 != Config.RPathsToUpdate.end())
-      return createStringError(
-          errc::invalid_argument,
-          "cannot specify both -rpath %s %s and -rpath %s %s",
-          It1->first.str().c_str(), It1->second.str().c_str(),
-          Old.str().c_str(), New.str().c_str());
+      return createStringError(errc::invalid_argument,
+                               "cannot specify both -rpath " + It1->getFirst() +
+                                   " " + It1->getSecond() + " and -rpath " +
+                                   Old + " " + New);
 
     // Cannot specify the same rpath under both -delete_rpath and -rpath
     auto It2 = find_if(Config.RPathsToRemove, Match);
     if (It2 != Config.RPathsToRemove.end())
-      return createStringError(
-          errc::invalid_argument,
-          "cannot specify both -delete_rpath %s and -rpath %s %s",
-          It2->str().c_str(), Old.str().c_str(), New.str().c_str());
+      return createStringError(errc::invalid_argument,
+                               "cannot specify both -delete_rpath " + *It2 +
+                                   " and -rpath " + Old + " " + New);
 
     // Cannot specify the same rpath under both -add_rpath and -rpath
     auto It3 = find_if(Config.RPathToAdd, Match);
     if (It3 != Config.RPathToAdd.end())
-      return createStringError(
-          errc::invalid_argument,
-          "cannot specify both -add_rpath %s and -rpath %s %s",
-          It3->str().c_str(), Old.str().c_str(), New.str().c_str());
+      return createStringError(errc::invalid_argument,
+                               "cannot specify both -add_rpath " + *It3 +
+                                   " and -rpath " + Old + " " + New);
 
-    Config.RPathsToUpdate.emplace_back(Old, New);
+    Config.RPathsToUpdate.insert({Old, New});
   }
 
   if (auto *Arg = InputArgs.getLastArg(INSTALL_NAME_TOOL_id))
     Config.SharedLibId = Arg->getValue();
 
   for (auto *Arg : InputArgs.filtered(INSTALL_NAME_TOOL_change)) {
-    Config.InstallNamesToUpdate.emplace_back(Arg->getValue(0),
-                                             Arg->getValue(1));
+    Config.InstallNamesToUpdate.insert({Arg->getValue(0), Arg->getValue(1)});
   }
 
   SmallVector<StringRef, 2> Positional;
diff --git a/llvm/tools/llvm-objcopy/CopyConfig.h b/llvm/tools/llvm-objcopy/CopyConfig.h
index ce119de..1341dd6 100644
--- a/llvm/tools/llvm-objcopy/CopyConfig.h
+++ b/llvm/tools/llvm-objcopy/CopyConfig.h
@@ -178,8 +178,8 @@ struct CopyConfig {
   std::vector<StringRef> DumpSection;
   std::vector<StringRef> SymbolsToAdd;
   std::vector<StringRef> RPathToAdd;
-  std::vector<std::pair<StringRef, StringRef>> RPathsToUpdate;
-  std::vector<std::pair<StringRef, StringRef>> InstallNamesToUpdate;
+  DenseMap<StringRef, StringRef> RPathsToUpdate;
+  DenseMap<StringRef, StringRef> InstallNamesToUpdate;
   DenseSet<StringRef> RPathsToRemove;
 
   // install-name-tool's id option
diff --git a/llvm/tools/llvm-objcopy/MachO/MachOObjcopy.cpp b/llvm/tools/llvm-objcopy/MachO/MachOObjcopy.cpp
index 3844b6f..9d0c366 100644
--- a/llvm/tools/llvm-objcopy/MachO/MachOObjcopy.cpp
+++ b/llvm/tools/llvm-objcopy/MachO/MachOObjcopy.cpp
@@ -42,35 +42,6 @@ static StringRef getPayloadString(const LoadCommand &LC) {
       .rtrim('\0');
 }
 
-static Error removeLoadCommands(const CopyConfig &Config, Object &Obj) {
-  DenseSet<StringRef> RPathsToRemove(Config.RPathsToRemove.begin(),
-                                     Config.RPathsToRemove.end());
-
-  LoadCommandPred RemovePred = [&RPathsToRemove](const LoadCommand &LC) {
-    if (LC.MachOLoadCommand.load_command_data.cmd == MachO::LC_RPATH) {
-      StringRef RPath = getPayloadString(LC);
-      if (RPathsToRemove.count(RPath)) {
-        RPathsToRemove.erase(RPath);
-        return true;
-      }
-    }
-    return false;
-  };
-
-  if (Error E = Obj.removeLoadCommands(RemovePred))
-    return E;
-
-  // Emit an error if the Mach-O binary does not contain an rpath path name
-  // specified in -delete_rpath.
-  for (StringRef RPath : Config.RPathsToRemove) {
-    if (RPathsToRemove.count(RPath))
-      return createStringError(errc::invalid_argument,
-                               "no LC_RPATH load command with path: %s",
-                               RPath.str().c_str());
-  }
-  return Error::success();
-}
-
 static Error removeSections(const CopyConfig &Config, Object &Obj) {
   SectionPred RemovePred = [](const std::unique_ptr<Section> &) {
     return false;
@@ -157,6 +128,103 @@ static LoadCommand buildRPathLoadCommand(StringRef Path) {
   return LC;
 }
 
+static Error processLoadCommands(const CopyConfig &Config, Object &Obj) {
+  // Remove RPaths.
+  DenseSet<StringRef> RPathsToRemove(Config.RPathsToRemove.begin(),
+                                     Config.RPathsToRemove.end());
+
+  LoadCommandPred RemovePred = [&RPathsToRemove](const LoadCommand &LC) {
+    if (LC.MachOLoadCommand.load_command_data.cmd == MachO::LC_RPATH) {
+      StringRef RPath = getPayloadString(LC);
+      if (RPathsToRemove.count(RPath)) {
+        RPathsToRemove.erase(RPath);
+        return true;
+      }
+    }
+    return false;
+  };
+
+  if (Error E = Obj.removeLoadCommands(RemovePred))
+    return E;
+
+  // Emit an error if the Mach-O binary does not contain an rpath path name
+  // specified in -delete_rpath.
+  for (StringRef RPath : Config.RPathsToRemove) {
+    if (RPathsToRemove.count(RPath))
+      return createStringError(errc::invalid_argument,
+                               "no LC_RPATH load command with path: %s",
+                               RPath.str().c_str());
+  }
+
+  DenseSet<StringRef> RPaths;
+
+  // Get all existing RPaths.
+  for (LoadCommand &LC : Obj.LoadCommands) {
+    if (LC.MachOLoadCommand.load_command_data.cmd == MachO::LC_RPATH)
+      RPaths.insert(getPayloadString(LC));
+  }
+
+  // Throw errors for invalid RPaths.
+  for (const auto &OldNew : Config.RPathsToUpdate) {
+    StringRef Old, New;
+    std::tie(Old, New) = OldNew;
+    if (RPaths.count(Old) == 0)
+      return createStringError(errc::invalid_argument,
+                               "no LC_RPATH load command with path: " + Old);
+    if (RPaths.count(New) != 0)
+      return createStringError(errc::invalid_argument,
+                               "rpath " + New +
+                                   " would create a duplicate load command");
+  }
+
+  // Update load commands.
+  for (LoadCommand &LC : Obj.LoadCommands) {
+    switch (LC.MachOLoadCommand.load_command_data.cmd) {
+    case MachO::LC_ID_DYLIB:
+      if (Config.SharedLibId) {
+        StringRef Id = Config.SharedLibId.getValue();
+        if (Id.empty())
+          return createStringError(errc::invalid_argument,
+                                   "cannot specify an empty id");
+        updateLoadCommandPayloadString<MachO::dylib_command>(LC, Id);
+      }
+      break;
+
+    case MachO::LC_RPATH: {
+      StringRef RPath = getPayloadString(LC);
+      StringRef NewRPath = Config.RPathsToUpdate.lookup(RPath);
+      if (!NewRPath.empty())
+        updateLoadCommandPayloadString<MachO::rpath_command>(LC, NewRPath);
+      break;
+    }
+
+    // TODO: Add LC_REEXPORT_DYLIB, LC_LAZY_LOAD_DYLIB, and LC_LOAD_UPWARD_DYLIB
+    // here once llvm-objcopy supports them.
+    case MachO::LC_LOAD_DYLIB:
+    case MachO::LC_LOAD_WEAK_DYLIB:
+      StringRef InstallName = getPayloadString(LC);
+      StringRef NewInstallName =
+          Config.InstallNamesToUpdate.lookup(InstallName);
+      if (!NewInstallName.empty())
+        updateLoadCommandPayloadString<MachO::dylib_command>(LC,
+                                                             NewInstallName);
+      break;
+    }
+  }
+
+  // Add new RPaths.
+  for (StringRef RPath : Config.RPathToAdd) {
+    if (RPaths.count(RPath) != 0)
+      return createStringError(errc::invalid_argument,
+                               "rpath " + RPath +
+                                   " would create a duplicate load command");
+    RPaths.insert(RPath);
+    Obj.addLoadCommand(buildRPathLoadCommand(RPath));
+  }
+
+  return Error::success();
+}
+
 static Error dumpSectionToFile(StringRef SecName, StringRef Filename,
                                Object &Obj) {
   for (LoadCommand &LC : Obj.LoadCommands)
@@ -273,34 +341,6 @@ static Error handleArgs(const CopyConfig &Config, Object &Obj) {
       for (std::unique_ptr<Section> &Sec : LC.Sections)
         Sec->Relocations.clear();
 
-  for (LoadCommand &LC : Obj.LoadCommands) {
-    switch (LC.MachOLoadCommand.load_command_data.cmd) {
-    case MachO::LC_ID_DYLIB:
-      if (Config.SharedLibId) {
-        StringRef Id = Config.SharedLibId.getValue();
-        if (Id.empty())
-          return createStringError(errc::invalid_argument,
-                                   "cannot specify an empty id");
-        updateLoadCommandPayloadString<MachO::dylib_command>(LC, Id);
-      }
-      break;
-
-    // TODO: Add LC_REEXPORT_DYLIB, LC_LAZY_LOAD_DYLIB, and LC_LOAD_UPWARD_DYLIB
-    // here once llvm-objcopy supports them.
-    case MachO::LC_LOAD_DYLIB:
-    case MachO::LC_LOAD_WEAK_DYLIB:
-      StringRef Old, New;
-      StringRef CurrentInstallName = getPayloadString(LC);
-      for (const auto &InstallNamePair : Config.InstallNamesToUpdate) {
-        std::tie(Old, New) = InstallNamePair;
-        if (CurrentInstallName == Old) {
-          updateLoadCommandPayloadString<MachO::dylib_command>(LC, New);
-          break;
-        }
-      }
-    }
-  }
-
   for (const auto &Flag : Config.AddSection) {
     std::pair<StringRef, StringRef> SecPair = Flag.split("=");
     StringRef SecName = SecPair.first;
@@ -311,45 +351,9 @@ static Error handleArgs(const CopyConfig &Config, Object &Obj) {
       return E;
   }
 
-  if (Error E = removeLoadCommands(Config, Obj))
+  if (Error E = processLoadCommands(Config, Obj))
     return E;
 
-  StringRef Old, New;
-  for (const auto &OldNew : Config.RPathsToUpdate) {
-    std::tie(Old, New) = OldNew;
-
-    auto FindRPathLC = [&Obj](StringRef RPath) {
-      return find_if(Obj.LoadCommands, [=](const LoadCommand &LC) {
-        return LC.MachOLoadCommand.load_command_data.cmd == MachO::LC_RPATH &&
-               getPayloadString(LC) == RPath;
-      });
-    };
-
-    auto NewIt = FindRPathLC(New);
-    if (NewIt != Obj.LoadCommands.end())
-      return createStringError(errc::invalid_argument,
-                               "rpath " + New +
-                                   " would create a duplicate load command");
-
-    auto OldIt = FindRPathLC(Old);
-    if (OldIt == Obj.LoadCommands.end())
-      return createStringError(errc::invalid_argument,
-                               "no LC_RPATH load command with path: " + Old);
-
-    updateLoadCommandPayloadString<MachO::rpath_command>(*OldIt, New);
-  }
-
-  for (StringRef RPath : Config.RPathToAdd) {
-    for (LoadCommand &LC : Obj.LoadCommands) {
-      if (LC.MachOLoadCommand.load_command_data.cmd == MachO::LC_RPATH &&
-          RPath == getPayloadString(LC)) {
-        return createStringError(errc::invalid_argument,
-                                 "rpath " + RPath +
-                                     " would create a duplicate load command");
-      }
-    }
-    Obj.addLoadCommand(buildRPathLoadCommand(RPath));
-  }
   return Error::success();
 }