Line data Source code
1 : //---------------------------------------------------------
2 : // Copyright 2015 Ontario Institute for Cancer Research
3 : // Written by Matei David (matei@cs.toronto.edu)
4 : //---------------------------------------------------------
5 :
6 : // Reference:
7 : // http://stackoverflow.com/questions/14086417/how-to-write-custom-input-stream-in-c
8 :
9 : #pragma once
10 :
11 : #include <cassert>
12 : #include <fstream>
13 : #include <sstream>
14 : #include <zlib.h>
15 : #include <memory>
16 : #include <iostream>
17 : #include "strict_fstream.hpp"
18 :
19 : #if defined(__GNUC__) && !defined(__clang__)
20 : #if (__GNUC__ > 5) || (__GNUC__ == 5 && __GNUC_MINOR__>0)
21 : #define CAN_MOVE_IOSTREAM
22 : #endif
23 : #else
24 : #define CAN_MOVE_IOSTREAM
25 : #endif
26 :
27 : namespace zstr
28 : {
29 :
30 : static const std::size_t default_buff_size = static_cast<std::size_t>(1 << 20);
31 :
32 : /// Exception class thrown by failed zlib operations.
33 : class Exception
34 : : public std::ios_base::failure
35 : {
36 : public:
37 0 : static std::string error_to_message(z_stream * zstrm_p, int ret)
38 : {
39 0 : std::string msg = "zlib: ";
40 0 : switch (ret)
41 : {
42 : case Z_STREAM_ERROR:
43 : msg += "Z_STREAM_ERROR: ";
44 : break;
45 : case Z_DATA_ERROR:
46 : msg += "Z_DATA_ERROR: ";
47 : break;
48 : case Z_MEM_ERROR:
49 : msg += "Z_MEM_ERROR: ";
50 : break;
51 : case Z_VERSION_ERROR:
52 : msg += "Z_VERSION_ERROR: ";
53 : break;
54 : case Z_BUF_ERROR:
55 : msg += "Z_BUF_ERROR: ";
56 : break;
57 0 : default:
58 0 : std::ostringstream oss;
59 0 : oss << ret;
60 0 : msg += "[" + oss.str() + "]: ";
61 : break;
62 : }
63 0 : if (zstrm_p->msg) {
64 : msg += zstrm_p->msg;
65 : }
66 : msg += " ("
67 0 : "next_in: " +
68 0 : std::to_string(uintptr_t(zstrm_p->next_in)) +
69 0 : ", avail_in: " +
70 0 : std::to_string(uintptr_t(zstrm_p->avail_in)) +
71 0 : ", next_out: " +
72 0 : std::to_string(uintptr_t(zstrm_p->next_out)) +
73 0 : ", avail_out: " +
74 0 : std::to_string(uintptr_t(zstrm_p->avail_out)) +
75 : ")";
76 0 : return msg;
77 : }
78 :
79 0 : Exception(z_stream * zstrm_p, int ret)
80 0 : : std::ios_base::failure(error_to_message(zstrm_p, ret))
81 : {
82 0 : }
83 : }; // class Exception
84 :
85 : namespace detail
86 : {
87 :
88 : class z_stream_wrapper
89 : : public z_stream
90 : {
91 : public:
92 4579 : z_stream_wrapper(bool _is_input, int _level, int _window_bits)
93 4579 : : is_input(_is_input)
94 : {
95 4579 : this->zalloc = nullptr;//Z_NULL
96 4579 : this->zfree = nullptr;//Z_NULL
97 4579 : this->opaque = nullptr;//Z_NULL
98 : int ret;
99 4579 : if (is_input)
100 : {
101 4516 : this->avail_in = 0;
102 4516 : this->next_in = nullptr;//Z_NULL
103 9032 : ret = inflateInit2(this, _window_bits ? _window_bits : 15+32);
104 : }
105 : else
106 : {
107 126 : ret = deflateInit2(this, _level, Z_DEFLATED, _window_bits ? _window_bits : 15+16, 8, Z_DEFAULT_STRATEGY);
108 : }
109 4579 : if (ret != Z_OK) throw Exception(this, ret);
110 4579 : }
111 4579 : ~z_stream_wrapper()
112 : {
113 4579 : if (is_input)
114 : {
115 4516 : inflateEnd(this);
116 : }
117 : else
118 : {
119 63 : deflateEnd(this);
120 : }
121 4579 : }
122 : private:
123 : bool is_input;
124 : }; // class z_stream_wrapper
125 :
126 : } // namespace detail
127 :
128 : class istreambuf
129 : : public std::streambuf
130 : {
131 : public:
132 134724 : istreambuf(std::streambuf * _sbuf_p,
133 : std::size_t _buff_size = default_buff_size, bool _auto_detect = true, int _window_bits = 0)
134 134724 : : sbuf_p(_sbuf_p),
135 : in_buff(),
136 134724 : in_buff_start(nullptr),
137 134724 : in_buff_end(nullptr),
138 : out_buff(),
139 134724 : zstrm_p(nullptr),
140 134724 : buff_size(_buff_size),
141 134724 : auto_detect(_auto_detect),
142 134724 : auto_detect_run(false),
143 134724 : is_text(false),
144 134724 : window_bits(_window_bits)
145 : {
146 : assert(sbuf_p);
147 134724 : in_buff = std::unique_ptr<char[]>(new char[buff_size]);
148 134724 : in_buff_start = in_buff.get();
149 134724 : in_buff_end = in_buff.get();
150 134724 : out_buff = std::unique_ptr<char[]>(new char[buff_size]);
151 : setg(out_buff.get(), out_buff.get(), out_buff.get());
152 134724 : }
153 :
154 : istreambuf(const istreambuf &) = delete;
155 : istreambuf & operator = (const istreambuf &) = delete;
156 :
157 0 : pos_type seekoff(off_type off, std::ios_base::seekdir dir,
158 : std::ios_base::openmode which) override
159 : {
160 0 : if (off != 0 || dir != std::ios_base::cur) {
161 : return std::streambuf::seekoff(off, dir, which);
162 : }
163 :
164 0 : if (!zstrm_p) {
165 0 : return 0;
166 : }
167 :
168 0 : return static_cast<long int>(zstrm_p->total_out - static_cast<uLong>(in_avail()));
169 : }
170 :
171 274134 : std::streambuf::int_type underflow() override
172 : {
173 274134 : if (this->gptr() == this->egptr())
174 : {
175 : // pointers for free region in output buffer
176 : char * out_buff_free_start = out_buff.get();
177 : int tries = 0;
178 : do
179 : {
180 274134 : if (++tries > 1000) {
181 0 : throw std::ios_base::failure("Failed to fill buffer after 1000 tries");
182 : }
183 :
184 : // read more input if none available
185 274134 : if (in_buff_start == in_buff_end)
186 : {
187 : // empty input buffer: refill from the start
188 269605 : in_buff_start = in_buff.get();
189 269605 : std::streamsize sz = sbuf_p->sgetn(in_buff.get(), static_cast<std::streamsize>(buff_size));
190 269605 : in_buff_end = in_buff_start + sz;
191 269605 : if (in_buff_end == in_buff_start) break; // end of input
192 : }
193 : // auto detect if the stream contains text or deflate data
194 139424 : if (auto_detect && ! auto_detect_run)
195 : {
196 134710 : auto_detect_run = true;
197 134710 : unsigned char b0 = *reinterpret_cast< unsigned char * >(in_buff_start);
198 134710 : unsigned char b1 = *reinterpret_cast< unsigned char * >(in_buff_start + 1);
199 : // Ref:
200 : // http://en.wikipedia.org/wiki/Gzip
201 : // http://stackoverflow.com/questions/9050260/what-does-a-zlib-header-look-like
202 269420 : is_text = ! (in_buff_start + 2 <= in_buff_end
203 134710 : && ((b0 == 0x1F && b1 == 0x8B) // gzip header
204 134557 : || (b0 == 0x78 && (b1 == 0x01 // zlib header
205 0 : || b1 == 0x9C
206 0 : || b1 == 0xDA))));
207 : }
208 139424 : if (is_text)
209 : {
210 : // simply swap in_buff and out_buff, and adjust pointers
211 : assert(in_buff_start == in_buff.get());
212 : std::swap(in_buff, out_buff);
213 134724 : out_buff_free_start = in_buff_end;
214 134724 : in_buff_start = in_buff.get();
215 134724 : in_buff_end = in_buff.get();
216 : }
217 : else
218 : {
219 : // run inflate() on input
220 9216 : if (! zstrm_p) zstrm_p = std::unique_ptr<detail::z_stream_wrapper>(new detail::z_stream_wrapper(true, Z_DEFAULT_COMPRESSION, window_bits));
221 4700 : zstrm_p->next_in = reinterpret_cast< decltype(zstrm_p->next_in) >(in_buff_start);
222 4700 : zstrm_p->avail_in = uint32_t(in_buff_end - in_buff_start);
223 4700 : zstrm_p->next_out = reinterpret_cast< decltype(zstrm_p->next_out) >(out_buff_free_start);
224 4700 : zstrm_p->avail_out = uint32_t((out_buff.get() + buff_size) - out_buff_free_start);
225 4700 : int ret = inflate(zstrm_p.get(), Z_NO_FLUSH);
226 : // process return code
227 4700 : if (ret != Z_OK && ret != Z_STREAM_END) throw Exception(zstrm_p.get(), ret);
228 : // update in&out pointers following inflate()
229 4700 : in_buff_start = reinterpret_cast< decltype(in_buff_start) >(zstrm_p->next_in);
230 4700 : in_buff_end = in_buff_start + zstrm_p->avail_in;
231 4700 : out_buff_free_start = reinterpret_cast< decltype(out_buff_free_start) >(zstrm_p->next_out);
232 : assert(out_buff_free_start + zstrm_p->avail_out == out_buff.get() + buff_size);
233 :
234 4700 : if (ret == Z_STREAM_END) {
235 : // if stream ended, deallocate inflator
236 : zstrm_p.reset();
237 : }
238 : }
239 139424 : } while (out_buff_free_start == out_buff.get());
240 : // 2 exit conditions:
241 : // - end of input: there might or might not be output available
242 : // - out_buff_free_start != out_buff: output available
243 : this->setg(out_buff.get(), out_buff.get(), out_buff_free_start);
244 : }
245 : return this->gptr() == this->egptr()
246 274134 : ? traits_type::eof()
247 274134 : : traits_type::to_int_type(*this->gptr());
248 : }
249 : private:
250 : std::streambuf * sbuf_p;
251 : std::unique_ptr<char[]> in_buff;
252 : char * in_buff_start;
253 : char * in_buff_end;
254 : std::unique_ptr<char[]> out_buff;
255 : std::unique_ptr<detail::z_stream_wrapper> zstrm_p;
256 : std::size_t buff_size;
257 : bool auto_detect;
258 : bool auto_detect_run;
259 : bool is_text;
260 : int window_bits;
261 :
262 : }; // class istreambuf
263 :
264 : class ostreambuf
265 : : public std::streambuf
266 : {
267 : public:
268 63 : ostreambuf(std::streambuf * _sbuf_p,
269 : std::size_t _buff_size = default_buff_size, int _level = Z_DEFAULT_COMPRESSION, int _window_bits = 0)
270 63 : : sbuf_p(_sbuf_p),
271 : in_buff(),
272 : out_buff(),
273 63 : zstrm_p(new detail::z_stream_wrapper(false, _level, _window_bits)),
274 63 : buff_size(_buff_size)
275 : {
276 : assert(sbuf_p);
277 63 : in_buff = std::unique_ptr<char[]>(new char[buff_size]);
278 63 : out_buff = std::unique_ptr<char[]>(new char[buff_size]);
279 63 : setp(in_buff.get(), in_buff.get() + buff_size);
280 63 : }
281 :
282 : ostreambuf(const ostreambuf &) = delete;
283 : ostreambuf & operator = (const ostreambuf &) = delete;
284 :
285 3861 : int deflate_loop(int flush)
286 : {
287 : while (true)
288 : {
289 5768 : zstrm_p->next_out = reinterpret_cast< decltype(zstrm_p->next_out) >(out_buff.get());
290 5768 : zstrm_p->avail_out = uint32_t(buff_size);
291 5768 : int ret = deflate(zstrm_p.get(), flush);
292 5768 : if (ret != Z_OK && ret != Z_STREAM_END && ret != Z_BUF_ERROR) {
293 0 : failed = true;
294 0 : throw Exception(zstrm_p.get(), ret);
295 : }
296 5768 : std::streamsize sz = sbuf_p->sputn(out_buff.get(), reinterpret_cast< decltype(out_buff.get()) >(zstrm_p->next_out) - out_buff.get());
297 5768 : if (sz != reinterpret_cast< decltype(out_buff.get()) >(zstrm_p->next_out) - out_buff.get())
298 : {
299 : // there was an error in the sink stream
300 : return -1;
301 : }
302 5705 : if (ret == Z_STREAM_END || ret == Z_BUF_ERROR || sz == 0)
303 : {
304 : break;
305 : }
306 : }
307 : return 0;
308 : }
309 :
310 126 : virtual ~ostreambuf()
311 63 : {
312 : // flush the zlib stream
313 : //
314 : // NOTE: Errors here (sync() return value not 0) are ignored, because we
315 : // cannot throw in a destructor. This mirrors the behaviour of
316 : // std::basic_filebuf::~basic_filebuf(). To see an exception on error,
317 : // close the ofstream with an explicit call to close(), and do not rely
318 : // on the implicit call in the destructor.
319 : //
320 63 : if (!failed) try {
321 63 : sync();
322 0 : } catch (...) {}
323 189 : }
324 1972 : std::streambuf::int_type overflow(std::streambuf::int_type c = traits_type::eof()) override
325 : {
326 1972 : zstrm_p->next_in = reinterpret_cast< decltype(zstrm_p->next_in) >(pbase());
327 1972 : zstrm_p->avail_in = uint32_t(pptr() - pbase());
328 3881 : while (zstrm_p->avail_in > 0)
329 : {
330 1909 : int r = deflate_loop(Z_NO_FLUSH);
331 1909 : if (r != 0)
332 : {
333 : setp(nullptr, nullptr);
334 0 : return traits_type::eof();
335 : }
336 : }
337 1972 : setp(in_buff.get(), in_buff.get() + buff_size);
338 1972 : return traits_type::eq_int_type(c, traits_type::eof()) ? traits_type::eof() : sputc(char_type(c));
339 : }
340 1952 : int sync() override
341 : {
342 : // first, call overflow to clear in_buff
343 1952 : overflow();
344 1952 : if (! pptr()) return -1;
345 : // then, call deflate asking to finish the zlib stream
346 1952 : zstrm_p->next_in = nullptr;
347 1952 : zstrm_p->avail_in = 0;
348 1952 : if (deflate_loop(Z_FINISH) != 0) return -1;
349 1889 : deflateReset(zstrm_p.get());
350 1889 : return 0;
351 : }
352 : private:
353 : std::streambuf * sbuf_p = nullptr;
354 : std::unique_ptr<char[]> in_buff;
355 : std::unique_ptr<char[]> out_buff;
356 : std::unique_ptr<detail::z_stream_wrapper> zstrm_p;
357 : std::size_t buff_size;
358 : bool failed = false;
359 :
360 : }; // class ostreambuf
361 :
362 : class istream
363 : : public std::istream
364 : {
365 : public:
366 : istream(std::istream & is,
367 : std::size_t _buff_size = default_buff_size, bool _auto_detect = true, int _window_bits = 0)
368 : : std::istream(new istreambuf(is.rdbuf(), _buff_size, _auto_detect, _window_bits))
369 : {
370 : exceptions(std::ios_base::badbit);
371 : }
372 : explicit istream(std::streambuf * sbuf_p)
373 : : std::istream(new istreambuf(sbuf_p))
374 : {
375 : exceptions(std::ios_base::badbit);
376 : }
377 0 : virtual ~istream()
378 0 : {
379 0 : delete rdbuf();
380 0 : }
381 : }; // class istream
382 :
383 : class ostream
384 : : public std::ostream
385 : {
386 : public:
387 : ostream(std::ostream & os,
388 : std::size_t _buff_size = default_buff_size, int _level = Z_DEFAULT_COMPRESSION, int _window_bits = 0)
389 : : std::ostream(new ostreambuf(os.rdbuf(), _buff_size, _level, _window_bits))
390 : {
391 : exceptions(std::ios_base::badbit);
392 : }
393 : explicit ostream(std::streambuf * sbuf_p)
394 : : std::ostream(new ostreambuf(sbuf_p))
395 : {
396 : exceptions(std::ios_base::badbit);
397 : }
398 0 : virtual ~ostream()
399 0 : {
400 0 : delete rdbuf();
401 0 : }
402 : }; // class ostream
403 :
404 : namespace detail
405 : {
406 :
407 : template < typename FStream_Type >
408 134332 : struct strict_fstream_holder
409 : {
410 134788 : strict_fstream_holder(const std::string& filename, std::ios_base::openmode mode = std::ios_base::in)
411 134788 : : _fs(filename, mode)
412 134787 : {}
413 : strict_fstream_holder() = default;
414 : FStream_Type _fs {};
415 : }; // class strict_fstream_holder
416 :
417 : } // namespace detail
418 :
419 : class ifstream
420 : : private detail::strict_fstream_holder< strict_fstream::ifstream >,
421 : public std::istream
422 : {
423 : public:
424 134724 : explicit ifstream(const std::string filename, std::ios_base::openmode mode = std::ios_base::in, size_t buff_size = default_buff_size)
425 134724 : : detail::strict_fstream_holder< strict_fstream::ifstream >(filename, mode),
426 269448 : std::istream(new istreambuf(_fs.rdbuf(), buff_size))
427 : {
428 134724 : exceptions(std::ios_base::badbit);
429 134724 : }
430 : explicit ifstream(): detail::strict_fstream_holder< strict_fstream::ifstream >(), std::istream(new istreambuf(_fs.rdbuf())){}
431 : void close() {
432 134269 : _fs.close();
433 134269 : }
434 : #ifdef CAN_MOVE_IOSTREAM
435 : void open(const std::string filename, std::ios_base::openmode mode = std::ios_base::in) {
436 : _fs.open(filename, mode);
437 : std::istream::operator=(std::istream(new istreambuf(_fs.rdbuf())));
438 : }
439 : #endif
440 : bool is_open() const {
441 : return _fs.is_open();
442 : }
443 163261 : virtual ~ifstream()
444 134269 : {
445 134269 : if (_fs.is_open()) close();
446 134269 : if (rdbuf()) delete rdbuf();
447 163261 : }
448 :
449 : /// Return the position within the compressed file (wrapped filestream)
450 : std::streampos compressed_tellg()
451 : {
452 : return _fs.tellg();
453 : }
454 : }; // class ifstream
455 :
456 : class ofstream
457 : : private detail::strict_fstream_holder< strict_fstream::ofstream >,
458 : public std::ostream
459 : {
460 : public:
461 64 : explicit ofstream(const std::string filename, std::ios_base::openmode mode = std::ios_base::out,
462 : int level = Z_DEFAULT_COMPRESSION, size_t buff_size = default_buff_size)
463 64 : : detail::strict_fstream_holder< strict_fstream::ofstream >(filename, mode | std::ios_base::binary),
464 127 : std::ostream(new ostreambuf(_fs.rdbuf(), buff_size, level))
465 : {
466 63 : exceptions(std::ios_base::badbit);
467 63 : }
468 : explicit ofstream(): detail::strict_fstream_holder< strict_fstream::ofstream >(), std::ostream(new ostreambuf(_fs.rdbuf())){}
469 : void close() {
470 63 : std::ostream::flush();
471 63 : _fs.close();
472 63 : }
473 : #ifdef CAN_MOVE_IOSTREAM
474 : void open(const std::string filename, std::ios_base::openmode mode = std::ios_base::out, int level = Z_DEFAULT_COMPRESSION) {
475 : flush();
476 : _fs.open(filename, mode | std::ios_base::binary);
477 : std::ostream::operator=(std::ostream(new ostreambuf(_fs.rdbuf(), default_buff_size, level)));
478 : }
479 : #endif
480 : bool is_open() const {
481 : return _fs.is_open();
482 : }
483 : ofstream& flush() {
484 : std::ostream::flush();
485 : _fs.flush();
486 : return *this;
487 : }
488 126 : virtual ~ofstream()
489 63 : {
490 63 : if (_fs.is_open()) close();
491 63 : if (rdbuf()) delete rdbuf();
492 126 : }
493 :
494 : // Return the position within the compressed file (wrapped filestream)
495 : std::streampos compressed_tellp()
496 : {
497 : return _fs.tellp();
498 : }
499 : }; // class ofstream
500 :
501 : } // namespace zstr
|
