Make Sink/Source::buffered no-op for buffered sinks and sources

fzhinkin · fzhinkin · commit 98762c8a8622 · 2023-06-23T12:07:06.000+02:00
diff --git a/core/common/src/Core.kt b/core/common/src/Core.kt
@@ -24,14 +24,24 @@ package kotlinx.io
  * Returns a new source that buffers reads from the source. The returned source will perform bulk
  * reads into its in-memory buffer. Use this wherever you read a source to get ergonomic and
  * efficient access to data.
+ *
+ * This method returns the receiver if it is already an instance of [Source].
  */
-public fun RawSource.buffered(): Source = RealSource(this)
+public fun RawSource.buffered(): Source = when (this) {
+  is Source -> this
+  else -> RealSource(this)
+}
 
 /**
  * Returns a new sink that buffers writes to the sink. The returned sink will batch writes to the sink.
  * Use this wherever you write to a sink to get ergonomic and efficient access to data.
+ *
+ * This method returns the receiver if it is already an instance of [Sink].
  */
-public fun RawSink.buffered(): Sink = RealSink(this)
+public fun RawSink.buffered(): Sink = when (this) {
+  is Sink -> this
+  else -> RealSink(this)
+}
 
 /**
  * Returns a sink that writes nowhere.
diff --git a/core/common/test/CommonPlatformTest.kt b/core/common/test/CommonPlatformTest.kt
@@ -33,12 +33,12 @@ class CommonPlatformTest {
   }
 
   @Test fun sinkBuffer() {
-    val sink = Buffer()
-    val buffered = (sink as RawSink).buffered()
+    val sink = BufferBackedSink()
+    val buffered = sink.buffered()
     buffered.writeUtf8("a")
-    assertEquals(sink.size, 0L)
+    assertEquals(sink.buffer.size, 0L)
     buffered.flush()
-    assertEquals(sink.size, 1L)
+    assertEquals(sink.buffer.size, 1L)
   }
 
   @Test fun blackhole() {
diff --git a/core/common/test/CommonRealSinkTest.kt b/core/common/test/CommonRealSinkTest.kt
@@ -32,68 +32,68 @@ import kotlin.test.assertFailsWith
 @OptIn(InternalIoApi::class)
 class CommonRealSinkTest {
   @Test fun bufferedSinkEmitsTailWhenItIsComplete() {
-    val sink = Buffer()
-    val bufferedSink = (sink as RawSink).buffered()
+    val sink = BufferBackedSink()
+    val bufferedSink = sink.buffered()
     bufferedSink.writeUtf8("a".repeat(Segment.SIZE - 1))
-    assertEquals(0, sink.size)
+    assertEquals(0, sink.buffer.size)
     bufferedSink.writeByte(0)
-    assertEquals(Segment.SIZE.toLong(), sink.size)
+    assertEquals(Segment.SIZE.toLong(), sink.buffer.size)
     assertEquals(0, bufferedSink.buffer.size)
   }
 
   @Test fun bufferedSinkEmitMultipleSegments() {
-    val sink = Buffer()
-    val bufferedSink = (sink as RawSink).buffered()
+    val sink = BufferBackedSink()
+    val bufferedSink = sink.buffered()
     bufferedSink.writeUtf8("a".repeat(Segment.SIZE * 4 - 1))
-    assertEquals(Segment.SIZE.toLong() * 3L, sink.size)
+    assertEquals(Segment.SIZE.toLong() * 3L, sink.buffer.size)
     assertEquals(Segment.SIZE.toLong() - 1L, bufferedSink.buffer.size)
   }
 
   @Test fun bufferedSinkFlush() {
-    val sink = Buffer()
-    val bufferedSink = (sink as RawSink).buffered()
+    val sink = BufferBackedSink()
+    val bufferedSink = sink.buffered()
     bufferedSink.writeByte('a'.code.toByte())
-    assertEquals(0, sink.size)
+    assertEquals(0, sink.buffer.size)
     bufferedSink.flush()
     assertEquals(0, bufferedSink.buffer.size)
-    assertEquals(1, sink.size)
+    assertEquals(1, sink.buffer.size)
   }
 
   @Test fun bytesEmittedToSinkWithFlush() {
-    val sink = Buffer()
-    val bufferedSink = (sink as RawSink).buffered()
+    val sink = BufferBackedSink()
+    val bufferedSink = sink.buffered()
     bufferedSink.writeUtf8("abc")
     bufferedSink.flush()
-    assertEquals(3, sink.size)
+    assertEquals(3, sink.buffer.size)
   }
 
   @Test fun bytesNotEmittedToSinkWithoutFlush() {
-    val sink = Buffer()
-    val bufferedSink = (sink as RawSink).buffered()
+    val sink = BufferBackedSink()
+    val bufferedSink = sink.buffered()
     bufferedSink.writeUtf8("abc")
-    assertEquals(0, sink.size)
+    assertEquals(0, sink.buffer.size)
   }
 
   @Test fun bytesEmittedToSinkWithEmit() {
-    val sink = Buffer()
-    val bufferedSink = (sink as RawSink).buffered()
+    val sink = BufferBackedSink()
+    val bufferedSink = sink.buffered()
     bufferedSink.writeUtf8("abc")
     bufferedSink.emit()
-    assertEquals(3, sink.size)
+    assertEquals(3, sink.buffer.size)
   }
 
   @Test fun completeSegmentsEmitted() {
-    val sink = Buffer()
-    val bufferedSink = (sink as RawSink).buffered()
+    val sink = BufferBackedSink()
+    val bufferedSink = sink.buffered()
     bufferedSink.writeUtf8("a".repeat(Segment.SIZE * 3))
-    assertEquals(Segment.SIZE.toLong() * 3L, sink.size)
+    assertEquals(Segment.SIZE.toLong() * 3L, sink.buffer.size)
   }
 
   @Test fun incompleteSegmentsNotEmitted() {
-    val sink = Buffer()
-    val bufferedSink = (sink as RawSink).buffered()
+    val sink = BufferBackedSink()
+    val bufferedSink = sink.buffered()
     bufferedSink.writeUtf8("a".repeat(Segment.SIZE * 3 - 1))
-    assertEquals(Segment.SIZE.toLong() * 2L, sink.size)
+    assertEquals(Segment.SIZE.toLong() * 2L, sink.buffer.size)
   }
 
   @Test fun closeWithExceptionWhenWriting() {
diff --git a/core/common/test/CommonRealSourceTest.kt b/core/common/test/CommonRealSourceTest.kt
@@ -47,76 +47,79 @@ class CommonRealSourceTest {
   }
 
   @Test fun requireTracksBufferFirst() {
-    val source = Buffer()
-    source.writeUtf8("bb")
+    val source = BufferBackedSource()
+    source.buffer.writeUtf8("bb")
 
-    val bufferedSource = (source as RawSource).buffered()
+    val bufferedSource = source.buffered()
     bufferedSource.buffer.writeUtf8("aa")
 
     bufferedSource.require(2)
     assertEquals(2, bufferedSource.buffer.size)
-    assertEquals(2, source.size)
+    assertEquals(2, source.buffer.size)
   }
 
   @Test fun requireIncludesBufferBytes() {
-    val source = Buffer()
-    source.writeUtf8("b")
+    val source = BufferBackedSource()
+    source.buffer.writeUtf8("b")
 
-    val bufferedSource = (source as RawSource).buffered()
+    val bufferedSource = source.buffered()
     bufferedSource.buffer.writeUtf8("a")
 
     bufferedSource.require(2)
     assertEquals("ab", bufferedSource.buffer.readUtf8(2))
   }
 
   @Test fun requireInsufficientData() {
-    val source = Buffer()
-    source.writeUtf8("a")
+    val source = BufferBackedSource()
+    source.buffer.writeUtf8("a")
 
-    val bufferedSource = (source as RawSource).buffered()
+    val bufferedSource = source.buffered()
 
     assertFailsWith<EOFException> {
       bufferedSource.require(2)
     }
   }
 
   @Test fun requireReadsOneSegmentAtATime() {
-    val source = Buffer()
-    source.writeUtf8("a".repeat(Segment.SIZE))
-    source.writeUtf8("b".repeat(Segment.SIZE))
+    val source = BufferBackedSource()
+    source.buffer.writeUtf8("a".repeat(Segment.SIZE))
+    source.buffer.writeUtf8("b".repeat(Segment.SIZE))
 
-    val bufferedSource = (source as RawSource).buffered()
+    val bufferedSource = source.buffered()
 
     bufferedSource.require(2)
-    assertEquals(Segment.SIZE.toLong(), source.size)
+    assertEquals(Segment.SIZE.toLong(), source.buffer.size)
     assertEquals(Segment.SIZE.toLong(), bufferedSource.buffer.size)
   }
 
   @Test fun skipReadsOneSegmentAtATime() {
-    val source = Buffer()
-    source.writeUtf8("a".repeat(Segment.SIZE))
-    source.writeUtf8("b".repeat(Segment.SIZE))
-    val bufferedSource = (source as RawSource).buffered()
+    val source = BufferBackedSource()
+    source.buffer.writeUtf8("a".repeat(Segment.SIZE))
+    source.buffer.writeUtf8("b".repeat(Segment.SIZE))
+    val bufferedSource = source.buffered()
     bufferedSource.skip(2)
-    assertEquals(Segment.SIZE.toLong(), source.size)
+    assertEquals(Segment.SIZE.toLong(), source.buffer.size)
     assertEquals(Segment.SIZE.toLong() - 2L, bufferedSource.buffer.size)
   }
 
   @Test fun skipTracksBufferFirst() {
-    val source = Buffer()
-    source.writeUtf8("bb")
+    val source = BufferBackedSource()
+    source.buffer.writeUtf8("bb")
 
-    val bufferedSource = (source as RawSource).buffered()
+    val bufferedSource = source.buffered()
     bufferedSource.buffer.writeUtf8("aa")
 
     bufferedSource.skip(2)
     assertEquals(0, bufferedSource.buffer.size)
-    assertEquals(2, source.size)
+    assertEquals(2, source.buffer.size)
   }
 
   @Test fun operationsAfterClose() {
-    val source = Buffer()
-    val bufferedSource = (source as RawSource).buffered()
+    val source = object : RawSource {
+      override fun readAtMostTo(sink: Buffer, byteCount: Long): Long = TODO()
+      override fun close() = Unit
+    }
+    val bufferedSource = source.buffered()
     bufferedSource.close()
 
     // Test a sample set of methods.
@@ -138,12 +141,12 @@ class CommonRealSourceTest {
     val write2 = Buffer().also { it.writeUtf8("b".repeat(Segment.SIZE)) }
     val write3 = Buffer().also { it.writeUtf8("c".repeat(Segment.SIZE)) }
 
-    val source = Buffer()
-    source.writeUtf8(
+    val source = BufferBackedSource()
+    source.buffer.writeUtf8(
       "${"a".repeat(Segment.SIZE)}${"b".repeat(Segment.SIZE)}${"c".repeat(Segment.SIZE)}")
 
     val mockSink = MockSink()
-    val bufferedSource = (source as RawSource).buffered()
+    val bufferedSource = source.buffered()
     assertEquals(Segment.SIZE.toLong() * 3L, bufferedSource.transferTo(mockSink))
     mockSink.assertLog(
       "write($write1, ${write1.size})",
diff --git a/core/common/test/DelicateApiTest.kt b/core/common/test/DelicateApiTest.kt
@@ -13,15 +13,14 @@ class DelicateApiTest {
     @Test
     @OptIn(InternalIoApi::class)
     fun testWriteIntoBuffer() {
-        val sink = Buffer()
-        val rawSink = sink as RawSink
+        val rawSink = BufferBackedSink()
         val bufferedSink = rawSink.buffered()
 
         bufferedSink.writeToInternalBuffer {
             it.writeByte(42)
         }
 
-        assertEquals(0, sink.size)
+        assertEquals(0, rawSink.buffer.size)
         assertEquals(1, bufferedSink.buffer.size)
 
         bufferedSink.writeToInternalBuffer {
@@ -35,7 +34,7 @@ class DelicateApiTest {
             it.writeByte(0x34)
         }
 
-        assertArrayEquals(byteArrayOf(0x0, 0x12), sink.readByteArray())
+        assertArrayEquals(byteArrayOf(0x0, 0x12), rawSink.buffer.readByteArray())
         assertArrayEquals(byteArrayOf(0x34), bufferedSink.buffer.readByteArray())
     }
 }
diff --git a/core/common/test/util.kt b/core/common/test/util.kt
@@ -67,3 +67,24 @@ fun Char.repeat(count: Int): String {
 fun assertArrayEquals(a: ByteArray, b: ByteArray) {
   assertEquals(a.contentToString(), b.contentToString())
 }
+
+class BufferBackedSink : RawSink {
+  val buffer = Buffer()
+
+  override fun write(source: Buffer, byteCount: Long) {
+    source.readTo(buffer, byteCount)
+  }
+
+  override fun flush() = Unit
+  override fun close() = Unit
+}
+
+class BufferBackedSource : RawSource {
+  val buffer = Buffer()
+
+  override fun readAtMostTo(sink: Buffer, byteCount: Long): Long {
+    return buffer.readAtMostTo(sink, byteCount)
+  }
+
+  override fun close() = Unit
+}

Original file line number	Diff line number	Diff line change
`@@ -13,15 +13,14 @@ class DelicateApiTest {`
`13`	`13`	`@Test`
`14`	`14`	`@OptIn(InternalIoApi::class)`
`15`	`15`	`fun testWriteIntoBuffer() {`
`16`		`- val sink = Buffer()`
`17`		`- val rawSink = sink as RawSink`
	`16`	`+ val rawSink = BufferBackedSink()`
`18`	`17`	`val bufferedSink = rawSink.buffered()`
`19`	`18`
`20`	`19`	`bufferedSink.writeToInternalBuffer {`
`21`	`20`	`it.writeByte(42)`
`22`	`21`	`}`
`23`	`22`
`24`		`- assertEquals(0, sink.size)`
	`23`	`+ assertEquals(0, rawSink.buffer.size)`
`25`	`24`	`assertEquals(1, bufferedSink.buffer.size)`
`26`	`25`
`27`	`26`	`bufferedSink.writeToInternalBuffer {`
`@@ -35,7 +34,7 @@ class DelicateApiTest {`
`35`	`34`	`it.writeByte(0x34)`
`36`	`35`	`}`
`37`	`36`
`38`		`- assertArrayEquals(byteArrayOf(0x0, 0x12), sink.readByteArray())`
	`37`	`+ assertArrayEquals(byteArrayOf(0x0, 0x12), rawSink.buffer.readByteArray())`
`39`	`38`	`assertArrayEquals(byteArrayOf(0x34), bufferedSink.buffer.readByteArray())`
`40`	`39`	`}`
`41`	`40`	`}`