[Support] mmap when possible in getSTDIN.

llvm/include/llvm/Support/MemoryBuffer.h

@@ -97,7 +97,7 @@ class LLVM_ABI MemoryBuffer {
   /// least the specified alignment.
   static ErrorOr<std::unique_ptr<MemoryBuffer>>
   getFile(const Twine &Filename, bool IsText = false,
-          bool RequiresNullTerminator = true, bool IsVolatile = false,
+          bool RequiresNullTerminator = false, bool IsVolatile = false,
           std::optional<Align> Alignment = std::nullopt);
 
   /// Read all of the specified file into a MemoryBuffer as a stream
@@ -125,31 +125,32 @@ class LLVM_ABI MemoryBuffer {
   /// least the specified alignment.
   static ErrorOr<std::unique_ptr<MemoryBuffer>>
   getOpenFile(sys::fs::file_t FD, const Twine &Filename, uint64_t FileSize,
-              bool RequiresNullTerminator = true, bool IsVolatile = false,
+              bool RequiresNullTerminator = false, bool IsVolatile = false,
               std::optional<Align> Alignment = std::nullopt);
 
   /// Open the specified memory range as a MemoryBuffer. Note that InputData
   /// must be null terminated if RequiresNullTerminator is true.
   static std::unique_ptr<MemoryBuffer>
   getMemBuffer(StringRef InputData, StringRef BufferName = "",
-               bool RequiresNullTerminator = true);
+               bool RequiresNullTerminator = false);
 
   static std::unique_ptr<MemoryBuffer>
-  getMemBuffer(MemoryBufferRef Ref, bool RequiresNullTerminator = true);
+  getMemBuffer(MemoryBufferRef Ref, bool RequiresNullTerminator = false);
 
   /// Open the specified memory range as a MemoryBuffer, copying the contents
   /// and taking ownership of it. InputData does not have to be null terminated.
   static std::unique_ptr<MemoryBuffer>
   getMemBufferCopy(StringRef InputData, const Twine &BufferName = "");
 
   /// Read all of stdin into a file buffer, and return it.
-  static ErrorOr<std::unique_ptr<MemoryBuffer>> getSTDIN();
+  static ErrorOr<std::unique_ptr<MemoryBuffer>>
+  getSTDIN(bool RequiresNullTerminator = false);
 
   /// Open the specified file as a MemoryBuffer, or open stdin if the Filename
   /// is "-".
   static ErrorOr<std::unique_ptr<MemoryBuffer>>
   getFileOrSTDIN(const Twine &Filename, bool IsText = false,
-                 bool RequiresNullTerminator = true,
+                 bool RequiresNullTerminator = false,
                  std::optional<Align> Alignment = std::nullopt);
 
   /// Map a subrange of the specified file as a MemoryBuffer.

llvm/lib/AsmParser/LLLexer.cpp

@@ -175,19 +175,10 @@ LLLexer::LLLexer(StringRef StartBuf, SourceMgr &SM, SMDiagnostic &Err,
 }
 
 int LLLexer::getNextChar() {
-  char CurChar = *CurPtr++;
-  switch (CurChar) {
-  default: return (unsigned char)CurChar;
-  case 0:
-    // A nul character in the stream is either the end of the current buffer or
-    // a random nul in the file.  Disambiguate that here.
-    if (CurPtr-1 != CurBuf.end())
-      return 0;  // Just whitespace.
-
-    // Otherwise, return end of file.
-    --CurPtr;  // Another call to lex will return EOF again.
+  if (CurPtr == CurBuf.end())
     return EOF;
-  }
+
+  return *CurPtr++;
 }
 
 lltok::Kind LLLexer::LexToken() {

llvm/lib/Support/MemoryBuffer.cpp

@@ -165,7 +165,7 @@ MemoryBuffer::getFileOrSTDIN(const Twine &Filename, bool IsText,
   StringRef NameRef = Filename.toStringRef(NameBuf);
 
   if (NameRef == "-")
-    return getSTDIN();
+    return getSTDIN(RequiresNullTerminator);
   return getFile(Filename, IsText, RequiresNullTerminator,
                  /*IsVolatile=*/false, Alignment);
 }
@@ -567,12 +567,33 @@ ErrorOr<std::unique_ptr<MemoryBuffer>> MemoryBuffer::getOpenFileSlice(
                                        IsVolatile, Alignment);
 }
 
-ErrorOr<std::unique_ptr<MemoryBuffer>> MemoryBuffer::getSTDIN() {
+ErrorOr<std::unique_ptr<MemoryBuffer>>
+MemoryBuffer::getSTDIN(bool RequiresNullTerminator) {
   // Read in all of the data from stdin, we cannot mmap stdin.
-  //
-  // FIXME: That isn't necessarily true, we should try to mmap stdin and
-  // fallback if it fails.
   sys::ChangeStdinMode(sys::fs::OF_Text);
+  std::error_code EC;
+  sys::fs::file_type Type;
+  sys::fs::file_status Status;
+  EC = sys::fs::status(sys::fs::getStdinHandle(), Status);
+  if (EC)
+    return EC;
+
+  Type = Status.type();
+  // If the FD is regular file or block file,
+  // we try to create a mmap buffer first.
+  // If failed, rollback to read and copy.
+  if ((Type == sys::fs::file_type::regular_file ||
+       Type == sys::fs::file_type::block_file) &&
+      shouldUseMmap(sys::fs::getStdinHandle(), Status.getSize(),
+                    Status.getSize(), 0, RequiresNullTerminator,
+                    sys::Process::getPageSizeEstimate(), false)) {
+    std::unique_ptr<MemoryBuffer> Result(
+        new (NamedBufferAlloc("<stdin>")) MemoryBufferMMapFile<MemoryBuffer>(
+            RequiresNullTerminator, sys::fs::getStdinHandle(), Status.getSize(),
+            0, EC));
+    if (!EC && (!RequiresNullTerminator || *Result->getBufferEnd() == '\0'))
+      return std::move(Result);
+  }
 
   return getMemoryBufferForStream(sys::fs::getStdinHandle(), "<stdin>");
 }

llvm/lib/TableGen/TGLexer.cpp

@@ -138,18 +138,14 @@ bool TGLexer::processEOF() {
 }
 
 int TGLexer::getNextChar() {
+  if (CurPtr == CurBuf.end())
+    return EOF;
   char CurChar = *CurPtr++;
   switch (CurChar) {
   default:
     return (unsigned char)CurChar;
 
   case 0: {
-    // A NUL character in the stream is either the end of the current buffer or
-    // a spurious NUL in the file. Disambiguate that here.
-    if (CurPtr - 1 == CurBuf.end()) {
-      --CurPtr; // Arrange for another call to return EOF again.
-      return EOF;
-    }
     PrintError(getLoc(),
                "NUL character is invalid in source; treated as space");
     return ' ';
@@ -160,7 +156,8 @@ int TGLexer::getNextChar() {
     // Handle the newline character by ignoring it and incrementing the line
     // count. However, be careful about 'dos style' files with \n\r in them.
     // Only treat a \n\r or \r\n as a single line.
-    if ((*CurPtr == '\n' || (*CurPtr == '\r')) && *CurPtr != CurChar)
+    if (CurPtr != CurBuf.end() && (*CurPtr == '\n' || (*CurPtr == '\r')) &&
+        *CurPtr != CurChar)
       ++CurPtr; // Eat the two char newline sequence.
     return '\n';
   }

llvm/unittests/AsmParser/AsmParserTest.cpp

@@ -39,10 +39,10 @@ TEST(AsmParserTest, NonNullTerminatedInput) {
   LLVMContext Ctx;
   StringRef Source = "; Empty module \n\1\2";
   SMDiagnostic Error;
-  std::unique_ptr<Module> Mod;
-  EXPECT_DEATH(Mod = parseAssemblyString(Source.substr(0, Source.size() - 2),
-                                         Error, Ctx),
-               "Buffer is not null terminated!");
+  std::unique_ptr<Module> Mod =
+      parseAssemblyString(Source.substr(0, Source.size() - 2), Error, Ctx);
+  EXPECT_TRUE(Mod != nullptr);
+  EXPECT_TRUE(Error.getMessage().empty());
 }
 
 #endif

mlir/lib/Tools/PDLL/Parser/Lexer.cpp

@@ -132,26 +132,19 @@ Token Lexer::emitError(const char *loc, const Twine &msg) {
 }
 
 int Lexer::getNextChar() {
+  if (curPtr == curBuffer.end())
+    return EOF;
   char curChar = *curPtr++;
   switch (curChar) {
   default:
     return static_cast<unsigned char>(curChar);
-  case 0: {
-    // A nul character in the stream is either the end of the current buffer
-    // or a random nul in the file. Disambiguate that here.
-    if (curPtr - 1 != curBuffer.end())
-      return 0;
-
-    // Otherwise, return end of file.
-    --curPtr;
-    return EOF;
-  }
   case '\n':
   case '\r':
     // Handle the newline character by ignoring it and incrementing the line
     // count. However, be careful about 'dos style' files with \n\r in them.
     // Only treat a \n\r or \r\n as a single line.
-    if ((*curPtr == '\n' || (*curPtr == '\r')) && *curPtr != curChar)
+    if (curPtr != curBuffer.end() && (*curPtr == '\n' || (*curPtr == '\r')) &&
+        *curPtr != curChar)
       ++curPtr;
     return '\n';
   }

mlir/tools/mlir-tblgen/FormatGen.cpp

@@ -53,26 +53,19 @@ FormatToken FormatLexer::emitErrorAndNote(SMLoc loc, const Twine &msg,
 }
 
 int FormatLexer::getNextChar() {
+  if (curPtr == curBuffer.end())
+    return EOF;
   char curChar = *curPtr++;
   switch (curChar) {
   default:
     return (unsigned char)curChar;
-  case 0: {
-    // A nul character in the stream is either the end of the current buffer or
-    // a random nul in the file. Disambiguate that here.
-    if (curPtr - 1 != curBuffer.end())
-      return 0;
-
-    // Otherwise, return end of file.
-    --curPtr;
-    return EOF;
-  }
   case '\n':
   case '\r':
     // Handle the newline character by ignoring it and incrementing the line
     // count. However, be careful about 'dos style' files with \n\r in them.
     // Only treat a \n\r or \r\n as a single line.
-    if ((*curPtr == '\n' || (*curPtr == '\r')) && *curPtr != curChar)
+    if (curPtr != curBuffer.end() && (*curPtr == '\n' || (*curPtr == '\r')) &&
+        *curPtr != curChar)
       ++curPtr;
     return '\n';
   }