util/slicesx: add Partition function

Signed-off-by: Andrew Dunham <andrew@du.nham.ca>
Change-Id: If97995ca9ee9fad40f327420dcb1857dd7ea2315

util/slicesx/slicesx.go

@@ -42,3 +42,18 @@ func Shuffle[S ~[]T, T any](s S) {
 		s[i], s[j] = s[j], s[i]
 	}
 }
+
+// Partition returns two slices, the first containing the elements of the input
+// slice for which the callback evaluates to true, the second containing the rest.
+//
+// This function does not mutate s.
+func Partition[S ~[]T, T any](s S, cb func(T) bool) (trues, falses S) {
+	for _, elem := range s {
+		if cb(elem) {
+			trues = append(trues, elem)
+		} else {
+			falses = append(falses, elem)
+		}
+	}
+	return
+}

util/slicesx/slicesx_test.go

@@ -44,6 +44,7 @@ func BenchmarkInterleave(b *testing.B) {
 		)
 	}
 }
+
 func TestShuffle(t *testing.T) {
 	var sl []int
 	for i := 0; i < 100; i++ {
@@ -64,3 +65,23 @@ func TestShuffle(t *testing.T) {
 		t.Errorf("expected shuffle after 10 tries")
 	}
 }
+
+func TestPartition(t *testing.T) {
+	var sl []int
+	for i := 1; i <= 10; i++ {
+		sl = append(sl, i)
+	}
+
+	evens, odds := Partition(sl, func(elem int) bool {
+		return elem%2 == 0
+	})
+
+	wantEvens := []int{2, 4, 6, 8, 10}
+	wantOdds := []int{1, 3, 5, 7, 9}
+	if !reflect.DeepEqual(evens, wantEvens) {
+		t.Errorf("evens: got %v, want %v", evens, wantEvens)
+	}
+	if !reflect.DeepEqual(odds, wantOdds) {
+		t.Errorf("odds: got %v, want %v", odds, wantOdds)
+	}
+}