Correctly split head segment (#350)

* Fixed how buffer's head is updated after the split

Author: fzhinkin

core/common/src/Buffer.kt

@@ -488,7 +488,7 @@ public class Buffer : Source, Sink {
 
         while (remainingByteCount > 0L) {
             // Is a prefix of the source's head segment all that we need to move?
-            if (remainingByteCount < source.head!!.limit - source.head!!.pos) {
+            if (remainingByteCount < source.head!!.size) {
                 val tail = tail
                 if (tail != null && tail.owner &&
                     remainingByteCount + tail.limit - (if (tail.shared) 0 else tail.pos) <= Segment.SIZE
@@ -501,17 +501,13 @@ public class Buffer : Source, Sink {
                 } else {
                     // We're going to need another segment. Split the source's head
                     // segment in two, then move the first of those two to this buffer.
-                    val newHead = source.head!!.split(remainingByteCount.toInt())
-                    if (source.head == source.tail) {
-                        source.tail = newHead
-                    }
-                    source.head = newHead
+                    source.head = source.head!!.split(remainingByteCount.toInt())
                 }
             }
 
             // Remove the source's head segment and append it to our tail.
-            val segmentToMove = source.head
-            val movedByteCount = (segmentToMove!!.limit - segmentToMove.pos).toLong()
+            val segmentToMove = source.head!!
+            val movedByteCount = segmentToMove.size.toLong()
             source.head = segmentToMove.pop()
             if (source.head == null) {
                 source.tail = null

core/common/test/CommonBufferTest.kt

@@ -601,4 +601,21 @@ class CommonBufferTest {
         buffer.clear()
         assertEquals(ByteString(), buffer.snapshot())
     }
+
+    @Test
+    fun splitHead() {
+        val dst = Buffer()
+        val src = Buffer().also { it.write(ByteArray(SEGMENT_SIZE)) }
+
+        dst.write(src, src.size / 2)
+        assertEquals(src.size, dst.size)
+
+        assertEquals(src.head, src.tail)
+        assertEquals(null, src.head?.prev)
+        assertEquals(null, src.tail?.next)
+
+        assertEquals(dst.head, dst.tail)
+        assertEquals(null, dst.head?.prev)
+        assertEquals(null, dst.tail?.next)
+    }
 }