lib/bindata/io.rb in bindata-2.3.1 vs lib/bindata/io.rb in bindata-2.3.2

@@ -16,11 +16,11 @@
         if io.respond_to?(:to_str)
           io = BinData::IO.create_string_io(io.to_str)
         end
 
         @raw_io = io
-        @buffer_end_pos = nil
+        @buffer_end_points = nil
 
         extend seekable? ? SeekableStream : UnSeekableStream
         stream_init
       end
 
@@ -36,49 +36,68 @@
       def seek(n)
         seek_raw(buffer_limited_n(n))
       end
 
       def buffer_limited_n(n)
-        if @buffer_end_pos
+        if @buffer_end_points
           if n.nil? or n > 0
-            max = @buffer_end_pos[1] - offset
+            max = @buffer_end_points[1] - offset
             n = max if n.nil? or n > max
           else
-            min = @buffer_end_pos[0] - offset
+            min = @buffer_end_points[0] - offset
             n = min if n < min
           end
         end
 
         n
       end
 
       def with_buffer_common(n, &block)
-        prev = @buffer_end_pos
+        prev = @buffer_end_points
         if prev
           avail = prev[1] - offset
           n = avail if n > avail
         end
-        @buffer_end_pos = [offset, offset + n]
+        @buffer_end_points = [offset, offset + n]
         begin
-          block.call(*@buffer_end_pos)
+          block.call(*@buffer_end_points)
         ensure
-          @buffer_end_pos = prev
+          @buffer_end_points = prev
         end
       end
 
       # Use #seek and #pos on seekable streams
       module SeekableStream
         # The number of bytes remaining in the input stream.
         def num_bytes_remaining
-          mark = @raw_io.pos
+          start_mark = @raw_io.pos
           @raw_io.seek(0, ::IO::SEEK_END)
-          bytes_remaining = @raw_io.pos - mark
-          @raw_io.seek(mark, ::IO::SEEK_SET)
+          end_mark = @raw_io.pos
 
+          if @buffer_end_points
+            if @buffer_end_points[1] < end_mark
+              end_mark = @buffer_end_points[1]
+            end
+          end
+
+          bytes_remaining = end_mark - start_mark
+          @raw_io.seek(start_mark, ::IO::SEEK_SET)
+
           bytes_remaining
         end
 
+        # All io calls in +block+ are rolled back after this
+        # method completes.
+        def with_readahead(&block)
+          mark = @raw_io.pos
+          begin
+            block.call
+          ensure
+            @raw_io.seek(mark, ::IO::SEEK_SET)
+          end
+        end
+
         #-----------
         private
 
         def stream_init
           @initial_pos = @raw_io.pos
@@ -110,19 +129,66 @@
         # The number of bytes remaining in the input stream.
         def num_bytes_remaining
           raise IOError, "stream is unseekable"
         end
 
+        # All io calls in +block+ are rolled back after this
+        # method completes.
+        def with_readahead(&block)
+          mark = @offset
+          @read_data = ""
+          @in_readahead = true
+
+          class << self
+            alias_method :read_raw_without_readahead, :read_raw
+            alias_method :read_raw, :read_raw_with_readahead
+          end
+
+          begin
+            block.call
+          ensure
+            @offset = mark
+            @in_readahead = false
+          end
+        end
+
         #-----------
         private
 
         def stream_init
           @offset = 0
         end
 
         def read_raw(n)
           data = @raw_io.read(n)
           @offset += data.size if data
+          data
+        end
+
+        def read_raw_with_readahead(n)
+          data = ""
+
+          if @read_data.length > 0 and not @in_readahead
+            bytes_to_consume = [n, @read_data.length].min
+            data << @read_data.slice!(0, bytes_to_consume)
+            n -= bytes_to_consume
+
+            if @read_data.length == 0
+              class << self
+                alias_method :read_raw, :read_raw_without_readahead
+              end
+            end
+          end
+
+          raw_data = @raw_io.read(n)
+          data << raw_data if raw_data
+
+          if @in_readahead
+            @read_data << data
+          end
+
+          @offset += data.size
+
           data
         end
 
         def write_raw(data)
           @offset += data.size