Repository: Kotlin/kotlinx-io
Pull request #383 (2024-09-04 17:35:59)

UnsafeBufferOperations.forEachSegment implementation
More concise replacement for 'iterate' function

qwwdfsad: This is a draft, please take a look.

If you are good with the idea, I'll deprecate `iterate` functions, write tests for `forEach` and re-tweak samples.


Patch:
diff --git a/core/api/kotlinx-io-core.api b/core/api/kotlinx-io-core.api
index 35fc21e7..3d47cde1 100644
--- a/core/api/kotlinx-io-core.api
+++ b/core/api/kotlinx-io-core.api
@@ -305,6 +305,7 @@ public abstract interface class kotlinx/io/unsafe/SegmentWriteContext {
 
 public final class kotlinx/io/unsafe/UnsafeBufferOperations {
 	public static final field INSTANCE Lkotlinx/io/unsafe/UnsafeBufferOperations;
+	public final fun forEachSegment (Lkotlinx/io/Buffer;Lkotlin/jvm/functions/Function2;)V
 	public final fun getMaxSafeWriteCapacity ()I
 	public final fun iterate (Lkotlinx/io/Buffer;JLkotlin/jvm/functions/Function3;)V
 	public final fun iterate (Lkotlinx/io/Buffer;Lkotlin/jvm/functions/Function2;)V
diff --git a/core/api/kotlinx-io-core.klib.api b/core/api/kotlinx-io-core.klib.api
index 06a3309e..2db224bc 100644
--- a/core/api/kotlinx-io-core.klib.api
+++ b/core/api/kotlinx-io-core.klib.api
@@ -211,6 +211,7 @@ final object kotlinx.io.unsafe/UnsafeBufferOperations { // kotlinx.io.unsafe/Uns
         final fun <get-maxSafeWriteCapacity>(): kotlin/Int // kotlinx.io.unsafe/UnsafeBufferOperations.maxSafeWriteCapacity.<get-maxSafeWriteCapacity>|<get-maxSafeWriteCapacity>(){}[0]
 
     final fun moveToTail(kotlinx.io/Buffer, kotlin/ByteArray, kotlin/Int = ..., kotlin/Int = ...) // kotlinx.io.unsafe/UnsafeBufferOperations.moveToTail|moveToTail(kotlinx.io.Buffer;kotlin.ByteArray;kotlin.Int;kotlin.Int){}[0]
+    final inline fun forEachSegment(kotlinx.io/Buffer, kotlin/Function2<kotlinx.io.unsafe/SegmentReadContext, kotlinx.io/Segment, kotlin/Unit>) // kotlinx.io.unsafe/UnsafeBufferOperations.forEachSegment|forEachSegment(kotlinx.io.Buffer;kotlin.Function2<kotlinx.io.unsafe.SegmentReadContext,kotlinx.io.Segment,kotlin.Unit>){}[0]
     final inline fun iterate(kotlinx.io/Buffer, kotlin/Function2<kotlinx.io.unsafe/BufferIterationContext, kotlinx.io/Segment?, kotlin/Unit>) // kotlinx.io.unsafe/UnsafeBufferOperations.iterate|iterate(kotlinx.io.Buffer;kotlin.Function2<kotlinx.io.unsafe.BufferIterationContext,kotlinx.io.Segment?,kotlin.Unit>){}[0]
     final inline fun iterate(kotlinx.io/Buffer, kotlin/Long, kotlin/Function3<kotlinx.io.unsafe/BufferIterationContext, kotlinx.io/Segment?, kotlin/Long, kotlin/Unit>) // kotlinx.io.unsafe/UnsafeBufferOperations.iterate|iterate(kotlinx.io.Buffer;kotlin.Long;kotlin.Function3<kotlinx.io.unsafe.BufferIterationContext,kotlinx.io.Segment?,kotlin.Long,kotlin.Unit>){}[0]
     final inline fun readFromHead(kotlinx.io/Buffer, kotlin/Function2<kotlinx.io.unsafe/SegmentReadContext, kotlinx.io/Segment, kotlin/Int>): kotlin/Int // kotlinx.io.unsafe/UnsafeBufferOperations.readFromHead|readFromHead(kotlinx.io.Buffer;kotlin.Function2<kotlinx.io.unsafe.SegmentReadContext,kotlinx.io.Segment,kotlin.Int>){}[0]
diff --git a/core/apple/src/BuffersApple.kt b/core/apple/src/BuffersApple.kt
index 98e2c0ea..29f698c1 100644
--- a/core/apple/src/BuffersApple.kt
+++ b/core/apple/src/BuffersApple.kt
@@ -58,19 +58,14 @@ internal fun Buffer.snapshotAsNSData(): NSData {
     val bytes = malloc(size.convert())?.reinterpret<uint8_tVar>()
         ?: throw Error("malloc failed: ${strerror(errno)?.toKString()}")
 
-    UnsafeBufferOperations.iterate(this) { ctx, head ->
-        var curr: Segment? = head
-        var index = 0
-        while (curr != null) {
-            val segment: Segment = curr
-            ctx.withData(segment) { data, pos, limit ->
-                val length = limit - pos
-                data.usePinned {
-                    memcpy(bytes + index, it.addressOf(pos), length.convert())
-                }
-                index += length
+    var index = 0
+    UnsafeBufferOperations.forEachSegment(this) { ctx, segment ->
+        ctx.withData(segment) { data, pos, limit ->
+            val length = limit - pos
+            data.usePinned {
+                memcpy(bytes + index, it.addressOf(pos), length.convert())
             }
-            curr = ctx.next(segment)
+            index += length
         }
     }
     return NSData.create(bytesNoCopy = bytes, length = size.convert())
diff --git a/core/common/src/Buffer.kt b/core/common/src/Buffer.kt
index 0dcc7d5b..d2c5dc96 100644
--- a/core/common/src/Buffer.kt
+++ b/core/common/src/Buffer.kt
@@ -69,6 +69,7 @@ public class Buffer : Source, Sink {
      */
     @InternalIoApi
     override val buffer: Buffer get() = this
+
     override fun exhausted(): Boolean = size == 0L
 
     override fun require(byteCount: Long) {
@@ -179,7 +180,7 @@ public class Buffer : Source, Sink {
     override fun flush(): Unit = Unit
 
     /**
-     * Copy bytes from this buffer's subrange starting at [startIndex] and ending at [endIndex], to [out] buffer.
+     * Copies bytes from this buffer's subrange starting at [startIndex] and ending at [endIndex], to [out] buffer.
      * This method does not consume data from the buffer.
      *
      * @param out the destination buffer to copy data into.
@@ -553,21 +554,16 @@ public class Buffer : Source, Sink {
         val len = minOf(maxPrintableBytes, size).toInt()
 
         val builder = StringBuilder(len * 2 + if (size > maxPrintableBytes) 1 else 0)
-
-        UnsafeBufferOperations.iterate(this) { ctx, head ->
-            var bytesWritten = 0
-            var seg: Segment? = head
-            do {
-                seg!!
-                var idx = 0
-                while (bytesWritten < len && idx < seg.size) {
-                    val b = ctx.getUnchecked(seg, idx++)
-                    bytesWritten++
-                    builder.append(HEX_DIGIT_CHARS[(b shr 4) and 0xf])
-                        .append(HEX_DIGIT_CHARS[b and 0xf])
-                }
-                seg = ctx.next(seg)
-            } while (seg != null)
+        var bytesWritten = 0
+        UnsafeBufferOperations.forEachSegment(this) { ctx, segment ->
+            var idx = 0
+            while (bytesWritten < len && idx < segment.size) {
+                val b = ctx.getUnchecked(segment, idx++)
+                bytesWritten++
+                builder
+                    .append(HEX_DIGIT_CHARS[(b shr 4) and 0xf])
+                    .append(HEX_DIGIT_CHARS[b and 0xf])
+            }
         }
 
         if (size > maxPrintableBytes) {
diff --git a/core/common/src/Buffers.kt b/core/common/src/Buffers.kt
index c25b0a2f..479e04cf 100644
--- a/core/common/src/Buffers.kt
+++ b/core/common/src/Buffers.kt
@@ -22,12 +22,8 @@ public fun Buffer.snapshot(): ByteString {
     check(size <= Int.MAX_VALUE) { "Buffer is too long ($size) to be converted into a byte string." }
 
     return buildByteString(size.toInt()) {
-        UnsafeBufferOperations.iterate(this@snapshot) { ctx, head ->
-            var curr = head
-            while (curr != null) {
-                ctx.withData(curr, this::append)
-                curr = ctx.next(curr)
-            }
+        UnsafeBufferOperations.forEachSegment(this@snapshot) { ctx, segment ->
+            ctx.withData(segment, this::append)
         }
     }
 }
diff --git a/core/common/src/Utf8.kt b/core/common/src/Utf8.kt
index f455fbce..698faa21 100644
--- a/core/common/src/Utf8.kt
+++ b/core/common/src/Utf8.kt
@@ -607,17 +607,17 @@ private fun Buffer.commonReadUtf8(byteCount: Long): String {
     // Invariant: byteCount was request()'ed into this buffer beforehand
     if (byteCount == 0L) return ""
 
-    UnsafeBufferOperations.iterate(this) { ctx, head ->
-        head!!
-        if (head.size >= byteCount) {
+    UnsafeBufferOperations.forEachSegment(this) { ctx, segment ->
+        if (segment.size >= byteCount) {
             var result = ""
-            ctx.withData(head) { data, pos, limit ->
+            ctx.withData(segment) { data, pos, limit ->
                 result = data.commonToUtf8String(pos, min(limit, pos + byteCount.toInt()))
                 skip(byteCount)
                 return result
             }
         }
-    }
         // If the string spans multiple segments, delegate to readBytes()
         return readByteArray(byteCount.toInt()).commonToUtf8String()
     }
+    error("Unreacheable")
+}
diff --git a/core/common/src/unsafe/UnsafeBufferOperations.kt b/core/common/src/unsafe/UnsafeBufferOperations.kt
index 1f0ce7f3..f6cceceb 100644
--- a/core/common/src/unsafe/UnsafeBufferOperations.kt
+++ b/core/common/src/unsafe/UnsafeBufferOperations.kt
@@ -295,13 +295,14 @@ public object UnsafeBufferOperations {
      * the [iterationAction].
      *
      * Both [iterationAction] arguments are valid only within [iterationAction] scope,
-     * it's an error to store and reuse it later.
+     * it is an error to store and reuse it later.
+     *
+     * For a full iteration over buffer's segments, see [forEachSegment].
      *
      * @param buffer a buffer to iterate over
      * @param iterationAction a callback to invoke with the head reference and an iteration context instance
      *
-     * @sample kotlinx.io.samples.unsafe.UnsafeReadWriteSamplesJvm.messageDigest
-     * @sample kotlinx.io.samples.unsafe.UnsafeBufferOperationsSamples.crc32Unsafe
+     * @sample kotlinx.io.samples.unsafe.UnsafeBufferOperationsSamples.crc32GetUnchecked
      */
     public inline fun iterate(
         buffer: Buffer,
@@ -352,6 +353,34 @@ public object UnsafeBufferOperations {
             iterationAction(BufferIterationContextImpl, s, o)
         }
     }
+
+    /**
+     * Iterates over [buffer] segments starting from the head.
+     *
+     * [action] is invoked with an instance of [SegmentReadContext]
+     * allowing to read and write in an unchecked manner from [buffer]'s segments
+     *
+     * It is considered an error to use a [SegmentReadContext] or a [Segment] instances outside the scope of
+     * the [action].
+     *
+     * Both [action] arguments are valid only within [action] scope, it is an error to store and reuse it later.
+     * The action might never be invoked if the given [buffer] is empty.
+     *
+     * @param buffer a buffer to iterate over
+     * @param action a callback to invoke with the head reference and an iteration context instance
+     * @sample kotlinx.io.samples.unsafe.UnsafeReadWriteSamplesJvm.messageDigest
+     * @sample kotlinx.io.samples.unsafe.UnsafeBufferOperationsSamples.crc32Unsafe
+     */
+    public inline fun forEachSegment(
+        buffer: Buffer,
+        action: (context: SegmentReadContext, segment: Segment) -> Unit
+    ) {
+        var curr: Segment? = buffer.head
+        while (curr != null) {
+            action(SegmentReadContextImpl, curr)
+            curr = curr.next
+        }
+    }
 }
 
 /**
@@ -494,8 +523,7 @@ public interface BufferIterationContext : SegmentReadContext {
      *
      * @param segment a segment for which a successor needs to be found
      *
-     * @sample kotlinx.io.samples.unsafe.UnsafeReadWriteSamplesJvm.messageDigest
-     * @sample kotlinx.io.samples.unsafe.UnsafeBufferOperationsSamples.crc32Unsafe
+     * @sample kotlinx.io.samples.unsafe.UnsafeBufferOperationsSamples.crc32GetUnchecked
      */
     public fun next(segment: Segment): Segment?
 }
diff --git a/core/jvm/src/BuffersJvm.kt b/core/jvm/src/BuffersJvm.kt
index 6b488a53..4b1364a7 100644
--- a/core/jvm/src/BuffersJvm.kt
+++ b/core/jvm/src/BuffersJvm.kt
@@ -30,7 +30,7 @@ import java.nio.ByteBuffer
 import java.nio.channels.ByteChannel
 
 /**
- * Read and exhaust bytes from [input] into this buffer. Stops reading data on [input] exhaustion.
+ * Reads and exhausts bytes from [input] into this buffer. Stops reading data on [input] exhaustion.
  *
  * @param input the stream to read data from.
  *
@@ -42,7 +42,7 @@ public fun Buffer.transferFrom(input: InputStream): Buffer {
 }
 
 /**
- * Read [byteCount] bytes from [input] into this buffer. Throws an exception when [input] is
+ * Reads [byteCount] bytes from [input] into this buffer. Throws an exception when [input] is
  * exhausted before reading [byteCount] bytes.
  *
  * @param input the stream to read data from.
@@ -107,7 +107,7 @@ public fun Buffer.readTo(out: OutputStream, byteCount: Long = size) {
 }
 
 /**
- * Copy bytes from this buffer's subrange, starting at [startIndex] and ending at [endIndex], to [out]. This method
+ * Copies bytes from this buffer's subrange, starting at [startIndex] and ending at [endIndex], to [out]. This method
  * does not consume data from the buffer.
  *
  * @param out the destination to copy data into.
diff --git a/core/jvm/src/RawSink.kt b/core/jvm/src/RawSink.kt
index 59b15b51..326a2e28 100644
--- a/core/jvm/src/RawSink.kt
+++ b/core/jvm/src/RawSink.kt
@@ -21,6 +21,7 @@
 package kotlinx.io
 
 import java.io.Flushable
+
 public actual interface RawSink : AutoCloseable, Flushable {
     public actual fun write(source: Buffer, byteCount: Long)
 


Test patch:
diff --git a/core/common/test/samples/unsafe/unsafeSamples.kt b/core/common/test/samples/unsafe/unsafeSamples.kt
index 5f300755..59617c21 100644
--- a/core/common/test/samples/unsafe/unsafeSamples.kt
+++ b/core/common/test/samples/unsafe/unsafeSamples.kt
@@ -267,20 +267,15 @@ class UnsafeBufferOperationsSamples {
         fun Buffer.crc32(): UInt {
             var crc32 = 0xffffffffU
             // Iterate over segments starting from buffer's head
-            UnsafeBufferOperations.iterate(this) { ctx, head ->
-                var currentSegment = head
-                // If a current segment is null, it means we ran out of segments.
-                while (currentSegment != null) {
-                    // Get data from a segment
-                    ctx.withData(currentSegment) { data, startIndex, endIndex ->
-                        for (idx in startIndex..<endIndex) {
-                            // Update crc32
-                            val index = data[idx].xor(crc32.toByte()).toUByte()
-                            crc32 = crc32Table[index.toInt()].xor(crc32.shr(8))
-                        }
+            UnsafeBufferOperations.forEachSegment(this) { ctx, segment ->
+                var currentSegment = segment
+                // Get data from a segment
+                ctx.withData(currentSegment) { data, startIndex, endIndex ->
+                    for (idx in startIndex..<endIndex) {
+                        // Update crc32
+                        val index = data[idx].toUInt().xor(crc32).toUByte()
+                        crc32 = crc32Table[index.toInt()].xor(crc32.shr(8))
                     }
-                    // Advance to the next segment
-                    currentSegment = ctx.next(currentSegment)
                 }
             }
             return crc32.xor(0xffffffffU)
diff --git a/core/common/test/unsafe/UnsafeBufferOperationsForEachTest.kt b/core/common/test/unsafe/UnsafeBufferOperationsForEachTest.kt
new file mode 100644
index 00000000..68ba2b5b
--- /dev/null
+++ b/core/common/test/unsafe/UnsafeBufferOperationsForEachTest.kt
@@ -0,0 +1,73 @@
+/*
+ * Copyright 2010-2024 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE.txt file.
+ */
+@file:OptIn(UnsafeIoApi::class)
+
+package kotlinx.io.unsafe
+
+import kotlinx.io.Buffer
+import kotlinx.io.UnsafeIoApi
+import kotlinx.io.assertArrayEquals
+import kotlinx.io.writeString
+import kotlin.test.Test
+import kotlin.test.assertEquals
+import kotlin.test.assertFailsWith
+import kotlin.test.assertNotNull
+import kotlin.test.assertNull
+import kotlin.test.assertTrue
+import kotlin.test.fail
+
+class UnsafeBufferOperationsForEachTest {
+
+    @Test
+    fun emptyBuffer() {
+        UnsafeBufferOperations.forEachSegment(Buffer()) { _, head ->
+            fail()
+        }
+    }
+
+    @Test
+    fun singleSegment() {
+        var counter = 0
+        UnsafeBufferOperations.forEachSegment(Buffer().also { it.writeByte(1) }) { ctx, segment ->
+            ++counter
+            assertEquals(1, segment.size)
+        }
+        assertEquals(1, counter)
+    }
+
+    @Test
+    fun multipleSegments() {
+        val buffer = Buffer()
+
+        val expectedSegments = 10
+        for (i in 0 ..< expectedSegments) {
+            UnsafeBufferOperations.moveToTail(buffer, byteArrayOf(i.toByte()))
+        }
+
+        val storedBytes = ByteArray(expectedSegments)
+        var idx = 0
+        UnsafeBufferOperations.forEachSegment(buffer) { ctx, segment ->
+            assertTrue(idx < expectedSegments)
+            storedBytes[idx++] = ctx.getUnchecked(segment, 0)
+        }
+
+        assertArrayEquals(ByteArray(expectedSegments) { it.toByte() }, storedBytes)
+    }
+
+    @Test
+    fun acquireDataDuringIteration() {
+        val buffer = Buffer().also { it.writeString("hello buffer") }
+
+        val expectedSize = buffer.size
+
+        UnsafeBufferOperations.forEachSegment(buffer) { ctx, segment ->
+            ctx.withData(segment) { data, startIndex, endIndex ->
+                assertEquals("hello buffer", data.decodeToString(startIndex, endIndex))
+            }
+        }
+
+        assertEquals(expectedSize, buffer.size)
+    }
+}
diff --git a/core/common/test/unsafe/UnsafeBufferOperationsIterationTest.kt b/core/common/test/unsafe/UnsafeBufferOperationsIterationTest.kt
index 867a8539..16b07c12 100644
--- a/core/common/test/unsafe/UnsafeBufferOperationsIterationTest.kt
+++ b/core/common/test/unsafe/UnsafeBufferOperationsIterationTest.kt
@@ -2,6 +2,7 @@
  * Copyright 2010-2024 JetBrains s.r.o. and Kotlin Programming Language contributors.
  * Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE.txt file.
  */
+@file:OptIn(UnsafeIoApi::class)
 
 package kotlinx.io.unsafe
 
@@ -11,8 +12,6 @@ import kotlinx.io.assertArrayEquals
 import kotlinx.io.writeString
 import kotlin.test.*
 
-@OptIn(UnsafeIoApi::class)
-
 class UnsafeBufferOperationsIterationTest {
     @Test
     fun callsInPlaceContract() {
diff --git a/core/jvm/test/samples/unsafeAccessSamplesJvm.kt b/core/jvm/test/samples/unsafeAccessSamplesJvm.kt
index 8255df72..0a633c9f 100644
--- a/core/jvm/test/samples/unsafeAccessSamplesJvm.kt
+++ b/core/jvm/test/samples/unsafeAccessSamplesJvm.kt
@@ -108,17 +108,13 @@ class UnsafeReadWriteSamplesJvm {
         fun Buffer.digest(algorithm: String): ByteString {
             val md = MessageDigest.getInstance(algorithm)
             // iterate over all segment and update data
-            UnsafeBufferOperations.iterate(this) { ctx, head ->
-                var segment = head
+            UnsafeBufferOperations.forEachSegment(this) { ctx, segment ->
                 // when segment is null, we reached the end of a buffer
-                while (segment != null) {
-                    // access segment data without copying it
-                    ctx.withData(segment) { data, startIndex, endIndex ->
-                        md.update(data, startIndex, endIndex - startIndex)
-                    }
-                    // advance to the next segment
-                    segment = ctx.next(segment)
+                // access segment data without copying it
+                ctx.withData(segment) { data, startIndex, endIndex ->
+                    md.update(data, startIndex, endIndex - startIndex)
                 }
+                // advance to the next segment
             }
             return UnsafeByteStringOperations.wrapUnsafe(md.digest())
         }