Skip to content

Commit 8a86a68

Browse files
authored
[Unsafe API 5/5] Hide Segment::data (#339)
* Hide Segment::data * Improve throwsEof implementation
1 parent e066a6f commit 8a86a68

File tree

4 files changed

+182
-160
lines changed

4 files changed

+182
-160
lines changed

core/common/src/Buffer.kt

Lines changed: 83 additions & 152 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* Copyright 2017-2023 JetBrains s.r.o. and respective authors and developers.
2+
* Copyright 2017-2024 JetBrains s.r.o. and respective authors and developers.
33
* Use of this source code is governed by the Apache 2.0 license that can be found in the LICENCE file.
44
*/
55

@@ -23,6 +23,7 @@ package kotlinx.io
2323
import kotlin.contracts.ExperimentalContracts
2424
import kotlin.contracts.InvocationKind.EXACTLY_ONCE
2525
import kotlin.contracts.contract
26+
import kotlinx.io.unsafe.UnsafeBufferOperations
2627
import kotlin.jvm.JvmSynthetic
2728

2829
/**
@@ -83,119 +84,84 @@ public class Buffer : Source, Sink {
8384
}
8485

8586
override fun readByte(): Byte {
86-
require(1)
87-
val segment = head!!
88-
var pos = segment.pos
89-
val limit = segment.limit
90-
val data = segment.data
91-
val b = data[pos++]
87+
val segment = head ?: throwEof(1)
88+
val segmentSize = segment.size
89+
if (segmentSize == 0) {
90+
recycleHead()
91+
return readByte()
92+
}
93+
val v = segment.readByte()
9294
sizeMut -= 1L
93-
if (pos == limit) {
95+
if (segmentSize == 1) {
9496
recycleHead()
95-
} else {
96-
segment.pos = pos
9797
}
98-
return b
98+
return v
9999
}
100100

101101
override fun readShort(): Short {
102-
require(2)
103-
104-
val segment = head!!
105-
var pos = segment.pos
106-
val limit = segment.limit
107-
108-
// If the short is split across multiple segments, delegate to readByte().
109-
if (limit - pos < 2) {
110-
val s = readByte() and 0xff shl 8 or (readByte() and 0xff)
111-
return s.toShort()
102+
val segment = head ?: throwEof(2)
103+
val segmentSize = segment.size
104+
if (segmentSize < 2) {
105+
// If the short is split across multiple segments, delegate to readByte().
106+
require(2)
107+
if (segmentSize == 0) {
108+
recycleHead()
109+
return readShort()
110+
}
111+
return (readByte() and 0xff shl 8 or (readByte() and 0xff)).toShort()
112112
}
113-
114-
val data = segment.data
115-
val s = data[pos++] and 0xff shl 8 or (data[pos++] and 0xff)
113+
val v = segment.readShort()
116114
sizeMut -= 2L
117-
118-
if (pos == limit) {
115+
if (segmentSize == 2) {
119116
recycleHead()
120-
} else {
121-
segment.pos = pos
122117
}
123-
124-
return s.toShort()
118+
return v
125119
}
126120

127121
override fun readInt(): Int {
128-
require(4)
129-
130-
val segment = head!!
131-
var pos = segment.pos
132-
val limit = segment.limit
133-
134-
// If the int is split across multiple segments, delegate to readByte().
135-
if (limit - pos < 4L) {
136-
return (
137-
readByte() and 0xff shl 24
138-
or (readByte() and 0xff shl 16)
139-
or (readByte() and 0xff shl 8)
140-
or (readByte() and 0xff)
141-
)
142-
}
143-
144-
val data = segment.data
145-
val i = (
146-
data[pos++] and 0xff shl 24
147-
or (data[pos++] and 0xff shl 16)
148-
or (data[pos++] and 0xff shl 8)
149-
or (data[pos++] and 0xff)
150-
)
122+
val segment = head ?: throwEof(4)
123+
val segmentSize = segment.size
124+
if (segmentSize < 4) {
125+
// If the short is split across multiple segments, delegate to readShort().
126+
require(4)
127+
if (segmentSize == 0) {
128+
recycleHead()
129+
return readInt()
130+
}
131+
return (readShort().toInt() shl 16 or (readShort().toInt() and 0xffff))
132+
}
133+
val v = segment.readInt()
151134
sizeMut -= 4L
152-
153-
if (pos == limit) {
135+
if (segmentSize == 4) {
154136
recycleHead()
155-
} else {
156-
segment.pos = pos
157137
}
158-
159-
return i
138+
return v
160139
}
161140

162141
override fun readLong(): Long {
163-
require(8)
164-
165-
val segment = head!!
166-
var pos = segment.pos
167-
val limit = segment.limit
168-
169-
// If the long is split across multiple segments, delegate to readInt().
170-
if (limit - pos < 8L) {
171-
return (
172-
readInt() and 0xffffffffL shl 32
173-
or (readInt() and 0xffffffffL)
174-
)
175-
}
176-
177-
val data = segment.data
178-
val v = (
179-
data[pos++] and 0xffL shl 56
180-
or (data[pos++] and 0xffL shl 48)
181-
or (data[pos++] and 0xffL shl 40)
182-
or (data[pos++] and 0xffL shl 32)
183-
or (data[pos++] and 0xffL shl 24)
184-
or (data[pos++] and 0xffL shl 16)
185-
or (data[pos++] and 0xffL shl 8)
186-
or (data[pos++] and 0xffL)
187-
)
188-
this.sizeMut -= 8L
189-
190-
if (pos == limit) {
142+
val segment = head ?: throwEof(8)
143+
val segmentSize = segment.size
144+
if (segmentSize < 8) {
145+
// If the short is split across multiple segments, delegate to readInt().
146+
require(8)
147+
if (segmentSize == 0) {
148+
recycleHead()
149+
return readLong()
150+
}
151+
return (readInt().toLong() shl 32 or (readInt().toLong() and 0xffffffffL))
152+
}
153+
val v = segment.readLong()
154+
sizeMut -= 8L
155+
if (segmentSize == 8) {
191156
recycleHead()
192-
} else {
193-
segment.pos = pos
194157
}
195-
196158
return v
197159
}
198160

161+
private fun throwEof(byteCount: Long): Nothing {
162+
throw EOFException("Buffer doesn't contain required number of bytes (size: $size, required: $byteCount)")
163+
}
164+
199165
/**
200166
* This method does not affect the buffer's content as there is no upstream to write data to.
201167
*/
@@ -295,7 +261,7 @@ public class Buffer : Source, Sink {
295261
return head!!.getUnchecked(0)
296262
}
297263
seek(position) { s, offset ->
298-
return s!!.data[(s.pos + position - offset).toInt()]
264+
return s!!.getUnchecked((position - offset).toInt())
299265
}
300266
}
301267

