gzfstream.h

Go to the documentation of this file.

/*
 * Souffle - A Datalog Compiler
 * Copyright (c) 2013, 2014, Oracle and/or its affiliates. All rights reserved
 * Licensed under the Universal Permissive License v 1.0 as shown at:
 * - https://opensource.org/licenses/UPL
 * - <souffle root>/licenses/SOUFFLE-UPL.txt
 */
 
/************************************************************************
 *
 * @file gzfstream.h
 * A simple zlib wrapper to provide gzip file streams.
 *
 ***********************************************************************/
 
#pragma once
 
#include <cstdio>
#include <cstring>
#include <iostream>
#include <string>
#include <zlib.h>
 
namespace souffle {
 
namespace gzfstream {
 
namespace internal {
 
class gzfstreambuf : public std::streambuf {
public:
    gzfstreambuf() {
        setp(buffer, buffer + (bufferSize - 1));
        setg(buffer + reserveSize, buffer + reserveSize, buffer + reserveSize);
    }
 
    gzfstreambuf(const gzfstreambuf&) = delete;
 
    gzfstreambuf(gzfstreambuf&& old) = default;
 
    gzfstreambuf* open(const std::string& filename, std::ios_base::openmode mode) {
        if (is_open()) {
            return nullptr;
        }
        if ((mode ^ std::ios::in ^ std::ios::out) == 0) {
            return nullptr;
        }
 
        this->mode = mode;
        std::string gzmode((mode & std::ios::in) != 0 ? "rb" : "wb");
        fileHandle = gzopen(filename.c_str(), gzmode.c_str());
 
        if (fileHandle == nullptr) {
            return nullptr;
        }
        isOpen = true;
 
        return this;
    }
 
    gzfstreambuf* close() {
        if (is_open()) {
            sync();
            isOpen = false;
            if (gzclose(fileHandle) == Z_OK) {
                return this;
            }
        }
        return nullptr;
    }
 
    bool is_open() const {
        return isOpen;
    }
 
    ~gzfstreambuf() override {
        try {
            close();
        } catch (...) {
            // Don't throw exceptions.
        }
    }
 
protected:
    int_type overflow(int c = EOF) override {
        if (((mode & std::ios::out) == 0) || !isOpen) {
            return EOF;
        }
 
        if (c != EOF) {
            *pptr() = c;
            pbump(1);
        }
        int toWrite = pptr() - pbase();
        if (gzwrite(fileHandle, pbase(), toWrite) != toWrite) {
            return EOF;
        }
        pbump(-toWrite);
 
        return c;
    }
 
    int_type underflow() override {
        if (((mode & std::ios::in) == 0) || !isOpen) {
            return EOF;
        }
        if ((gptr() != nullptr) && (gptr() < egptr())) {
            return traits_type::to_int_type(*gptr());
        }
 
        unsigned charsPutBack = gptr() - eback();
        if (charsPutBack > reserveSize) {
            charsPutBack = reserveSize;
        }
        memcpy(buffer + reserveSize - charsPutBack, gptr() - charsPutBack, charsPutBack);
 
        int charsRead = gzread(fileHandle, buffer + reserveSize, bufferSize - reserveSize);
        if (charsRead <= 0) {
            return EOF;
        }
 
        setg(buffer + reserveSize - charsPutBack, buffer + reserveSize, buffer + reserveSize + charsRead);
 
        return traits_type::to_int_type(*gptr());
    }
 
    int sync() override {
        if ((pptr() != nullptr) && pptr() > pbase()) {
            int toWrite = pptr() - pbase();
            if (gzwrite(fileHandle, pbase(), toWrite) != toWrite) {
                return -1;
            }
            pbump(-toWrite);
        }
        return 0;
    }
 
private:
    static constexpr unsigned int bufferSize = 65536;
    static constexpr unsigned int reserveSize = 16;
 
    char buffer[bufferSize] = {};
    gzFile fileHandle = {};
    bool isOpen = false;
    std::ios_base::openmode mode = std::ios_base::in;
};
 
class gzfstream : virtual public std::ios {
public:
    gzfstream() {
        init(&buf);
    }
 
    gzfstream(const std::string& filename, std::ios_base::openmode mode) {
        init(&buf);
        open(filename, mode);
    }
 
    gzfstream(const gzfstream&) = delete;
 
    gzfstream(gzfstream&&) = delete;
 
    ~gzfstream() override = default;
 
    void open(const std::string& filename, std::ios_base::openmode mode) {
        if (buf.open(filename, mode) == nullptr) {
            clear(rdstate() | std::ios::badbit);
        }
    }
 
    bool is_open() {
        return buf.is_open();
    }
 
    void close() {
        if (buf.is_open()) {
            if (buf.close() == nullptr) {
                clear(rdstate() | std::ios::badbit);
            }
        }
    }
 
    gzfstreambuf* rdbuf() const {
        return &buf;
    }
 
protected:
    mutable gzfstreambuf buf;
};
 
}  // namespace internal
 
class igzfstream : public internal::gzfstream, public std::istream {
public:
    igzfstream() : internal::gzfstream(), std::istream(&buf) {}
 
    explicit igzfstream(const std::string& filename, std::ios_base::openmode mode = std::ios::in)
            : internal::gzfstream(filename, mode), std::istream(&buf) {}
 
    igzfstream(const igzfstream&) = delete;
 
    igzfstream(igzfstream&&) = delete;
 
    internal::gzfstreambuf* rdbuf() const {
        return internal::gzfstream::rdbuf();
    }
 
    void open(const std::string& filename, std::ios_base::openmode mode = std::ios::in) {
        internal::gzfstream::open(filename, mode);
    }
};
 
class ogzfstream : public internal::gzfstream, public std::ostream {
public:
    ogzfstream() : std::ostream(&buf) {}
 
    explicit ogzfstream(const std::string& filename, std::ios_base::openmode mode = std::ios::out)
            : internal::gzfstream(filename, mode), std::ostream(&buf) {}
 
    ogzfstream(const ogzfstream&) = delete;
 
    ogzfstream(ogzfstream&&) = delete;
 
    internal::gzfstreambuf* rdbuf() const {
        return internal::gzfstream::rdbuf();
    }
 
    void open(const std::string& filename, std::ios_base::openmode mode = std::ios::out) {
        internal::gzfstream::open(filename, mode);
    }
};
 
} /* namespace gzfstream */
 
} /* namespace souffle */