Widen signature of readbytes! methods to contiguous SubArrays

This is essential for performance.
caleb-allen · Jul 10, 2020 · 4bc1885 · 4bc1885
1 parent c14460f
commit 4bc1885
Show file tree

Hide file tree

Showing 2 changed files with 47 additions and 4 deletions.
diff --git a/base/iostream.jl b/base/iostream.jl
@@ -445,7 +445,9 @@ function readline(s::IOStream; keep::Bool=false)
     @_lock_ios s ccall(:jl_readuntil, Ref{String}, (Ptr{Cvoid}, UInt8, UInt8, UInt8), s.ios, '\n', 1, keep ? 0 : 2)
 end
 
-function readbytes_all!(s::IOStream, b::Array{UInt8}, nb)
+function readbytes_all!(s::IOStream,
+                        b::Union{Array{UInt8}, FastContiguousSubArray{UInt8,<:Any,<:Array{UInt8}}},
+                        nb::Integer)
     olb = lb = length(b)
     nr = 0
     @_lock_ios s begin
@@ -466,7 +468,9 @@ function readbytes_all!(s::IOStream, b::Array{UInt8}, nb)
     return nr
 end
 
-function readbytes_some!(s::IOStream, b::Array{UInt8}, nb)
+function readbytes_some!(s::IOStream,
+                         b::Union{Array{UInt8}, FastContiguousSubArray{UInt8,<:Any,<:Array{UInt8}}},
+                         nb::Integer)
     olb = length(b)
     if nb > olb
         resize!(b, nb)
@@ -495,7 +499,10 @@ requested bytes, until an error or end-of-file occurs. If `all` is `false`, at m
 `read` call is performed, and the amount of data returned is device-dependent. Note that not
 all stream types support the `all` option.
 """
-function readbytes!(s::IOStream, b::Array{UInt8}, nb=length(b); all::Bool=true)
+function readbytes!(s::IOStream,
+                    b::Union{Array{UInt8}, FastContiguousSubArray{UInt8,<:Any,<:Array{UInt8}}},
+                    nb=length(b);
+                    all::Bool=true)
     return all ? readbytes_all!(s, b, nb) : readbytes_some!(s, b, nb)
 end
 

diff --git a/test/iostream.jl b/test/iostream.jl
@@ -39,15 +39,17 @@
     end
 end
 
-@testset "readbytes_some! via readbytes!" begin
+@testset "readbytes!" begin
     mktemp() do path, file
         function append_to_file(str)
             mark(file)
             print(file, str)
             flush(file)
             reset(file)
         end
+        # Array
         append_to_file("aaaaaaaaaaaaaaaaa")
+        # readbytes_some
         b = UInt8[0]
         readbytes!(file, b, all=false)
         @test String(b) == "a"
@@ -59,8 +61,42 @@ end
         b = UInt8[]
         readbytes!(file, b, 15)
         @test String(b) == "aaaaaaaaaaaaaa"
+
+        # SubArray
+        append_to_file("aaaaaaaaaaaaaaaaa")
+        # readbytes_some
+        b = view(UInt8[0, 0, 0], 2:2)
+        readbytes!(file, b, all=false)
+        @test String(b) == "a"
+        b = view(UInt8[0, 0, 0], 2:3)
+        readbytes!(file, b, 2, all=false)
+        @test String(b) == "aa"
+        b = view(UInt8[0, 0, 0], 1:3)
+        readbytes!(file, b, 2, all=false)
+        @test b == UInt8['a', 'a', 0]
+        @test String(b[1:2]) == "aa"
+        # with resizing of b
+        b = view(UInt8[0, 0, 0], 1:0)
+        @test_throws MethodError readbytes!(file, b, 2, all=false)
+        @test isempty(b)
+        # readbytes_all
+        b = view(UInt8[0, 0, 0], 2:2)
+        readbytes!(file, b)
+        @test String(b) == "a"
+        b = view(UInt8[0, 0, 0], 2:3)
+        readbytes!(file, b, 2)
+        @test String(b) == "aa"
+        b = view(UInt8[0, 0, 0], 1:3)
+        readbytes!(file, b, 2)
+        @test b == UInt8['a', 'a', 0]
+        @test String(b[1:2]) == "aa"
+        #  with resizing of b
+        b = view(UInt8[0, 0, 0], 1:0)
+        @test_throws MethodError readbytes!(file, b, 2)
+        @test isempty(b)
     end
 end
+
 @testset "issue #18755" begin
     mktemp() do path, io
         write(io, zeros(UInt8, 131073))