diff --git a/dali/operators/reader/loader/indexed_file_loader.h b/dali/operators/reader/loader/indexed_file_loader.h index 387e33de46..bfb755fe86 100755 --- a/dali/operators/reader/loader/indexed_file_loader.h +++ b/dali/operators/reader/loader/indexed_file_loader.h @@ -30,6 +30,7 @@ #include "dali/util/uri.h" #include "dali/util/file.h" #include "dali/util/odirect_file.h" +#include "dali/core/call_at_exit.h" namespace dali { @@ -192,13 +193,21 @@ class IndexedFileLoader : public Loader, true> { DALI_ENFORCE(index_uris.size() == paths_.size(), "Number of index files needs to match the number of data files"); for (size_t i = 0; i < index_uris.size(); ++i) { - std::ifstream fin(index_uris[i]); - DALI_ENFORCE(fin.good(), "Failed to open file " + index_uris[i]); + const auto& path = index_uris[i]; + auto uri = URI::Parse(path); + auto index_file = FileStream::Open(path); + auto index_file_cleanup = AtScopeExit([&index_file] { + if (index_file) + index_file->Close(); + }); + + FileStreamBuf<> stream_buf(index_file.get()); + std::istream fin(&stream_buf); + DALI_ENFORCE(fin.good(), "Failed to open file " + path); int64 pos, size; while (fin >> pos >> size) { indices_.emplace_back(pos, size, i); } - fin.close(); } } diff --git a/dali/util/file.h b/dali/util/file.h index 8cf74e6893..4869e205fd 100644 --- a/dali/util/file.h +++ b/dali/util/file.h @@ -16,6 +16,7 @@ #define DALI_UTIL_FILE_H_ #include +#include #include #include #include @@ -107,6 +108,33 @@ class DLL_PUBLIC FileStream : public InputStream { std::string path_; }; +/** + * @brief Custom streambuf implementation that reads from FileStream. + * @remarks It is useful to be used together with std::istream + */ +template +class FileStreamBuf : public std::streambuf { + public: + explicit FileStreamBuf(FileStream *reader) : reader_(reader) { + setg(buffer_, buffer_, buffer_); // Initialize get area pointers + } + + protected: + int_type underflow() override { + if (gptr() == egptr()) { // get area is exhausted + size_t nbytes = reader_->Read(buffer_, BufferSize); + if (nbytes == 0) + return traits_type::eof(); + setg(buffer_, buffer_, buffer_ + nbytes); + } + return traits_type::to_int_type(*gptr()); + } + + private: + FileStream *reader_; + char buffer_[BufferSize]; +}; + } // namespace dali #endif // DALI_UTIL_FILE_H_