From dee56964d94d28f3f609bc4f62042dd6d1e153ae Mon Sep 17 00:00:00 2001
From: Aleksandr Nogikh <nogikh@google.com>
Date: Thu, 22 Aug 2024 17:48:30 +0200
Subject: 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.
---
 pkg/bisect/minimize/slice.go | 45 ++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 45 insertions(+)

(limited to 'pkg/bisect/minimize/slice.go')

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]
-- 
cgit mrf-deployment