[llvm-rc] Use proper search algorithm for finding resources.

Previously we would only look in the current directory for a
resource, which might not be the same as the directory of the
rc file.  Furthermore, MSVC rc supports a /I option, and can
also look in the system environment.  This patch adds support
for this search algorithm.

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

git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@315499 91177308-0d34-0410-b5e6-96231b3b80d8

5 files changed, 85 insertions, 29 deletions

diff --git a/tools/llvm-rc/ResourceFileWriter.cpp b/tools/llvm-rc/ResourceFileWriter.cpp
index 8534b6c576a..c43f128eec2 100644
--- a/tools/llvm-rc/ResourceFileWriter.cpp
+++ b/tools/llvm-rc/ResourceFileWriter.cpp
@@ -17,6 +17,8 @@
 #include "llvm/Support/ConvertUTF.h"
 #include "llvm/Support/Endian.h"
 #include "llvm/Support/EndianStream.h"
+#include "llvm/Support/Path.h"
+#include "llvm/Support/Process.h"
 
 using namespace llvm::support;
 
@@ -41,12 +43,13 @@ public:
   ~ContextKeeper() { FileWriter->ObjectData = SavedInfo; }
 };
 
-static Error createError(Twine Message,
+static Error createError(const Twine &Message,
                          std::errc Type = std::errc::invalid_argument) {
   return make_error<StringError>(Message, std::make_error_code(Type));
 }
 
-static Error checkNumberFits(uint32_t Number, size_t MaxBits, Twine FieldName) {
+static Error checkNumberFits(uint32_t Number, size_t MaxBits,
+                             const Twine &FieldName) {
   assert(1 <= MaxBits && MaxBits <= 32);
   if (!(Number >> MaxBits))
     return Error::success();
@@ -56,13 +59,13 @@ static Error checkNumberFits(uint32_t Number, size_t MaxBits, Twine FieldName) {
 }
 
 template <typename FitType>
-static Error checkNumberFits(uint32_t Number, Twine FieldName) {
+static Error checkNumberFits(uint32_t Number, const Twine &FieldName) {
   return checkNumberFits(Number, sizeof(FitType) * 8, FieldName);
 }
 
 // A similar function for signed integers.
 template <typename FitType>
-static Error checkSignedNumberFits(uint32_t Number, Twine FieldName,
+static Error checkSignedNumberFits(uint32_t Number, const Twine &FieldName,
                                    bool CanBeNegative) {
   int32_t SignedNum = Number;
   if (SignedNum < std::numeric_limits<FitType>::min() ||
@@ -79,13 +82,13 @@ static Error checkSignedNumberFits(uint32_t Number, Twine FieldName,
   return Error::success();
 }
 
-static Error checkRCInt(RCInt Number, Twine FieldName) {
+static Error checkRCInt(RCInt Number, const Twine &FieldName) {
   if (Number.isLong())
     return Error::success();
   return checkNumberFits<uint16_t>(Number, FieldName);
 }
 
-static Error checkIntOrString(IntOrString Value, Twine FieldName) {
+static Error checkIntOrString(IntOrString Value, const Twine &FieldName) {
   if (!Value.isInt())
     return Error::success();
   return checkNumberFits<uint16_t>(Value.getInt(), FieldName);
@@ -356,15 +359,10 @@ Error ResourceFileWriter::appendFile(StringRef Filename) {
   bool IsLong;
   stripQuotes(Filename, IsLong);
 
-  // Filename path should be relative to the current working directory.
-  // FIXME: docs say so, but reality is more complicated, script
-  // location and include paths must be taken into account.
-  ErrorOr<std::unique_ptr<MemoryBuffer>> File =
-      MemoryBuffer::getFile(Filename, -1, false);
+  auto File = loadFile(Filename);
   if (!File)
-    return make_error<StringError>("Error opening file '" + Filename +
-                                       "': " + File.getError().message(),
-                                   File.getError());
+    return File.takeError();
+
   *FS << (*File)->getBuffer();
   return Error::success();
 }
@@ -805,15 +803,10 @@ Error ResourceFileWriter::visitIconOrCursorResource(const RCResource *Base) {
 
   bool IsLong;
   stripQuotes(FileStr, IsLong);
-  ErrorOr<std::unique_ptr<MemoryBuffer>> File =
-      MemoryBuffer::getFile(FileStr, -1, false);
+  auto File = loadFile(FileStr);
 
   if (!File)
-    return make_error<StringError>(
-        "Error opening " +
-            Twine(Type == IconCursorGroupType::Icon ? "icon" : "cursor") +
-            " '" + FileStr + "': " + File.getError().message(),
-        File.getError());
+    return File.takeError();
 
   BinaryStreamReader Reader((*File)->getBuffer(), support::little);
 
@@ -1413,5 +1406,43 @@ Error ResourceFileWriter::writeVersionInfoBody(const RCResource *Base) {
   return Error::success();
 }
 
+Expected<std::unique_ptr<MemoryBuffer>>
+ResourceFileWriter::loadFile(StringRef File) const {
+  SmallString<128> Path;
+  SmallString<128> Cwd;
+  std::unique_ptr<MemoryBuffer> Result;
+
+  // 1. The current working directory.
+  sys::fs::current_path(Cwd);
+  Path.assign(Cwd.begin(), Cwd.end());
+  sys::path::append(Path, File);
+  if (sys::fs::exists(Path))
+    return errorOrToExpected(MemoryBuffer::getFile(Path, -1i64, false));
+
+  // 2. The directory of the input resource file, if it is different from the
+  // current
+  //    working directory.
+  StringRef InputFileDir = sys::path::parent_path(Params.InputFilePath);
+  Path.assign(InputFileDir.begin(), InputFileDir.end());
+  sys::path::append(Path, File);
+  if (sys::fs::exists(Path))
+    return errorOrToExpected(MemoryBuffer::getFile(Path, -1i64, false));
+
+  // 3. All of the include directories specified on the command line.
+  for (StringRef ForceInclude : Params.Include) {
+    Path.assign(ForceInclude.begin(), ForceInclude.end());
+    sys::path::append(Path, File);
+    if (sys::fs::exists(Path))
+      return errorOrToExpected(MemoryBuffer::getFile(Path, -1i64, false));
+  }
+
+  if (auto Result =
+          llvm::sys::Process::FindInEnvPath("INCLUDE", File, Params.NoInclude))
+    return errorOrToExpected(MemoryBuffer::getFile(*Result, -1i64, false));
+
+  return make_error<StringError>("error : file not found : " + Twine(File),
+                                 inconvertibleErrorCode());
+}
+
 } // namespace rc
 } // namespace llvm
diff --git a/tools/llvm-rc/ResourceFileWriter.h b/tools/llvm-rc/ResourceFileWriter.h
index 8d193d6a948..b06b8cf8a6f 100644
--- a/tools/llvm-rc/ResourceFileWriter.h
+++ b/tools/llvm-rc/ResourceFileWriter.h
@@ -22,10 +22,17 @@
 namespace llvm {
 namespace rc {
 
+struct SearchParams {
+  std::vector<std::string> Include;   // Additional folders to search for files.
+  std::vector<std::string> NoInclude; // Folders to exclude from file search.
+  StringRef InputFilePath;            // The full path of the input file.
+};
+
 class ResourceFileWriter : public Visitor {
 public:
-  ResourceFileWriter(std::unique_ptr<raw_fd_ostream> Stream)
-      : FS(std::move(Stream)), IconCursorID(1) {
+  ResourceFileWriter(const SearchParams &Params,
+                     std::unique_ptr<raw_fd_ostream> Stream)
+      : Params(Params), FS(std::move(Stream)), IconCursorID(1) {
     assert(FS && "Output stream needs to be provided to the serializator");
   }
 
@@ -136,6 +143,8 @@ private:
   Error writeVersionInfoBlock(const VersionInfoBlock &);
   Error writeVersionInfoValue(const VersionInfoValue &);
 
+  const SearchParams &Params;
+
   // Output stream handling.
   std::unique_ptr<raw_fd_ostream> FS;
 
@@ -170,6 +179,8 @@ private:
 
   void padStream(uint64_t Length);
 
+  Expected<std::unique_ptr<MemoryBuffer>> loadFile(StringRef File) const;
+
   // Icon and cursor IDs are allocated starting from 1 and increasing for
   // each icon/cursor dumped. This maintains the current ID to be allocated.
   uint16_t IconCursorID;
diff --git a/tools/llvm-rc/ResourceScriptParser.cpp b/tools/llvm-rc/ResourceScriptParser.cpp
index 4acae313558..769b47a20bd 100644
--- a/tools/llvm-rc/ResourceScriptParser.cpp
+++ b/tools/llvm-rc/ResourceScriptParser.cpp
@@ -12,6 +12,10 @@
 //===---------------------------------------------------------------------===//
 
 #include "ResourceScriptParser.h"
+#include "llvm/Option/ArgList.h"
+#include "llvm/Support/FileSystem.h"
+#include "llvm/Support/Path.h"
+#include "llvm/Support/Process.h"
 
 // Take an expression returning llvm::Error and forward the error if it exists.
 #define RETURN_IF_ERROR(Expr)                                                  \
diff --git a/tools/llvm-rc/ResourceScriptParser.h b/tools/llvm-rc/ResourceScriptParser.h
index 1a124d4ee2e..84fdfd5a586 100644
--- a/tools/llvm-rc/ResourceScriptParser.h
+++ b/tools/llvm-rc/ResourceScriptParser.h
@@ -25,6 +25,9 @@
 #include <vector>
 
 namespace llvm {
+namespace opt {
+class InputArgList;
+}
 namespace rc {
 
 class RCParser {
@@ -51,7 +54,7 @@ public:
     LocIter ErrorLoc, FileEnd;
   };
 
-  RCParser(std::vector<RCToken> TokenList);
+  explicit RCParser(std::vector<RCToken> TokenList);
 
   // Reads and returns a single resource definition, or error message if any
   // occurred.
diff --git a/tools/llvm-rc/llvm-rc.cpp b/tools/llvm-rc/llvm-rc.cpp
index 2a4faeb2d2e..f82a0dbe0e3 100644
--- a/tools/llvm-rc/llvm-rc.cpp
+++ b/tools/llvm-rc/llvm-rc.cpp
@@ -68,7 +68,7 @@ public:
 
 static ExitOnError ExitOnErr;
 
-LLVM_ATTRIBUTE_NORETURN static void fatalError(Twine Message) {
+LLVM_ATTRIBUTE_NORETURN static void fatalError(const Twine &Message) {
   errs() << Message << "\n";
   exit(1);
 }
@@ -107,10 +107,10 @@ int main(int argc_, const char *argv_[]) {
   }
 
   // Read and tokenize the input file.
-  const Twine &Filename = InArgsInfo[0];
-  ErrorOr<std::unique_ptr<MemoryBuffer>> File = MemoryBuffer::getFile(Filename);
+  ErrorOr<std::unique_ptr<MemoryBuffer>> File =
+      MemoryBuffer::getFile(InArgsInfo[0]);
   if (!File) {
-    fatalError("Error opening file '" + Filename +
+    fatalError("Error opening file '" + Twine(InArgsInfo[0]) +
                "': " + File.getError().message());
   }
 
@@ -138,6 +138,13 @@ int main(int argc_, const char *argv_[]) {
     }
   }
 
+  SearchParams Params;
+  SmallString<128> InputFile(InArgsInfo[0]);
+  llvm::sys::fs::make_absolute(InputFile);
+  Params.InputFilePath = InputFile;
+  Params.Include = InputArgs.getAllArgValues(OPT_INCLUDE);
+  Params.NoInclude = InputArgs.getAllArgValues(OPT_NOINCLUDE);
+
   std::unique_ptr<ResourceFileWriter> Visitor;
   bool IsDryRun = InputArgs.hasArg(OPT_DRY_RUN);
 
@@ -153,7 +160,7 @@ int main(int argc_, const char *argv_[]) {
     if (EC)
       fatalError("Error opening output file '" + OutArgsInfo[0] +
                  "': " + EC.message());
-    Visitor = llvm::make_unique<ResourceFileWriter>(std::move(FOut));
+    Visitor = llvm::make_unique<ResourceFileWriter>(Params, std::move(FOut));
     Visitor->AppendNull = InputArgs.hasArg(OPT_ADD_NULL);
 
     ExitOnErr(NullResource().visit(Visitor.get()));