FollyIobufAsStdStreambuf.h

gistfile1.txt

#pragma once

#include <exception>
#include <streambuf>
#include <iostream>
#include <memory>
#include <string>
#include <unordered_map>
#include <folly/io/IOBuf.h>
#include <folly/io/Cursor.h>

namespace fs {
    class IOBufStdWrap : public std::streambuf {
    public:
        explicit IOBufStdWrap(folly::IOBuf* ioBuf)
            : iobuf_{ ioBuf }, rchain_{ iobuf_ }, wchain_{ iobuf_ }
        {
            reload();
        }

        IOBufStdWrap* reload() {
            rchain_ = wchain_ = iobuf_;

            while (rchain_->length() == 0 && rchain_->next() != iobuf_) {
                rchain_ = rchain_->next();
            }

            if (wchain_->prev() != iobuf_) {
                wchain_ = wchain_->prev();
            }
            
            setg(
                const_cast<char*>(reinterpret_cast<const char*>(rchain_->data())),
                const_cast<char*>(reinterpret_cast<const char*>(rchain_->data())),
                const_cast<char*>(reinterpret_cast<const char*>(rchain_->tail()))
            );

            setp(
                reinterpret_cast<char*>(wchain_->writableTail()),
                const_cast<char*>(reinterpret_cast<const char*>(wchain_->bufferEnd()))
            );

            return this;
        }

    protected:
        std::streamsize xsgetn(char* s, std::streamsize n) override {
            auto numRead = 0;
            while (numRead < n) {
                auto size = std::min(n - numRead, egptr() - gptr());
                memcpy(s + numRead, gptr(), size);
                numRead += size;
                gbump(size);

                // no more data
                if (egptr() == gptr() && underflow() == EOF) {
                    break;
                }
            }
            return numRead;
        }

        std::streamsize xsputn(const char* s, std::streamsize n) override {
            auto numWritten = 0;
            while (numWritten < n) {
                auto size = std::min(n, epptr() - pptr());
                memcpy(pptr(), s + numWritten, size);
                numWritten += size;
                pbump(size);
                wchain_->append(size);

                // cannot write any more
                if (epptr() == pptr() && overflow() == EOF) {
                    break;
                }
            }
            return numWritten;
        }

        int_type underflow() override {
            // has data on current chain
            if (gptr() < egptr()) {
                return *gptr();
            }

            // has data on some next chain
            while (rchain_->next() != iobuf_) {
                rchain_ = rchain_->next();
                if (rchain_->length() > 0) {
                    setg(
                        const_cast<char*>(reinterpret_cast<const char*>(rchain_->data())),
                        const_cast<char*>(reinterpret_cast<const char*>(rchain_->data())),
                        const_cast<char*>(reinterpret_cast<const char*>(rchain_->tail()))
                    );
                    return *gptr();
                }
            }

            return EOF;
        }

        int_type overflow(int_type ch = EOF) override {
            // has room on current chain
            if (pptr() < epptr()) {
                return 0;
            }

            // has room on some next chain (unlikely)
            if (wchain_->next() != iobuf_) {
                wchain_ = wchain_->next();
                if (wchain_->length() > 0) {
                    setp(
                        reinterpret_cast<char*>(wchain_->writableTail()),
                        const_cast<char*>(reinterpret_cast<const char*>(wchain_->bufferEnd()))
                    );
                    return 0;
                }
            }

            // create new chain
            wchain_->appendChain(folly::IOBuf::create(8192));
            wchain_ = wchain_->next();
            setp(
                reinterpret_cast<char*>(wchain_->writableTail()),
                const_cast<char*>(reinterpret_cast<const char*>(wchain_->bufferEnd()))
            );

            return 0;
        }

    private:
        folly::IOBuf* iobuf_;
        folly::IOBuf* rchain_;
        folly::IOBuf* wchain_;
    };
    
    
    class InMemoryFile {
    public:
        InMemoryFile() : buf_{ folly::IOBuf::create(8192) }, bufWrap_{ std::make_unique<IOBufStdWrap>(buf_.get()) } {}
        InMemoryFile(const InMemoryFile&) = delete;
        InMemoryFile(InMemoryFile&&) = default;

        std::unique_ptr<std::istream> openRead() { return std::make_unique<std::istream>(bufWrap_->reload()); }
        std::unique_ptr<std::ostream> openWrite() { return std::make_unique<std::ostream>(bufWrap_->reload()); }
    private:
        std::unique_ptr<folly::IOBuf> buf_;
        std::unique_ptr<IOBufStdWrap> bufWrap_;
    };

    class InMemoryFileSystem : public IFileSystem {
    public:
        ~InMemoryFileSystem() override = default;
        
        std::unique_ptr<std::istream> openInputStream(const std::filesystem::path& path) override {
            auto iter = files_.find(path.string());
            if (iter == files_.end()) {
                return nullptr;
            }
            return iter->second->openRead();
        }
        
        std::unique_ptr<std::ostream> openOutputStream(const std::filesystem::path& path) override {
            auto iter = files_.find(path.string());
            if (iter == files_.end()) {
                auto created = files_.insert_or_assign(path.string(), std::make_unique<InMemoryFile>());
                return created.first->second->openWrite();
            }
            return iter->second->openWrite();
        }

    private:
        std::unordered_map<std::string, std::unique_ptr<InMemoryFile>> files_;
    };

}