reader.h

#ifndef SINGA_IO_READER_H_
#define SINGA_IO_READER_H_

#include <cstring>
#include <fstream>
#include <string>
#include "singa/singa_config.h"

#ifdef USE_LMDB
#include <lmdb.h>
#include <sys/stat.h>
#include <vector>
#endif  // USE_LMDB

namespace singa {
namespace io {

using std::string;

class Reader {
 public:
  virtual ~Reader() {}

  virtual bool Open(const std::string& path) = 0;

  virtual void Close() = 0;

  virtual bool Read(std::string* key, std::string* value) = 0;

  virtual int Count() = 0;

  virtual void SeekToFirst() = 0;
};

class BinFileReader : public Reader {
 public:
  ~BinFileReader() { Close(); }
  bool Open(const std::string& path) override;
  bool Open(const std::string& path, int capacity);
  void Close() override;
  bool Read(std::string* key, std::string* value) override;
  int Count() override;
  void SeekToFirst() override;
  inline std::string path() { return path_; }

 protected:
  bool OpenFile();
  bool ReadField(std::string* content);
  bool PrepareNextField(int size);

 private:
  std::string path_ = "";
  std::ifstream fdat_;
  char* buf_ = nullptr;
  int offset_ = 0;
  int capacity_ = 10485760;
  int bufsize_ = 0;
  const char kMagicWord[2] = {'s', 'g'};
};

class TextFileReader : public Reader {
 public:
  ~TextFileReader() { Close(); }
  bool Open(const std::string& path) override;
  void Close() override;
  bool Read(std::string* key, std::string* value) override;
  int Count() override;
  void SeekToFirst() override;
  inline std::string path() { return path_; }

 private:
  std::string path_ = "";
  std::ifstream fdat_;
  int lineNo_ = 0;
};

#ifdef USE_LMDB
class LMDBReader : public Reader {
 public:
  ~LMDBReader() { Close(); }
  bool Open(const std::string& path) override;
  void Close() override;
  bool Read(std::string* key, std::string* value) override;
  int Count() override;
  void SeekToFirst() override;
  inline std::string path() { return path_; }
  inline bool valid() { return valid_; }

 protected:
  void Seek(MDB_cursor_op op);
  inline void MDB_CHECK(int mdb_status);

 private:
  std::string path_ = "";
  MDB_env* mdb_env_ = nullptr;
  MDB_dbi mdb_dbi_;
  MDB_txn* mdb_txn_ = nullptr;
  MDB_cursor* mdb_cursor_ = nullptr;
  MDB_val mdb_key_, mdb_value_;
  bool valid_;
  bool first_;
};
#endif  // USE_LMDB
}  // namespace io
}  // namespace singa

#endif  // SINGA_IO_READER_H_