00001
00002
00003
00004 #ifndef UTIL_FAKE_OFSTREAM_H
00005 #define UTIL_FAKE_OFSTREAM_H
00006
00007 #include "util/file.hh"
00008 #include "util/float_to_string.hh"
00009 #include "util/integer_to_string.hh"
00010 #include "util/scoped.hh"
00011 #include "util/string_piece.hh"
00012
00013 #include <cassert>
00014 #include <cstring>
00015
00016 #include <stdint.h>
00017
00018 namespace util {
00019 class FakeOFStream {
00020 public:
00021
00022
00023
00024
00025
00026
00027 explicit FakeOFStream(int out = -1, std::size_t buffer_size = 1048576)
00028 : buf_(util::MallocOrThrow(std::max(buffer_size, (size_t)20))),
00029 current_(static_cast<char*>(buf_.get())),
00030 end_(current_ + std::max(buffer_size, (size_t)20)),
00031 fd_(out) {}
00032
00033 ~FakeOFStream() {
00034
00035 flush();
00036 }
00037
00038 void SetFD(int to) {
00039 flush();
00040 fd_ = to;
00041 }
00042
00043 FakeOFStream &write(const void *data, std::size_t length) {
00044 if (UTIL_LIKELY(current_ + length <= end_)) {
00045 std::memcpy(current_, data, length);
00046 current_ += length;
00047 return *this;
00048 }
00049 flush();
00050 if (current_ + length <= end_) {
00051 std::memcpy(current_, data, length);
00052 current_ += length;
00053 } else {
00054 util::WriteOrThrow(fd_, data, length);
00055 }
00056 return *this;
00057 }
00058
00059
00060 FakeOFStream &operator<<(StringPiece str) {
00061 return write(str.data(), str.size());
00062 }
00063
00064
00065
00066
00067 private:
00068 template <int Arg> struct EnableIfKludge {
00069 typedef FakeOFStream type;
00070 };
00071 public:
00072 template <class T> typename EnableIfKludge<ToStringBuf<T>::kBytes>::type &operator<<(const T value) {
00073 EnsureRemaining(ToStringBuf<T>::kBytes);
00074 current_ = ToString(value, current_);
00075 assert(current_ <= end_);
00076 return *this;
00077 }
00078
00079 FakeOFStream &operator<<(char c) {
00080 EnsureRemaining(1);
00081 *current_++ = c;
00082 return *this;
00083 }
00084
00085 FakeOFStream &operator<<(unsigned char c) {
00086 EnsureRemaining(1);
00087 *current_++ = static_cast<char>(c);
00088 return *this;
00089 }
00090
00091
00092
00093
00094
00095
00096
00097
00098
00099 FakeOFStream &operator<<(std::size_t value) {
00100 EnsureRemaining(ToStringBuf<uint64_t>::kBytes);
00101 current_ = ToString(static_cast<uint64_t>(value), current_);
00102 return *this;
00103 }
00104
00105
00106 void flush() {
00107 if (current_ != buf_.get()) {
00108 util::WriteOrThrow(fd_, buf_.get(), current_ - (char*)buf_.get());
00109 current_ = static_cast<char*>(buf_.get());
00110 }
00111 }
00112
00113
00114 void Finish() {
00115 flush();
00116 buf_.reset();
00117 current_ = NULL;
00118 util::FSyncOrThrow(fd_);
00119 }
00120
00121 private:
00122 void EnsureRemaining(std::size_t amount) {
00123 if (UTIL_UNLIKELY(current_ + amount > end_)) {
00124 flush();
00125 assert(current_ + amount <= end_);
00126 }
00127 }
00128
00129 util::scoped_malloc buf_;
00130 char *current_, *end_;
00131
00132 int fd_;
00133 };
00134
00135 }
00136
00137 #endif