lib/simd/src/sort/key.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2
3 pub const Direction = enum {
4 ascending,
5 descending,
6
7 pub fn reverse(self: Direction) Direction {
8 return switch (self) {
9 .ascending => .descending,
10 .descending => .ascending,
11 };
12 }
13 };
14
15 pub const Ascending = struct {};
16 pub const Descending = struct {};
17
18 pub const K32V32 = extern struct {
19 value: u32 align(8),
20 key: u32,
21
22 pub fn init(key_value: u32, value_value: u32) K32V32 {
23 return .{ .value = value_value, .key = key_value };
24 }
25 };
26
27 pub const K64V64 = extern struct {
28 value: u64 align(16),
29 key: u64,
30
31 pub fn init(key_value: u64, value_value: u64) K64V64 {
32 return .{ .value = value_value, .key = key_value };
33 }
34 };
35
36 pub fn supported(comptime Key: type) bool {
37 if (Key == K32V32 or Key == K64V64 or Key == u128) return true;
38 return switch (@typeInfo(Key)) {
39 .int => |info| info.bits == 16 or info.bits == 32 or info.bits == 64,
40 .float => |info| info.bits == 16 or info.bits == 32 or info.bits == 64,
41 else => false,
42 };
43 }
44
45 pub inline fn before(
46 comptime Key: type,
47 comptime direction: Direction,
48 a: Key,
49 b: Key,
50 ) bool {
51 requireSupported(Key);
52 return if (direction == .ascending)
53 lessKeyAscending(Key, a, b)
54 else
55 lessKeyAscending(Key, b, a);
56 }
57
58 pub inline fn networkBefore(
59 comptime Key: type,
60 comptime direction: Direction,
61 a: Key,
62 b: Key,
63 ) bool {
64 requireSupported(Key);
65 return if (direction == .ascending)
66 lessFullAscending(Key, a, b)
67 else
68 lessFullAscending(Key, b, a);
69 }
70
71 pub inline fn equivalent(comptime Key: type, a: Key, b: Key) bool {
72 requireSupported(Key);
73 if (Key == K32V32 or Key == K64V64) return a.key == b.key;
74 return a == b;
75 }
76
77 pub inline fn first(
78 comptime Key: type,
79 comptime direction: Direction,
80 a: Key,
81 b: Key,
82 ) Key {
83 return if (before(Key, direction, a, b)) a else b;
84 }
85
86 pub inline fn last(
87 comptime Key: type,
88 comptime direction: Direction,
89 a: Key,
90 b: Key,
91 ) Key {
92 return if (before(Key, direction, a, b)) b else a;
93 }
94
95 pub fn medianOf3(
96 comptime Key: type,
97 comptime direction: Direction,
98 a: Key,
99 b: Key,
100 c: Key,
101 ) Key {
102 var low = a;
103 var high = c;
104 if (networkBefore(Key, direction, high, low)) std.mem.swap(Key, &low, &high);
105 const middle = if (networkBefore(Key, direction, low, b)) b else low;
106 return if (networkBefore(Key, direction, middle, high)) middle else high;
107 }
108
109 pub fn firstValue(comptime Key: type, comptime direction: Direction) Key {
110 return if (direction == .ascending) smallestValue(Key) else largestValue(Key);
111 }
112
113 pub fn lastValue(comptime Key: type, comptime direction: Direction) Key {
114 return if (direction == .ascending) largestValue(Key) else smallestValue(Key);
115 }
116
117 pub fn previous(
118 comptime Key: type,
119 comptime direction: Direction,
120 value: Key,
121 ) Key {
122 requireSupported(Key);
123 if (Key == K32V32) {
124 return .{
125 .value = value.value,
126 .key = if (direction == .ascending) value.key -% 1 else value.key +% 1,
127 };
128 }
129 if (Key == K64V64) {
130 return .{
131 .value = value.value,
132 .key = if (direction == .ascending) value.key -% 1 else value.key +% 1,
133 };
134 }
135 return switch (@typeInfo(Key)) {
136 .int => if (direction == .ascending) value -% 1 else value +% 1,
137 .float => if (direction == .ascending)
138 smallerFloat(Key, value)
139 else
140 largerFloat(Key, value),
141 else => unreachable,
142 };
143 }
144
145 pub fn validForNetwork(comptime Key: type, value: Key) bool {
146 requireSupported(Key);
147 return switch (@typeInfo(Key)) {
148 .float => !std.math.isNan(value),
149 else => true,
150 };
151 }
152
153 fn lessKeyAscending(comptime Key: type, a: Key, b: Key) bool {
154 if (Key == K32V32 or Key == K64V64) return a.key < b.key;
155 return a < b;
156 }
157
158 fn lessFullAscending(comptime Key: type, a: Key, b: Key) bool {
159 if (Key == K32V32 or Key == K64V64) {
160 return a.key < b.key or (a.key == b.key and a.value < b.value);
161 }
162 return a < b;
163 }
164
165 fn smallestValue(comptime Key: type) Key {
166 requireSupported(Key);
167 if (Key == K32V32) return .{ .value = 0, .key = 0 };
168 if (Key == K64V64) return .{ .value = 0, .key = 0 };
169 return switch (@typeInfo(Key)) {
170 .int => std.math.minInt(Key),
171 .float => -std.math.inf(Key),
172 else => unreachable,
173 };
174 }
175
176 fn largestValue(comptime Key: type) Key {
177 requireSupported(Key);
178 if (Key == K32V32) {
179 return .{ .value = std.math.maxInt(u32), .key = std.math.maxInt(u32) };
180 }
181 if (Key == K64V64) {
182 return .{ .value = std.math.maxInt(u64), .key = std.math.maxInt(u64) };
183 }
184 return switch (@typeInfo(Key)) {
185 .int => std.math.maxInt(Key),
186 .float => std.math.inf(Key),
187 else => unreachable,
188 };
189 }
190
191 fn largerFloat(comptime Float: type, value: Float) Float {
192 const UInt = @Int(.unsigned, @bitSizeOf(Float));
193 const sign_bit = @as(UInt, 1) << (@bitSizeOf(Float) - 1);
194 const bits: UInt = @bitCast(value);
195 const magnitude = bits & ~sign_bit;
196 const was_positive = bits <= sign_bit;
197 const add: UInt = if (was_positive) 1 else std.math.maxInt(UInt);
198 var result: Float = @bitCast(magnitude +% add);
199 if (!std.math.isFinite(result)) {
200 result = if (was_positive) std.math.inf(Float) else std.math.floatMax(Float);
201 }
202 return if (was_positive) result else -result;
203 }
204
205 fn smallerFloat(comptime Float: type, value: Float) Float {
206 const UInt = @Int(.unsigned, @bitSizeOf(Float));
207 const sign_bit = @as(UInt, 1) << (@bitSizeOf(Float) - 1);
208 const bits: UInt = @bitCast(value);
209 const magnitude = bits & ~sign_bit;
210 const was_positive = value > 0;
211 const add: UInt = if (was_positive) std.math.maxInt(UInt) else 1;
212 var result: Float = @bitCast(magnitude +% add);
213 if (!std.math.isFinite(result)) {
214 result = if (was_positive) std.math.floatMax(Float) else std.math.inf(Float);
215 }
216 return if (was_positive) result else -result;
217 }
218
219 fn requireSupported(comptime Key: type) void {
220 if (comptime !supported(Key)) @compileError("VQSort key type is not supported");
221 }
222
223 test "Highway sort key records preserve native field layout" {
224 try std.testing.expectEqual(@as(usize, 8), @sizeOf(K32V32));
225 try std.testing.expectEqual(@as(usize, 8), @alignOf(K32V32));
226 try std.testing.expectEqual(@as(usize, 0), @offsetOf(K32V32, "value"));
227 try std.testing.expectEqual(@as(usize, 4), @offsetOf(K32V32, "key"));
228 try std.testing.expectEqual(@as(usize, 16), @sizeOf(K64V64));
229 try std.testing.expectEqual(@as(usize, 16), @alignOf(K64V64));
230 try std.testing.expectEqual(@as(usize, 0), @offsetOf(K64V64, "value"));
231 try std.testing.expectEqual(@as(usize, 8), @offsetOf(K64V64, "key"));
232 }
233
234 test "Highway sort key order separates equivalence from network ties" {
235 const low = K32V32.init(7, 1);
236 const high = K32V32.init(7, 9);
237 try std.testing.expect(equivalent(K32V32, low, high));
238 try std.testing.expect(!before(K32V32, .ascending, low, high));
239 try std.testing.expect(networkBefore(K32V32, .ascending, low, high));
240 try std.testing.expect(networkBefore(K32V32, .descending, high, low));
241 }
242
243 fn verifyFloatNeighbors(comptime Float: type) !void {
244 const UInt = @Int(.unsigned, @bitSizeOf(Float));
245 const positive_subnormal: Float = @bitCast(@as(UInt, 1));
246 const negative_subnormal: Float = @bitCast(
247 (@as(UInt, 1) << (@bitSizeOf(Float) - 1)) | 1,
248 );
249 try std.testing.expectEqual(positive_subnormal, previous(Float, .descending, 0));
250 try std.testing.expectEqual(negative_subnormal, previous(Float, .ascending, 0));
251 try std.testing.expectEqual(
252 std.math.inf(Float),
253 previous(Float, .descending, std.math.inf(Float)),
254 );
255 try std.testing.expectEqual(
256 -std.math.inf(Float),
257 previous(Float, .ascending, -std.math.inf(Float)),
258 );
259 try std.testing.expectEqual(
260 std.math.floatMax(Float),
261 previous(Float, .ascending, std.math.inf(Float)),
262 );
263 try std.testing.expectEqual(
264 -std.math.floatMax(Float),
265 previous(Float, .descending, -std.math.inf(Float)),
266 );
267 }
268
269 test "Highway floating previous values cover zeros finite limits and infinities" {
270 try verifyFloatNeighbors(f16);
271 try verifyFloatNeighbors(f32);
272 try verifyFloatNeighbors(f64);
273 }
274
275 test "Highway median of three follows key order" {
276 try std.testing.expectEqual(@as(u64, 4), medianOf3(u64, .ascending, 9, 1, 4));
277 try std.testing.expectEqual(@as(u64, 4), medianOf3(u64, .descending, 1, 9, 4));
278 for (0..8) |bits| {
279 const a: u64 = @intFromBool(bits & 1 != 0);
280 const b: u64 = @intFromBool(bits & 2 != 0);
281 const c: u64 = @intFromBool(bits & 4 != 0);
282 const expected: u64 = @intFromBool(@popCount(bits) >= 2);
283 try std.testing.expectEqual(expected, medianOf3(u64, .ascending, a, b, c));
284 }
285 const a = K64V64.init(7, 4);
286 const b = K64V64.init(7, 1);
287 const c = K64V64.init(7, 9);
288 try std.testing.expectEqual(a, medianOf3(K64V64, .ascending, a, b, c));
289 try std.testing.expectEqual(a, medianOf3(K64V64, .descending, a, b, c));
290 }