1 files changed, 169 insertions, 1 deletions
diff --git a/src/streams.h b/src/streams.h
index ed2e985d76..d4a309be3c 100644
--- a/src/streams.h
+++ b/src/streams.h
@@ -1,5 +1,5 @@
 // Copyright (c) 2009-2010 Satoshi Nakamoto
-// Copyright (c) 2009-2017 The Bitcoin Core developers
+// Copyright (c) 2009-2018 The Bitcoin Core developers
 // Distributed under the MIT software license, see the accompanying
 // file COPYING or http://www.opensource.org/licenses/mit-license.php.
 
@@ -61,6 +61,7 @@ public:
 
     int GetVersion() const { return nVersion; }
     int GetType() const { return nType; }
+    size_t size() const { return stream->size(); }
 };
 
 template<typename S>
@@ -138,6 +139,80 @@ private:
     size_t nPos;
 };
 
+/** Minimal stream for reading from an existing vector by reference
+ */
+class VectorReader
+{
+private:
+    const int m_type;
+    const int m_version;
+    const std::vector<unsigned char>& m_data;
+    size_t m_pos = 0;
+
+public:
+
+    /*
+     * @param[in]  type Serialization Type
+     * @param[in]  version Serialization Version (including any flags)
+     * @param[in]  data Referenced byte vector to overwrite/append
+     * @param[in]  pos Starting position. Vector index where reads should start.
+     */
+    VectorReader(int type, int version, const std::vector<unsigned char>& data, size_t pos)
+        : m_type(type), m_version(version), m_data(data)
+    {
+        seek(pos);
+    }
+
+    /*
+     * (other params same as above)
+     * @param[in]  args  A list of items to deserialize starting at pos.
+     */
+    template <typename... Args>
+    VectorReader(int type, int version, const std::vector<unsigned char>& data, size_t pos,
+                  Args&&... args)
+        : VectorReader(type, version, data, pos)
+    {
+        ::UnserializeMany(*this, std::forward<Args>(args)...);
+    }
+
+    template<typename T>
+    VectorReader& operator>>(T& obj)
+    {
+        // Unserialize from this stream
+        ::Unserialize(*this, obj);
+        return (*this);
+    }
+
+    int GetVersion() const { return m_version; }
+    int GetType() const { return m_type; }
+
+    size_t size() const { return m_data.size() - m_pos; }
+    bool empty() const { return m_data.size() == m_pos; }
+
+    void read(char* dst, size_t n)
+    {
+        if (n == 0) {
+            return;
+        }
+
+        // Read from the beginning of the buffer
+        size_t pos_next = m_pos + n;
+        if (pos_next > m_data.size()) {
+            throw std::ios_base::failure("VectorReader::read(): end of data");
+        }
+        memcpy(dst, m_data.data() + m_pos, n);
+        m_pos = pos_next;
+    }
+
+    void seek(size_t n)
+    {
+        m_pos += n;
+        if (m_pos > m_data.size()) {
+            throw std::ios_base::failure("VectorReader::seek(): end of data");
+        }
+    }
+};
+
 /** Double ended buffer combining vector and stream-like interfaces.
  *
  * >> and << read and write unformatted data using the above serialization templates.
@@ -435,12 +510,105 @@ public:
     }
 };
 
+template <typename IStream>
+class BitStreamReader
+{
+private:
+    IStream& m_istream;
+
+    /// Buffered byte read in from the input stream. A new byte is read into the
+    /// buffer when m_offset reaches 8.
+    uint8_t m_buffer{0};
+
+    /// Number of high order bits in m_buffer already returned by previous
+    /// Read() calls. The next bit to be returned is at this offset from the
+    /// most significant bit position.
+    int m_offset{8};
+
+public:
+    explicit BitStreamReader(IStream& istream) : m_istream(istream) {}
+
+    /** Read the specified number of bits from the stream. The data is returned
+     * in the nbits least signficant bits of a 64-bit uint.
+     */
+    uint64_t Read(int nbits) {
+        if (nbits < 0 || nbits > 64) {
+            throw std::out_of_range("nbits must be between 0 and 64");
+        }
+
+        uint64_t data = 0;
+        while (nbits > 0) {
+            if (m_offset == 8) {
+                m_istream >> m_buffer;
+                m_offset = 0;
+            }
 
+            int bits = std::min(8 - m_offset, nbits);
+            data <<= bits;
+            data |= static_cast<uint8_t>(m_buffer << m_offset) >> (8 - bits);
+            m_offset += bits;
+            nbits -= bits;
+        }
+        return data;
+    }
+};
 
+template <typename OStream>
+class BitStreamWriter
+{
+private:
+    OStream& m_ostream;
 
+    /// Buffered byte waiting to be written to the output stream. The byte is
+    /// written buffer when m_offset reaches 8 or Flush() is called.
+    uint8_t m_buffer{0};
 
+    /// Number of high order bits in m_buffer already written by previous
+    /// Write() calls and not yet flushed to the stream. The next bit to be
+    /// written to is at this offset from the most significant bit position.
+    int m_offset{0};
 
+public:
+    explicit BitStreamWriter(OStream& ostream) : m_ostream(ostream) {}
 
+    ~BitStreamWriter()
+    {
+        Flush();
+    }
+
+    /** Write the nbits least significant bits of a 64-bit int to the output
+     * stream. Data is buffered until it completes an octet.
+     */
+    void Write(uint64_t data, int nbits) {
+        if (nbits < 0 || nbits > 64) {
+            throw std::out_of_range("nbits must be between 0 and 64");
+        }
+
+        while (nbits > 0) {
+            int bits = std::min(8 - m_offset, nbits);
+            m_buffer |= (data << (64 - nbits)) >> (64 - 8 + m_offset);
+            m_offset += bits;
+            nbits -= bits;
+
+            if (m_offset == 8) {
+                Flush();
+            }
+        }
+    }
+
+    /** Flush any unwritten bits to the output stream, padding with 0's to the
+     * next byte boundary.
+     */
+    void Flush() {
+        if (m_offset == 0) {
+            return;
+        }
+
+        m_ostream << m_buffer;
+        m_buffer = 0;
+        m_offset = 0;
+    }
+};