Skip to content

Commit 2e22c89

Browse files
MahdiBMLukasa
andauthored
Add InlineArray helpers to ByteBuffer (#3252)
### Motivation: Useful for parsing packets, for example `ipv4: InlineArray<4, UInt8>` and `ipv6: InlineArray<16, UInt8>`. ### Modifications: For now I've only added a `readInlineArray` function. I know some other functions are missing, such as `writeInlineArray`. I wanted to first open up a discussion and see if these changes are acceptable. I can add those functions too if required, in this PR or other PRs. ### Result: Users can read `ByteBuffer` into stack-allocated memory, which can be more performant than the other alternatives like `Array`, or more convenient than reading as a tuple like `(UInt8, UInt8, UInt8, UInt8)`. ### Caveats: Swift 6.2 is required so I've used `#if compiler(>=6.2)`. Furthermore, `InlineArray` is marked as available on `macOS(9999)` since the Swift team have yet to update that mark, although `InlineArray` is planned for Swift 6.2 per the [proposal](https://github.com/swiftlang/swift-evolution/blob/main/proposals/0453-vector.md). Edit: to be clear things work fine on Linux and that's where I've been using this same `readInlineArray` function that I've proposed. --------- Co-authored-by: Cory Benfield <lukasa@apple.com>
1 parent a18bddb commit 2e22c89

File tree

2 files changed

+98
-0
lines changed

2 files changed

+98
-0
lines changed

Sources/NIOCore/ByteBuffer-aux.swift

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -58,6 +58,43 @@ extension ByteBuffer {
5858
return result
5959
}
6060

61+
#if compiler(>=6.2)
62+
@inlinable
63+
@available(macOS 26, iOS 26, tvOS 26, watchOS 26, visionOS 26, *)
64+
public mutating func readInlineArray<
65+
let count: Int,
66+
IntegerType: FixedWidthInteger
67+
>(
68+
endianness: Endianness = .big,
69+
as: InlineArray<count, IntegerType>.Type = InlineArray<count, IntegerType>.self
70+
) -> InlineArray<count, IntegerType>? {
71+
// use stride to account for padding bytes
72+
let stride = MemoryLayout<IntegerType>.stride
73+
let bytesRequired = stride * count
74+
75+
guard self.readableBytes >= bytesRequired else {
76+
return nil
77+
}
78+
79+
let inlineArray = InlineArray<count, IntegerType> { (outputSpan: inout OutputSpan<IntegerType>) in
80+
for index in 0..<count {
81+
// already made sure of 'self.readableBytes >= bytesRequired' above,
82+
// so this is safe to force-unwrap as it's guaranteed to exist
83+
let integer = self.getInteger(
84+
// this is less than 'bytesRequired' so is safe to multiply
85+
at: stride &* index,
86+
endianness: endianness,
87+
as: IntegerType.self
88+
)!
89+
outputSpan.append(integer)
90+
}
91+
}
92+
// already made sure of 'self.readableBytes >= bytesRequired' above
93+
self._moveReaderIndex(forwardBy: bytesRequired)
94+
return inlineArray
95+
}
96+
#endif
97+
6198
/// Returns the Bytes at the current reader index without advancing it.
6299
///
63100
/// This method is equivalent to calling `getBytes(at: readerIndex, ...)`

Tests/NIOCoreTests/ByteBufferTest.swift

Lines changed: 61 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4461,5 +4461,66 @@ extension ByteBufferTest {
44614461
let result = self.buf.readBytes(length: 4)
44624462
XCTAssertEqual(Array(0..<4), result!)
44634463
}
4464+
4465+
@available(macOS 26, iOS 26, tvOS 26, watchOS 26, visionOS 26, *)
4466+
func testReadInlineArrayOfUInt8() throws {
4467+
let bytes = (0..<10).map { _ in UInt8.random(in: .min ... .max) }
4468+
4469+
let startWriterIndex = self.buf.writerIndex
4470+
let written = self.buf.writeBytes(bytes)
4471+
XCTAssertEqual(startWriterIndex + written, self.buf.writerIndex)
4472+
XCTAssertEqual(written, self.buf.readableBytes)
4473+
4474+
let result = try XCTUnwrap(
4475+
self.buf.readInlineArray(as: InlineArray<10, UInt8>.self)
4476+
)
4477+
XCTAssertEqual(10, result.count)
4478+
for idx in result.indices {
4479+
XCTAssertEqual(bytes[idx], result[idx])
4480+
}
4481+
XCTAssertEqual(0, self.buf.readableBytes)
4482+
XCTAssertEqual(10, self.buf.readerIndex)
4483+
}
4484+
4485+
@available(macOS 26, iOS 26, tvOS 26, watchOS 26, visionOS 26, *)
4486+
func testReadInlineArrayOfUInt64() throws {
4487+
let bytes = (0..<15).map { _ in UInt64.random(in: .min ... .max) }
4488+
4489+
let startWriterIndex = self.buf.writerIndex
4490+
var written = 0
4491+
for byte in bytes {
4492+
written += self.buf.writeInteger(byte)
4493+
}
4494+
XCTAssertEqual(startWriterIndex + written, self.buf.writerIndex)
4495+
XCTAssertEqual(written, self.buf.readableBytes)
4496+
4497+
let result = try XCTUnwrap(
4498+
self.buf.readInlineArray(as: InlineArray<15, UInt64>.self)
4499+
)
4500+
XCTAssertEqual(15, result.count)
4501+
for idx in result.indices {
4502+
XCTAssertEqual(bytes[idx], result[idx])
4503+
}
4504+
XCTAssertEqual(0, self.buf.readableBytes)
4505+
XCTAssertEqual(120, self.buf.readerIndex)
4506+
}
4507+
4508+
@available(macOS 26, iOS 26, tvOS 26, watchOS 26, visionOS 26, *)
4509+
func testNotEnoughBytesToReadInlineArrayOfInt32() throws {
4510+
let startWriterIndex = self.buf.writerIndex
4511+
var written = 0
4512+
/// Write 15 bytes. This won't be enough to read an `InlineArray<5, Int32>`.
4513+
for _ in 0..<15 {
4514+
written += self.buf.writeInteger(UInt8.random(in: .min ... .max))
4515+
}
4516+
XCTAssertEqual(startWriterIndex + written, self.buf.writerIndex)
4517+
XCTAssertEqual(written, self.buf.readableBytes)
4518+
4519+
let result = self.buf.readInlineArray(as: InlineArray<5, Int32>.self)
4520+
4521+
XCTAssertNil(result)
4522+
XCTAssertEqual(written, self.buf.readableBytes)
4523+
XCTAssertEqual(0, self.buf.readerIndex)
4524+
}
44644525
}
44654526
#endif

0 commit comments

Comments
 (0)