pkg/corpus: upgrade corpus minimization algorithm

3 files changed, 49 insertions, 4 deletions

diff --git a/pkg/corpus/fenwik.go b/pkg/corpus/fenwik.go
index 3a49219ac..dd6bfdbef 100644
--- a/pkg/corpus/fenwik.go
+++ b/pkg/corpus/fenwik.go
@@ -43,6 +43,11 @@ func (t *fenwikTree) prefixSumTo(ind int) (sum int64) {
 	return
 }
 
+func (t *fenwikTree) reserve(cap int) {
+	t.capacity = cap
+	t.extendAndRebuild()
+}
+
 func (t *fenwikTree) extendAndRebuild() {
 	if t.capacity == 0 {
 		t.capacity = 16
diff --git a/pkg/corpus/minimize.go b/pkg/corpus/minimize.go
index d01e06f57..d3d68ee27 100644
--- a/pkg/corpus/minimize.go
+++ b/pkg/corpus/minimize.go
@@ -41,11 +41,30 @@ func (corpus *Corpus) Minimize(cover bool) {
 
 	// Overwrite the program lists.
 	corpus.ProgramsList = &ProgramsList{}
+	corpus.ProgramsList.prios.reserve(len(inputs))
 	for _, area := range corpus.focusAreas {
 		area.ProgramsList = &ProgramsList{}
 	}
-	for _, ctx := range signal.Minimize(inputs) {
+	progs := signal.Minimize(inputs)
+
+	sort.SliceStable(progs, func(i, j int) bool {
+		first := progs[i].(*Item)
+		second := progs[j].(*Item)
+		return first.ExecLast > second.ExecLast
+	})
+
+	inputs = make([]signal.Context, 0, len(progs))
+	for _, inp := range progs {
+		inputs = append(inputs, signal.Context{
+			Signal:  inp.(*Item).StableSignal,
+			Context: inp,
+		})
+	}
+	progs = signal.Minimize(inputs)
+
+	for _, ctx := range progs {
 		inp := ctx.(*Item)
+
 		corpus.progsMap[inp.Sig] = inp
 		corpus.saveProgram(inp.Prog, inp.Signal)
 		for area := range inp.areas {
diff --git a/pkg/signal/signal.go b/pkg/signal/signal.go
index a22d19f67..ffa27ce97 100644
--- a/pkg/signal/signal.go
+++ b/pkg/signal/signal.go
@@ -146,7 +146,7 @@ type Context struct {
 func Minimize(corpus []Context) []any {
 	type ContextPrio struct {
 		prio prioType
-		idx  int
+		idx  []int
 	}
 	covered := make(map[elemType]ContextPrio)
 	for i, inp := range corpus {
@@ -154,14 +154,35 @@ func Minimize(corpus []Context) []any {
 			if prev, ok := covered[e]; !ok || p > prev.prio {
 				covered[e] = ContextPrio{
 					prio: p,
-					idx:  i,
+					idx:  []int{i},
+				}
+			} else if prev := covered[e]; p == prev.prio {
+				covered[e] = ContextPrio{
+					prio: p,
+					idx:  append(prev.idx, i),
 				}
 			}
 		}
 	}
 	indices := make(map[int]struct{}, len(corpus))
 	for _, cp := range covered {
-		indices[cp.idx] = struct{}{}
+		if len(cp.idx) == 1 {
+			indices[cp.idx[0]] = struct{}{}
+		}
+	}
+	for _, cp := range covered {
+		if len(cp.idx) != 1 {
+			alreadyCovered := false
+			for pc := range cp.idx {
+				if _, ok := indices[pc]; ok {
+					alreadyCovered = true
+					break
+				}
+			}
+			if !alreadyCovered {
+				indices[cp.idx[0]] = struct{}{}
+			}
+		}
 	}
 	result := make([]any, 0, len(indices))
 	for idx := range indices {