IO: improve inferrability of write and unsafe_write (part of #37081)

JuliaLang · Aug 25, 2020 · 7c9a2a5 · 7c9a2a5
1 parent 71404fc
commit 7c9a2a5
Show file tree

Hide file tree

Showing 5 changed files with 13 additions and 6 deletions.
diff --git a/base/io.jl b/base/io.jl
@@ -358,7 +358,7 @@ function pipe_reader end
 function pipe_writer end
 
 write(io::AbstractPipe, byte::UInt8) = write(pipe_writer(io)::IO, byte)
-unsafe_write(io::AbstractPipe, p::Ptr{UInt8}, nb::UInt) = unsafe_write(pipe_writer(io)::IO, p, nb)
+unsafe_write(io::AbstractPipe, p::Ptr{UInt8}, nb::UInt) = unsafe_write(pipe_writer(io)::IO, p, nb)::Union{Int,UInt}
 buffer_writes(io::AbstractPipe, args...) = buffer_writes(pipe_writer(io)::IO, args...)
 flush(io::AbstractPipe) = flush(pipe_writer(io)::IO)
 

diff --git a/base/iobuffer.jl b/base/iobuffer.jl
@@ -418,7 +418,7 @@ end
 function unsafe_write(to::GenericIOBuffer, p::Ptr{UInt8}, nb::UInt)
     ensureroom(to, nb)
     ptr = (to.append ? to.size+1 : to.ptr)
-    written = Int(min(nb, length(to.data) - ptr + 1))
+    written = Int(min(nb, Int(length(to.data))::Int - ptr + 1))
     towrite = written
     d = to.data
     while towrite > 0

diff --git a/base/strings/io.jl b/base/strings/io.jl
@@ -178,12 +178,12 @@ string(a::Symbol) = String(a)
 # note: print uses an encoding determined by `io` (defaults to UTF-8), whereas
 #       write uses an encoding determined by `s` (UTF-8 for `String`)
 print(io::IO, s::AbstractString) = for c in s; print(io, c); end
-write(io::IO, s::AbstractString) = (len = 0; for c in s; len += write(io, c); end; len)
+write(io::IO, s::AbstractString) = (len = 0; for c in s; len += Int(write(io, c))::Int; end; len)
 show(io::IO, s::AbstractString) = print_quoted(io, s)
 
 # optimized methods to avoid iterating over chars
 write(io::IO, s::Union{String,SubString{String}}) =
-    GC.@preserve s unsafe_write(io, pointer(s), reinterpret(UInt, sizeof(s)))
+    GC.@preserve s Int(unsafe_write(io, pointer(s), reinterpret(UInt, sizeof(s))))::Int
 print(io::IO, s::Union{String,SubString{String}}) = (write(io, s); nothing)
 
 ## printing literal quoted string data ##

diff --git a/test/iostream.jl b/test/iostream.jl
@@ -166,3 +166,8 @@ end
         @test length(readavailable(io)) > 0
     end
 end
+
+@testset "inference" begin
+    @test all(T -> T <: Union{UInt, Int}, Base.return_types(unsafe_write, (IO, Ptr{UInt8}, UInt)))
+    @test all(T -> T === Bool, Base.return_types(eof, (IO,)))
+end
diff --git a/test/strings/basic.jl b/test/strings/basic.jl
@@ -10,7 +10,7 @@ using Random
 
     # Check that resizing empty source vector does not corrupt string
     b = IOBuffer()
-    write(b, "ab")
+    @inferred write(b, "ab")
     x = take!(b)
     s = String(x)
     resize!(x, 0)
@@ -263,8 +263,10 @@ end
     for c in x
         nb += write(f, c)
     end
-    @test nb == 3
+    @test nb === 3
     @test String(take!(f)) == "123"
+
+    @test all(T -> T <: Union{Union{}, Int}, Base.return_types(write, (IO, AbstractString)))
 end
 
 @testset "issue #7248" begin