@@ -333,16 +299,12 @@ public class Buffer : Source, Sink {
333299
override fun readAtMostTo(sink: ByteArray, startIndex: Int, endIndex: Int): Int {
334300
checkBounds(sink.size, startIndex, endIndex)
335301

336-
val s = head ?: return -1
337-
val toCopy = minOf(endIndex - startIndex, s.limit - s.pos)
338-
s.data.copyInto(
339-
destination = sink, destinationOffset = startIndex, startIndex = s.pos, endIndex = s.pos + toCopy
340-
)
341-
342-
s.pos += toCopy
302+
val s = this.head ?: return -1
303+
val toCopy = minOf(endIndex - startIndex, s.size)
304+
s.readTo(sink, startIndex, startIndex + toCopy)
343305
sizeMut -= toCopy.toLong()
344306

345-
if (s.pos == s.limit) {
307+
if (s.isEmpty()) {
346308
recycleHead()
347309
}
348310

@@ -406,19 +368,11 @@ public class Buffer : Source, Sink {
406368
var currentOffset = startIndex
407369
while (currentOffset < endIndex) {
408370
val tail = writableSegment(1)
409-
410-
val toCopy = minOf(endIndex - currentOffset, Segment.SIZE - tail.limit)
411-
source.copyInto(
412-
destination = tail.data,
413-
destinationOffset = tail.limit,
414-
startIndex = currentOffset,
415-
endIndex = currentOffset + toCopy
416-
)
417-
371+
val toCopy = minOf(endIndex - currentOffset, tail.remainingCapacity)
372+
tail.write(source, currentOffset, currentOffset + toCopy)
418373
currentOffset += toCopy
419-
tail.limit += toCopy
420374
}
421-
sizeMut += endIndex - startIndex
375+
sizeMut += endIndex - startIndex
422376
}
423377

424378
override fun write(source: RawSource, byteCount: Long) {
@@ -536,46 +490,22 @@ public class Buffer : Source, Sink {
536490
}
537491

538492
override fun writeByte(byte: Byte) {
539-
val tail = writableSegment(1)
540-
tail.data[tail.limit++] = byte
493+
writableSegment(1).writeByte(byte)
541494
sizeMut += 1L
542495
}
543496

544497
override fun writeShort(short: Short) {
545-
val tail = writableSegment(2)
546-
val data = tail.data
547-
var limit = tail.limit
548-
data[limit++] = (short.toInt() ushr 8 and 0xff).toByte()
549-
data[limit++] = (short.toInt() and 0xff).toByte()
550-
tail.limit = limit
498+
writableSegment(2).writeShort(short)
551499
sizeMut += 2L
552500
}
553501

554502
override fun writeInt(int: Int) {
555-
val tail = writableSegment(4)
556-
val data = tail.data
557-
var limit = tail.limit
558-
data[limit++] = (int ushr 24 and 0xff).toByte()
559-
data[limit++] = (int ushr 16 and 0xff).toByte()
560-
data[limit++] = (int ushr 8 and 0xff).toByte()
561-
data[limit++] = (int and 0xff).toByte()
562-
tail.limit = limit
503+
writableSegment(4).writeInt(int)
563504
sizeMut += 4L
564505
}
565506

566507
override fun writeLong(long: Long) {
567-
val tail = writableSegment(8)
568-
val data = tail.data
569-
var limit = tail.limit
570-
data[limit++] = (long ushr 56 and 0xffL).toByte()
571-
data[limit++] = (long ushr 48 and 0xffL).toByte()
572-
data[limit++] = (long ushr 40 and 0xffL).toByte()
573-
data[limit++] = (long ushr 32 and 0xffL).toByte()
574-
data[limit++] = (long ushr 24 and 0xffL).toByte()
575-
data[limit++] = (long ushr 16 and 0xffL).toByte()
576-
data[limit++] = (long ushr 8 and 0xffL).toByte()
577-
data[limit++] = (long and 0xffL).toByte()
578-
tail.limit = limit
508+
writableSegment(8).writeLong(long)
579509
sizeMut += 8L
580510
}
581511

@@ -615,6 +545,7 @@ public class Buffer : Source, Sink {
615545
*
616546
* @sample kotlinx.io.samples.KotlinxIoCoreCommonSamples.bufferToString
617547
*/
548+
@OptIn(UnsafeIoApi::class)
618549
override fun toString(): String {
619550
if (size == 0L) return "Buffer(size=0)"
620551

@@ -623,20 +554,20 @@ public class Buffer : Source, Sink {
623554

624555
val builder = StringBuilder(len * 2 + if (size > maxPrintableBytes) 1 else 0)
625556

626-
var curr = head!!
627-
var bytesWritten = 0
628-
var pos = curr.pos
629-
while (bytesWritten < len) {
630-
if (pos == curr.limit) {
631-
curr = curr.next!!
632-
pos = curr.pos
633-
}
634-
635-
val b = curr.data[pos++].toInt()
636-
bytesWritten++
637-
638-
builder.append(HEX_DIGIT_CHARS[(b shr 4) and 0xf])
639-
.append(HEX_DIGIT_CHARS[b and 0xf])
557+
UnsafeBufferOperations.iterate(this) { ctx, head ->
558+
var bytesWritten = 0
559+
var seg: Segment? = head
560+
do {
561+
seg!!
562+
var idx = 0
563+
while (bytesWritten < len && idx < seg.size) {
564+
val b = ctx.getUnchecked(seg, idx++)
565+
bytesWritten++
566+
builder.append(HEX_DIGIT_CHARS[(b shr 4) and 0xf])
567+
.append(HEX_DIGIT_CHARS[b and 0xf])
568+
}
569+
seg = ctx.next(seg)
570+
} while (seg != null)
640571
}
641572

642573
if (size > maxPrintableBytes) {

0 commit comments

Comments
 (0)