Skip to content

Commit a028b77

Browse files
authored
Optimize the varint size calculations to be branchless. (#1906)
The logic for this comes from the C++ code: https://github.com/protocolbuffers/protobuf/blob/195da42b90bd39aa8c0b6551188fccf77a83f506/src/google/protobuf/io/coded_stream.h#L1754-L1826 I've always thought Swift didn't have a way to get to the optimal instructions, but I was wrong, it has all the building blocks needed. Looking at new code in godbolt, in `-O`, without or without the arithmetic overflow protection, the code is the same (tiny), but by turning off the arithmetic overflow checks, we also keep the debug code a bit smaller.
1 parent 01c4572 commit a028b77

File tree

1 file changed

+15
-44
lines changed

1 file changed

+15
-44
lines changed

Sources/SwiftProtobuf/Varint.swift

Lines changed: 15 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -19,74 +19,45 @@ package enum Varint {
1919
///
2020
/// - Parameter value: The number whose varint size should be calculated.
2121
/// - Returns: The size, in bytes, of the 32-bit varint.
22+
@usableFromInline
2223
package static func encodedSize(of value: UInt32) -> Int {
23-
if (value & (~0 << 7)) == 0 {
24-
return 1
25-
}
26-
if (value & (~0 << 14)) == 0 {
27-
return 2
28-
}
29-
if (value & (~0 << 21)) == 0 {
30-
return 3
31-
}
32-
if (value & (~0 << 28)) == 0 {
33-
return 4
34-
}
35-
return 5
24+
// This logic comes from the upstream C++ for CodedOutputStream::VarintSize32(uint32_t),
25+
// it provides a branchless calculation of the size.
26+
let clz = value.leadingZeroBitCount
27+
return ((UInt32.bitWidth &* 9 &+ 64) &- (clz &* 9)) / 64
3628
}
3729

3830
/// Computes the number of bytes that would be needed to store a signed 32-bit varint, if it were
3931
/// treated as an unsigned integer with the same bit pattern.
4032
///
4133
/// - Parameter value: The number whose varint size should be calculated.
4234
/// - Returns: The size, in bytes, of the 32-bit varint.
35+
@inline(__always)
4336
package static func encodedSize(of value: Int32) -> Int {
44-
if value >= 0 {
45-
return encodedSize(of: UInt32(bitPattern: value))
46-
} else {
47-
// Must sign-extend.
48-
return encodedSize(of: Int64(value))
49-
}
37+
// Must sign-extend.
38+
encodedSize(of: Int64(value))
5039
}
5140

5241
/// Computes the number of bytes that would be needed to store a 64-bit varint.
5342
///
5443
/// - Parameter value: The number whose varint size should be calculated.
5544
/// - Returns: The size, in bytes, of the 64-bit varint.
45+
@inline(__always)
5646
static func encodedSize(of value: Int64) -> Int {
57-
// Handle two common special cases up front.
58-
if (value & (~0 << 7)) == 0 {
59-
return 1
60-
}
61-
if value < 0 {
62-
return 10
63-
}
64-
65-
// Divide and conquer the remaining eight cases.
66-
var value = value
67-
var n = 2
68-
69-
if (value & (~0 << 35)) != 0 {
70-
n &+= 4
71-
value >>= 28
72-
}
73-
if (value & (~0 << 21)) != 0 {
74-
n &+= 2
75-
value >>= 14
76-
}
77-
if (value & (~0 << 14)) != 0 {
78-
n &+= 1
79-
}
80-
return n
47+
encodedSize(of: UInt64(bitPattern: value))
8148
}
8249

8350
/// Computes the number of bytes that would be needed to store an unsigned 64-bit varint, if it
8451
/// were treated as a signed integer with the same bit pattern.
8552
///
8653
/// - Parameter value: The number whose varint size should be calculated.
8754
/// - Returns: The size, in bytes, of the 64-bit varint.
55+
@usableFromInline
8856
static func encodedSize(of value: UInt64) -> Int {
89-
encodedSize(of: Int64(bitPattern: value))
57+
// This logic comes from the upstream C++ for CodedOutputStream::VarintSize64(uint64_t),
58+
// it provides a branchless calculation of the size.
59+
let clz = value.leadingZeroBitCount
60+
return ((UInt64.bitWidth &* 9 &+ 64) &- (clz &* 9)) / 64
9061
}
9162

9263
/// Counts the number of distinct varints in a packed byte buffer.

0 commit comments

Comments
 (0)