OR-Tools  8.2
recordio.h
Go to the documentation of this file.
1 // Copyright 2010-2018 Google LLC
2 // Licensed under the Apache License, Version 2.0 (the "License");
3 // you may not use this file except in compliance with the License.
4 // You may obtain a copy of the License at
5 //
6 // http://www.apache.org/licenses/LICENSE-2.0
7 //
8 // Unless required by applicable law or agreed to in writing, software
9 // distributed under the License is distributed on an "AS IS" BASIS,
10 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
11 // See the License for the specific language governing permissions and
12 // limitations under the License.
13 
14 #ifndef OR_TOOLS_BASE_RECORDIO_H_
15 #define OR_TOOLS_BASE_RECORDIO_H_
16 
17 #include <memory>
18 #include <string>
19 
20 #include "ortools/base/file.h"
21 
22 // This file defines some IO interfaces to compatible with Google
23 // IO specifications.
24 namespace recordio {
25 // This class appends a protocol buffer to a file in a binary format.
26 // The data written in the file follows the following format (sequentially):
27 // - MagicNumber (32 bits) to recognize this format.
28 // - Uncompressed data payload size (64 bits).
29 // - Compressed data payload size (64 bits), or 0 if the
30 // data is not compressed.
31 // - Payload, possibly compressed. See RecordWriter::Compress()
32 // and RecordReader::Uncompress.
33 class RecordWriter {
34  public:
35  // Magic number when reading and writing protocol buffers.
36  static const int kMagicNumber;
37 
38  explicit RecordWriter(File* const file);
39 
40  template <class P>
41  bool WriteProtocolMessage(const P& proto) {
42  std::string uncompressed_buffer;
43  proto.SerializeToString(&uncompressed_buffer);
44  const uint64 uncompressed_size = uncompressed_buffer.size();
45  const std::string compressed_buffer =
46  use_compression_ ? Compress(uncompressed_buffer) : "";
47  const uint64 compressed_size = compressed_buffer.size();
48  if (file_->Write(&kMagicNumber, sizeof(kMagicNumber)) !=
49  sizeof(kMagicNumber)) {
50  return false;
51  }
52  if (file_->Write(&uncompressed_size, sizeof(uncompressed_size)) !=
53  sizeof(uncompressed_size)) {
54  return false;
55  }
56  if (file_->Write(&compressed_size, sizeof(compressed_size)) !=
57  sizeof(compressed_size)) {
58  return false;
59  }
60  if (use_compression_) {
61  if (file_->Write(compressed_buffer.c_str(), compressed_size) !=
62  compressed_size) {
63  return false;
64  }
65  } else {
66  if (file_->Write(uncompressed_buffer.c_str(), uncompressed_size) !=
67  uncompressed_size) {
68  return false;
69  }
70  }
71  return true;
72  }
73  // Closes the underlying file.
74  bool Close();
75 
76  void set_use_compression(bool use_compression);
77 
78  private:
79  std::string Compress(const std::string& input) const;
80  File* const file_;
81  bool use_compression_;
82 };
83 
84 // This class reads a protocol buffer from a file.
85 // The format must be the one described in RecordWriter, above.
86 class RecordReader {
87  public:
88  explicit RecordReader(File* const file);
89 
90  template <class P>
91  bool ReadProtocolMessage(P* const proto) {
92  uint64 usize = 0;
93  uint64 csize = 0;
94  int magic_number = 0;
95  if (file_->Read(&magic_number, sizeof(magic_number)) !=
96  sizeof(magic_number)) {
97  return false;
98  }
99  if (magic_number != RecordWriter::kMagicNumber) {
100  return false;
101  }
102  if (file_->Read(&usize, sizeof(usize)) != sizeof(usize)) {
103  return false;
104  }
105  if (file_->Read(&csize, sizeof(csize)) != sizeof(csize)) {
106  return false;
107  }
108  std::unique_ptr<char[]> buffer(new char[usize + 1]);
109  if (csize != 0) { // The data is compressed.
110  std::unique_ptr<char[]> compressed_buffer(new char[csize + 1]);
111  if (file_->Read(compressed_buffer.get(), csize) != csize) {
112  return false;
113  }
114  compressed_buffer[csize] = '\0';
115  Uncompress(compressed_buffer.get(), csize, buffer.get(), usize);
116  } else {
117  if (file_->Read(buffer.get(), usize) != usize) {
118  return false;
119  }
120  }
121  proto->ParseFromArray(buffer.get(), usize);
122  return true;
123  }
124 
125  // Closes the underlying file.
126  bool Close();
127 
128  private:
129  void Uncompress(const char* const source, uint64 source_size,
130  char* const output_buffer, uint64 output_size) const;
131 
132  File* const file_;
133 };
134 } // namespace recordio
135 
136 #endif // OR_TOOLS_BASE_RECORDIO_H_
Definition: base/file.h:32
size_t Write(const void *const buff, size_t size)
Definition: file.cc:77
size_t Read(void *const buff, size_t size)
Definition: file.cc:70
RecordReader(File *const file)
Definition: recordio.cc:52
bool ReadProtocolMessage(P *const proto)
Definition: recordio.h:91
static const int kMagicNumber
Definition: recordio.h:36
RecordWriter(File *const file)
Definition: recordio.cc:26
void set_use_compression(bool use_compression)
Definition: recordio.cc:31
bool WriteProtocolMessage(const P &proto)
Definition: recordio.h:41
CpModelProto proto
uint64_t uint64
Definition: file.cc:141
static int input(yyscan_t yyscanner)