io.h

Go to the documentation of this file.

 
#ifndef DMLC_IO_H_
#define DMLC_IO_H_
#include <cstdio>
#include <string>
#include <cstring>
#include <vector>
#include <istream>
#include <ostream>
#include <streambuf>
#include "./logging.h"
 
// include uint64_t only to make io standalone
#ifdef _MSC_VER
typedef unsigned __int64 uint64_t;
#else
#include <inttypes.h>
#endif
 
namespace dmlc {
class Stream {  // NOLINT(*)
 public:
  virtual size_t Read(void *ptr, size_t size) = 0;
  virtual void Write(const void *ptr, size_t size) = 0;
  virtual ~Stream(void) {}
  static Stream *Create(const char *uri,
                        const char* const flag,
                        bool allow_null = false);
  // helper functions to write/read different data structures
  template<typename T>
  inline void Write(const T &data);
  template<typename T>
  inline bool Read(T *out_data);
  template<typename T>
  inline void WriteArray(const T* data, size_t num_elems);
  template<typename T>
  inline bool ReadArray(T* data, size_t num_elems);
};
 
class SeekStream: public Stream {
 public:
  // virtual destructor
  virtual ~SeekStream(void) {}
  virtual void Seek(size_t pos) = 0;
  virtual size_t Tell(void) = 0;
  static SeekStream *CreateForRead(const char *uri,
                                   bool allow_null = false);
};
 
class Serializable {
 public:
  virtual ~Serializable() {}
  virtual void Load(Stream *fi) = 0;
  virtual void Save(Stream *fo) const = 0;
};
 
class InputSplit {
 public:
  struct Blob {
    void *dptr;
    size_t size;
  };
  virtual void HintChunkSize(size_t /*chunk_size*/) {}
  virtual size_t GetTotalSize(void) = 0;
  virtual void BeforeFirst(void) = 0;
  virtual bool NextRecord(Blob *out_rec) = 0;
  virtual bool NextChunk(Blob *out_chunk) = 0;
  virtual bool NextBatch(Blob *out_chunk, size_t /*n_records*/) { return NextChunk(out_chunk); }
  virtual ~InputSplit(void) DMLC_THROW_EXCEPTION {}
  virtual void ResetPartition(unsigned part_index, unsigned num_parts) = 0;
  static InputSplit* Create(const char *uri,
                            unsigned part_index,
                            unsigned num_parts,
                            const char *type);
  static InputSplit* Create(const char *uri,
                            const char *index_uri,
                            unsigned part_index,
                            unsigned num_parts,
                            const char *type,
                            const bool shuffle = false,
                            const int seed = 0,
                            const size_t batch_size = 256,
                            const bool recurse_directories = false);
};
 
#ifndef _LIBCPP_SGX_NO_IOSTREAMS
class ostream : public std::basic_ostream<char> {
 public:
  explicit ostream(Stream *stream,
                   size_t buffer_size = (1 << 10))
      : std::basic_ostream<char>(NULL), buf_(buffer_size) {
    this->set_stream(stream);
  }
  // explictly synchronize the buffer
  virtual ~ostream() DMLC_NO_EXCEPTION {
    buf_.pubsync();
  }
  inline void set_stream(Stream *stream) {
    buf_.set_stream(stream);
    this->rdbuf(&buf_);
  }
 
  inline size_t bytes_written(void) const {
    return buf_.bytes_out();
  }
 
 private:
  // internal streambuf
  class OutBuf : public std::streambuf {
   public:
    explicit OutBuf(size_t buffer_size)
        : stream_(NULL), buffer_(buffer_size), bytes_out_(0) {
      if (buffer_size == 0) buffer_.resize(2);
    }
    // set stream to the buffer
    inline void set_stream(Stream *stream);
 
    inline size_t bytes_out() const { return bytes_out_; }
   private:
    Stream *stream_;
    std::vector<char> buffer_;
    size_t bytes_out_;
    // override sync
    inline int_type sync(void);
    // override overflow
    inline int_type overflow(int c);
  };
  OutBuf buf_;
};
 
class istream : public std::basic_istream<char> {
 public:
  explicit istream(Stream *stream,
                   size_t buffer_size = (1 << 10))
      : std::basic_istream<char>(NULL), buf_(buffer_size) {
    this->set_stream(stream);
  }
  virtual ~istream() DMLC_NO_EXCEPTION {}
  inline void set_stream(Stream *stream) {
    buf_.set_stream(stream);
    this->rdbuf(&buf_);
  }
  inline size_t bytes_read(void) const {
    return buf_.bytes_read();
  }
 
