do lazy reallocation after take!(iobuffer) (#48676)

closes #27741. closes #48651

Author: stevengj

base/iobuffer.jl

@@ -5,6 +5,7 @@
 # Stateful string
 mutable struct GenericIOBuffer{T<:AbstractVector{UInt8}} <: IO
     data::T # T should support: getindex, setindex!, length, copyto!, and resize!
+    reinit::Bool # if true, data needs to be re-allocated (after take!)
     readable::Bool
     writable::Bool
     seekable::Bool # if not seekable, implementation is free to destroy (compact) past read data
@@ -17,7 +18,7 @@ mutable struct GenericIOBuffer{T<:AbstractVector{UInt8}} <: IO
     function GenericIOBuffer{T}(data::T, readable::Bool, writable::Bool, seekable::Bool, append::Bool,
                                 maxsize::Integer) where T<:AbstractVector{UInt8}
         require_one_based_indexing(data)
-        new(data,readable,writable,seekable,append,length(data),maxsize,1,-1)
+        new(data,false,readable,writable,seekable,append,length(data),maxsize,1,-1)
     end
 end
 const IOBuffer = GenericIOBuffer{Vector{UInt8}}
@@ -137,8 +138,12 @@ PipeBuffer(data::Vector{UInt8}=UInt8[]; maxsize::Int = typemax(Int)) =
     GenericIOBuffer(data,true,true,false,true,maxsize)
 PipeBuffer(maxsize::Integer) = (x = PipeBuffer(StringVector(maxsize), maxsize = maxsize); x.size=0; x)
 
+_similar_data(b::GenericIOBuffer, len::Int) = similar(b.data, len)
+_similar_data(b::IOBuffer, len::Int) = StringVector(len)
+
 function copy(b::GenericIOBuffer)
-    ret = typeof(b)(b.writable ? copy(b.data) : b.data,
+    ret = typeof(b)(b.reinit ? _similar_data(b, 0) : b.writable ?
+                    copyto!(_similar_data(b, length(b.data)), b.data) : b.data,
                     b.readable, b.writable, b.seekable, b.append, b.maxsize)
     ret.size = b.size
     ret.ptr  = b.ptr
@@ -270,7 +275,10 @@ function truncate(io::GenericIOBuffer, n::Integer)
     io.seekable || throw(ArgumentError("truncate failed, IOBuffer is not seekable"))
     n < 0 && throw(ArgumentError("truncate failed, n bytes must be ≥ 0, got $n"))
     n > io.maxsize && throw(ArgumentError("truncate failed, $(n) bytes is exceeds IOBuffer maxsize $(io.maxsize)"))
-    if n > length(io.data)
+    if io.reinit
+        io.data = _similar_data(io, n)
+        io.reinit = false
+    elseif n > length(io.data)
         resize!(io.data, n)
     end
     io.data[io.size+1:n] .= 0
@@ -325,9 +333,14 @@ end
         ensureroom_slowpath(io, nshort)
     end
     n = min((nshort % Int) + (io.append ? io.size : io.ptr-1), io.maxsize)
-    l = length(io.data)
-    if n > l
-        _growend!(io.data, (n - l) % UInt)
+    if io.reinit
+        io.data = _similar_data(io, n)
+        io.reinit = false
+    else
+        l = length(io.data)
+        if n > l
+            _growend!(io.data, (n - l) % UInt)
+        end
     end
     return io
 end
@@ -390,18 +403,26 @@ end
 function take!(io::IOBuffer)
     ismarked(io) && unmark(io)
     if io.seekable
-        data = io.data
         if io.writable
-            maxsize = (io.maxsize == typemax(Int) ? 0 : min(length(io.data),io.maxsize))
-            io.data = StringVector(maxsize)
+            if io.reinit
+                data = StringVector(0)
+            else
+                data = resize!(io.data, io.size)
+                io.reinit = true
+            end
         else
-            data = copy(data)
+            data = copyto!(StringVector(io.size), 1, io.data, 1, io.size)
         end
-        resize!(data,io.size)
     else
         nbytes = bytesavailable(io)
-        a = StringVector(nbytes)
-        data = read!(io, a)
+        if io.writable
+            data = io.data
+            io.reinit = true
+            _deletebeg!(data, io.ptr-1)
+            resize!(data, nbytes)
+        else
+            data = read!(io, StringVector(nbytes))
+        end
     end
     if io.writable
         io.ptr = 1
@@ -410,6 +431,19 @@ function take!(io::IOBuffer)
     return data
 end
 
+"""
+    _unsafe_take!(io::IOBuffer)
+
+This simply returns the raw resized `io.data`, with no checks to be
+sure that `io` is readable etcetera, and leaves `io` in an inconsistent
+state.  This should only be used internally for performance-critical
+`String` routines that immediately discard `io` afterwards, and it
+*assumes* that `io` is writable and seekable.
+
+It saves no allocations compared to `take!`, it just omits some checks.
+"""
+_unsafe_take!(io::IOBuffer) = resize!(io.data, io.size)
+
 function write(to::IO, from::GenericIOBuffer)
     if to === from
         from.ptr = from.size + 1

base/strings/basic.jl

@@ -636,7 +636,7 @@ function filter(f, s::AbstractString)
     for c in s
         f(c) && write(out, c)
     end
-    String(take!(out))
+    String(_unsafe_take!(out))
 end
 
 ## string first and last ##

base/strings/io.jl

@@ -113,7 +113,7 @@ function sprint(f::Function, args...; context=nothing, sizehint::Integer=0)
     else
         f(s, args...)
     end
-    String(resize!(s.data, s.size))
+    String(_unsafe_take!(s))
 end
 
 function _str_sizehint(x)
@@ -147,7 +147,7 @@ function print_to_string(xs...)
     for x in xs
         print(s, x)
     end
-    String(resize!(s.data, s.size))
+    String(_unsafe_take!(s))
 end
 
 function string_with_env(env, xs...)
@@ -164,7 +164,7 @@ function string_with_env(env, xs...)
     for x in xs
         print(env_io, x)
     end
-    String(resize!(s.data, s.size))
+    String(_unsafe_take!(s))
 end
 
 """

stdlib/REPL/src/LineEdit.jl

@@ -761,10 +761,11 @@ function edit_splice!(s::BufferLike, r::Region=region(s), ins::String = ""; rigi
     elseif buf.mark >= B
         buf.mark += sizeof(ins) - B + A
     end
+    ensureroom(buf, B) # handle !buf.reinit from take!
     ret = splice!(buf.data, A+1:B, codeunits(String(ins))) # position(), etc, are 0-indexed
     buf.size = buf.size + sizeof(ins) - B + A
     adjust_pos && seek(buf, position(buf) + sizeof(ins))
-    return String(ret)
+    return String(copy(ret))
 end
 
 edit_splice!(s::MIState, ins::AbstractString) = edit_splice!(s, region(s), ins)
@@ -1281,7 +1282,7 @@ end
 # compute the number of spaces from b till the next non-space on the right
 # (which can also be "end of line" or "end of buffer")
 function leadingspaces(buf::IOBuffer, b::Int)
-    ls = something(findnext(_notspace, buf.data, b+1), 0)-1
+    @views ls = something(findnext(_notspace, buf.data[1:buf.size], b+1), 0)-1
     ls == -1 && (ls = buf.size)
     ls -= b
     return ls
@@ -2238,7 +2239,7 @@ end
 
 function move_line_end(buf::IOBuffer)
     eof(buf) && return
-    pos = findnext(isequal(UInt8('\n')), buf.data, position(buf)+1)
+    @views pos = findnext(isequal(UInt8('\n')), buf.data[1:buf.size], position(buf)+1)
     if pos === nothing
         move_input_end(buf)
         return

stdlib/REPL/test/lineedit.jl

@@ -306,21 +306,21 @@ seek(buf,0)
 
 ## edit_delete_prev_word ##
 
-buf = IOBuffer("type X\n ")
+buf = IOBuffer(Vector{UInt8}("type X\n "), read=true, write=true)
 seekend(buf)
 @test !isempty(@inferred(LineEdit.edit_delete_prev_word(buf)))
 @test position(buf) == 5
 @test buf.size == 5
 @test content(buf) == "type "
 
-buf = IOBuffer("4 +aaa+ x")
+buf = IOBuffer(Vector{UInt8}("4 +aaa+ x"), read=true, write=true)
 seek(buf,8)
 @test !isempty(LineEdit.edit_delete_prev_word(buf))
 @test position(buf) == 3
 @test buf.size == 4
 @test content(buf) == "4 +x"
 
-buf = IOBuffer("x = func(arg1,arg2 , arg3)")
+buf = IOBuffer(Vector{UInt8}("x = func(arg1,arg2 , arg3)"), read=true, write=true)
 seekend(buf)
 LineEdit.char_move_word_left(buf)
 @test position(buf) == 21