OILS / vendor / souffle / io / gzfstream.h View on Github | oilshell.org

236 lines, 131 significant
1/*
2 * Souffle - A Datalog Compiler
3 * Copyright (c) 2021, The Souffle Developers. All rights reserved
4 * Licensed under the Universal Permissive License v 1.0 as shown at:
5 * - https://opensource.org/licenses/UPL
6 * - <souffle root>/licenses/SOUFFLE-UPL.txt
7 */
8
9/************************************************************************
10 *
11 * @file gzfstream.h
12 * A simple zlib wrapper to provide gzip file streams.
13 *
14 ***********************************************************************/
15
16#pragma once
17
18#include <cstdio>
19#include <cstring>
20#include <iostream>
21#include <string>
22#include <zlib.h>
23
24namespace souffle {
25
26namespace gzfstream {
27
28namespace internal {
29
30class gzfstreambuf : public std::streambuf {
31public:
32 gzfstreambuf() {
33 setp(buffer, buffer + (bufferSize - 1));
34 setg(buffer + reserveSize, buffer + reserveSize, buffer + reserveSize);
35 }
36
37 gzfstreambuf(const gzfstreambuf&) = delete;
38
39 gzfstreambuf(gzfstreambuf&& old) = default;
40
41 gzfstreambuf* open(const std::string& filename, std::ios_base::openmode mode) {
42 if (is_open()) {
43 return nullptr;
44 }
45 if ((mode ^ std::ios::in ^ std::ios::out) == 0) {
46 return nullptr;
47 }
48
49 this->mode = mode;
50 std::string gzmode((mode & std::ios::in) != 0 ? "rb" : "wb");
51 fileHandle = gzopen(filename.c_str(), gzmode.c_str());
52
53 if (fileHandle == nullptr) {
54 return nullptr;
55 }
56 isOpen = true;
57
58 return this;
59 }
60
61 gzfstreambuf* close() {
62 if (is_open()) {
63 sync();
64 isOpen = false;
65 if (gzclose(fileHandle) == Z_OK) {
66 return this;
67 }
68 }
69 return nullptr;
70 }
71
72 bool is_open() const {
73 return isOpen;
74 }
75
76 ~gzfstreambuf() override {
77 try {
78 close();
79 } catch (...) {
80 // Don't throw exceptions.
81 }
82 }
83
84protected:
85 int_type overflow(int c = EOF) override {
86 if (((mode & std::ios::out) == 0) || !isOpen) {
87 return EOF;
88 }
89
90 if (c != EOF) {
91 *pptr() = c;
92 pbump(1);
93 }
94 const int toWrite = static_cast<int>(pptr() - pbase());
95 if (gzwrite(fileHandle, pbase(), static_cast<unsigned int>(toWrite)) != toWrite) {
96 return EOF;
97 }
98 pbump(-toWrite);
99
100 return c;
101 }
102
103 int_type underflow() override {
104 if (((mode & std::ios::in) == 0) || !isOpen) {
105 return EOF;
106 }
107 if ((gptr() != nullptr) && (gptr() < egptr())) {
108 return traits_type::to_int_type(*gptr());
109 }
110
111 std::size_t charsPutBack = gptr() - eback();
112 if (charsPutBack > reserveSize) {
113 charsPutBack = reserveSize;
114 }
115 memcpy(buffer + reserveSize - charsPutBack, gptr() - charsPutBack, charsPutBack);
116
117 int charsRead =
118 gzread(fileHandle, buffer + reserveSize, static_cast<unsigned int>(bufferSize - reserveSize));
119 if (charsRead <= 0) {
120 return EOF;
121 }
122
123 setg(buffer + reserveSize - charsPutBack, buffer + reserveSize, buffer + reserveSize + charsRead);
124
125 return traits_type::to_int_type(*gptr());
126 }
127
128 int sync() override {
129 if ((pptr() != nullptr) && pptr() > pbase()) {
130 const int toWrite = static_cast<int>(pptr() - pbase());
131 if (gzwrite(fileHandle, pbase(), static_cast<unsigned int>(toWrite)) != toWrite) {
132 return -1;
133 }
134 pbump(-toWrite);
135 }
136 return 0;
137 }
138
139private:
140 static constexpr std::size_t bufferSize = 65536;
141 static constexpr std::size_t reserveSize = 16;
142
143 char buffer[bufferSize] = {};
144 gzFile fileHandle = {};
145 bool isOpen = false;
146 std::ios_base::openmode mode = std::ios_base::in;
147};
148
149class gzfstream : virtual public std::ios {
150public:
151 gzfstream() {
152 init(&buf);
153 }
154
155 gzfstream(const std::string& filename, std::ios_base::openmode mode) {
156 init(&buf);
157 open(filename, mode);
158 }
159
160 gzfstream(const gzfstream&) = delete;
161
162 gzfstream(gzfstream&&) = delete;
163
164 ~gzfstream() override = default;
165
166 void open(const std::string& filename, std::ios_base::openmode mode) {
167 if (buf.open(filename, mode) == nullptr) {
168 clear(rdstate() | std::ios::badbit);
169 }
170 }
171
172 bool is_open() {
173 return buf.is_open();
174 }
175
176 void close() {
177 if (buf.is_open()) {
178 if (buf.close() == nullptr) {
179 clear(rdstate() | std::ios::badbit);
180 }
181 }
182 }
183
184 gzfstreambuf* rdbuf() const {
185 return &buf;
186 }
187
188protected:
189 mutable gzfstreambuf buf;
190};
191
192} // namespace internal
193
194class igzfstream : public internal::gzfstream, public std::istream {
195public:
196 igzfstream() : internal::gzfstream(), std::istream(&buf) {}
197
198 explicit igzfstream(const std::string& filename, std::ios_base::openmode mode = std::ios::in)
199 : internal::gzfstream(filename, mode), std::istream(&buf) {}
200
201 igzfstream(const igzfstream&) = delete;
202
203 igzfstream(igzfstream&&) = delete;
204
205 internal::gzfstreambuf* rdbuf() const {
206 return internal::gzfstream::rdbuf();
207 }
208
209 void open(const std::string& filename, std::ios_base::openmode mode = std::ios::in) {
210 internal::gzfstream::open(filename, mode);
211 }
212};
213
214class ogzfstream : public internal::gzfstream, public std::ostream {
215public:
216 ogzfstream() : std::ostream(&buf) {}
217
218 explicit ogzfstream(const std::string& filename, std::ios_base::openmode mode = std::ios::out)
219 : internal::gzfstream(filename, mode), std::ostream(&buf) {}
220
221 ogzfstream(const ogzfstream&) = delete;
222
223 ogzfstream(ogzfstream&&) = delete;
224
225 internal::gzfstreambuf* rdbuf() const {
226 return internal::gzfstream::rdbuf();
227 }
228
229 void open(const std::string& filename, std::ios_base::openmode mode = std::ios::out) {
230 internal::gzfstream::open(filename, mode);
231 }
232};
233
234} /* namespace gzfstream */
235
236} /* namespace souffle */