Created
June 16, 2023 19:58
-
-
Save stormouse/b4b4aa16ee18d95fcc574e5357af3eaa to your computer and use it in GitHub Desktop.
FollyIobufAsStdStreambuf.h
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
#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_; | |
}; | |
} |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment