Align StreamReader with Ruby's IO conventions

lib/zstd-ruby/stream_reader.rb

@@ -1,22 +1,98 @@
 module Zstd
-  # @todo Exprimental
+  # @todo Experimental
+  #
+  # Wraps a readable IO containing zstd-compressed data and exposes the
+  # decompressed bytes through a standard read interface.
+  #
+  #   File.open('data.zst') do |file|
+  #     reader = Zstd::StreamReader.new(file)
+  #     while (chunk = reader.read(16 * 1024))
+  #       # ...
+  #     end
+  #   end
+  #
+  # +length+ is a number of *decompressed* bytes, and +read+ returns +nil+ once
+  # the stream is exhausted, so the reader can drive any consumer that expects
+  # an IO-like object.
   class StreamReader
-    def initialize(io)
+    # Number of compressed bytes pulled from the underlying IO per refill.
+    DEFAULT_CHUNK_SIZE = 64 * 1024
+
+    def initialize(io, chunk_size: DEFAULT_CHUNK_SIZE)
+      raise ArgumentError, "chunk_size must be positive" unless chunk_size.to_int > 0
+
       @io = io
+      @chunk_size = chunk_size.to_int
       @stream = Zstd::StreamingDecompress.new
+      @buffer = +''.b
+      @pending = +''.b
+      @source_eof = false
     end
 
-    def read(length)
-      if @io.eof?
-        raise StandardError, "EOF"
+    # Reads and returns up to +length+ decompressed bytes, or every remaining
+    # byte when +length+ is nil.
+    #
+    # Returns nil at the end of the stream, or an empty String when +length+ is
+    # zero, mirroring IO#read.
+    def read(length = nil, outbuf = nil)
+      if length.nil?
+        fill_until(Float::INFINITY)
+        data = @buffer
+        @buffer = +''.b
+        return outbuf ? outbuf.replace(data) : data
+      end
+
+      length = length.to_int
+      raise ArgumentError, "negative length #{length} given" if length < 0
+      return outbuf ? outbuf.replace(+''.b) : +''.b if length == 0
+
+      fill_until(length)
+      if @buffer.empty?
+        outbuf&.replace(+''.b)
+        return nil
       end
-      data = @io.read(length)
-      @stream.decompress(data)
+
+      data = @buffer.byteslice(0, length)
+      @buffer = @buffer.byteslice(data.bytesize..) || +''.b
+      outbuf ? outbuf.replace(data) : data
+    end
+
+    def eof?
+      fill_until(1)
+      @buffer.empty?
     end
+    alias_method :eof, :eof?
 
     def close
-      @io.write(@stream.finish)
       @io.close
+      nil
+    end
+
+    private
+
+    def fill_until(n)
+      while @buffer.bytesize < n
+        if @pending.empty?
+          break if @source_eof
+
+          chunk = @io.read(@chunk_size)
+          if chunk.nil? || chunk.empty?
+            @source_eof = true
+            break
+          end
+          @pending = chunk
+        end
+
+        # decompress_with_pos writes at most ZSTD_DStreamOutSize bytes per call
+        # and reports how much input it consumed, so a high compression ratio
+        # cannot balloon the buffer the way decompress would.
+        decompressed, consumed = @stream.decompress_with_pos(@pending)
+        @pending = @pending.byteslice(consumed..) || +''.b
+        @buffer << decompressed
+
+        # No forward progress: the frame is truncated or complete.
+        break if consumed == 0 && decompressed.empty?
+      end
     end
   end
 end

lib/zstd-ruby/stream_writer.rb

@@ -1,5 +1,5 @@
 module Zstd
-  # @todo Exprimental
+  # @todo Experimental
   class StreamWriter
     def initialize(io, level: nil)
       @io = io

spec/zstd-ruby-stream_reader_spec.rb

@@ -1,21 +1,202 @@
 require "spec_helper"
 require 'zstd-ruby'
-require 'pry'
+require 'stringio'
 
 RSpec.describe Zstd::StreamReader do
-  describe 'read' do
-    it 'shoud work' do
+  # Compressible payload large enough to span several internal zstd blocks.
+  let(:compressible) { (1..20_000).map { |n| %({"name":"pkg-#{n}","version":"1.0.#{n}"}) }.join("\n") << "\n" }
+  # Incompressible payload: compressed size ends up close to the original size.
+  let(:incompressible) { Random.new(42).bytes(512 * 1024) }
+
+  def reader_for(data, **opts)
+    described_class.new(StringIO.new(Zstd.compress(data)), **opts)
+  end
+
+  def read_all(reader, length)
+    +''.b.tap do |out|
+      while (chunk = reader.read(length))
+        out << chunk
+      end
+    end
+  end
+
+  describe '#read' do
+    it 'reads the data written by StreamWriter' do
       io = StringIO.new
       writer = Zstd::StreamWriter.new(io)
       writer.write("abc")
       writer.write("def")
       writer.finish
       io.rewind
 
