diff --git a/core/slice/sort.odin b/core/slice/sort.odin index 779ba37fd..8020cede7 100644 --- a/core/slice/sort.odin +++ b/core/slice/sort.odin @@ -38,18 +38,50 @@ sort :: proc(data: $T/[]$E) where ORD(E) { } } -// sort sorts a slice + +sort_by_indices :: proc{ sort_by_indices_allocate, _sort_by_indices} + +sort_by_indices_allocate :: proc(data: $T/[]$E, indices: []int, allocator := context.allocator) -> (sorted: T) { + assert(len(data) == len(indices)) + sorted = make([]int, len(data), allocator) + for v, i in indices { + sorted[i] = data[v] + } + return +} + +_sort_by_indices :: proc(data, sorted: $T/[]$E, indices: []int) { + assert(len(data) == len(indices)) + assert(len(data) == len(sorted)) + for v, i in indices { + sorted[i] = data[v] + } +} + +sort_by_indices_overwrite :: proc(data: $T/[]$E, indices: []int) { + assert(len(data) == len(indices)) + temp := make([]int, len(data), context.allocator) + defer delete(temp) + for v, i in indices { + temp[i] = data[v] + } + swap_with_slice(data, temp) +} + +// sort sorts a slice and returns a slice of the original indices // This sort is not guaranteed to be stable -sort_with_indices :: proc(data: $T/[]$E, indices: []int) where ORD(E) { +sort_with_indices :: proc(data: $T/[]$E, allocator := context.allocator) -> (indices: []int) where ORD(E) { + indices = make([]int, len(data), allocator) when size_of(E) != 0 { if n := len(data); n > 1 { - assert(len(data) == len(indices)) for _, idx in indices { indices[idx] = idx } _quick_sort_general_with_indices(data, indices, 0, n, _max_depth(n), struct{}{}, .Ordered) } + return indices } + return indices } // sort_by sorts a slice with a given procedure to test whether two values are ordered "i < j" @@ -64,16 +96,18 @@ sort_by :: proc(data: $T/[]$E, less: proc(i, j: E) -> bool) { // sort_by sorts a slice with a given procedure to test whether two values are ordered "i < j" // This sort is not guaranteed to be stable -sort_by_with_indices :: proc(data: $T/[]$E, indices: []int, less: proc(i, j: E) -> bool) { +sort_by_with_indices :: proc(data: $T/[]$E, less: proc(i, j: E) -> bool, allocator := context.allocator) -> (indices : []int) { when size_of(E) != 0 { if n := len(data); n > 1 { - assert(len(data) == len(indices)) + indices = make([]int, len(data), allocator) for _, idx in indices { indices[idx] = idx } _quick_sort_general_with_indices(data, indices, 0, n, _max_depth(n), less, .Less) + return indices } } + return nil } sort_by_cmp :: proc(data: $T/[]$E, cmp: proc(i, j: E) -> Ordering) { diff --git a/tests/core/slice/test_core_slice.odin b/tests/core/slice/test_core_slice.odin index 12577ab20..05dfb58ff 100644 --- a/tests/core/slice/test_core_slice.odin +++ b/tests/core/slice/test_core_slice.odin @@ -51,23 +51,20 @@ test_sort_with_indices :: proc(t: ^testing.T) { r := rand.create(seed) vals := make([]u64, test_size) - f_idx := make([]int, test_size) // Forward index, will be sorted r_idx := make([]int, test_size) // Reverse index - defer { delete(vals) - delete(f_idx) delete(r_idx) } // Set up test values for _, i in vals { vals[i] = rand.uint64(&r) - f_idx[i] = i } // Sort - slice.sort_with_indices(vals, f_idx) + f_idx := slice.sort_with_indices(vals) + defer delete(f_idx) // Verify sorted test values rand.init(&r, seed) @@ -94,4 +91,92 @@ test_sort_with_indices :: proc(t: ^testing.T) { last = v } } -} \ No newline at end of file +} + +@test +test_sort_by_indices :: proc(t: ^testing.T) { + seed := rand.uint64() + fmt.printf("Random seed: %v\n", seed) + + // Test sizes are all prime. + test_sizes :: []int{7, 13, 347, 1031, 10111, 100003} + + for test_size in test_sizes { + fmt.printf("Sorting %v random u64 values along with index.\n", test_size) + + r := rand.create(seed) + + vals := make([]u64, test_size) + r_idx := make([]int, test_size) // Reverse index + defer { + delete(vals) + delete(r_idx) + } + + // Set up test values + for _, i in vals { + vals[i] = rand.uint64(&r) + } + + // Sort + f_idx := slice.sort_with_indices(vals) + defer delete(f_idx) + + // Verify sorted test values + rand.init(&r, seed) + + { + indices := make([]int, test_size) + defer delete(indices) + for _, i in indices { + indices[i] = i + } + + sorted_indices := slice.sort_by_indices(indices, f_idx) + defer delete(sorted_indices) + for v, i in sorted_indices { + idx_pass := v == f_idx[i] + expect(t, idx_pass, "Expected the sorted index to be the same as the result from sort_with_indices") + if !idx_pass { + break + } + } + } + { + indices := make([]int, test_size) + defer delete(indices) + for _, i in indices { + indices[i] = i + } + + slice.sort_by_indices_overwrite(indices, f_idx) + for v, i in indices { + idx_pass := v == f_idx[i] + expect(t, idx_pass, "Expected the sorted index to be the same as the result from sort_with_indices") + if !idx_pass { + break + } + } + } + { + indices := make([]int, test_size) + swap := make([]int, test_size) + defer { + delete(indices) + delete(swap) + } + for _, i in indices { + indices[i] = i + } + + slice.sort_by_indices(indices, swap, f_idx) + for v, i in swap { + idx_pass := v == f_idx[i] + expect(t, idx_pass, "Expected the sorted index to be the same as the result from sort_with_indices") + if !idx_pass { + break + } + } + } + } +}