lib/simd/src/robust.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2
3 pub const Error = error{InsufficientScratch};
4 pub const max_array_scratch_bytes: usize = 16 * 1024;
5
6 pub fn CountingSortEntry(comptime T: type) type {
7 comptime requireInteger(T);
8 return struct {
9 value: T,
10 count: usize,
11 };
12 }
13
14 pub fn countingSort(
15 comptime T: type,
16 values: []T,
17 scratch: []CountingSortEntry(T),
18 ) Error!void {
19 comptime requireInteger(T);
20 const Entry = CountingSortEntry(T);
21 var unique_count: usize = 0;
22 for (values) |value| {
23 var unique_index: usize = 0;
24 while (unique_index < unique_count and scratch[unique_index].value != value) {
25 unique_index += 1;
26 }
27 if (unique_index == unique_count) {
28 if (unique_count == scratch.len) return error.InsufficientScratch;
29 scratch[unique_count] = .{ .value = value, .count = 1 };
30 unique_count += 1;
31 } else {
32 scratch[unique_index].count += 1;
33 }
34 }
35
36 std.mem.sort(Entry, scratch[0..unique_count], {}, EntryAscending(Entry).lessThan);
37
38 var output_index: usize = 0;
39 for (scratch[0..unique_count]) |entry| {
40 const end = output_index + entry.count;
41 @memset(values[output_index..end], entry.value);
42 output_index = end;
43 }
44 std.debug.assert(output_index == values.len);
45 }
46
47 pub fn countingSortArray(values: anytype) void {
48 const T = arrayElement(@TypeOf(values));
49 const count = comptime arrayLength(@TypeOf(values));
50 comptime requireArrayScratch(CountingSortEntry(T), count);
51 var scratch: [count]CountingSortEntry(T) = undefined;
52 countingSort(T, values[0..], &scratch) catch unreachable;
53 }
54
55 pub fn minRange(
56 comptime T: type,
57 sorted: []const T,
58 index_begin: usize,
59 half_count: usize,
60 ) usize {
61 comptime requireUnsignedInteger(T);
62 std.debug.assert(half_count != 0);
63 std.debug.assert(index_begin <= sorted.len);
64 std.debug.assert(half_count <= (sorted.len - index_begin) / 2);
65 var min_range: T = std.math.maxInt(T);
66 var min_index = index_begin;
67 for (index_begin..index_begin + half_count) |index| {
68 std.debug.assert(sorted[index] <= sorted[index + half_count]);
69 const range = sorted[index + half_count] - sorted[index];
70 if (range < min_range) {
71 min_range = range;
72 min_index = index;
73 }
74 }
75 return min_index;
76 }
77
78 pub fn modeOfSorted(comptime T: type, sorted: []const T) T {
79 comptime requireUnsignedInteger(T);
80 std.debug.assert(sorted.len != 0);
81 var index_begin: usize = 0;
82 var half_count = sorted.len / 2;
83 while (half_count > 1) {
84 index_begin = minRange(T, sorted, index_begin, half_count);
85 half_count >>= 1;
86 }
87 const lower = sorted[index_begin];
88 if (half_count == 0) return lower;
89 std.debug.assert(half_count == 1);
90 return roundedUnsignedAverage(lower, sorted[index_begin + 1]);
91 }
92
93 pub fn mode(
94 comptime T: type,
95 values: []T,
96 scratch: []CountingSortEntry(T),
97 ) Error!T {
98 comptime requireUnsignedInteger(T);
99 std.debug.assert(values.len != 0);
100 try countingSort(T, values, scratch);
101 return modeOfSorted(T, values);
102 }
103
104 pub fn modeArray(values: anytype) arrayElement(@TypeOf(values)) {
105 const T = arrayElement(@TypeOf(values));
106 const count = comptime arrayLength(@TypeOf(values));
107 comptime if (count == 0) @compileError("mode requires at least one value");
108 comptime requireArrayScratch(CountingSortEntry(T), count);
109 var scratch: [count]CountingSortEntry(T) = undefined;
110 return mode(T, values[0..], &scratch) catch unreachable;
111 }
112
113 pub fn median(comptime T: type, values: []T) T {
114 comptime requireOrderType(T);
115 std.debug.assert(values.len != 0);
116 std.mem.sort(T, values, {}, std.sort.asc(T));
117 const half = values.len / 2;
118 if (values.len % 2 != 0) return values[half];
119 return switch (@typeInfo(T)) {
120 .int => roundedIntegerMedian(T, values[half - 1], values[half]),
121 .float => (values[half - 1] + values[half]) / 2,
122 else => unreachable,
123 };
124 }
125
126 pub fn medianArray(values: anytype) arrayElement(@TypeOf(values)) {
127 const T = arrayElement(@TypeOf(values));
128 const count = comptime arrayLength(@TypeOf(values));
129 comptime if (count == 0) @compileError("median requires at least one value");
130 return median(T, values[0..]);
131 }
132
133 pub fn medianAbsoluteDeviation(
134 comptime T: type,
135 values: []const T,
136 median_value: T,
137 scratch: []T,
138 ) Error!T {
139 comptime requireOrderType(T);
140 std.debug.assert(values.len != 0);
141 if (scratch.len < values.len) return error.InsufficientScratch;
142 for (values, scratch[0..values.len]) |value, *deviation| {
143 deviation.* = absoluteDeviation(T, value, median_value);
144 }
145 return median(T, scratch[0..values.len]);
146 }
147
148 pub fn medianAbsoluteDeviationArray(
149 values: anytype,
150 median_value: arrayElement(@TypeOf(values)),
151 ) arrayElement(@TypeOf(values)) {
152 const T = arrayElement(@TypeOf(values));
153 const count = comptime arrayLength(@TypeOf(values));
154 comptime if (count == 0) @compileError("median absolute deviation requires at least one value");
155 comptime requireArrayScratch(T, count);
156 var scratch: [count]T = undefined;
157 return medianAbsoluteDeviation(T, values[0..], median_value, &scratch) catch unreachable;
158 }
159
160 fn roundedUnsignedAverage(lower: anytype, upper: @TypeOf(lower)) @TypeOf(lower) {
161 std.debug.assert(lower <= upper);
162 const difference = upper - lower;
163 return lower + difference / 2 + difference % 2;
164 }
165
166 fn EntryAscending(comptime Entry: type) type {
167 return struct {
168 fn lessThan(_: void, lhs: Entry, rhs: Entry) bool {
169 return lhs.value < rhs.value;
170 }
171 };
172 }
173
174 fn roundedIntegerMedian(comptime T: type, lower: T, upper: T) T {
175 std.debug.assert(lower <= upper);
176 const info = @typeInfo(T).int;
177 if (info.signedness == .unsigned) return roundedUnsignedAverage(lower, upper);
178 if (info.bits > 64) @compileError("signed integer medians support up to 64 bits");
179 const sum = @as(i128, lower) + @as(i128, upper) + 1;
180 return @intCast(@divTrunc(sum, 2));
181 }
182
183 fn absoluteDeviation(comptime T: type, value: T, median_value: T) T {
184 const value_i64 = asI64(T, value);
185 const median_i64 = asI64(T, median_value);
186 const difference = @as(i128, value_i64) - @as(i128, median_i64);
187 const magnitude: u128 = @intCast(if (difference < 0) -difference else difference);
188 return switch (@typeInfo(T)) {
189 .int => |info| blk: {
190 if (info.signedness == .signed) {
191 std.debug.assert(magnitude <= std.math.maxInt(T));
192 }
193 break :blk @intCast(magnitude);
194 },
195 .float => @floatFromInt(magnitude),
196 else => unreachable,
197 };
198 }
199
200 fn asI64(comptime T: type, value: T) i64 {
201 return switch (@typeInfo(T)) {
202 .int => |info| blk: {
203 if (info.bits > 64) @compileError("absolute deviations support up to 64-bit integers");
204 if (info.signedness == .unsigned) {
205 std.debug.assert(value <= std.math.maxInt(i64));
206 }
207 break :blk @intCast(value);
208 },
209 .float => blk: {
210 std.debug.assert(std.math.isFinite(value));
211 break :blk @intFromFloat(value);
212 },
213 else => unreachable,
214 };
215 }
216
217 fn arrayElement(comptime Pointer: type) type {
218 return switch (@typeInfo(Pointer)) {
219 .pointer => |pointer| switch (@typeInfo(pointer.child)) {
220 .array => |array| array.child,
221 else => @compileError("expected a pointer to an array"),
222 },
223 else => @compileError("expected a pointer to an array"),
224 };
225 }
226
227 fn arrayLength(comptime Pointer: type) usize {
228 return switch (@typeInfo(Pointer)) {
229 .pointer => |pointer| switch (@typeInfo(pointer.child)) {
230 .array => |array| array.len,
231 else => @compileError("expected a pointer to an array"),
232 },
233 else => @compileError("expected a pointer to an array"),
234 };
235 }
236
237 fn requireInteger(comptime T: type) void {
238 switch (@typeInfo(T)) {
239 .int => {},
240 else => @compileError("counting sort requires an integer type"),
241 }
242 }
243
244 fn requireUnsignedInteger(comptime T: type) void {
245 switch (@typeInfo(T)) {
246 .int => |info| if (info.signedness != .unsigned) {
247 @compileError("half-sample mode requires an unsigned integer type");
248 },
249 else => @compileError("half-sample mode requires an unsigned integer type"),
250 }
251 }
252
253 fn requireOrderType(comptime T: type) void {
254 switch (@typeInfo(T)) {
255 .int, .float => {},
256 else => @compileError("robust order statistics require an integer or float type"),
257 }
258 }
259
260 fn requireArrayScratch(comptime T: type, comptime count: usize) void {
261 if (@sizeOf(T) == 0) return;
262 if (count > max_array_scratch_bytes / @sizeOf(T)) {
263 @compileError("fixed-array robust scratch exceeds max_array_scratch_bytes");
264 }
265 }
266
267 test "Highway robust counting sort orders signed values and reports scratch exhaustion" {
268 var empty = [0]u8{};
269 countingSortArray(&empty);
270
271 var values = [_]i16{ 7, -2, 7, 0, -2, 4, 7 };
272 var scratch: [4]CountingSortEntry(i16) = undefined;
273 try countingSort(i16, &values, &scratch);
274 try std.testing.expectEqualSlices(i16, &.{ -2, -2, 0, 4, 7, 7, 7 }, &values);
275
276 var unique = [_]u8{ 3, 1, 2 };
277 const original = unique;
278 var short: [2]CountingSortEntry(u8) = undefined;
279 try std.testing.expectError(error.InsufficientScratch, countingSort(u8, &unique, &short));
280 try std.testing.expectEqualSlices(u8, &original, &unique);
281 }
282
283 test "Highway robust half-sample mode chooses first narrow ranges and rounded pairs" {
284 const sorted = [_]u64{ 0, 10, 11, 20, 21, 30 };
285 try std.testing.expectEqual(@as(usize, 1), minRange(u64, &sorted, 1, 2));
286 try std.testing.expectEqual(@as(u64, 7), modeOfSorted(u64, &.{7}));
287 try std.testing.expectEqual(@as(u64, 5), modeOfSorted(u64, &.{ 1, 9 }));
288 try std.testing.expectEqual(
289 std.math.maxInt(u64),
290 modeOfSorted(u64, &.{ std.math.maxInt(u64) - 1, std.math.maxInt(u64) }),
291 );
292
293 var values = [_]u64{ 100, 2, 3, 2, 1 };
294 try std.testing.expectEqual(@as(u64, 2), modeArray(&values));
295 try std.testing.expectEqualSlices(u64, &.{ 1, 2, 2, 3, 100 }, &values);
296 }
297
298 test "Highway robust medians preserve integer rounding and floating averages" {
299 var odd = [_]u32{ 9, 1, 4 };
300 try std.testing.expectEqual(@as(u32, 4), medianArray(&odd));
301 var even = [_]u32{ 3, 2 };
302 try std.testing.expectEqual(@as(u32, 3), medianArray(&even));
303 var negative = [_]i32{ -3, -1 };
304 try std.testing.expectEqual(@as(i32, -1), medianArray(&negative));
305 var floating = [_]f64{ 20, 1, 10, 2 };
306 try std.testing.expectEqual(@as(f64, 6), medianArray(&floating));
307 }
308
309 test "Highway robust median absolute deviation uses integer-cast deviations" {
310 const values = [_]u64{ 1, 1, 2, 2, 4, 6, 9 };
311 try std.testing.expectEqual(@as(u64, 1), medianAbsoluteDeviationArray(&values, 2));
312 const fractional = [_]f64{ 1.2, 2.8, 4.9 };
313 try std.testing.expectEqual(@as(f64, 1), medianAbsoluteDeviationArray(&fractional, 2.8));
314
315 const constant = [_]u32{ 17, 17, 17, 17, 17 };
316 var constant_mode = constant;
317 try std.testing.expectEqual(@as(u32, 17), modeArray(&constant_mode));
318 try std.testing.expectEqual(@as(u32, 0), medianAbsoluteDeviationArray(&constant, 17));
319
320 var short: [2]u64 = undefined;
321 try std.testing.expectError(
322 error.InsufficientScratch,
323 medianAbsoluteDeviation(u64, &values, 2, &short),
324 );
325 }