From 700e51047c4517ea5deda744a14acf89dba9aa3b Mon Sep 17 00:00:00 2001 From: fumoboy007 <2100868+fumoboy007@users.noreply.github.com> Date: Tue, 18 Jun 2024 17:45:53 -0700 Subject: [PATCH] Change `MessageWriter` to manually manage an unsafe mutable buffer pointer instead of using `Data`. `Data` is slow to append single bytes because it always calls `memmove` rather than just storing the byte. --- Sources/MessagePack/MessagePackEncoder.swift | 4 +- Sources/MessagePack/MessageWriter.swift | 94 +++++++++++++------ .../MessagePackTests/MessageWriterTests.swift | 27 ++++-- 3 files changed, 85 insertions(+), 40 deletions(-) diff --git a/Sources/MessagePack/MessagePackEncoder.swift b/Sources/MessagePack/MessagePackEncoder.swift index e3a6dee..30a6e99 100644 --- a/Sources/MessagePack/MessagePackEncoder.swift +++ b/Sources/MessagePack/MessagePackEncoder.swift @@ -1,6 +1,6 @@ // MIT License // -// Copyright © 2023 Darren Mo. +// Copyright © 2023–2024 Darren Mo. // // Permission is hereby granted, free of charge, to any person obtaining a copy // of this software and associated documentation files (the "Software"), to deal @@ -44,6 +44,6 @@ public struct MessagePackEncoder { var messageWriter = MessageWriter() try messagePackValue.encode(to: &messageWriter) - return messageWriter.message + return messageWriter.finish() } } diff --git a/Sources/MessagePack/MessageWriter.swift b/Sources/MessagePack/MessageWriter.swift index a67a73d..46ab913 100644 --- a/Sources/MessagePack/MessageWriter.swift +++ b/Sources/MessagePack/MessageWriter.swift @@ -1,6 +1,6 @@ // MIT License // -// Copyright © 2023 Darren Mo. +// Copyright © 2023–2024 Darren Mo. // // Permission is hereby granted, free of charge, to any person obtaining a copy // of this software and associated documentation files (the "Software"), to deal @@ -23,45 +23,83 @@ import Foundation public struct MessageWriter: ~Copyable { - // TODO: Manually manage an unsafe mutable buffer pointer instead of using `Data` - // after Swift 5.8 support is dropped. `Data` is slow to append single bytes - // because it always calls `memmove` rather than just storing the byte. - // - // We need to wait until Swift 5.8 support is dropped because we need the - // noncopyable struct feature. A class is another alternative but it is slow - // because it needs to perform runtime checks to enforce exclusivity. - private(set) var message = Data() - - // Hack: This is a workaround for an issue (https://github.com/fumoboy007/msgpack-swift/issues/4) - // related to Whole Module Optimization in Swift 5.9+. Although the issue only - // appeared on Linux so far, add the workaround for every platform to be safe. - // - // TODO: Remove this workaround after the root cause has been fixed. See - // https://github.com/apple/swift/issues/70979 for more details. -#if swift(>=5.9) - @inline(never) -#endif + static let initialCapacity = NSPageSize() + + // Manually manage an unsafe mutable buffer pointer instead of using `Data`. + // `Data` is slow to append single bytes because it always calls `memmove` + // rather than just storing the byte. + private var buffer: UnsafeMutableBufferPointer! = .allocate(capacity: initialCapacity) + private var totalByteCount = 0 + + deinit { + buffer?.deallocate() + } + + // MARK: - Writing Bytes + public mutating func write(byte: UInt8) { - withUnsafePointer(to: byte) { - message.append($0, count: 1) - } + let writeIndex = totalByteCount + + totalByteCount += 1 + increaseCapacityIfNeeded() + + buffer.initializeElement(at: writeIndex, to: byte) } public mutating func write(_ bytes: UnsafeRawBufferPointer) { + let writeStartIndex = totalByteCount + + totalByteCount += bytes.count + increaseCapacityIfNeeded() + bytes.withMemoryRebound(to: UInt8.self) { bytes in - guard let baseAddress = bytes.baseAddress else { - return - } - message.append(baseAddress, count: bytes.count) + let writeEndIndex = buffer[writeStartIndex.. Void) { - let byteCountBeforeWrites = message.count + let byteCountBeforeWrites = totalByteCount writeBytes(&self) - let writtenByteCount = message.count - byteCountBeforeWrites + let writtenByteCount = totalByteCount - byteCountBeforeWrites precondition(writtenByteCount == byteCount, "Expected \(byteCount) byte(s) to be written but found \(writtenByteCount).") } + + private mutating func increaseCapacityIfNeeded() { + var capacity = buffer.count + guard totalByteCount > capacity else { + return + } + + let pageSize = NSPageSize() + var newCapacityInPages = totalByteCount / pageSize + if totalByteCount > newCapacityInPages * pageSize { + newCapacityInPages += 1 + } + + capacity = newCapacityInPages * pageSize + precondition(totalByteCount <= capacity) + + let newBaseAddress = realloc(buffer.baseAddress, capacity)!.assumingMemoryBound(to: UInt8.self) + buffer = UnsafeMutableBufferPointer(start: newBaseAddress, + count: capacity) + } + + // MARK: - Getting the Message + + consuming func finish() -> Data { + guard let baseAddress = buffer.baseAddress else { + return Data() + } + + // Set to `nil` so that `deinit` does not prematurely deallocate the buffer. + // The buffer’s lifetime will be managed by the `Data` instance. + buffer = nil + + return Data(bytesNoCopy: baseAddress, + count: totalByteCount, + deallocator: .custom({ (baseAddress, _) in baseAddress.deallocate() })) + } } diff --git a/Tests/MessagePackTests/MessageWriterTests.swift b/Tests/MessagePackTests/MessageWriterTests.swift index 735f77d..8d4c0cb 100644 --- a/Tests/MessagePackTests/MessageWriterTests.swift +++ b/Tests/MessagePackTests/MessageWriterTests.swift @@ -28,7 +28,8 @@ class MessageWriterTests: XCTestCase { func testNoWrites() { let writer = MessageWriter() - XCTAssertEqual(writer.message, Data()) + let writtenBytes = writer.finish() + XCTAssertEqual(writtenBytes, Data()) } func testWriteByte() { @@ -37,14 +38,15 @@ class MessageWriterTests: XCTestCase { let byte = UInt8.random(in: .min...(.max)) Self.write(byte, to: &writer) - XCTAssertEqual(writer.message, Data([byte])) + let writtenBytes = writer.finish() + XCTAssertEqual(writtenBytes, Data([byte])) } func testWriteByte_multiple() { var writer = MessageWriter() var bytes = [UInt8]() - let byteCount = Int.random(in: 1...Int(UInt8.max)) + let byteCount = MessageWriter.initialCapacity + 1 for _ in 0..