00001
00002
00003
00004
00005
00006
00007
00008
00009
00010
00011
00012
00013
00014 #ifndef OR_TOOLS_BASE_RECORDIO_H_
00015 #define OR_TOOLS_BASE_RECORDIO_H_
00016
00017 #include <string>
00018 #include "base/file.h"
00019 #include "base/scoped_ptr.h"
00020
00021
00022
00023 namespace operations_research {
00024
00025
00026
00027
00028
00029
00030
00031
00032 class RecordWriter {
00033 public:
00034
00035 static const int kMagicNumber;
00036
00037 explicit RecordWriter(File* const file);
00038
00039 template <class P> bool WriteProtocolMessage(const P& proto) {
00040 std::string uncompressed_buffer;
00041 proto.SerializeToString(&uncompressed_buffer);
00042 const uint64 uncompressed_size = uncompressed_buffer.size();
00043 const std::string compressed_buffer =
00044 use_compression_ ? Compress(uncompressed_buffer) : "";
00045 const uint64 compressed_size = compressed_buffer.size();
00046 if (file_->Write(&kMagicNumber, sizeof(kMagicNumber)) !=
00047 sizeof(kMagicNumber)) {
00048 return false;
00049 }
00050 if (file_->Write(&uncompressed_size, sizeof(uncompressed_size)) !=
00051 sizeof(uncompressed_size)) {
00052 return false;
00053 }
00054 if (file_->Write(&compressed_size, sizeof(compressed_size)) !=
00055 sizeof(compressed_size)) {
00056 return false;
00057 }
00058 if (use_compression_) {
00059 if (file_->Write(compressed_buffer.c_str(), compressed_size) !=
00060 compressed_size) {
00061 return false;
00062 }
00063 } else {
00064 if (file_->Write(uncompressed_buffer.c_str(), uncompressed_size) !=
00065 uncompressed_size) {
00066 return false;
00067 }
00068 }
00069 return true;
00070 }
00071
00072 bool Close();
00073
00074 void set_use_compression(bool use_compression);
00075
00076 private:
00077 std::string Compress(const std::string& input) const;
00078 File* const file_;
00079 bool use_compression_;
00080 };
00081
00082
00083
00084 class RecordReader {
00085 public:
00086 explicit RecordReader(File* const file);
00087
00088 template <class P> bool ReadProtocolMessage(P* const proto) {
00089 uint64 usize = 0;
00090 uint64 csize = 0;
00091 int magic_number = 0;
00092 if (file_->Read(&magic_number, sizeof(magic_number)) !=
00093 sizeof(magic_number)) {
00094 return false;
00095 }
00096 if (magic_number != RecordWriter::kMagicNumber) {
00097 return false;
00098 }
00099 if (file_->Read(&usize, sizeof(usize)) != sizeof(usize)) {
00100 return false;
00101 }
00102 if (file_->Read(&csize, sizeof(csize)) != sizeof(csize)) {
00103 return false;
00104 }
00105 scoped_array<char> buffer(new char[usize + 1]);
00106 if (csize != 0) {
00107 scoped_array<char> compressed_buffer(new char[csize + 1]);
00108 if (file_->Read(compressed_buffer.get(), csize) != csize) {
00109 return false;
00110 }
00111 compressed_buffer[csize] = '\0';
00112 Uncompress(compressed_buffer.get(), usize, buffer.get(), usize);
00113 } else {
00114 if (file_->Read(buffer.get(), usize) != usize) {
00115 return false;
00116 }
00117 }
00118 proto->ParseFromArray(buffer.get(), usize);
00119 return true;
00120 }
00121
00122
00123 bool Close();
00124
00125 private:
00126 void Uncompress(const char* const source,
00127 uint64 source_size,
00128 char* const output_buffer,
00129 uint64 output_size) const;
00130
00131 File* const file_;
00132 };
00133 }
00134
00135 #endif // OR_TOOLS_BASE_RECORDIO_H_