Stream.hpp

#ifndef STREAM_HPP_
#define STREAM_HPP_

#include <algorithm>
#include <cstring>

#include "Util.hpp"

class Tellable {
public:
  virtual uint64_t tell() const {
    unimplementedError(__FUNCTION__);
    return 0;
  }
};

class Stream : public Tellable {
public:
  virtual int get() = 0;
  virtual void put(int c) = 0;
  virtual size_t read(uint8_t* buf, size_t n) {
    size_t count;
    for (count = 0; count < n; ++count) {
      auto c = get();
      if (c == EOF) {
        break;
      }
      buf[count] = c;
    }
    return count;
  }
  // Not thread safe by default.
  virtual size_t readat(uint64_t pos, uint8_t* buf, size_t n) {
    seek(pos);
    return read(buf, n);
  }
  virtual void write(const uint8_t* buf, size_t n) {
    for (;n; --n) {
      put(*(buf++));
    }
  }
  virtual void writeat(uint64_t pos, const uint8_t* buf, size_t n) {
    seek(pos);
    write(buf, n);
  }
  virtual void seek(uint64_t pos) {
    unimplementedError(__FUNCTION__);
  }
  virtual ~Stream() {
  }
  // Helper
  void put16(uint16_t n) {
    put(static_cast<uint8_t>(n >> 8));
    put(static_cast<uint8_t>(n >> 0));
  }
  uint16_t get16() {
    uint16_t ret = 0;
    ret = (ret << 8) | static_cast<uint16_t>(get());
    ret = (ret << 8) | static_cast<uint16_t>(get());
    return ret;
  }
#if 0
  inline void leb128Encode(int64_t n) {
    bool neg = n < 0;
    if (neg) n = -n;
    leb128Encode(static_cast<uint64_t>((n << 1) | (neg ? 1u : 0)));
  }
#endif
  inline void leb128Encode(uint64_t n) {
    while (n >= 0x80) {
      auto c = static_cast<uint8_t>(0x80 | (n & 0x7F));
      put(c);
      n >>= 7;
    }
    put(static_cast<uint8_t>(n));
  }
  uint64_t leb128Decode() {
    uint64_t ret = 0;
    uint64_t shift = 0;
    while (true) {
      const uint8_t c = get();
      ret |= static_cast<uint64_t>(c & 0x7F) << shift;
      shift += 7;
      if ((c & 0x80) == 0) break;
    }
    return ret;
  }
  void writeString(const char* str, char terminator) {
    while (*str != '\0') {
      put(*(str++));
    }
    put(terminator);
  }
  std::string readString() {
    std::string s;
    for (;;) {
      int c = get();
      if (c == EOF || c == '\0') break;
      s.push_back(static_cast<char>(c));
    }
    return s;
  }
};

class WriteStream : public Stream {
public:
  virtual int get() {
    unimplementedError(__FUNCTION__);
    return 0;
  }
  virtual void put(int c) = 0;
  virtual ~WriteStream() {}
};

class VoidWriteStream : public WriteStream {
  uint64_t pos_;
public:
  VoidWriteStream() : pos_(0) {}
  virtual ~VoidWriteStream() {}
  virtual void write(const uint8_t*, size_t n) {
    pos_ += n;
  }
  virtual void put(int) {
    ++pos_;
  }
  virtual uint64_t tell() const {
    return pos_;
  }
  virtual void seek(uint64_t pos) {
    pos_ = pos;
  }
};

class ReadStream : public Stream {
public:
  virtual int get() = 0;
  virtual void put(int c) {
    unimplementedError(__FUNCTION__);
  }
  virtual ~ReadStream() {}
};

