1 files changed, 35 insertions, 11 deletions

diff --git a/util/stream/rewindable_stream.hh b/util/stream/rewindable_stream.hh
index 9ee637c99..560825cde 100644
--- a/util/stream/rewindable_stream.hh
+++ b/util/stream/rewindable_stream.hh
@@ -5,6 +5,8 @@
 
 #include <boost/noncopyable.hpp>
 
+#include <deque>
+
 namespace util {
 namespace stream {
 
@@ -23,6 +25,10 @@ class RewindableStream : boost::noncopyable {
      */
     RewindableStream();
 
+    ~RewindableStream() {
+      Poison();
+    }
+
     /**
      * Initializes an existing RewindableStream at a specific position in
      * a Chain.
@@ -38,21 +44,32 @@ class RewindableStream : boost::noncopyable {
      *
      * Equivalent to RewindableStream a(); a.Init(....);
      */
-    explicit RewindableStream(const ChainPosition &position);
+    explicit RewindableStream(const ChainPosition &position)
+      : in_(NULL) {
+      Init(position);
+    }
 
     /**
      * Gets the record at the current stream position. Const version.
      */
-    const void *Get() const;
+    const void *Get() const {
+      assert(!poisoned_);
+      assert(current_);
+      return current_;
+    }
 
     /**
      * Gets the record at the current stream position.
      */
-    void *Get();
+    void *Get() {
+      assert(!poisoned_);
+      assert(current_);
+      return current_;
+    }
 
-    operator bool() const { return current_; }
+    operator bool() const { return !poisoned_; }
 
-    bool operator!() const { return !(*this); }
+    bool operator!() const { return poisoned_; }
 
     /**
      * Marks the current position in the stream to be rewound to later.
@@ -80,19 +97,26 @@ class RewindableStream : boost::noncopyable {
     void Poison();
 
   private:
-    void FetchBlock();
+    void AppendBlock();
+
+    void Flush(std::deque<Block>::iterator to);
+
+    std::deque<Block> blocks_;
+    // current_ is in blocks_[blocks_it_] unless poisoned_.
+    std::size_t blocks_it_;
 
     std::size_t entry_size_;
     std::size_t block_size_;
+    std::size_t block_count_;
 
-    uint8_t *marked_, *current_, *end_;
-
-    Block first_bl_;
-    Block second_bl_;
-    Block* current_bl_;
+    uint8_t *marked_, *current_;
+    const uint8_t *block_end_;
 
     PCQueue<Block> *in_, *out_;
 
+    // Have we hit poison at the end of the stream, even if rewinding?
+    bool hit_poison_;
+    // Is the curren position poison?
     bool poisoned_;
 
     WorkerProgress progress_;