lib/simd/src/sort/partition.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const key = @import("key.zig");
3
4 pub const Bounds = struct {
5 before: usize,
6 after: usize,
7 };
8
9 pub fn around(
10 comptime Key: type,
11 comptime direction: key.Direction,
12 values: []Key,
13 pivot: Key,
14 ) Bounds {
15 if (comptime !key.supported(Key)) @compileError("VQSort key type is not supported");
16 var before: usize = 0;
17 var scan: usize = 0;
18 var after = values.len;
19 var iterations: usize = 0;
20 while (scan < after) : (iterations += 1) {
21 std.debug.assert(iterations < values.len);
22 if (key.before(Key, direction, values[scan], pivot)) {
23 std.mem.swap(Key, &values[before], &values[scan]);
24 before += 1;
25 scan += 1;
26 } else if (key.before(Key, direction, pivot, values[scan])) {
27 after -= 1;
28 std.mem.swap(Key, &values[scan], &values[after]);
29 } else {
30 scan += 1;
31 }
32 }
33 std.debug.assert(before <= after);
34 std.debug.assert(after <= values.len);
35 return .{ .before = before, .after = after };
36 }
37
38 pub fn nanToBack(comptime Key: type, values: []Key) usize {
39 if (comptime !key.supported(Key)) @compileError("VQSort key type is not supported");
40 if (comptime floating(Key)) {
41 var write: usize = 0;
42 for (values) |value| {
43 if (!std.math.isNan(value)) {
44 values[write] = value;
45 write += 1;
46 }
47 }
48 @memset(values[write..], canonicalNaN(Key));
49 return values.len - write;
50 }
51 return 0;
52 }
53
54 pub fn canonicalNaN(comptime Float: type) Float {
55 if (comptime !floating(Float)) @compileError("canonical NaN requires a floating type");
56 const UInt = @Int(.unsigned, @bitSizeOf(Float));
57 const bits: UInt = std.math.maxInt(UInt) >> 1;
58 return @bitCast(bits);
59 }
60
61 fn floating(comptime Key: type) bool {
62 return switch (@typeInfo(Key)) {
63 .float => true,
64 else => false,
65 };
66 }
67
68 fn Order(comptime Key: type, comptime direction: key.Direction) type {
69 return struct {
70 fn lessThan(_: void, a: Key, b: Key) bool {
71 return key.networkBefore(Key, direction, a, b);
72 }
73 };
74 }
75
76 test "Highway VQSort partition preserves red zones counts and boundaries" {
77 var storage = [_]u32{
78 0xa55a_5aa5,
79 7,
80 2,
81 9,
82 4,
83 4,
84 1,
85 8,
86 4,
87 3,
88 0x5aa5_a55a,
89 };
90 const original = storage;
91 const bounds = around(u32, .ascending, storage[1 .. storage.len - 1], 4);
92 try std.testing.expectEqual(@as(usize, 3), bounds.before);
93 try std.testing.expectEqual(@as(usize, 6), bounds.after);
94 for (storage[1 .. bounds.before + 1]) |value| try std.testing.expect(value < 4);
95 for (storage[bounds.before + 1 .. bounds.after + 1]) |value| {
96 try std.testing.expectEqual(@as(u32, 4), value);
97 }
98 for (storage[bounds.after + 1 .. storage.len - 1]) |value| try std.testing.expect(value > 4);
99 try std.testing.expectEqual(original[0], storage[0]);
100 try std.testing.expectEqual(original[original.len - 1], storage[storage.len - 1]);
101 var expected = original[1 .. original.len - 1].*;
102 var actual = storage[1 .. storage.len - 1].*;
103 std.mem.sort(u32, &expected, {}, Order(u32, .ascending).lessThan);
104 std.mem.sort(u32, &actual, {}, Order(u32, .ascending).lessThan);
105 try std.testing.expectEqualSlices(u32, &expected, &actual);
106 }
107
108 test "Highway VQSort partition treats record values as equivalent keys" {
109 var values = [_]key.K32V32{
110 key.K32V32.init(7, 9),
111 key.K32V32.init(3, 8),
112 key.K32V32.init(7, 1),
113 key.K32V32.init(11, 4),
114 key.K32V32.init(7, 5),
115 };
116 const bounds = around(key.K32V32, .ascending, &values, key.K32V32.init(7, 0));
117 try std.testing.expectEqual(@as(usize, 1), bounds.before);
118 try std.testing.expectEqual(@as(usize, 4), bounds.after);
119 for (values[bounds.before..bounds.after]) |value| try std.testing.expectEqual(@as(u32, 7), value.key);
120 }
121
122 test "Highway VQSort NaN compaction retains infinities and canonicalizes the tail" {
123 const payload_nan: f32 = @bitCast(@as(u32, 0x7fc0_1234));
124 var values = [_]f32{ 3, payload_nan, std.math.inf(f32), -4, std.math.nan(f32), -std.math.inf(f32) };
125 try std.testing.expectEqual(@as(usize, 2), nanToBack(f32, &values));
126 try std.testing.expectEqualSlices(f32, &[_]f32{ 3, std.math.inf(f32), -4, -std.math.inf(f32) }, values[0..4]);
127 const canonical_bits: u32 = @bitCast(canonicalNaN(f32));
128 for (values[4..]) |value| try std.testing.expectEqual(canonical_bits, @as(u32, @bitCast(value)));
129
130 var integers = [_]u64{ 3, 1, 2 };
131 const expected = integers;
132 try std.testing.expectEqual(@as(usize, 0), nanToBack(u64, &integers));
133 try std.testing.expectEqualSlices(u64, &expected, &integers);
134 }