authorgravatar for andrew@ziglang.orgAndrew Kelley <andrew@ziglang.org> 2022-03-09 22:50:27-07:00
committergravatar for andrew@ziglang.orgAndrew Kelley <andrew@ziglang.org> 2022-03-10 13:13:17-05:00
log6e49ba77f3001fe04d4b2177a5c37d28bd422758
treecfd55e3566fbf651e339529ff03693c75484907b
parentf736cde397a6abb1399827ed5988c43001706580

std: add sort method to ArrayHashMap and MultiArrayList

This also adds `std.sort.sortContext` and `std.sort.insertionSortContext` which are more advanced methods that allow overriding the `swap` method. The former calls the latter for now because reworking the main sort implementation is a big task that can be done later without any changes to the API.

3 files changed, 124 insertions(+), 9 deletions(-)

lib/std/array_hash_map.zig+57
......@@ -408,6 +408,13 @@ pub fn ArrayHashMap(
408408 return self.unmanaged.reIndexContext(self.allocator, self.ctx);
409409 }
410410
411 /// Sorts the entries and then rebuilds the index.
412 /// `sort_ctx` must have this method:
413 /// `fn lessThan(ctx: @TypeOf(ctx), a_index: usize, b_index: usize) bool`
414 pub fn sort(self: *Self, sort_ctx: anytype) void {
415 return self.unmanaged.sortContext(sort_ctx, self.ctx);
416 }
417
411418 /// Shrinks the underlying `Entry` array to `new_len` elements and discards any associated
412419 /// index entries. Keeps capacity the same.
413420 pub fn shrinkRetainingCapacity(self: *Self, new_len: usize) void {
......@@ -1169,6 +1176,22 @@ pub fn ArrayHashMapUnmanaged(
11691176 self.index_header = new_header;
11701177 }
11711178
1179 /// Sorts the entries and then rebuilds the index.
1180 /// `sort_ctx` must have this method:
1181 /// `fn lessThan(ctx: @TypeOf(ctx), a_index: usize, b_index: usize) bool`
1182 pub inline fn sort(self: *Self, sort_ctx: anytype) void {
1183 if (@sizeOf(ByIndexContext) != 0)
1184 @compileError("Cannot infer context " ++ @typeName(Context) ++ ", call sortContext instead.");
1185 return self.sortContext(sort_ctx, undefined);
1186 }
1187
1188 pub fn sortContext(self: *Self, sort_ctx: anytype, ctx: Context) void {
1189 self.entries.sort(sort_ctx);
1190 const header = self.index_header orelse return;
1191 header.reset();
1192 self.insertAllEntriesIntoNewHeader(if (store_hash) {} else ctx, header);
1193 }
1194
11721195 /// Shrinks the underlying `Entry` array to `new_len` elements and discards any associated
11731196 /// index entries. Keeps capacity the same.
11741197 pub fn shrinkRetainingCapacity(self: *Self, new_len: usize) void {
......@@ -1868,6 +1891,14 @@ const IndexHeader = struct {
18681891 allocator.free(slice);
18691892 }
18701893
1894 /// Puts an IndexHeader into the state that it would be in after being freshly allocated.
1895 fn reset(header: *IndexHeader) void {
1896 const index_size = hash_map.capacityIndexSize(header.bit_index);
1897 const ptr = @ptrCast([*]align(@alignOf(IndexHeader)) u8, header);
1898 const nbytes = @sizeOf(IndexHeader) + header.length() * index_size;
1899 @memset(ptr + @sizeOf(IndexHeader), 0xff, nbytes - @sizeOf(IndexHeader));
1900 }
1901
18711902 // Verify that the header has sufficient alignment to produce aligned arrays.
18721903 comptime {
18731904 if (@alignOf(u32) > @alignOf(IndexHeader))
......@@ -2218,6 +2249,32 @@ test "auto store_hash" {
22182249 try testing.expect(meta.fieldInfo(HasExpensiveEqlUn.Data, .hash).field_type != void);
22192250}
22202251
2252test "sort" {
2253 var map = AutoArrayHashMap(i32, i32).init(std.testing.allocator);
2254 defer map.deinit();
2255
2256 for ([_]i32{ 8, 3, 12, 10, 2, 4, 9, 5, 6, 13, 14, 15, 16, 1, 11, 17, 7 }) |x| {
2257 try map.put(x, x * 3);
2258 }
2259
2260 const C = struct {
2261 keys: []i32,
2262
2263 pub fn lessThan(ctx: @This(), a_index: usize, b_index: usize) bool {
2264 return ctx.keys[a_index] < ctx.keys[b_index];
2265 }
2266 };
2267
2268 map.sort(C{ .keys = map.keys() });
2269
2270 var x: i32 = 1;
2271 for (map.keys()) |key, i| {
2272 try testing.expect(key == x);
2273 try testing.expect(map.values()[i] == x * 3);
2274 x += 1;
2275 }
2276}
2277
22212278pub fn getHashPtrAddrFn(comptime K: type, comptime Context: type) (fn (Context, K) u32) {
22222279 return struct {
22232280 fn hash(ctx: Context, key: K) u32 {
lib/std/multi_array_list.zig+28
......@@ -392,6 +392,34 @@ pub fn MultiArrayList(comptime S: type) type {
392392 return result;
393393 }
394394
395 /// `ctx` has the following method:
396 /// `fn lessThan(ctx: @TypeOf(ctx), a_index: usize, b_index: usize) bool`
397 pub fn sort(self: Self, ctx: anytype) void {
398 const SortContext = struct {
399 sub_ctx: @TypeOf(ctx),
400 slice: Slice,
401
402 pub fn swap(sc: @This(), a_index: usize, b_index: usize) void {
403 inline for (fields) |field_info, i| {
404 if (@sizeOf(field_info.field_type) != 0) {
405 const field = @intToEnum(Field, i);
406 const ptr = sc.slice.items(field);
407 mem.swap(field_info.field_type, &ptr[a_index], &ptr[b_index]);
408 }
409 }
410 }
411
412 pub fn lessThan(sc: @This(), a_index: usize, b_index: usize) bool {
413 return sc.sub_ctx.lessThan(a_index, b_index);
414 }
415 };
416
417 std.sort.sortContext(self.len, SortContext{
418 .sub_ctx = ctx,
419 .slice = self.slice(),
420 });
421 }
422
395423 fn capacityInBytes(capacity: usize) usize {
396424 const sizes_vector: std.meta.Vector(sizes.bytes.len, usize) = sizes.bytes;
397425 const capacity_vector = @splat(sizes.bytes.len, capacity);
lib/std/sort.zig+39-9
......@@ -73,7 +73,10 @@ test "binarySearch" {
7373 );
7474}
7575
76/// Stable in-place sort. O(n) best case, O(pow(n, 2)) worst case. O(1) memory (no allocator required).
76/// Stable in-place sort. O(n) best case, O(pow(n, 2)) worst case.
77/// O(1) memory (no allocator required).
78/// This can be expressed in terms of `insertionSortContext` but the glue
79/// code is slightly longer than the direct implementation.
7780pub fn insertionSort(
7881 comptime T: type,
7982 items: []T,
......@@ -91,6 +94,18 @@ pub fn insertionSort(
9194 }
9295}
9396
97/// Stable in-place sort. O(n) best case, O(pow(n, 2)) worst case.
98/// O(1) memory (no allocator required).
99pub fn insertionSortContext(len: usize, context: anytype) void {
100 var i: usize = 1;
101 while (i < len) : (i += 1) {
102 var j: usize = i;
103 while (j > 0 and context.lessThan(j, j - 1)) : (j -= 1) {
104 context.swap(j, j - 1);
105 }
106 }
107}
108
94109const Range = struct {
95110 start: usize,
96111 end: usize,
......@@ -178,7 +193,8 @@ const Pull = struct {
178193 range: Range,
179194};
180195
181/// Stable in-place sort. O(n) best case, O(n*log(n)) worst case and average case. O(1) memory (no allocator required).
196/// Stable in-place sort. O(n) best case, O(n*log(n)) worst case and average case.
197/// O(1) memory (no allocator required).
182198/// Currently implemented as block sort.
183199pub fn sort(
184200 comptime T: type,
......@@ -186,6 +202,7 @@ pub fn sort(
186202 context: anytype,
187203 comptime lessThan: fn (context: @TypeOf(context), lhs: T, rhs: T) bool,
188204) void {
205
189206 // Implementation ported from https://github.com/BonzaiThePenguin/WikiSort/blob/master/WikiSort.c
190207 var cache: [512]T = undefined;
191208
......@@ -291,10 +308,13 @@ pub fn sort(
291308
292309 // then merge sort the higher levels, which can be 8-15, 16-31, 32-63, 64-127, etc.
293310 while (true) {
294 // if every A and B block will fit into the cache, use a special branch specifically for merging with the cache
295 // (we use < rather than <= since the block size might be one more than iterator.length())
311 // if every A and B block will fit into the cache, use a special branch
312 // specifically for merging with the cache
313 // (we use < rather than <= since the block size might be one more than
314 // iterator.length())
296315 if (iterator.length() < cache.len) {
297 // if four subarrays fit into the cache, it's faster to merge both pairs of subarrays into the cache,
316 // if four subarrays fit into the cache, it's faster to merge both
317 // pairs of subarrays into the cache,
298318 // then merge the two merged subarrays from the cache back into the original array
299319 if ((iterator.length() + 1) * 4 <= cache.len and iterator.length() * 4 <= items.len) {
300320 iterator.begin();
......@@ -767,11 +787,15 @@ pub fn sort(
767787 }
768788 }
769789
770 // when we're finished with this merge step we should have the one or two internal buffers left over, where the second buffer is all jumbled up
771 // insertion sort the second buffer, then redistribute the buffers back into the items using the opposite process used for creating the buffer
790 // when we're finished with this merge step we should have the one
791 // or two internal buffers left over, where the second buffer is all jumbled up
792 // insertion sort the second buffer, then redistribute the buffers
793 // back into the items using the opposite process used for creating the buffer
772794
773 // while an unstable sort like quicksort could be applied here, in benchmarks it was consistently slightly slower than a simple insertion sort,
774 // even for tens of millions of items. this may be because insertion sort is quite fast when the data is already somewhat sorted, like it is here
795 // while an unstable sort like quicksort could be applied here, in benchmarks
796 // it was consistently slightly slower than a simple insertion sort,
797 // even for tens of millions of items. this may be because insertion
798 // sort is quite fast when the data is already somewhat sorted, like it is here
775799 insertionSort(T, items[buffer2.start..buffer2.end], context, lessThan);
776800
777801 pull_index = 0;
......@@ -808,6 +832,12 @@ pub fn sort(
808832 }
809833}
810834
835/// TODO currently this just calls `insertionSortContext`. The block sort implementation
836/// in this file needs to be adapted to use the sort context.
837pub fn sortContext(len: usize, context: anytype) void {
838 return insertionSortContext(len, context);
839}
840
811841// merge operation without a buffer
812842fn mergeInPlace(
813843 comptime T: type,