diff options
Diffstat (limited to 'util/file_stream.hh')
-rw-r--r-- | util/file_stream.hh | 89 |
1 files changed, 89 insertions, 0 deletions
diff --git a/util/file_stream.hh b/util/file_stream.hh new file mode 100644 index 000000000..ae9ad5aa7 --- /dev/null +++ b/util/file_stream.hh @@ -0,0 +1,89 @@ +/* Like std::ofstream but without being incredibly slow. Backed by a raw fd. + * Supports most of the built-in types except for long double. + */ +#ifndef UTIL_FILE_STREAM_H +#define UTIL_FILE_STREAM_H + +#include "util/fake_ostream.hh" +#include "util/file.hh" +#include "util/scoped.hh" + +#include <cassert> +#include <cstring> + +#include <stdint.h> + +namespace util { + +class FileStream : public FakeOStream<FileStream> { + public: + FileStream(int out = -1, std::size_t buffer_size = 8192) + : buf_(util::MallocOrThrow(std::max<std::size_t>(buffer_size, kToStringMaxBytes))), + current_(static_cast<char*>(buf_.get())), + end_(current_ + std::max<std::size_t>(buffer_size, kToStringMaxBytes)), + fd_(out) {} + + ~FileStream() { + flush(); + } + + void SetFD(int to) { + flush(); + fd_ = to; + } + + FileStream &flush() { + if (current_ != buf_.get()) { + util::WriteOrThrow(fd_, buf_.get(), current_ - (char*)buf_.get()); + current_ = static_cast<char*>(buf_.get()); + } + return *this; + } + + // For writes of arbitrary size. + FileStream &write(const void *data, std::size_t length) { + if (UTIL_LIKELY(current_ + length <= end_)) { + std::memcpy(current_, data, length); + current_ += length; + return *this; + } + flush(); + if (current_ + length <= end_) { + std::memcpy(current_, data, length); + current_ += length; + } else { + util::WriteOrThrow(fd_, data, length); + } + return *this; + } + + FileStream &seekp(uint64_t to) { + util::SeekOrThrow(fd_, to); + return *this; + } + + protected: + friend class FakeOStream<FileStream>; + // For writes directly to buffer guaranteed to have amount < buffer size. + char *Ensure(std::size_t amount) { + if (UTIL_UNLIKELY(current_ + amount > end_)) { + flush(); + assert(current_ + amount <= end_); + } + return current_; + } + + void AdvanceTo(char *to) { + current_ = to; + assert(current_ <= end_); + } + + private: + util::scoped_malloc buf_; + char *current_, *end_; + int fd_; +}; + +} // namespace + +#endif |