pkg/bisect/minimize: introduce SliceWithFixed

The function is similar to minimize.Slice(), but it allows to designate
a subset of slice elements that must always remain.

2 files changed, 73 insertions, 0 deletions

diff --git a/pkg/bisect/minimize/slice.go b/pkg/bisect/minimize/slice.go
index 76fe82101..fee125542 100644
--- a/pkg/bisect/minimize/slice.go
+++ b/pkg/bisect/minimize/slice.go
@@ -7,6 +7,7 @@ import (
 	"errors"
 	"fmt"
 	"math"
+	"slices"
 	"strings"
 )
 
@@ -45,6 +46,50 @@ func Slice[T any](config Config[T], slice []T) ([]T, error) {
 	return ctx.bisect()
 }
 
+// SliceWithFixed behaves like Slice, but also allows to designate the elements that
+// must always remain (the "fixed" ones).
+func SliceWithFixed[T any](config Config[T], slice []T, fixed func(T) bool) ([]T, error) {
+	var freeIdx, fixedIdx []int
+	for i := 0; i < len(slice); i++ {
+		if fixed(slice[i]) {
+			fixedIdx = append(fixedIdx, i)
+		} else {
+			freeIdx = append(freeIdx, i)
+		}
+	}
+	if len(freeIdx) == 0 {
+		return slice, nil
+	}
+	convert := func(idx []int) []T {
+		ret := make([]T, 0, len(idx)+len(fixedIdx))
+		idx, fixedIdx := slices.Clone(idx), slices.Clone(fixedIdx)
+		for len(idx)+len(fixedIdx) > 0 {
+			if len(idx) > 0 && (len(fixedIdx) == 0 ||
+				len(fixedIdx) > 0 && idx[0] < fixedIdx[0]) {
+				ret = append(ret, slice[idx[0]])
+				idx = idx[1:]
+			} else {
+				ret = append(ret, slice[fixedIdx[0]])
+				fixedIdx = fixedIdx[1:]
+			}
+		}
+		return ret
+	}
+	newConfig := Config[int]{
+		MaxSteps:  config.MaxSteps,
+		MaxChunks: config.MaxChunks,
+		Pred: func(idx []int) (bool, error) {
+			return config.Pred(convert(idx))
+		},
+		Logf: config.Logf,
+	}
+	result, err := Slice[int](newConfig, freeIdx)
+	if err != nil {
+		return nil, err
+	}
+	return convert(result), nil
+}
+
 type sliceCtx[T any] struct {
 	Config[T]
 	chunks   []*arrayChunk[T]
diff --git a/pkg/bisect/minimize/slice_test.go b/pkg/bisect/minimize/slice_test.go
index 4e3e65202..cd6ec37c6 100644
--- a/pkg/bisect/minimize/slice_test.go
+++ b/pkg/bisect/minimize/slice_test.go
@@ -42,6 +42,34 @@ func TestBisectSliceFull(t *testing.T) {
 	assert.Equal(t, ret, array)
 }
 
+func TestBisectSliceWithFixed(t *testing.T) {
+	t.Parallel()
+	array := make([]int, 100)
+	for i := 0; i < 100; i++ {
+		array[i] = i
+	}
+	ret, err := SliceWithFixed(Config[int]{
+		Pred: func(arr []int) (bool, error) {
+			// Let's ensure that there are always fixed elements.
+			values := map[int]bool{}
+			for _, v := range arr {
+				values[v] = true
+			}
+			if !values[0] || !values[20] || !values[40] {
+				t.Fatal("predicate step without fixed elements")
+			}
+			// And also require the elements 25 and 50.
+			return values[25] && values[50], nil
+		},
+		Logf: t.Logf,
+	}, array, func(elem int) bool {
+		// Let's keep these 3 elements.
+		return elem == 0 || elem == 20 || elem == 40
+	})
+	assert.NoError(t, err)
+	assert.Equal(t, []int{0, 20, 25, 40, 50}, ret)
+}
+
 func TestBisectRandomSlice(t *testing.T) {
 	t.Parallel()
 	r := rand.New(testutil.RandSource(t))