[PECOFF] Support HEAPSIZE directive.

llvm-svn: 198994

4 files changed, 81 insertions, 3 deletions

diff --git a/lld/include/lld/Driver/WinLinkModuleDef.h b/lld/include/lld/Driver/WinLinkModuleDef.h
index 1cdf10eaa65..25e68591b72 100644
--- a/lld/include/lld/Driver/WinLinkModuleDef.h
+++ b/lld/include/lld/Driver/WinLinkModuleDef.h
@@ -27,9 +27,11 @@ enum class Kind {
   unknown,
   eof,
   identifier,
+  comma,
   equal,
   kw_data,
   kw_exports,
+  kw_heapsize,
   kw_noname,
 };
 
@@ -58,7 +60,7 @@ private:
 
 class Directive {
 public:
-  enum class Kind { exports };
+  enum class Kind { exports, heapsize };
 
   Kind getKind() const { return _kind; }
   virtual ~Directive() {}
@@ -87,6 +89,23 @@ private:
   const std::vector<PECOFFLinkingContext::ExportDesc> _exports;
 };
 
+class Heapsize : public Directive {
+public:
+  explicit Heapsize(uint64_t reserve, uint64_t commit)
+      : Directive(Kind::heapsize), _reserve(reserve), _commit(commit) {}
+
+  static bool classof(const Directive *dir) {
+    return dir->getKind() == Kind::heapsize;
+  }
+
+  uint64_t getReserve() const { return _reserve; }
+  uint64_t getCommit() const { return _commit; }
+
+private:
+  const uint64_t _reserve;
+  const uint64_t _commit;
+};
+
 class Parser {
 public:
   explicit Parser(Lexer &lex, llvm::BumpPtrAllocator &alloc)
@@ -100,6 +119,7 @@ private:
   void error(const Token &tok, Twine msg);
 
   bool parseExport(PECOFFLinkingContext::ExportDesc &result);
+  bool parseHeapsize(uint64_t &reserve, uint64_t &commit);
 
   Lexer &_lex;
   llvm::BumpPtrAllocator &_alloc;
diff --git a/lld/lib/Driver/WinLinkDriver.cpp b/lld/lib/Driver/WinLinkDriver.cpp
index ae324e54a5d..5224124a6d4 100644
--- a/lld/lib/Driver/WinLinkDriver.cpp
+++ b/lld/lib/Driver/WinLinkDriver.cpp
@@ -932,6 +932,9 @@ WinLinkDriver::parse(int argc, const char *argv[], PECOFFLinkingContext &ctx,
           desc.name = ctx.decorateSymbol(desc.name);
           ctx.addDllExport(desc);
         }
+      } else if (auto *hs = dyn_cast<moduledef::Heapsize>(dir.getValue())) {
+        ctx.setHeapReserve(hs->getReserve());
+        ctx.setHeapCommit(hs->getCommit());
       } else {
         llvm::dbgs() << static_cast<int>(dir.getValue()->getKind()) << "\n";
         llvm_unreachable("Unknown module-definition directive.\n");
diff --git a/lld/lib/Driver/WinLinkModuleDef.cpp b/lld/lib/Driver/WinLinkModuleDef.cpp
index 24bcd5c153a..3faa45518fd 100644
--- a/lld/lib/Driver/WinLinkModuleDef.cpp
+++ b/lld/lib/Driver/WinLinkModuleDef.cpp
@@ -27,6 +27,9 @@ Token Lexer::lex() {
   case '=':
     _buffer = _buffer.drop_front();
     return Token(Kind::equal, "=");
+  case ',':
+    _buffer = _buffer.drop_front();
+    return Token(Kind::comma, ",");
   case '"': {
     size_t end = _buffer.find('"', 1);
     Token ret;
@@ -42,11 +45,12 @@ Token Lexer::lex() {
   default: {
     size_t end = _buffer.find_first_not_of(
         "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ"
-        "0123456789_.*~+!@#$%^&*()");
+        "0123456789_.*~+!@#$%^&*()/");
     StringRef word = _buffer.substr(0, end);
     Kind kind = llvm::StringSwitch<Kind>(word)
                     .Case("DATA", Kind::kw_data)
                     .Case("EXPORTS", Kind::kw_exports)
+                    .Case("HEAPSIZE", Kind::kw_heapsize)
                     .Case("NONAME", Kind::kw_noname)
                     .Default(Kind::identifier);
     _buffer = (end == _buffer.npos) ? "" : _buffer.drop_front(end);
@@ -84,7 +88,13 @@ llvm::Optional<Directive *> Parser::parse() {
     }
     return new (_alloc) Exports(exports);
   }
-  error(_tok, "Expected EXPORTS");
+  if (_tok._kind == Kind::kw_heapsize) {
+    uint64_t reserve, commit;
+    if (!parseHeapsize(reserve, commit))
+      return llvm::None;
+    return new (_alloc) Heapsize(reserve, commit);
+  }
+  error(_tok, "Expected EXPORTS or HEAPSIZE");
   return llvm::None;
 }
 
@@ -117,5 +127,30 @@ bool Parser::parseExport(PECOFFLinkingContext::ExportDesc &result) {
   }
 }
 
+bool Parser::parseHeapsize(uint64_t &reserve, uint64_t &commit) {
+  consumeToken();
+  if (_tok._kind != Kind::identifier) {
+    ungetToken();
+    return false;
+  }
+  if (_tok._range.getAsInteger(0, reserve))
+    return false;
+
+  consumeToken();
+  if (_tok._kind != Kind::comma) {
+    ungetToken();
+    commit = 0;
+    return true;
+  }
+  consumeToken();
+  if (_tok._kind != Kind::identifier) {
+    ungetToken();
+    return false;
+  }
+  if (_tok._range.getAsInteger(0, commit))
+    return false;
+  return true;
+}
+
 } // moddef
 } // namespace lld
diff --git a/lld/unittests/DriverTests/WinLinkModuleDefTest.cpp b/lld/unittests/DriverTests/WinLinkModuleDefTest.cpp
index 5f3b079c8ea..b4f1b5263c4 100644
--- a/lld/unittests/DriverTests/WinLinkModuleDefTest.cpp
+++ b/lld/unittests/DriverTests/WinLinkModuleDefTest.cpp
@@ -67,3 +67,23 @@ TEST_F(ParserTest, Exports) {
   EXPECT_EQ(exports[4].noname, true);
   EXPECT_EQ(exports[4].isData, true);
 }
+
+TEST_F(ParserTest, Heapsize1) {
+  llvm::BumpPtrAllocator alloc;
+  llvm::Optional<moduledef::Directive *> dir = parse("HEAPSIZE 65536", alloc);
+  EXPECT_TRUE(dir.hasValue());
+  auto *heapsize = dyn_cast<moduledef::Heapsize>(dir.getValue());
+  EXPECT_TRUE(heapsize != nullptr);
+  EXPECT_EQ(65536U, heapsize->getReserve());
+  EXPECT_EQ(0U, heapsize->getCommit());
+}
+
+TEST_F(ParserTest, Heapsize2) {
+  llvm::BumpPtrAllocator alloc;
+  llvm::Optional<moduledef::Directive *> dir = parse("HEAPSIZE 65536, 8192", alloc);
+  EXPECT_TRUE(dir.hasValue());
+  auto *heapsize = dyn_cast<moduledef::Heapsize>(dir.getValue());
+  EXPECT_TRUE(heapsize != nullptr);
+  EXPECT_EQ(65536U, heapsize->getReserve());
+  EXPECT_EQ(8192U, heapsize->getCommit());
+}