From be1e354b5d736460ddb66b1ca4ce84a4e98fd4fb Mon Sep 17 00:00:00 2001
From: Jakob Nybo Nissen <jakobnybonissen@gmail.com>
Date: Wed, 24 Jul 2024 10:27:58 +0200
Subject: [PATCH] Address reviewer comments

* Do not assume io.seekable in unsafe_takestring!(::IOBuffer)
* Add fallback definitions unsafe_takestring!(::GenericIOBuffer) and
  takestring!(::GenericIOBuffer)
* Test takestring! with a nonzero offset in buffer
---
 base/iobuffer.jl | 21 +++++++++++----------
 test/iobuffer.jl |  9 +++++++++
 2 files changed, 20 insertions(+), 10 deletions(-)

diff --git a/base/iobuffer.jl b/base/iobuffer.jl
index 3052cb56d5bf9..48b3e2ff00512 100644
--- a/base/iobuffer.jl
+++ b/base/iobuffer.jl
@@ -489,23 +489,23 @@ function take!(io::IOBuffer)
     return data
 end
 
-"Internal method. This method can be faster than takestring!, because it assumes
-the buffer is seekable, and because it does not reset the buffer to a usable state.
+"Internal method. This method can be faster than takestring!, because it does not
+reset the buffer to a usable state, and it does not check for io.reinit.
 Using the buffer after calling unsafe_takestring! may cause undefined behaviour.
 This function is meant to be used when the buffer is only used as a temporary
 string builder, which is discarded after the string is built."
 function unsafe_takestring!(io::IOBuffer)
-    off = io.offset
-    nbytes = io.size - off
+    start = io.seekable ? io.offset + 1 : io.ptr
+    nbytes = io.size - start + 1
     iszero(nbytes) && return ""
     # The C function can only copy from the start of the memory.
     # Fortunately, in most cases, the offset will be zero.
-    return if iszero(off)
+    return if isone(start)
         ccall(:jl_genericmemory_to_string, Ref{String}, (Any, Int), io.data, nbytes)
     else
         mem = StringMemory(nbytes)
-        unsafe_copyto!(mem, 1, io.data, off + 1, nbytes)
-        unsafe_takestring!(mem)
+        unsafe_copyto!(mem, 1, io.data, start, nbytes)
+        unsafe_takestring(mem)
     end
 end
 
@@ -538,9 +538,6 @@ function takestring!(io::IOBuffer)
     # we can return an empty string without interacting with the buffer at all.
     io.reinit && return ""
 
-    # Currently, `IOBuffer`s (as opposed to `GenericIOBuffer`) is always seekable, since all
-    # IOBuffer constructors return seekable buffers, and GenericIOBuffer is internal.
-    # This makes `unsafe_takestring!` valid.
     s = unsafe_takestring!(io)
 
     # Restore the buffer to a usable state, making it no longer undefined behaviour to
@@ -558,6 +555,10 @@ function takestring!(io::IOBuffer)
     s
 end
 
+# Fallback methods
+unsafe_takestring!(io::GenericIOBuffer) = takestring!(io)
+takestring!(io::GenericIOBuffer) = String(take!(io))
+
 """
     _unsafe_take!(io::IOBuffer)
 
diff --git a/test/iobuffer.jl b/test/iobuffer.jl
index 8fddf07915fde..0509db746c766 100644
--- a/test/iobuffer.jl
+++ b/test/iobuffer.jl
@@ -63,6 +63,15 @@ end
     write(buf, "xyz")
     @test takestring!(buf) == "xyz"
     buf = IOBuffer()
+
+    # Test with a nonzero offset in the buffer
+    v = rand(UInt8, 8)
+    for i in 1:8
+        pushfirst!(v, rand(UInt8))
+    end
+    buf = IOBuffer(v)
+    s = String(copy(v))
+    @test takestring!(buf) == s
 end
 
 @testset "Read/write readonly IOBuffer" begin