-      reader = Zstd::StreamReader.new(io)
-      expect(reader.read(10)).to eq('a')
-      expect(reader.read(10)).to eq('bcdef')
-      expect(reader.read(10)).to eq('')
+      reader = described_class.new(io)
+      expect(reader.read(10)).to eq('abcdef')
+      expect(reader.read(10)).to be_nil
+    end
+
+    it 'returns exactly the requested number of decompressed bytes' do
+      reader = reader_for(compressible)
+
+      5.times { expect(reader.read(512).bytesize).to eq(512) }
+    end
+
+    it 'returns the requested length regardless of the compression ratio' do
+      # A tiny compressed frame expands to far more than the requested length;
+      # the reader must not hand back the whole frame at once.
+      reader = reader_for("hello world\n" * 20_000)
+
+      expect(reader.read(512).bytesize).to eq(512)
+    end
+
+    it 'round-trips compressible data' do
+      expect(read_all(reader_for(compressible), 512)).to eq(compressible)
+    end
+
+    it 'round-trips incompressible data' do
+      expect(read_all(reader_for(incompressible), 4096)).to eq(incompressible)
+    end
+
+    it 'round-trips data smaller than a single read' do
+      expect(read_all(reader_for('sample data'), 4096)).to eq('sample data')
+    end
+
+    it 'returns a short final read rather than padding' do
+      reader = reader_for('0123456789')
+
+      expect(reader.read(4)).to eq('0123')
+      expect(reader.read(100)).to eq('456789')
+      expect(reader.read(100)).to be_nil
+    end
+
+    it 'returns nil once the stream is exhausted' do
+      reader = reader_for('sample data')
+      reader.read(4096)
+
+      expect(reader.read(4096)).to be_nil
+      expect(reader.read(4096)).to be_nil
+    end
+
+    it 'reads the whole stream when no length is given' do
+      expect(reader_for(compressible).read).to eq(compressible)
+    end
+
+    it 'returns an empty String at EOF when no length is given' do
+      reader = reader_for('sample data')
+      reader.read
+
+      expect(reader.read).to eq('')
+    end
+
+    it 'returns an empty String for a zero length' do
+      expect(reader_for('sample data').read(0)).to eq('')
+    end
+
+    it 'raises ArgumentError for a negative length' do
+      expect { reader_for('sample data').read(-1) }.to raise_error(ArgumentError)
+    end
+
+    it 'writes into outbuf when given' do
+      reader = reader_for(compressible)
+      outbuf = +'previous contents'
+
+      result = reader.read(512, outbuf)
+
+      expect(result).to equal(outbuf)
+      expect(outbuf.bytesize).to eq(512)
+    end
+
+    it 'clears outbuf and returns nil at EOF' do
+      reader = reader_for('sample data')
+      reader.read(4096)
+      outbuf = +'previous contents'
+
+      expect(reader.read(4096, outbuf)).to be_nil
+      expect(outbuf).to be_empty
+    end
+
+    it 'honours a custom chunk_size' do
+      reader = reader_for(compressible, chunk_size: 1024)
+
+      expect(read_all(reader, 512)).to eq(compressible)
+    end
+
+    it 'rejects a non-positive chunk_size' do
+      expect { reader_for('sample data', chunk_size: 0) }.to raise_error(ArgumentError)
+    end
+  end
+
+  describe '#eof?' do
+    it 'is false while data remains and true once drained' do
+      reader = reader_for('sample data')
+
+      expect(reader.eof?).to eq(false)
+      expect(reader.read(4096)).to eq('sample data')
+      expect(reader.eof?).to eq(true)
+    end
+
+    it 'does not consume data' do
+      reader = reader_for('sample data')
+
+      expect(reader.eof?).to eq(false)
+      expect(reader.read(4096)).to eq('sample data')
+    end
+  end
+
+  describe '#close' do
+    it 'closes the underlying IO' do
+      io = StringIO.new(Zstd.compress('sample data'))
+      reader = described_class.new(io)
+
+      expect { reader.close }.not_to raise_error
+      expect(io).to be_closed
+    end
+  end
+
+  describe 'driving an IO consumer' do
+    it 'can be wrapped to feed Gem::Package::TarReader' do
+      require 'rubygems/package'
+
+      tar = StringIO.new
+      Gem::Package::TarWriter.new(tar) do |writer|
+        3.times { |i| writer.add_file("0#{i}.txt", 0644) { |f| f.write(compressible) } }
+      end
+
+      io_like = Class.new do
+        def initialize(reader)
+          @reader = reader
+          @pos = 0
+        end
+
+        attr_reader :pos
+
+        def read(length = nil, outbuf = nil)
+          data = @reader.read(length, outbuf)
+          @pos += data.bytesize if data
+          data
+        end
+
+        def eof? = @reader.eof?
+
+        def seek(amount, whence = IO::SEEK_SET)
+          raise Errno::EINVAL unless whence == IO::SEEK_CUR
+
+          while amount > 0
+            chunk = read(amount < 65_536 ? amount : 65_536)
+            break unless chunk
+
+            amount -= chunk.bytesize
+          end
+          0
+        end
+      end
+
+      entries = []
+      source = StringIO.new(Zstd.compress(tar.string))
+      Gem::Package::TarReader.new(io_like.new(described_class.new(source))) do |reader|
+        reader.each { |entry| entries << [entry.full_name, entry.read.bytesize] }
+      end
+
+      expect(entries).to eq(3.times.map { |i| ["0#{i}.txt", compressible.bytesize] })
     end
   end
 end