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(...@@ -408,6 +408,13 @@ pub fn ArrayHashMap(
408 return self.unmanaged.reIndexContext(self.allocator, self.ctx);408 return self.unmanaged.reIndexContext(self.allocator, self.ctx);
409 }409 }
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
411 /// Shrinks the underlying `Entry` array to `new_len` elements and discards any associated418 /// Shrinks the underlying `Entry` array to `new_len` elements and discards any associated
412 /// index entries. Keeps capacity the same.419 /// index entries. Keeps capacity the same.
413 pub fn shrinkRetainingCapacity(self: *Self, new_len: usize) void {420 pub fn shrinkRetainingCapacity(self: *Self, new_len: usize) void {
...@@ -1169,6 +1176,22 @@ pub fn ArrayHashMapUnmanaged(...@@ -1169,6 +1176,22 @@ pub fn ArrayHashMapUnmanaged(
1169 self.index_header = new_header;1176 self.index_header = new_header;
1170 }1177 }
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
1172 /// Shrinks the underlying `Entry` array to `new_len` elements and discards any associated1195 /// Shrinks the underlying `Entry` array to `new_len` elements and discards any associated
1173 /// index entries. Keeps capacity the same.1196 /// index entries. Keeps capacity the same.
1174 pub fn shrinkRetainingCapacity(self: *Self, new_len: usize) void {1197 pub fn shrinkRetainingCapacity(self: *Self, new_len: usize) void {
...@@ -1868,6 +1891,14 @@ const IndexHeader = struct {...@@ -1868,6 +1891,14 @@ const IndexHeader = struct {
1868 allocator.free(slice);1891 allocator.free(slice);
1869 }1892 }
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
1871 // Verify that the header has sufficient alignment to produce aligned arrays.1902 // Verify that the header has sufficient alignment to produce aligned arrays.
1872 comptime {1903 comptime {
1873 if (@alignOf(u32) > @alignOf(IndexHeader))1904 if (@alignOf(u32) > @alignOf(IndexHeader))
...@@ -2218,6 +2249,32 @@ test "auto store_hash" {...@@ -2218,6 +2249,32 @@ test "auto store_hash" {
2218 try testing.expect(meta.fieldInfo(HasExpensiveEqlUn.Data, .hash).field_type != void);2249 try testing.expect(meta.fieldInfo(HasExpensiveEqlUn.Data, .hash).field_type != void);
2219}2250}
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
2221pub fn getHashPtrAddrFn(comptime K: type, comptime Context: type) (fn (Context, K) u32) {2278pub fn getHashPtrAddrFn(comptime K: type, comptime Context: type) (fn (Context, K) u32) {
2222 return struct {2279 return struct {
2223 fn hash(ctx: Context, key: K) u32 {2280 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 {...@@ -392,6 +392,34 @@ pub fn MultiArrayList(comptime S: type) type {
392 return result;392 return result;
393 }393 }
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
395 fn capacityInBytes(capacity: usize) usize {423 fn capacityInBytes(capacity: usize) usize {
396 const sizes_vector: std.meta.Vector(sizes.bytes.len, usize) = sizes.bytes;424 const sizes_vector: std.meta.Vector(sizes.bytes.len, usize) = sizes.bytes;
397 const capacity_vector = @splat(sizes.bytes.len, capacity);425 const capacity_vector = @splat(sizes.bytes.len, capacity);
lib/std/sort.zig+39-9
...@@ -73,7 +73,10 @@ test "binarySearch" {...@@ -73,7 +73,10 @@ test "binarySearch" {
73 );73 );
74}74}
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.
77pub fn insertionSort(80pub fn insertionSort(
78 comptime T: type,81 comptime T: type,
79 items: []T,82 items: []T,
...@@ -91,6 +94,18 @@ pub fn insertionSort(...@@ -91,6 +94,18 @@ pub fn insertionSort(
91 }94 }
92}95}
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
94const Range = struct {109const Range = struct {
95 start: usize,110 start: usize,
96 end: usize,111 end: usize,
...@@ -178,7 +193,8 @@ const Pull = struct {...@@ -178,7 +193,8 @@ const Pull = struct {
178 range: Range,193 range: Range,
179};194};
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).
182/// Currently implemented as block sort.198/// Currently implemented as block sort.
183pub fn sort(199pub fn sort(
184 comptime T: type,200 comptime T: type,
...@@ -186,6 +202,7 @@ pub fn sort(...@@ -186,6 +202,7 @@ pub fn sort(
186 context: anytype,202 context: anytype,
187 comptime lessThan: fn (context: @TypeOf(context), lhs: T, rhs: T) bool,203 comptime lessThan: fn (context: @TypeOf(context), lhs: T, rhs: T) bool,
188) void {204) void {
205
189 // Implementation ported from https://github.com/BonzaiThePenguin/WikiSort/blob/master/WikiSort.c206 // Implementation ported from https://github.com/BonzaiThePenguin/WikiSort/blob/master/WikiSort.c
190 var cache: [512]T = undefined;207 var cache: [512]T = undefined;
191208
...@@ -291,10 +308,13 @@ pub fn sort(...@@ -291,10 +308,13 @@ pub fn sort(
291308
292 // then merge sort the higher levels, which can be 8-15, 16-31, 32-63, 64-127, etc.309 // then merge sort the higher levels, which can be 8-15, 16-31, 32-63, 64-127, etc.
293 while (true) {310 while (true) {
294 // if every A and B block will fit into the cache, use a special branch specifically for merging with the cache311 // if every A and B block will fit into the cache, use a special branch
295 // (we use < rather than <= since the block size might be one more than iterator.length())312 // specifically for merging with the cache
313 // (we use < rather than <= since the block size might be one more than
314 // iterator.length())
296 if (iterator.length() < cache.len) {315 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,
298 // then merge the two merged subarrays from the cache back into the original array318 // then merge the two merged subarrays from the cache back into the original array
299 if ((iterator.length() + 1) * 4 <= cache.len and iterator.length() * 4 <= items.len) {319 if ((iterator.length() + 1) * 4 <= cache.len and iterator.length() * 4 <= items.len) {
300 iterator.begin();320 iterator.begin();
...@@ -767,11 +787,15 @@ pub fn sort(...@@ -767,11 +787,15 @@ pub fn sort(
767 }787 }
768 }788 }
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 up790 // when we're finished with this merge step we should have the one
771 // insertion sort the second buffer, then redistribute the buffers back into the items using the opposite process used for creating the buffer791 // 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,795 // while an unstable sort like quicksort could be applied here, in benchmarks
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 here796 // 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
775 insertionSort(T, items[buffer2.start..buffer2.end], context, lessThan);799 insertionSort(T, items[buffer2.start..buffer2.end], context, lessThan);
776800
777 pull_index = 0;801 pull_index = 0;
...@@ -808,6 +832,12 @@ pub fn sort(...@@ -808,6 +832,12 @@ pub fn sort(
808 }832 }
809}833}
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
811// merge operation without a buffer841// merge operation without a buffer
812fn mergeInPlace(842fn mergeInPlace(
813 comptime T: type,843 comptime T: type,