class ReadMemoryStream : public ReadStream {
public:
  ReadMemoryStream(const std::vector<uint8_t>* buffer)
    : buffer_(buffer->data())
    , pos_(buffer->data())
    , limit_(buffer->data() + buffer->size()) {
  }
  ReadMemoryStream(const uint8_t* buffer, const uint8_t* limit) : buffer_(buffer), pos_(buffer), limit_(limit) {
  }
  virtual int get() {
    if (pos_ >= limit_) {
      return EOF;
    }
    return *pos_++;
  }
  virtual size_t read(uint8_t* buf, size_t n) {
    const size_t remain = limit_ - pos_;
    const size_t read_count = std::min(remain, n);
    std::copy(pos_, pos_ + read_count, buf);
    pos_ += read_count;
    return read_count;
  }
  virtual uint64_t tell() const {
    return pos_ - buffer_;
  }

private:
  const uint8_t* const buffer_;
  const uint8_t* pos_;
  const uint8_t* const limit_;
};

class WriteMemoryStream : public WriteStream {
public:
  explicit WriteMemoryStream(uint8_t* buffer) : buffer_(buffer), pos_(buffer) {
  }
  virtual void put(int c) {
    *pos_++ = static_cast<uint8_t>(static_cast<unsigned int>(c));
  }
  virtual void write(const uint8_t* data, uint32_t count) {
    memcpy(pos_, data, count);
    pos_ += count;
  }
  virtual uint64_t tell() const {
    return pos_ - buffer_;
  }

private:
  uint8_t* buffer_;
  uint8_t* pos_;
};

class WriteVectorStream : public WriteStream {
public:
  explicit WriteVectorStream(std::vector<uint8_t>* buffer) : buffer_(buffer) {
  }
  virtual void put(int c) {
    buffer_->push_back(c);
  }
  virtual void write(const uint8_t* data, uint32_t count) {
    buffer_->insert(buffer_->end(), data, data + count);
  }
  virtual uint64_t tell() const {
    return buffer_->size();
  }

private:
  std::vector<uint8_t>* const buffer_;
};

template <typename T>
class OStreamWrapper : public std::ostream {
  class StreamBuf : public std::streambuf {
  public:
  };
public:
};

template <const uint32_t buffer_size>
class BufferedStreamReader {
public:
  Stream* stream;
  size_t buffer_count, buffer_pos;
  uint8_t buffer[buffer_size];
  bool done_;

  bool done() const { return done_; }

  BufferedStreamReader(Stream* stream) {
    assert(stream != nullptr);
    init(stream);
  }
  virtual ~BufferedStreamReader() {
  }
  void init(Stream* new_stream) {
    stream = new_stream;
    buffer_pos = 0;
    buffer_count = 0;
    done_ = false;
  }
  ALWAYS_INLINE size_t remain() const {
    return buffer_count - buffer_pos;
  }
  ALWAYS_INLINE int get() {
    if (UNLIKELY(remain() == 0 && Refill() == false)) {
      return EOF;
    }
    return buffer[buffer_pos++];
  }
  void put(int c) {
    unimplementedError(__FUNCTION__);
  }
  uint64_t tell() const {
    return stream->tell() + buffer_pos;
  }
private:
  NO_INLINE bool Refill() {
    buffer_pos = 0;
    buffer_count = stream->read(buffer, buffer_size);
    if (UNLIKELY(buffer_count == 0)) {
      done_ = true;
      return false;
    }
    return true;
  }
};

template <const uint32_t kBufferSize>
class BufferedStreamWriter {
public:
  BufferedStreamWriter(Stream* stream) {
    assert(stream != nullptr);
    init(stream);
  }
  virtual ~BufferedStreamWriter() {
    flush();
    assert(ptr_ == buffer_);
  }
  void init(Stream* new_stream) {
    stream_ = new_stream;
    ptr_ = buffer_;
  }
  NO_INLINE void flush() {
    stream_->write(buffer_, ptr_ - buffer_);
    ptr_ = buffer_;
  }
  ALWAYS_INLINE void put(uint8_t c) {
    if (UNLIKELY(ptr_ >= end())) {
      flush();
    }
    *(ptr_++) = c;
  }
  int get() {
    unimplementedError(__FUNCTION__);
    return 0;
  }
  uint64_t tell() const {
    return stream_->tell() + (ptr_ - buffer_);
  }

private:
  const uint8_t* end() const {
    return &buffer_[kBufferSize];
  }
  Stream* stream_;
  uint8_t buffer_[kBufferSize];
  uint8_t* ptr_;
};

