[PECOFF][Driver] Add -heap command line option.

llvm-svn: 183622

lld/include/lld/ReaderWriter/PECOFFTargetInfo.h

@@ -25,6 +25,8 @@ public:
   PECOFFTargetInfo()
       : _stackReserve(1024 * 1024),
         _stackCommit(4096),
+        _heapReserve(1024 * 1024),
+        _heapCommit(4096),
         _subsystem(llvm::COFF::IMAGE_SUBSYSTEM_UNKNOWN),
         _minOSVersion(6, 0) {}
 
@@ -44,11 +46,15 @@ public:
   virtual void addPasses(PassManager &pm) const {}
 
   void setStackReserve(uint64_t size) { _stackReserve = size; }
-  uint64_t getStackReserve() const { return _stackReserve; }
-
   void setStackCommit(uint64_t size) { _stackCommit = size; }
+  uint64_t getStackReserve() const { return _stackReserve; }
   uint64_t getStackCommit() const { return _stackCommit; }
 
+  void setHeapReserve(uint64_t size) { _heapReserve = size; }
+  void setHeapCommit(uint64_t size) { _heapCommit = size; }
+  uint64_t getHeapReserve() const { return _heapReserve; }
+  uint64_t getHeapCommit() const { return _heapCommit; }
+
   void setSubsystem(llvm::COFF::WindowsSubsystem ss) { _subsystem = ss; }
   llvm::COFF::WindowsSubsystem getSubsystem() const { return _subsystem; }
 
@@ -68,6 +74,8 @@ public:
 private:
   uint64_t _stackReserve;
   uint64_t _stackCommit;
+  uint64_t _heapReserve;
+  uint64_t _heapCommit;
   llvm::COFF::WindowsSubsystem _subsystem;
   OSVersion _minOSVersion;
 

lld/lib/Driver/WinLinkDriver.cpp

@@ -78,20 +78,44 @@ bool checkNumber(StringRef version, const char *errorMessage,
   return true;
 }
 
-// Parse -stack command line option. The form of the option is
-// "-stack:stackReserveSize[,stackCommitSize]".
+// Parse an argument for -stack or -heap. The expected string is
+// "reserveSize[,stackCommitSize]".
+bool parseMemoryOption(const StringRef &arg, raw_ostream &diagnostics,
+                       uint64_t &reserve, uint64_t &commit) {
+  StringRef reserveStr, commitStr;
+  llvm::tie(reserveStr, commitStr) = arg.split(',');
+  if (!checkNumber(reserveStr, "invalid stack size: ", diagnostics))
+    return false;
+  reserve = atoi(reserveStr.str().c_str());
+  if (!commitStr.empty()) {
+    if (!checkNumber(commitStr, "invalid stack size: ", diagnostics))
+      return false;
+    commit = atoi(commitStr.str().c_str());
+  }
+  return true;
+}
+
+// Parse -stack command line option
 bool parseStackOption(PECOFFTargetInfo &info, const StringRef &arg,
                       raw_ostream &diagnostics) {
-  StringRef reserve, commit;
-  llvm::tie(reserve, commit) = arg.split(',');
-  if (!checkNumber(reserve, "invalid stack size: ", diagnostics))
+  uint64_t reserve;
+  uint64_t commit = info.getStackCommit();
+  if (!parseMemoryOption(arg, diagnostics, reserve, commit))
     return false;
-  info.setStackReserve(atoi(reserve.str().c_str()));
-  if (!commit.empty()) {
-    if (!checkNumber(commit, "invalid stack size: ", diagnostics))
-      return false;
-    info.setStackCommit(atoi(commit.str().c_str()));
-  }
+  info.setStackReserve(reserve);
+  info.setStackCommit(commit);
+  return true;
+}
+
+// Parse -heap command line option.
+bool parseHeapOption(PECOFFTargetInfo &info, const StringRef &arg,
+                     raw_ostream &diagnostics) {
+  uint64_t reserve;
+  uint64_t commit = info.getHeapCommit();
+  if (!parseMemoryOption(arg, diagnostics, reserve, commit))
+    return false;
+  info.setHeapReserve(reserve);
+  info.setHeapCommit(commit);
   return true;
 }
 
@@ -216,6 +240,11 @@ bool WinLinkDriver::parse(int argc, const char *argv[],
     if (!parseStackOption(info, arg->getValue(), diagnostics))
       return true;
 
+  // Handle -heap
+  if (llvm::opt::Arg *arg = parsedArgs->getLastArg(OPT_heap))
+    if (!parseHeapOption(info, arg->getValue(), diagnostics))
+      return true;
+
   // Handle -subsystem
   if (llvm::opt::Arg *arg = parsedArgs->getLastArg(OPT_subsystem))
     if (!parseSubsystemOption(info, arg->getValue(), diagnostics))

lld/lib/Driver/WinLinkOptions.td

@@ -13,6 +13,10 @@ def stack : Separate<["-", "/"], "stack">,
     HelpText<"Size of the stack">;
 def stack_c: Joined<["-", "/"], "stack:">, Alias<stack>;
 
+def heap : Separate<["-", "/"], "heap">,
+    HelpText<"Size of the heap">;
+def heap_c: Joined<["-", "/"], "heap:">, Alias<heap>;
+
 def subsystem : Separate<["-", "/"], "subsystem">,
     HelpText<"Specify subsystem">;
 def subsystem_c: Joined<["-", "/"], "subsystem:">, Alias<subsystem>;

lld/lib/ReaderWriter/PECOFF/PECOFFTargetInfo.cpp

@@ -30,6 +30,13 @@ bool PECOFFTargetInfo::validate(raw_ostream &diagnostics) {
     return true;
   }
 
+  if (_heapReserve < _heapCommit) {
+    diagnostics << "Invalid heap size: reserve size must be equal to or "
+                << "greater than commit size, but got "
+                << _heapCommit << " and " << _heapReserve << ".\n";
+    return true;
+  }
+
   _reader = createReaderPECOFF(*this);
   _writer = createWriterPECOFF(*this);
   return false;

lld/lib/ReaderWriter/PECOFF/WriterPECOFF.cpp

@@ -157,8 +157,8 @@ public:
 
     _peHeader.SizeOfStackReserve = targetInfo.getStackReserve();
     _peHeader.SizeOfStackCommit = targetInfo.getStackCommit();
-    _peHeader.SizeOfHeapReserve = 0x100000;
-    _peHeader.SizeOfHeapCommit = 0x1000;
+    _peHeader.SizeOfHeapReserve = targetInfo.getHeapReserve();
+    _peHeader.SizeOfHeapCommit = targetInfo.getHeapCommit();
 
     // The number of data directory entries. We always have 16 entries.
     _peHeader.NumberOfRvaAndSize = 16;

lld/unittests/DriverTests/WinLinkDriverTest.cpp

@@ -96,4 +96,16 @@ TEST_F(WinLinkParserTest, StackReserveAndCommit) {
   EXPECT_EQ(16384, info->getStackReserve());
   EXPECT_EQ(8192, info->getStackCommit());
 }
+
+TEST_F(WinLinkParserTest, HeapReserve) {
+  parse("link.exe", "-heap", "8192", nullptr);
+  EXPECT_EQ(8192, info->getHeapReserve());
+  EXPECT_EQ(4096, info->getHeapCommit());
+}
+
+TEST_F(WinLinkParserTest, HeapReserveAndCommit) {
+  parse("link.exe", "-heap", "16384,8192", nullptr);
+  EXPECT_EQ(16384, info->getHeapReserve());
+  EXPECT_EQ(8192, info->getHeapCommit());
+}
 }  // end anonymous namespace