[WinLink] Fix use-after-return.

llvm-svn: 183603

2 files changed, 16 insertions, 7 deletions

diff --git a/lld/include/lld/ReaderWriter/PECOFFTargetInfo.h b/lld/include/lld/ReaderWriter/PECOFFTargetInfo.h
index 012e81158a2..32c7f3a0100 100644
--- a/lld/include/lld/ReaderWriter/PECOFFTargetInfo.h
+++ b/lld/include/lld/ReaderWriter/PECOFFTargetInfo.h
@@ -14,6 +14,7 @@
 #include "lld/ReaderWriter/Reader.h"
 #include "lld/ReaderWriter/Writer.h"
 
+#include "llvm/Support/Allocator.h"
 #include "llvm/Support/COFF.h"
 #include "llvm/Support/ErrorHandling.h"
 
@@ -59,12 +60,20 @@ public:
     return _minOSVersion;
   }
 
+  StringRef allocateString(const StringRef &ref) {
+    char *x = _extraStrings.Allocate<char>(ref.size() + 1);
+    memcpy(x, ref.data(), ref.size());
+    x[ref.size()] = '\0';
+    return x;
+  }
+
 private:
   llvm::COFF::WindowsSubsystem _subsystem;
   OSVersion _minOSVersion;
 
   mutable std::unique_ptr<Reader> _reader;
   mutable std::unique_ptr<Writer> _writer;
+  llvm::BumpPtrAllocator _extraStrings;
 };
 
 } // end namespace lld
diff --git a/lld/lib/Driver/WinLinkDriver.cpp b/lld/lib/Driver/WinLinkDriver.cpp
index 7bdbd4928a0..8730a9cc22e 100644
--- a/lld/lib/Driver/WinLinkDriver.cpp
+++ b/lld/lib/Driver/WinLinkDriver.cpp
@@ -127,19 +127,19 @@ bool parseSubsystemOption(PECOFFTargetInfo &info, std::string arg,
 }
 
 // Add ".obj" extension if the given path name has no file extension.
-std::string canonicalizeInputFileName(std::string path) {
+StringRef canonicalizeInputFileName(PECOFFTargetInfo &info, std::string path) {
   if (llvm::sys::path::extension(path).empty())
-    return path.append(".obj");
-  return path;
+    path.append(".obj");
+  return info.allocateString(path);
 }
 
 // Replace a file extension with ".exe". If the given file has no
 // extension, just add ".exe".
-std::string getDefaultOutputFileName(std::string path) {
+StringRef getDefaultOutputFileName(PECOFFTargetInfo &info, std::string path) {
   StringRef ext = llvm::sys::path::extension(path);
   if (!ext.empty())
     path.erase(path.size() - ext.size());
-  return path.append(".exe");
+  return info.allocateString(path.append(".exe"));
 }
 
 } // namespace
@@ -223,12 +223,12 @@ bool WinLinkDriver::parse(int argc, const char *argv[],
 
   // Add ".obj" extension for those who have no file extension.
   for (const StringRef &path : inputPaths)
-    info.appendInputFile(canonicalizeInputFileName(path));
+    info.appendInputFile(canonicalizeInputFileName(info, path));
 
   // If -out option was not specified, the default output file name is
   // constructed by replacing an extension with ".exe".
   if (info.outputPath().empty() && !inputPaths.empty())
-    info.setOutputPath(getDefaultOutputFileName(inputPaths[0]));
+    info.setOutputPath(getDefaultOutputFileName(info, inputPaths[0]));
 
   // Validate the combination of options used.
   return info.validate(diagnostics);