2 files changed, 92 insertions, 0 deletions

diff --git a/llvm/lib/Support/Unix/Program.inc b/llvm/lib/Support/Unix/Program.inc
index 7bf6eceda73..4124340dfed 100644
--- a/llvm/lib/Support/Unix/Program.inc
+++ b/llvm/lib/Support/Unix/Program.inc
@@ -17,6 +17,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "Unix.h"
+#include "llvm/ADT/StringExtras.h"
 #include "llvm/Support/Compiler.h"
 #include "llvm/Support/FileSystem.h"
 #include "llvm/Support/raw_ostream.h"
@@ -100,6 +101,33 @@ sys::FindProgramByName(const std::string& progName) {
   return "";
 }
 
+ErrorOr<std::string> sys::findProgramByName(StringRef Name,
+                                            ArrayRef<StringRef> Paths) {
+  assert(!Name.empty() && "Must have a name!");
+  // Use the given path verbatim if it contains any slashes; this matches
+  // the behavior of sh(1) and friends.
+  if (Name.find('/') != StringRef::npos)
+    return std::string(Name);
+
+  if (Paths.empty()) {
+    SmallVector<StringRef, 16> SearchPaths;
+    SplitString(std::getenv("PATH"), SearchPaths, ":");
+    return findProgramByName(Name, SearchPaths);
+  }
+
+  for (auto Path : Paths) {
+    if (Path.empty())
+      continue;
+
+    // Check to see if this first directory contains the executable...
+    SmallString<128> FilePath(Path);
+    sys::path::append(FilePath, Name);
+    if (sys::fs::can_execute(FilePath.c_str()))
+      return std::string(FilePath.str()); // Found the executable!
+  }
+  return std::errc::no_such_file_or_directory;
+}
+
 static bool RedirectIO(const StringRef *Path, int FD, std::string* ErrMsg) {
   if (!Path) // Noop
     return false;
diff --git a/llvm/lib/Support/Windows/Program.inc b/llvm/lib/Support/Windows/Program.inc
index 995c9fedb96..5755b0ad035 100644
--- a/llvm/lib/Support/Windows/Program.inc
+++ b/llvm/lib/Support/Windows/Program.inc
@@ -12,9 +12,11 @@
 //===----------------------------------------------------------------------===//
 
 #include "WindowsSupport.h"
+#include "llvm/ADT/StringExtras.h"
 #include "llvm/Support/ConvertUTF.h"
 #include "llvm/Support/FileSystem.h"
 #include "llvm/Support/raw_ostream.h"
+#include "llvm/Support/WindowsError.h"
 #include <cstdio>
 #include <fcntl.h>
 #include <io.h>
@@ -69,6 +71,68 @@ std::string sys::FindProgramByName(const std::string &progName) {
   return std::string(result.data(), result.size());
 }
 
+ErrorOr<std::string> sys::findProgramByName(StringRef Name,
+                                            ArrayRef<StringRef> Paths) {
+  assert(!Name.empty() && "Must have a name!");
+
+  if (Name.find_first_of("/\\") != StringRef::npos)
+    return std::string(Name);
+
+  const char16_t *Path = nullptr;
+  std::u16string PathStorage;
+  if (!Paths.empty()) {
+    PathStorage.reserve(Paths.size() * MAX_PATH);
+    for (int i = 0; i < Paths.size(); ++i) {
+      if (i)
+        PathStorage.push_back(';');
+      StringRef P = Paths[i];
+      SmallVector<wchar_t, MAX_PATH> TmpPath;
+      if (std::error_code EC = windows::UTF8ToUTF16(P, TmpPath))
+        return EC;
+      PathStorage.append(TmpPath.begin(), TmpPath.end());
+    }
+    Path = PathStorage.c_str();
+  }
+
+  SmallVector<wchar_t, MAX_PATH> U16Name;
+  if (std::error_code EC = windows::UTF8ToUTF16(Name, U16Name))
+    return EC;
+
+  SmallVector<StringRef, 12> PathExts;
+  PathExts.push_back("");
+  SplitString(std::getenv("PATHEXT"), PathExts, ";");
+
+  SmallVector<wchar_t, MAX_PATH> U16Result;
+  DWORD Len = MAX_PATH;
+  for (StringRef Ext : PathExts) {
+    SmallVector<wchar_t, MAX_PATH> U16Ext;
+    if (std::error_code EC = windows::UTF8ToUTF16(Ext, U16Ext))
+      return EC;
+
+    do {
+      U16Result.reserve(Len);
+      Len = ::SearchPathW((const wchar_t *)Path, c_str(U16Name),
+                          U16Ext.empty() ? nullptr : c_str(U16Ext),
+                          U16Result.capacity(), U16Result.data(), nullptr);
+    } while (Len > U16Result.capacity());
+
+    if (Len != 0)
+      break; // Found it.
+  }
+
+  if (Len == 0)
+    return mapWindowsError(::GetLastError());
+
+  U16Result.set_size(Len);
+
+  SmallVector<char, MAX_PATH> U8Result;
+  if (std::error_code EC =
+          windows::UTF16ToUTF8(U16Result.data(), U16Result.size(), U8Result))
+    return EC;
+
+  return std::string(U8Result.begin(), U8Result.end());
+}
+
 static HANDLE RedirectIO(const StringRef *path, int fd, std::string* ErrMsg) {
   HANDLE h;
   if (path == 0) {