Skip to content

Commit a9d7f6c

Browse files
committed
[Support] mmap when possible in getSTDIN.
When stdin is executed like ./prog < file, we are possible to mmap the buffer so that we can reduce the total memory usage if we try to open a large file. For example, programs like llvm-strings iterates the buffer until the end of given size.
1 parent 5401210 commit a9d7f6c

File tree

3 files changed

+30
-8
lines changed

3 files changed

+30
-8
lines changed

llvm/include/llvm/Support/MemoryBuffer.h

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -75,7 +75,6 @@ class LLVM_ABI MemoryBuffer {
7575
/// Return an identifier for this buffer, typically the filename it was read
7676
/// from.
7777
virtual StringRef getBufferIdentifier() const { return "Unknown buffer"; }
78-
7978
/// For read-only MemoryBuffer_MMap, mark the buffer as unused in the near
8079
/// future and the kernel can free resources associated with it. Further
8180
/// access is supported but may be expensive. This calls
@@ -143,7 +142,8 @@ class LLVM_ABI MemoryBuffer {
143142
getMemBufferCopy(StringRef InputData, const Twine &BufferName = "");
144143

145144
/// Read all of stdin into a file buffer, and return it.
146-
static ErrorOr<std::unique_ptr<MemoryBuffer>> getSTDIN();
145+
static ErrorOr<std::unique_ptr<MemoryBuffer>>
146+
getSTDIN(bool RequiresNullTerminatorRequires = true);
147147

148148
/// Open the specified file as a MemoryBuffer, or open stdin if the Filename
149149
/// is "-".

llvm/lib/Support/MemoryBuffer.cpp

Lines changed: 26 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -165,7 +165,7 @@ MemoryBuffer::getFileOrSTDIN(const Twine &Filename, bool IsText,
165165
StringRef NameRef = Filename.toStringRef(NameBuf);
166166

167167
if (NameRef == "-")
168-
return getSTDIN();
168+
return getSTDIN(RequiresNullTerminator);
169169
return getFile(Filename, IsText, RequiresNullTerminator,
170170
/*IsVolatile=*/false, Alignment);
171171
}
@@ -567,12 +567,33 @@ ErrorOr<std::unique_ptr<MemoryBuffer>> MemoryBuffer::getOpenFileSlice(
567567
IsVolatile, Alignment);
568568
}
569569

570-
ErrorOr<std::unique_ptr<MemoryBuffer>> MemoryBuffer::getSTDIN() {
570+
ErrorOr<std::unique_ptr<MemoryBuffer>>
571+
MemoryBuffer::getSTDIN(bool RequiresNullTerminator) {
571572
// Read in all of the data from stdin, we cannot mmap stdin.
572-
//
573-
// FIXME: That isn't necessarily true, we should try to mmap stdin and
574-
// fallback if it fails.
575573
sys::ChangeStdinMode(sys::fs::OF_Text);
574+
std::error_code EC;
575+
sys::fs::file_type Type;
576+
sys::fs::file_status Status;
577+
EC = sys::fs::status(sys::fs::getStdinHandle(), Status);
578+
if (EC)
579+
return EC;
580+
581+
Type = Status.type();
582+
// If the FD is regular file or block file,
583+
// we try to create a mmap buffer first.
584+
// If failed, rollback to read and copy.
585+
if ((Type == sys::fs::file_type::regular_file ||
586+
Type == sys::fs::file_type::block_file) &&
587+
shouldUseMmap(sys::fs::getStdinHandle(), Status.getSize(),
588+
Status.getSize(), 0, RequiresNullTerminator,
589+
sys::Process::getPageSizeEstimate(), false)) {
590+
std::unique_ptr<MemoryBuffer> Result(
591+
new (NamedBufferAlloc("<stdin>")) MemoryBufferMMapFile<MemoryBuffer>(
592+
RequiresNullTerminator, sys::fs::getStdinHandle(), Status.getSize(),
593+
0, EC));
594+
if (!EC && (!RequiresNullTerminator || *Result->getBufferEnd() == '\0'))
595+
return std::move(Result);
596+
}
576597

577598
return getMemoryBufferForStream(sys::fs::getStdinHandle(), "<stdin>");
578599
}

llvm/tools/llvm-strings/llvm-strings.cpp

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -175,7 +175,8 @@ int main(int argc, char **argv) {
175175

176176
for (const auto &File : InputFileNames) {
177177
ErrorOr<std::unique_ptr<MemoryBuffer>> Buffer =
178-
MemoryBuffer::getFileOrSTDIN(File, /*IsText=*/true);
178+
MemoryBuffer::getFileOrSTDIN(File, /*IsText=*/true,
179+
/*RequiresNullTerminator*/ false);
179180
if (std::error_code EC = Buffer.getError())
180181
errs() << File << ": " << EC.message() << '\n';
181182
else

0 commit comments

Comments
 (0)