template <const bool kLazy = true>
class MemoryBitStream {
  uint8_t* no_alias data_;
  uint32_t buffer_;
  uint32_t bits_;
  static const uint32_t kBitsPerSizeT = sizeof(uint32_t) * kBitsPerByte;
public:
  ALWAYS_INLINE MemoryBitStream(uint8_t* data) : data_(data), buffer_(0), bits_(0) {
  }

  uint8_t* getData() {
    return data_;
  }

  ALWAYS_INLINE void tryReadByte() {
    if (bits_ <= kBitsPerSizeT - kBitsPerByte) {
      readByte();
    }
  }

  ALWAYS_INLINE void readByte() {
    buffer_ = (buffer_ << kBitsPerByte) | *data_++;
    bits_ += kBitsPerByte;
  }

  ALWAYS_INLINE uint32_t readBits(uint32_t bits) {
    if (kLazy) {
      while (bits_ < bits) {
        readByte();
      }
    } else {
      // This might be slower
      tryReadByte();
      tryReadByte();
      tryReadByte();
    }
    bits_ -= bits;
    uint32_t ret = buffer_ >> bits_;
    buffer_ -= ret << bits_;
    return ret;
  }

  ALWAYS_INLINE void flushByte() {
    bits_ -= kBitsPerByte;
    uint32_t byte = buffer_ >> bits_;
    buffer_ -= byte << bits_;
    *data_++ = byte;
  }

  void flush() {
    while (bits_ > kBitsPerByte) {
      flushByte();
    }
    *data_++ = buffer_;
  }

  ALWAYS_INLINE void writeBits(uint32_t data, uint32_t bits) {
    bits_ += bits;
    buffer_ = (buffer_ << bits) | data;
    while (bits_ >= kBitsPerByte) {
      flushByte();
    }
  }
};

class VerifyStream : public WriteStream {
public:
  Stream* const stream_;
  uint64_t differences_;
  uint64_t count_, ref_count_;

  VerifyStream(Stream* stream, size_t ref_count) : stream_(stream), count_(0), ref_count_(ref_count) {
    init();
  }
  uint64_t getCount() const {
    return count_;
  }
  void resetCount() {
    count_ = 0;
  }
  void init() {
    differences_ = 0;
  }
  void put(int c) {
    auto ref = stream_->get();
    if (c != ref) {
      difference(ref, c);
    }
    ++count_;
  }
  void seek(uint64_t pos) {
    stream_->seek(pos);
  }
  void write(const uint8_t* buf, size_t n) {
    uint8_t buffer[4 * KB];
    while (n != 0) {
      size_t count = stream_->read(buffer, std::min(static_cast<size_t>(4 * KB), n));
      for (size_t i = 0; i < count; ++i) {
        auto ref = buffer[i];
        if (buf[i] != buffer[i]) {
          difference(buffer[i], buf[i]);
        }
      }
      buf += count;
      n -= count;
      count_ += count;
    }
  }
  void difference(int ref, int c) {
    if (differences_ == 0) {
      std::cerr << "Difference found at byte! " << stream_->tell() << " b1: " << "ref: "
        << static_cast<int>(ref) << " new: " << static_cast<int>(c) << std::endl;
    }
    ++differences_;
  }
  virtual uint64_t tell() const {
    return count_;
  }
  void summary() {
    if (count_ != ref_count_) {
      std::cerr << "ERROR: Missing bytes " << count_ << "/" << ref_count_ << " differences=" << differences_ << std::endl;
    } else {
      if (differences_) {
        std::cerr << "ERROR: differences=" << differences_ << std::endl;
      } else {
        std::cout << "No differences found!" << std::endl;
      }
    }
  }
};

#endif