slices: optimize Compact and CompactFunc

src/slices/slices.go

@@ -346,40 +346,46 @@ func Clone[S ~[]E, E any](s S) S {
 // which may have a smaller length.
 // Compact zeroes the elements between the new length and the original length.
 func Compact[S ~[]E, E comparable](s S) S {
-	if len(s) < 2 {
-		return s
-	}
-	i := 1
-	for k := 1; k < len(s); k++ {
-		if s[k] != s[k-1] {
-			if i != k {
-				s[i] = s[k]
+	if len(s) > 1 {
+		for k := 1; k < len(s); k++ {
+			if s[k] == s[k-1] {
+				s2 := s[k:]
+				for k2 := 1; k2 < len(s2); k2++ {
+					if s2[k2] != s2[k2-1] {
+						s[k] = s2[k2]
+						k++
+					}
+				}
+
+				clear(s[k:]) // zero/nil out the obsolete elements, for GC
+				return s[:k]
 			}
-			i++
 		}
 	}
-	clear(s[i:]) // zero/nil out the obsolete elements, for GC
-	return s[:i]
+	return s
 }
 
 // CompactFunc is like [Compact] but uses an equality function to compare elements.
 // For runs of elements that compare equal, CompactFunc keeps the first one.
 // CompactFunc zeroes the elements between the new length and the original length.
 func CompactFunc[S ~[]E, E any](s S, eq func(E, E) bool) S {
-	if len(s) < 2 {
-		return s
-	}
-	i := 1
-	for k := 1; k < len(s); k++ {
-		if !eq(s[k], s[k-1]) {
-			if i != k {
-				s[i] = s[k]
+	if len(s) > 1 {
+		for k := 1; k < len(s); k++ {
+			if eq(s[k], s[k-1]) {
+				s2 := s[k:]
+				for k2 := 1; k2 < len(s2); k2++ {
+					if !eq(s2[k2], s2[k2-1]) {
+						s[k] = s2[k2]
+						k++
+					}
+				}
+
+				clear(s[k:]) // zero/nil out the obsolete elements, for GC
+				return s[:k]
 			}
-			i++
 		}
 	}
-	clear(s[i:]) // zero/nil out the obsolete elements, for GC
-	return s[:i]
+	return s
 }
 
 // Grow increases the slice's capacity, if necessary, to guarantee space for

src/slices/slices_test.go

@@ -761,7 +761,7 @@ var compactTests = []struct {
 		[]int{1, 2, 3},
 	},
 	{
-		"1 item",
+		"2 items",
 		[]int{1, 1, 2},
 		[]int{1, 2},
 	},
@@ -800,12 +800,26 @@ func BenchmarkCompact(b *testing.B) {
 }
 
 func BenchmarkCompact_Large(b *testing.B) {
-	type Large [4 * 1024]byte
-
-	ss := make([]Large, 1024)
-	for i := 0; i < b.N; i++ {
-		_ = Compact(ss)
-	}
+	type Large [16]int
+	const N = 1024
+
+	b.Run("all_dup", func(b *testing.B) {
+		ss := make([]Large, N)
+		b.ResetTimer()
+		for i := 0; i < b.N; i++ {
+			_ = Compact(ss)
+		}
+	})
+	b.Run("no_dup", func(b *testing.B) {
+		ss := make([]Large, N)
+		for i := range ss {
+			ss[i][0] = i
+		}
+		b.ResetTimer()
+		for i := 0; i < b.N; i++ {
+			_ = Compact(ss)
+		}
+	})
 }
 
 func TestCompactFunc(t *testing.T) {
@@ -871,15 +885,42 @@ func TestCompactFuncClearTail(t *testing.T) {
 	}
 }
 
-func BenchmarkCompactFunc_Large(b *testing.B) {
-	type Large [4 * 1024]byte
-
-	ss := make([]Large, 1024)
-	for i := 0; i < b.N; i++ {
-		_ = CompactFunc(ss, func(a, b Large) bool { return a == b })
+func BenchmarkCompactFunc(b *testing.B) {
+	for _, c := range compactTests {
+		b.Run(c.name, func(b *testing.B) {
+			ss := make([]int, 0, 64)
+			for k := 0; k < b.N; k++ {
+				ss = ss[:0]
+				ss = append(ss, c.s...)
+				_ = CompactFunc(ss, func(a, b int) bool { return a == b })
+			}
+		})
 	}
 }
 
+func BenchmarkCompactFunc_Large(b *testing.B) {
+	type Element = int
+	const N = 1024 * 1024
+
+	b.Run("all_dup", func(b *testing.B) {
+		ss := make([]Element, N)
+		b.ResetTimer()
+		for i := 0; i < b.N; i++ {
+			_ = CompactFunc(ss, func(a, b Element) bool { return a == b })
+		}
+	})
+	b.Run("no_dup", func(b *testing.B) {
+		ss := make([]Element, N)
+		for i := range ss {
+			ss[i] = i
+		}
+		b.ResetTimer()
+		for i := 0; i < b.N; i++ {
+			_ = CompactFunc(ss, func(a, b Element) bool { return a == b })
+		}
+	})
+}
+
 func TestGrow(t *testing.T) {
 	s1 := []int{1, 2, 3}