Merge pull request #74 from segmentio/sort-fix

Sort fix
segmentio · Apr 8, 2022 · 56bf31b · 56bf31b
2 parents 6453584 + 5b0f4f9
commit 56bf31b
Show file tree

Hide file tree

Showing 5 changed files with 148 additions and 23 deletions.
diff --git a/qsort/sort.go b/qsort/sort.go
@@ -20,7 +20,7 @@ func Sort(data []byte, size int, swap func(int, int)) {
 	}
 
 	// No specialization available. Use the slower generic sorting routine.
-	if purego || size%8 != 0 || size > 32 {
+	if size%8 != 0 || size > 32 {
 		sort.Sort(newGeneric(data, size, swap))
 		return
 	}
@@ -37,11 +37,11 @@ func Sort(data []byte, size int, swap func(int, int)) {
 	// If no indirect swapping is required, try to use the hybrid partitioning scheme from
 	// https://blog.reverberate.org/2020/05/29/hoares-rebuttal-bubble-sorts-comeback.html
 	switch {
-	case swap == nil && size == 8 && cpu.X86.Has(x86.CMOV):
+	case swap == nil && !purego && size == 8 && cpu.X86.Has(x86.CMOV):
 		hybridQuicksort64(unsafeBytesTo64(data))
-	case swap == nil && size == 16 && cpu.X86.Has(x86.AVX):
+	case swap == nil && !purego && size == 16 && cpu.X86.Has(x86.AVX):
 		hybridQuicksort128(unsafeBytesTo128(data))
-	case swap == nil && size == 32 && cpu.X86.Has(x86.AVX2):
+	case swap == nil && !purego && size == 32 && cpu.X86.Has(x86.AVX2):
 		hybridQuicksort256(unsafeBytesTo256(data))
 	case size == 8:
 		quicksort64(unsafeBytesTo64(data), 0, smallCutoff, insertionsort64, hoarePartition64, swap)

diff --git a/qsort/sort16.go b/qsort/sort16.go
@@ -112,5 +112,5 @@ func hybridPartition128(data, scratch []uint128) int {
 }
 
 func less128(a, b uint128) bool {
-	return a.hi < b.hi || (a.hi == b.hi && a.lo <= b.lo)
+	return a.hi < b.hi || (a.hi == b.hi && a.lo < b.lo)
 }
diff --git a/qsort/sort24.go b/qsort/sort24.go
@@ -83,5 +83,5 @@ func hoarePartition192(data []uint192, base int, swap func(int, int)) int {
 func less192(a, b uint192) bool {
 	return a.hi < b.hi ||
 		(a.hi == b.hi && a.mid < b.mid) ||
-		(a.hi == b.hi && a.mid == b.mid && a.lo <= b.lo)
+		(a.hi == b.hi && a.mid == b.mid && a.lo < b.lo)
 }
diff --git a/qsort/sort32.go b/qsort/sort32.go
@@ -112,5 +112,5 @@ func less256(a, b uint256) bool {
 	return a.a < b.a ||
 		(a.a == b.a && a.b < b.b) ||
 		(a.a == b.a && a.b == b.b && a.c < b.c) ||
-		(a.a == b.a && a.b == b.b && a.c == b.c && a.d <= b.d)
+		(a.a == b.a && a.b == b.b && a.c == b.c && a.d < b.d)
 }
diff --git a/qsort/sort_test.go b/qsort/sort_test.go
@@ -81,6 +81,74 @@ func testSort(t *testing.T, size int) {
 	}
 }
 
+func TestPivot8(t *testing.T) {
+	lo := uint64(1)
+	mid := uint64(2)
+	hi := uint64(3)
+
+	for i := 0; i < 1000; i++ {
+		input := []uint64{lo, mid, hi}
+		rand.Shuffle(3, func(i, j int) {
+			input[i], input[j] = input[j], input[i]
+		})
+		medianOfThree64(input, 3, nil)
+		if input[0] != mid {
+			t.Fatal("medianOfThree128 did not put pivot in first position")
+		}
+	}
+}
+
+func TestPivot16(t *testing.T) {
+	lo := uint128{lo: 1}
+	mid := uint128{lo: 2}
+	hi := uint128{lo: 3}
+
+	for i := 0; i < 1000; i++ {
+		input := []uint128{lo, mid, hi}
+		rand.Shuffle(3, func(i, j int) {
+			input[i], input[j] = input[j], input[i]
+		})
+		medianOfThree128(input, 3, nil)
+		if input[0] != mid {
+			t.Fatal("medianOfThree128 did not put pivot in first position")
+		}
+	}
+}
+
+func TestPivot24(t *testing.T) {
+	lo := uint192{lo: 1}
+	mid := uint192{lo: 2}
+	hi := uint192{lo: 3}
+
+	for i := 0; i < 1000; i++ {
+		input := []uint192{lo, mid, hi}
+		rand.Shuffle(3, func(i, j int) {
+			input[i], input[j] = input[j], input[i]
+		})
+		medianOfThree192(input, 3, nil)
+		if input[0] != mid {
+			t.Fatal("medianOfThree192 did not put pivot in first position")
+		}
+	}
+}
+
+func TestPivot32(t *testing.T) {
+	lo := uint256{d: 1}
+	mid := uint256{d: 2}
+	hi := uint256{d: 3}
+
+	for i := 0; i < 1000; i++ {
+		input := []uint256{lo, mid, hi}
+		rand.Shuffle(3, func(i, j int) {
+			input[i], input[j] = input[j], input[i]
+		})
+		medianOfThree256(input, 3, nil)
+		if input[0] != mid {
+			t.Fatal("medianOfThree256 did not put pivot in first position")
+		}
+	}
+}
+
 func randint(lo, hi int) int {
 	if hi == lo {
 		return lo
@@ -89,35 +157,92 @@ func randint(lo, hi int) int {
 }
 
 func BenchmarkSort8(b *testing.B) {
-	benchSort(b, 8)
+	for _, count := range []int{1e3, 1e4, 1e5, 1e6} {
+		b.Run(strconv.Itoa(count), benchSort(count, 8, random, nil))
+	}
+}
+
+func BenchmarkSort8Indirect(b *testing.B) {
+	swap := func(int, int) {}
+	const count = 100000
+	b.Run("random", benchSort(count, 8, random, swap))
+	b.Run("asc", benchSort(count, 8, asc, swap))
+	b.Run("desc", benchSort(count, 8, desc, swap))
 }
 
 func BenchmarkSort16(b *testing.B) {
-	benchSort(b, 16)
+	for _, count := range []int{1e3, 1e4, 1e5, 1e6} {
+		b.Run(strconv.Itoa(count), benchSort(count, 16, random, nil))
+	}
+}
+
+func BenchmarkSort16Indirect(b *testing.B) {
+	swap := func(int, int) {}
+	const count = 100000
+	b.Run("random", benchSort(count, 16, random, swap))
+	b.Run("asc", benchSort(count, 16, asc, swap))
+	b.Run("desc", benchSort(count, 16, desc, swap))
 }
 
 func BenchmarkSort24(b *testing.B) {
-	benchSort(b, 24)
+	for _, count := range []int{1e3, 1e4, 1e5, 1e6} {
+		b.Run(strconv.Itoa(count), benchSort(count, 24, random, nil))
+	}
 }
 
-func BenchmarkSort32(b *testing.B) {
-	benchSort(b, 32)
+func BenchmarkSort24Indirect(b *testing.B) {
+	swap := func(int, int) {}
+	const count = 100000
+	b.Run("random", benchSort(count, 24, random, swap))
+	b.Run("asc", benchSort(count, 24, asc, swap))
+	b.Run("desc", benchSort(count, 24, desc, swap))
 }
 
-func benchSort(b *testing.B, size int) {
+func BenchmarkSort32(b *testing.B) {
 	for _, count := range []int{1e3, 1e4, 1e5, 1e6} {
-		b.Run(strconv.Itoa(count), func(b *testing.B) {
-			buf := make([]byte, count*size)
-			unsorted := make([]byte, count*size)
-			prng.Read(unsorted)
+		b.Run(strconv.Itoa(count), benchSort(count, 32, random, nil))
+	}
+}
+
+func BenchmarkSort32Indirect(b *testing.B) {
+	swap := func(int, int) {}
+	const count = 100000
+	b.Run("random", benchSort(count, 32, random, swap))
+	b.Run("asc", benchSort(count, 32, asc, swap))
+	b.Run("desc", benchSort(count, 32, desc, swap))
+}
+
+type order int
+
+const (
+	random order = iota
+	asc
+	desc
+)
 
-			b.SetBytes(int64(len(buf)))
-			b.ResetTimer()
+func benchSort(count, size int, order order, indirect func(int, int)) func(*testing.B) {
+	return func(b *testing.B) {
+		buf := make([]byte, count*size)
+		unsorted := make([]byte, count*size)
+		prng.Read(unsorted)
 
-			for i := 0; i < b.N; i++ {
-				copy(buf, unsorted)
-				Sort(buf, size, nil)
+		if order == asc || order == desc {
+			sort.Sort(newGeneric(unsorted, size, nil))
+		}
+		if order == desc {
+			g := newGeneric(unsorted, size, nil)
+			items := g.Len()
+			for i := 0; i < items/2; i++ {
+				g.Swap(i, items-1-i)
 			}
-		})
+		}
+
+		b.SetBytes(int64(len(buf)))
+		b.ResetTimer()
+
+		for i := 0; i < b.N; i++ {
+			copy(buf, unsorted)
+			Sort(buf, size, indirect)
+		}
 	}
 }