Fix end gap processing in takeWhile (#23)

Author: Sergey Mashkov

kotlinx-io-js/src/main/kotlin/kotlinx/io/core/ByteReadPacket.kt

@@ -69,7 +69,12 @@ actual class ByteReadPacket
     internal actual constructor(head: IoBuffer, remaining: Long, pool: ObjectPool<IoBuffer>) : ByteReadPacketPlatformBase(head, remaining, pool), Input {
     actual constructor(head: IoBuffer, pool: ObjectPool<IoBuffer>) : this(head, head.remainingAll(), pool)
 
+    init {
+        markNoMoreChunksAvailable()
+    }
+
     final override fun fill() = null
+
     override fun closeSource() {
     }
 

kotlinx-io-jvm/src/main/kotlin/kotlinx/io/core/ByteReadPacket.kt

@@ -39,7 +39,12 @@ internal actual constructor(head: IoBuffer, remaining: Long, pool: ObjectPool<Io
         pool
     )
 
+    init {
+        markNoMoreChunksAvailable()
+    }
+
     final override fun fill() = null
+
     override fun closeSource() {
     }
 

kotlinx-io-native/src/main/kotlin/kotlinx/io/core/ByteReadPacket.kt

@@ -40,9 +40,7 @@ internal actual constructor(head: IoBuffer, remaining: Long, pool: ObjectPool<Io
     actual constructor(head: IoBuffer, pool: ObjectPool<IoBuffer>) : this(head, head.remainingAll(), pool)
 
     init {
-        if (remaining == 0L) {
-            doFill()
-        }
+        markNoMoreChunksAvailable()
     }
 
     final override fun fill() = null

src/main/kotlin/kotlinx/io/core/Packet.kt

@@ -619,6 +619,8 @@ abstract class ByteReadPacketBase(@PublishedApi internal var head: IoBuffer,
 
         val remaining = current.readRemaining
         val overrunSize = minOf(remaining, ReservedSize - current.endGap)
+        if (next.startGap < overrunSize) return fixGapAfterReadFallback(current)
+
         next.restoreStartGap(overrunSize)
 
         if (remaining > overrunSize) {
@@ -637,6 +639,12 @@ abstract class ByteReadPacketBase(@PublishedApi internal var head: IoBuffer,
     }
 
     private fun fixGapAfterReadFallback(current: IoBuffer) {
+        if (noMoreChunksAvailable) {
+            this.headRemaining = current.readRemaining
+            this.tailRemaining = 0
+            return
+        }
+
         val size = current.readRemaining
         val overrun = minOf(size, ReservedSize - current.endGap)
 
@@ -645,6 +653,7 @@ abstract class ByteReadPacketBase(@PublishedApi internal var head: IoBuffer,
         } else {
             val new = pool.borrow()
             new.reserveEndGap(ReservedSize)
+            new.next = current.next
 
             new.writeBufferAppend(current, size)
             this.head = new
@@ -665,6 +674,7 @@ abstract class ByteReadPacketBase(@PublishedApi internal var head: IoBuffer,
         chunk1.reserveEndGap(ByteReadPacketBase.ReservedSize)
         chunk2.reserveEndGap(ByteReadPacketBase.ReservedSize)
         chunk1.next = chunk2
+        chunk2.next = current.next
 
         chunk1.writeBufferAppend(current, size - overrun)
         chunk2.writeBufferAppend(current, overrun)
@@ -708,6 +718,12 @@ abstract class ByteReadPacketBase(@PublishedApi internal var head: IoBuffer,
 
     protected abstract fun closeSource()
 
+    internal fun markNoMoreChunksAvailable() {
+        if (!noMoreChunksAvailable) {
+            noMoreChunksAvailable = true
+        }
+    }
+
     protected fun doFill(): IoBuffer? {
         if (noMoreChunksAvailable) return null
         val chunk = fill()

src/test/kotlin/kotlinx/io/tests/ReadTextCommonTest.kt

@@ -321,6 +321,66 @@ class ReadTextCommonTest {
         assertEquals("\u0BF5", text)
     }
 
+    @Test
+    fun testReadTextAfterInt() {
+        buildPacket {
+            writeInt(5)
+            writeText("Hello")
+            writeInt(1)
+        }.use { packet ->
+            assertEquals(5, packet.readInt())
+            assertEquals("Hello", packet.readText(min = 5, max = 5))
+            assertEquals(1, packet.readInt())
+        }
+    }
+
+    @Test
+    fun testReadTextAfterIntFromInput() {
+        val content = buildPacket {
+            writeInt(5)
+            writeText("Hello")
+            writeInt(1)
+
+            writeFully(ByteArray(8192))
+        }.readBytes()
+
+        val input = object : AbstractInput() {
+            private var offset = 0
+
+            override fun fill(): IoBuffer? {
+                if (offset >= content.size) return null
+
+                val buffer = pool.borrow()
+                buffer.reserveEndGap(8)
+
+                val size = minOf(buffer.writeRemaining, content.size - offset)
+                buffer.writeFully(content, offset, size)
+                offset += size
+
+                return buffer
+            }
+
+            override fun closeSource() {
+                offset = Int.MAX_VALUE
+            }
+        }
+
+        try {
+            assertEquals(5, input.readInt())
+            assertEquals("Hello", input.readText(min = 5, max = 5))
+            assertEquals(1, input.readInt())
+        } finally {
+            input.close()
+        }
+    }
+
+    @Test
+    fun testReadTextFromPacketFromByteArray() {
+        val content = byteArrayOf(0x31, 0x32, 0x33, 0x00)
+        val packet = ByteReadPacket(content)
+        assertEquals("123", packet.readText(max = 3))
+    }
+
     private inline fun buildPacket(startGap: Int = 0, block: BytePacketBuilder.() -> Unit): ByteReadPacket {
         val builder = BytePacketBuilder(startGap, pool)
         try {

src/test/kotlin/kotlinx/io/tests/TakeWhileTest.kt

@@ -5,24 +5,24 @@ import kotlin.test.*
 
 class TakeWhileTest {
     private val pool = VerifyingObjectPool(IoBuffer.NoPool)
-    private val ch1 = pool.borrow()
-    private val ch2 = pool.borrow()
+    private val chunk1 = pool.borrow()
+    private val chunk2 = pool.borrow()
 
     private val chunks = ArrayList<IoBuffer>()
     private val packets = ArrayList<ByteReadPacket>()
 
     @BeforeTest
     fun prepare() {
-        ch1.resetForWrite()
-        ch2.resetForWrite()
+        chunk1.resetForWrite()
+        chunk2.resetForWrite()
 
-        ch1.reserveEndGap(8)
-        ch2.reserveEndGap(8)
+        chunk1.reserveEndGap(8)
+        chunk2.reserveEndGap(8)
 
-        ch1.next = ch2
+        chunk1.next = chunk2
 
-        chunks.add(ch1)
-        chunks.add(ch2)
+        chunks.add(chunk1)
+        chunks.add(chunk2)
     }
 
     @AfterTest
@@ -39,49 +39,112 @@ class TakeWhileTest {
 
     @Test
     fun takeWhileSizeAtEdgeNotConsumed() {
-        ch1.writeFully(cycle(ch1.writeRemaining))
-        ch2.writeFully(cycle(3))
+        chunk1.writeFully(cycle(chunk1.writeRemaining))
+        chunk2.writeFully(cycle(3))
 
-        val pkt = ByteReadPacket(ch1, pool)
+        val pkt = ByteReadPacket(chunk1, pool)
         packets.add(pkt)
-        chunks.remove(ch1)
-        chunks.remove(ch2)
+        chunks.remove(chunk1)
+        chunks.remove(chunk2)
 
-        pkt.discardExact(ch1.readRemaining - 1)
+        pkt.discardExact(chunk1.readRemaining - 1)
 
         assertEquals(4, pkt.remaining)
-        assertEquals(1, ch1.readRemaining)
-        assertEquals(3, ch2.readRemaining)
+        assertEquals(1, chunk1.readRemaining)
+        assertEquals(3, chunk2.readRemaining)
 
         pkt.takeWhileSize(4) { 0 }
         assertEquals(4, pkt.remaining)
     }
 
     @Test
     fun takeWhileSizeAtEdgeNotConsumed2() {
-        ch1.writeFully(cycle(ch1.writeRemaining))
-        ch2.writeFully(cycle(10))
+        chunk1.writeFully(cycle(chunk1.writeRemaining))
+        chunk2.writeFully(cycle(10))
 
-        val pkt = ByteReadPacket(ch1, pool)
+        val pkt = ByteReadPacket(chunk1, pool)
         packets.add(pkt)
-        chunks.remove(ch1)
-        chunks.remove(ch2)
+        chunks.remove(chunk1)
+        chunks.remove(chunk2)
 
-        pkt.discardExact(ch1.readRemaining - 1)
+        pkt.discardExact(chunk1.readRemaining - 1)
 
         assertEquals(11, pkt.remaining)
-        assertEquals(1, ch1.readRemaining)
-        assertEquals(10, ch2.readRemaining)
+        assertEquals(1, chunk1.readRemaining)
+        assertEquals(10, chunk2.readRemaining)
 
         pkt.takeWhileSize(8) { it.discard(7); 0 }
 
-        assertNotNull(ch1.next)
+        assertNotNull(chunk1.next)
         assertEquals(4, pkt.remaining)
 
         pkt.takeWhileSize(4) { 0 }
 
         assertEquals(4, pkt.remaining)
     }
 
+    @Test
+    fun testTakeWhileFromByteArrayPacket1() {
+        val content = byteArrayOf(0x31, 0x32, 0x33, 0x00)
+        val packet = ByteReadPacket(content)
+
+        packet.takeWhileSize { it.discardExact(3); 0 }
+        assertEquals(1, packet.remaining)
+        packet.discardExact(1)
+    }
+
+    @Test
+    fun testTakeWhileFromByteArrayPacket2() {
+        val content = byteArrayOf(0x31, 0x32, 0x33, 0x00)
+        val packet = ByteReadPacket(content)
+
+        packet.takeWhileSize { it.discardExact(4); 0 }
+        assertEquals(0, packet.remaining)
+    }
+
+    @Test
+    fun testTakeWhileFromByteArrayPacket3() {
+        val content = byteArrayOf(0x31, 0x32, 0x33, 0x00)
+        val packet = ByteReadPacket(content)
+
+        packet.takeWhileSize { 0 }
+        assertEquals(4, packet.remaining)
+        packet.discardExact(4)
+    }
+
+    @Test
+    @Ignore
+    fun testLong() {
+        val goldenCopy = cycle(8192)
+        for (discardBefore in 0..8192) {
+            buildPacket {
+                writeFully(goldenCopy)
+            }.use { pkt ->
+                pkt.discardExact(discardBefore)
+                val expected = goldenCopy.copyOfRange(discardBefore, goldenCopy.size)
+
+                for (initialSize in 1..8) {
+                    for (consume in 0..initialSize) {
+                        val actual = buildPacket {
+                            pkt.copy().use { copy ->
+                                copy.takeWhileSize(initialSize) { buffer ->
+                                    val consumed = buffer.readBytes(consume)
+                                    writeFully(consumed)
+                                    if (consume == 0) 0 else initialSize
+                                }
+
+                                val remainingBytes = copy.readBytes()
+                                writeFully(remainingBytes)
+                            }
+                        }.readBytes()
+
+                        val equals = actual.contentEquals(expected)
+                        assertTrue(equals)
+                    }
+                }
+            }
+        }
+    }
+
     private fun cycle(size: Int) = ByteArray(size) { (it and 0xff).toByte() }
 }