-
Notifications
You must be signed in to change notification settings - Fork 4
/
zstdstream.h
188 lines (157 loc) · 4.06 KB
/
zstdstream.h
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
/*
* Copyright (c) 2015-present, Facebook, Inc.
*
* This source code is licensed under the MIT license found in the
* LICENSE file in the root directory of this source tree.
*
* This source code is licensed under the BSD-style license found in the
* LICENSE file in the root directory of this source tree. An additional grant
* of patent rights can be found in the PATENTS file in the same directory.
*
* STL stream classes for Zstd compression and decompression.
* Partially inspired by https://github.com/mateidavid/zstr.
*/
#pragma once
#include <cassert>
#include <fstream>
#include <stdexcept>
#include <vector>
// Required to access ZSTD_isFrame()
#define ZSTD_STATIC_LINKING_ONLY
#include <zstd.h>
namespace zstd {
/**
* Custom exception for zstd error codes
*/
class exception : public std::exception {
public:
explicit exception(int code);
const char* what() const throw();
private:
std::string msg_;
};
inline size_t check(size_t code) {
if (ZSTD_isError(code)) {
throw exception(code);
}
return code;
}
/**
* Provides stream compression functionality
*/
class cstream {
public:
static constexpr int defaultLevel = 5;
cstream();
~cstream();
size_t init(int level = defaultLevel);
size_t compress(ZSTD_outBuffer* output, ZSTD_inBuffer* input);
size_t flush(ZSTD_outBuffer* output);
size_t end(ZSTD_outBuffer* output);
private:
ZSTD_CStream* cstrm_;
};
/**
* Provides stream decompression functionality
*/
class dstream {
public:
dstream();
~dstream();
size_t decompress(ZSTD_outBuffer* output, ZSTD_inBuffer* input);
private:
ZSTD_DStream* dstrm_;
};
/**
* Zstd stream buffer for compression. Data is written in a single big frame.
*/
class ostreambuf : public std::streambuf {
public:
explicit ostreambuf(std::streambuf* sbuf, int level = cstream::defaultLevel);
virtual ~ostreambuf();
using int_type = typename std::streambuf::int_type;
virtual int_type overflow(int_type ch = traits_type::eof());
virtual int sync();
private:
ssize_t compress(size_t pos);
std::streambuf* sbuf_;
int clevel_;
cstream strm_;
std::vector<char> inbuf_;
std::vector<char> outbuf_;
size_t inhint_;
bool strInit_;
};
/**
* Zstd stream buffer for decompression. If input data is not compressed, this
* stream will simply copy it.
*/
class istreambuf : public std::streambuf {
public:
explicit istreambuf(std::streambuf* sbuf);
virtual std::streambuf::int_type underflow();
private:
std::streambuf* sbuf_;
dstream strm_;
std::vector<char> inbuf_;
std::vector<char> outbuf_; // only needed if actually compressed
size_t inhint_;
size_t inpos_ = 0;
size_t inavail_ = 0;
bool detected_ = false;
bool compressed_ = false;
};
// Input stream for Zstd-compressed data
class istream : public std::istream {
public:
istream(std::streambuf* sbuf);
virtual ~istream();
};
/**
* Output stream for Zstd-compressed data
*/
class ostream : public std::ostream {
public:
ostream(std::streambuf* sbuf);
virtual ~ostream();
};
/**
* This class enables [io]fstream below to inherit from [io]stream (required for
* setting a custom streambuf) while still constructing a corresponding
* [io]fstream first (required for initializing the Zstd streambufs).
*/
template <typename T>
class fsholder {
public:
explicit fsholder(
const std::string& path,
std::ios_base::openmode mode = std::ios_base::out)
: fs_(path, mode) {}
protected:
T fs_;
};
/**
* Output file stream that writes Zstd-compressed data
*/
class ofstream : private fsholder<std::ofstream>, public std::ostream {
public:
explicit ofstream(
const std::string& path,
std::ios_base::openmode mode = std::ios_base::out);
virtual ~ofstream();
virtual operator bool() const;
void close();
};
/**
* Input file stream for Zstd-compressed data
*/
class ifstream : private fsholder<std::ifstream>, public std::istream {
public:
explicit ifstream(
const std::string& path,
std::ios_base::openmode mode = std::ios_base::in);
virtual ~ifstream();
operator bool() const;
void close();
};
} // namespace zstd