Merge pull request #2005 from thePHTest/sort_with_indices

have sort_with_indices allocate. Add a couple convenience procs for using the indices result to sort other slices.

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) {

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
 		}
 	}
-}
+}
+
+@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
+				}
+			}
+		}
+	}
+}