summaryrefslogtreecommitdiff
path: root/klm/util/read_compressed.cc
diff options
context:
space:
mode:
authorKenneth Heafield <github@kheafield.com>2013-01-22 21:37:49 +0000
committerKenneth Heafield <github@kheafield.com>2013-01-22 21:37:49 +0000
commitb35a7f3a96ff8ae42e15922dd6949bf9f5d15501 (patch)
tree26edbe78931ffc50864a899c087d851005fe560b /klm/util/read_compressed.cc
parent51a412aa7f5f50035cf28a274a70508c839f3d40 (diff)
KenLM 58da338b
Diffstat (limited to 'klm/util/read_compressed.cc')
-rw-r--r--klm/util/read_compressed.cc28
1 files changed, 27 insertions, 1 deletions
diff --git a/klm/util/read_compressed.cc b/klm/util/read_compressed.cc
index 7a1a8fb5..b81549e4 100644
--- a/klm/util/read_compressed.cc
+++ b/klm/util/read_compressed.cc
@@ -320,6 +320,23 @@ class XZip : public ReadBase {
};
#endif // HAVE_XZLIB
+class IStreamReader : public ReadBase {
+ public:
+ explicit IStreamReader(std::istream &stream) : stream_(stream) {}
+
+ std::size_t Read(void *to, std::size_t amount, ReadCompressed &thunk) {
+ if (!stream_.read(static_cast<char*>(to), amount)) {
+ UTIL_THROW_IF(!stream_.eof(), ErrnoException, "istream error");
+ amount = stream_.gcount();
+ }
+ ReadCount(thunk) += amount;
+ return amount;
+ }
+
+ private:
+ std::istream &stream_;
+};
+
enum MagicResult {
UNKNOWN, GZIP, BZIP, XZIP
};
@@ -329,7 +346,7 @@ MagicResult DetectMagic(const void *from_void) {
if (header[0] == 0x1f && header[1] == 0x8b) {
return GZIP;
}
- if (header[0] == 'B' && header[1] == 'Z') {
+ if (header[0] == 'B' && header[1] == 'Z' && header[2] == 'h') {
return BZIP;
}
const uint8_t xzmagic[6] = { 0xFD, '7', 'z', 'X', 'Z', 0x00 };
@@ -387,6 +404,10 @@ ReadCompressed::ReadCompressed(int fd) {
Reset(fd);
}
+ReadCompressed::ReadCompressed(std::istream &in) {
+ Reset(in);
+}
+
ReadCompressed::ReadCompressed() {}
ReadCompressed::~ReadCompressed() {}
@@ -396,6 +417,11 @@ void ReadCompressed::Reset(int fd) {
internal_.reset(ReadFactory(fd, raw_amount_));
}
+void ReadCompressed::Reset(std::istream &in) {
+ internal_.reset();
+ internal_.reset(new IStreamReader(in));
+}
+
std::size_t ReadCompressed::Read(void *to, std::size_t amount) {
return internal_->Read(to, amount, *this);
}