From b5720354efb61f1ee8784fa6809427b8a68cffa7 Mon Sep 17 00:00:00 2001
From: Alexander Shaposhnikov <alexshap@fb.com>
Date: Wed, 21 Apr 2021 08:18:20 -0700
Subject: [PATCH] [lld][MachO] Refactor findCommand

Refactor findCommand to allow passing multiple types. NFC.

Test plan: make check-lld-macho

Differential revision: https://reviews.llvm.org/D100954
---
 lld/MachO/InputFiles.cpp | 36 +++++++++++++++++++-----------------
 lld/MachO/InputFiles.h   |  8 +++++---
 lld/MachO/ObjC.cpp       |  4 ++--
 3 files changed, 26 insertions(+), 22 deletions(-)

diff --git a/lld/MachO/InputFiles.cpp b/lld/MachO/InputFiles.cpp
index 6edeb2ccc284..61f2840b2f83 100644
--- a/lld/MachO/InputFiles.cpp
+++ b/lld/MachO/InputFiles.cpp
@@ -109,29 +109,31 @@ static Optional<PlatformInfo> getPlatformInfo(const InputFile *input) {
 
   using Header = typename LP::mach_header;
   auto *hdr = reinterpret_cast<const Header *>(input->mb.getBufferStart());
+
   PlatformInfo platformInfo;
   if (const auto *cmd =
           findCommand<build_version_command>(hdr, LC_BUILD_VERSION)) {
     platformInfo.target.Platform = static_cast<PlatformKind>(cmd->platform);
     platformInfo.minimum = decodeVersion(cmd->minos);
     return platformInfo;
-  } else if (const auto *cmd =
-                 findCommand<version_min_command>(hdr, LC_VERSION_MIN_MACOSX)) {
-    platformInfo.target.Platform = PlatformKind::macOS;
-    platformInfo.minimum = decodeVersion(cmd->version);
-    return platformInfo;
-  } else if (const auto *cmd = findCommand<version_min_command>(
-                 hdr, LC_VERSION_MIN_IPHONEOS)) {
-    platformInfo.target.Platform = PlatformKind::iOS;
-    platformInfo.minimum = decodeVersion(cmd->version);
-    return platformInfo;
-  } else if (const auto *cmd =
-                 findCommand<version_min_command>(hdr, LC_VERSION_MIN_TVOS)) {
-    platformInfo.target.Platform = PlatformKind::tvOS;
-    platformInfo.minimum = decodeVersion(cmd->version);
-  } else if (const auto *cmd = findCommand<version_min_command>(
-                 hdr, LC_VERSION_MIN_WATCHOS)) {
-    platformInfo.target.Platform = PlatformKind::watchOS;
+  }
+  if (const auto *cmd = findCommand<version_min_command>(
+          hdr, LC_VERSION_MIN_MACOSX, LC_VERSION_MIN_IPHONEOS,
+          LC_VERSION_MIN_TVOS, LC_VERSION_MIN_WATCHOS)) {
+    switch (cmd->cmd) {
+    case LC_VERSION_MIN_MACOSX:
+      platformInfo.target.Platform = PlatformKind::macOS;
+      break;
+    case LC_VERSION_MIN_IPHONEOS:
+      platformInfo.target.Platform = PlatformKind::iOS;
+      break;
+    case LC_VERSION_MIN_TVOS:
+      platformInfo.target.Platform = PlatformKind::tvOS;
+      break;
+    case LC_VERSION_MIN_WATCHOS:
+      platformInfo.target.Platform = PlatformKind::watchOS;
+      break;
+    }
     platformInfo.minimum = decodeVersion(cmd->version);
     return platformInfo;
   }
diff --git a/lld/MachO/InputFiles.h b/lld/MachO/InputFiles.h
index 6ead6900e19c..4db93bec8611 100644
--- a/lld/MachO/InputFiles.h
+++ b/lld/MachO/InputFiles.h
@@ -189,13 +189,15 @@ extern llvm::SetVector<InputFile *> inputFiles;
 
 llvm::Optional<MemoryBufferRef> readFile(StringRef path);
 
-template <class CommandType = llvm::MachO::load_command, class Header>
-const CommandType *findCommand(const Header *hdr, uint32_t type) {
+template <class CommandType = llvm::MachO::load_command, class Header,
+          class... Types>
+const CommandType *findCommand(const Header *hdr, Types... types) {
+  std::initializer_list<uint32_t> typesList{types...};
   const uint8_t *p = reinterpret_cast<const uint8_t *>(hdr) + sizeof(Header);
 
   for (uint32_t i = 0, n = hdr->ncmds; i < n; ++i) {
     auto *cmd = reinterpret_cast<const CommandType *>(p);
-    if (cmd->cmd == type)
+    if (llvm::is_contained(typesList, cmd->cmd))
       return cmd;
     p += cmd->cmdsize;
   }
diff --git a/lld/MachO/ObjC.cpp b/lld/MachO/ObjC.cpp
index 6a341b280576..8fd1628435c4 100644
--- a/lld/MachO/ObjC.cpp
+++ b/lld/MachO/ObjC.cpp
@@ -23,8 +23,8 @@ template <class LP> static bool hasObjCSection(MemoryBufferRef mb) {
 
   auto *hdr =
       reinterpret_cast<const typename LP::mach_header *>(mb.getBufferStart());
-  if (const load_command *cmd = findCommand(hdr, LP::segmentLCType)) {
-    auto *c = reinterpret_cast<const typename LP::segment_command *>(cmd);
+  if (const auto *c =
+          findCommand<typename LP::segment_command>(hdr, LP::segmentLCType)) {
     auto sectionHeaders =
         ArrayRef<Section>{reinterpret_cast<const Section *>(c + 1), c->nsects};
     for (const Section &sec : sectionHeaders) {