 private:
  // internal streambuf
  class InBuf : public std::streambuf {
   public:
    explicit InBuf(size_t buffer_size)
        : stream_(NULL), bytes_read_(0),
          buffer_(buffer_size) {
      if (buffer_size == 0) buffer_.resize(2);
    }
    // set stream to the buffer
    inline void set_stream(Stream *stream);
    // return how many bytes read so far
    inline size_t bytes_read(void) const {
      return bytes_read_;
    }
   private:
    Stream *stream_;
    size_t bytes_read_;
    std::vector<char> buffer_;
    // override underflow
    inline int_type underflow();
  };
  InBuf buf_;
};
#endif
}  // namespace dmlc
 
#include "./serializer.h"
 
namespace dmlc {
// implementations of inline functions
template<typename T>
inline void Stream::Write(const T &data) {
  serializer::Handler<T>::Write(this, data);
}
template<typename T>
inline bool Stream::Read(T *out_data) {
  return serializer::Handler<T>::Read(this, out_data);
}
 
template<typename T>
inline void Stream::WriteArray(const T* data, size_t num_elems) {
  for (size_t i = 0; i < num_elems; ++i) {
    this->Write<T>(data[i]);
  }
}
 
template<typename T>
inline bool Stream::ReadArray(T* data, size_t num_elems) {
  for (size_t i = 0; i < num_elems; ++i) {
    if (!this->Read<T>(data + i)) return false;
  }
  return true;
}
 
#ifndef _LIBCPP_SGX_NO_IOSTREAMS
// implementations for ostream
inline void ostream::OutBuf::set_stream(Stream *stream) {
  if (stream_ != NULL) this->pubsync();
  this->stream_ = stream;
  this->setp(&buffer_[0], &buffer_[0] + buffer_.size() - 1);
}
inline int ostream::OutBuf::sync(void) {
  if (stream_ == NULL) return -1;
  std::ptrdiff_t n = pptr() - pbase();
  stream_->Write(pbase(), n);
  this->pbump(-static_cast<int>(n));
  bytes_out_ += n;
  return 0;
}
inline int ostream::OutBuf::overflow(int c) {
  *(this->pptr()) = c;
  std::ptrdiff_t n = pptr() - pbase();
  this->pbump(-static_cast<int>(n));
  if (c == EOF) {
    stream_->Write(pbase(), n);
    bytes_out_ += n;
  } else {
    stream_->Write(pbase(), n + 1);
    bytes_out_ += n + 1;
  }
  return c;
}
 
// implementations for istream
inline void istream::InBuf::set_stream(Stream *stream) {
  stream_ = stream;
  this->setg(&buffer_[0], &buffer_[0], &buffer_[0]);
}
inline int istream::InBuf::underflow() {
  char *bhead = &buffer_[0];
  if (this->gptr() == this->egptr()) {
    size_t sz = stream_->Read(bhead, buffer_.size());
    this->setg(bhead, bhead, bhead + sz);
    bytes_read_ += sz;
  }
  if (this->gptr() == this->egptr()) {
    return traits_type::eof();
  } else {
    return traits_type::to_int_type(*gptr());
  }
}
#endif
 
namespace io {
struct URI {
  std::string protocol;
  std::string host;
  std::string name;
  URI(void) {}
  explicit URI(const char *uri) {
    const char *p = std::strstr(uri, "://");
    if (p == NULL) {
      name = uri;
    } else {
      protocol = std::string(uri, p - uri + 3);
      uri = p + 3;
      p = std::strchr(uri, '/');
      if (p == NULL) {
        host = uri; name = '/';
      } else {
        host = std::string(uri, p - uri);
        name = p;
      }
    }
  }
  inline std::string str(void) const {
    return protocol + host + name;
  }
};
 
enum FileType {
  kFile,
  kDirectory
};
 
struct FileInfo {
  URI path;
  size_t size;
  FileType type;
  FileInfo() : size(0), type(kFile) {}
};
 
class FileSystem {
 public:
  static FileSystem *GetInstance(const URI &path);
  virtual ~FileSystem() {}
  virtual FileInfo GetPathInfo(const URI &path) = 0;
  virtual void ListDirectory(const URI &path, std::vector<FileInfo> *out_list) = 0;
  virtual void ListDirectoryRecursive(const URI &path,
                                      std::vector<FileInfo> *out_list);
  virtual Stream *Open(const URI &path,
                       const char* const flag,
                       bool allow_null = false) = 0;
  virtual SeekStream *OpenForRead(const URI &path,
                                  bool allow_null = false) = 0;
};
 
}  // namespace io
}  // namespace dmlc
#endif  // DMLC_IO_H_