add total_in and total_out

Author: bicycle1885

src/buffer.jl

@@ -17,17 +17,21 @@
 # `markpos` ≤ `bufferpos` ≤ `marginpos` must hold whenever possible.
 
 mutable struct Buffer
+    # data and positions (see above)
     data::Vector{UInt8}
     markpos::Int
     bufferpos::Int
     marginpos::Int
 
+    # the number of total bytes passed through this buffer
+    total::Int64
+
     function Buffer(size::Integer)
-        return new(Vector{UInt8}(size), 0, 1, 1)
+        return new(Vector{UInt8}(size), 0, 1, 1, 0)
     end
 
     function Buffer(data::Vector{UInt8})
-        return new(data, 0, 1, length(data)+1)
+        return new(data, 0, 1, length(data)+1, 0)
     end
 end
 
@@ -82,6 +86,7 @@ end
 function writebyte!(buf::Buffer, b::UInt8)
     buf.data[buf.marginpos] = b
     buf.marginpos += 1
+    buf.total += 1
     return 1
 end
 
@@ -167,6 +172,7 @@ end
 function initbuffer!(buf::Buffer)
     buf.markpos = 0
     buf.bufferpos = buf.marginpos = 1
+    buf.total = 0
     return buf
 end
 
@@ -189,6 +195,7 @@ function readdata!(input::IO, output::Buffer)
     Base.unsafe_read(input, marginptr(output), n)
     output.marginpos += n
     nread += n
+    output.total += nread
     return nread
 end
 

src/stream.jl

@@ -251,6 +251,7 @@ function Base.unsafe_write(stream::TranscodingStream, input::Ptr{UInt8}, nbytes:
         unsafe_copy!(marginptr(buffer1), p, m)
         p += m
         buffer1.marginpos += m
+        buffer1.total += m
     end
     return Int(p - input)
 end
@@ -322,6 +323,32 @@ function Base.transcode(codec::Codec, data::Vector{UInt8})
 end
 
 
+# Utils
+# -----
+
+function total_in(stream::TranscodingStream)::Int64
+    state = stream.state
+    if state.state == :read
+        return state.buffer2.total
+    elseif state.state == :write
+        return state.buffer1.total
+    else
+        return zero(Int64)
+    end
+end
+
+function total_out(stream::TranscodingStream)::Int64
+    state = stream.state
+    if state.state == :read
+        return state.buffer1.total
+    elseif state.state == :write
+        return state.buffer2.total
+    else
+        return zero(Int64)
+    end
+end
+
+
 # Buffering
 # ---------
 
@@ -347,6 +374,7 @@ function fillbuffer(stream::TranscodingStream)
         Δin, Δout, stream.state.code = process(stream.codec, buffermem(buffer2), marginmem(buffer1))
         buffer2.bufferpos += Δin
         buffer1.marginpos += Δout
+        buffer1.total += Δout
         nfilled += Δout
     end
     return nfilled
@@ -395,6 +423,7 @@ function process_to_write(stream::TranscodingStream)
     Δin, Δout, stream.state.code = process(stream.codec, buffermem(buffer1), marginmem(buffer2))
     buffer1.bufferpos += Δin
     buffer2.marginpos += Δout
+    buffer2.total += Δout
     makemargin!(buffer1, 0)
     return Δin
 end

test/runtests.jl

@@ -121,6 +121,21 @@ using Base.Test
     @test hash(read(stream)) == hash(data)
     @test hash(stream.state.buffer1.data[1:length(data)]) == hash(data)
 
+    stream = TranscodingStream(Identity(), IOBuffer(b"foobar"))
+    @test TranscodingStreams.total_in(stream) === Int64(0)
+    @test TranscodingStreams.total_out(stream) === Int64(0)
+    read(stream)
+    @test TranscodingStreams.total_in(stream) === Int64(6)
+    @test TranscodingStreams.total_out(stream) === Int64(6)
+
+    stream = TranscodingStream(Identity(), IOBuffer())
+    @test TranscodingStreams.total_in(stream) === Int64(0)
+    @test TranscodingStreams.total_out(stream) === Int64(0)
+    write(stream, b"foobar")
+    flush(stream)
+    @test TranscodingStreams.total_in(stream) === Int64(6)
+    @test TranscodingStreams.total_out(stream) === Int64(6)
+
     # transcode
     @test transcode(Identity(), b"") == b""
     @test transcode(Identity(), b"foo") == b"foo"