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