One Month In Minsk. Made in Belarus.

1 files changed, 75 insertions, 0 deletions

diff --git a/coding/var_record_reader.hpp b/coding/var_record_reader.hpp
new file mode 100644
index 0000000000..d6e0129d68
--- /dev/null
+++ b/coding/var_record_reader.hpp
@@ -0,0 +1,75 @@
+#pragma once
+#include "byte_stream.hpp"
+#include "reader.hpp"
+#include "varint.hpp"
+#include "../base/base.hpp"
+#include "../std/algorithm.hpp"
+#include "../std/vector.hpp"
+
+inline uint32_t VarRecordSizeReaderVarint(ArrayByteSource & source)
+{
+  return ReadVarUint<uint32_t>(source);
+}
+
+inline uint32_t VarRecordSizeReaderFixed(ArrayByteSource & source)
+{
+  return ReadPrimitiveFromSource<uint32_t>(source);
+}
+
+// Efficiently reads records, encoded as [VarUint size] [Data] .. [VarUint size] [Data].
+// If size of a record is less than expectedRecordSize, exactly 1 Reader.Read() call is made,
+// otherwise exactly 2 Reader.Read() calls are made.
+// Second template parameter is strategy for reading record size,
+// either &VarRecordSizeReaderVarint or &VarRecordSizeReaderFixed.
+template <class ReaderT, uint32_t (*VarRecordSizeReaderFn)(ArrayByteSource &)>
+class VarRecordReader
+{
+public:
+  VarRecordReader(ReaderT const & reader, uint32_t expectedRecordSize)
+  : m_Reader(reader), m_ReaderSize(m_Reader.Size()), m_ExpectedRecordSize(expectedRecordSize)
+  {
+    ASSERT_GREATER_OR_EQUAL(expectedRecordSize, 4, ());
+  }
+
+  uint64_t ReadRecord(uint64_t const pos, vector<char> & buffer, uint32_t & recordOffset) const
+  {
+    ASSERT_LESS(pos, m_ReaderSize, ());
+    uint32_t const initialSize = static_cast<uint32_t>(
+        min(static_cast<uint64_t>(m_ExpectedRecordSize), m_ReaderSize - pos));
+    buffer.resize(initialSize);
+    m_Reader.Read(pos, &buffer[0], initialSize);
+    ArrayByteSource source(&buffer[0]);
+    uint32_t const recordSize = VarRecordSizeReaderFn(source);
+    uint32_t const recordSizeSize = static_cast<char const *>(source.Ptr()) - &buffer[0];
+    uint32_t const fullSize = recordSize + recordSizeSize;
+    ASSERT_LESS_OR_EQUAL(pos + fullSize, m_ReaderSize, ());
+    buffer.resize(fullSize);
+    if (initialSize < fullSize)
+      m_Reader.Read(pos + initialSize, &buffer[initialSize], fullSize - initialSize);
+
+    recordOffset = recordSizeSize;
+    return pos + fullSize;
+  }
+
+  template <typename F>
+  void ForEachRecord(F const & f) const
+  {
+    uint64_t pos = 0;
+    vector<char> buffer;
+    while (pos < m_ReaderSize)
+    {
+      uint32_t offset;
+      uint64_t nextPos = ReadRecord(pos, buffer, offset);
+      f(pos, &buffer[offset], buffer.size() - offset);
+      pos = nextPos;
+    }
+    ASSERT_EQUAL(pos, m_ReaderSize, ());
+  }
+
+  bool IsEqual(string const & fName) const { return m_Reader.IsEqual(fName); }
+
+protected:
+  ReaderT m_Reader;
+  uint64_t m_ReaderSize;
+  uint32_t m_ExpectedRecordSize; // Expected size of a record.
+};