prog: restricts hints to at most 10 attempts per single kernel PC

We are getting too many generated candidates, the fuzzer may not keep up
with them at all (hints jobs keep growing infinitely). If a hint indeed came
from the input w/o transformation, then we should guess it on the first
attempt (or at least after few attempts). If it did not come from the input,
or came with a non-trivial transformation, then any number of attempts won't
help. So limit the total number of attempts (until the next restart).

3 files changed, 137 insertions, 20 deletions

diff --git a/prog/hints.go b/prog/hints.go
index ce83c009f..c054b9852 100644
--- a/prog/hints.go
+++ b/prog/hints.go
@@ -23,18 +23,13 @@ import (
 	"encoding/binary"
 	"fmt"
 	"sort"
+	"sync"
 
 	"github.com/google/syzkaller/pkg/image"
 )
 
-// Example: for comparisons {(op1, op2), (op1, op3), (op1, op4), (op2, op1)}
-// this map will store the following:
-//
-//	m = {
-//			op1: {map[op2]: true, map[op3]: true, map[op4]: true},
-//			op2: {map[op1]: true}
-//	}.
-type CompMap map[uint64]map[uint64]bool
+// CompMap maps comparison operand that could come from the input to the second operand to the PC.
+type CompMap map[uint64]map[uint64]map[uint64]bool
 
 const (
 	maxDataLength = 100
@@ -42,11 +37,18 @@ const (
 
 var specialIntsSet map[uint64]bool
 
-func (m CompMap) AddComp(arg1, arg2 uint64) {
+func (m CompMap) Add(pc, arg1, arg2 uint64, isConst bool) {
 	if _, ok := m[arg1]; !ok {
-		m[arg1] = make(map[uint64]bool)
+		m[arg1] = make(map[uint64]map[uint64]bool)
+	}
+	if _, ok := m[arg1][arg2]; !ok {
+		m[arg1][arg2] = make(map[uint64]bool)
+	}
+	m[arg1][arg2][pc] = true
+	if !isConst {
+		// Both operands could come from the input.
+		m.Add(pc, arg2, arg1, true)
 	}
-	m[arg1][arg2] = true
 }
 
 func (m CompMap) String() string {
@@ -66,8 +68,13 @@ func (m CompMap) String() string {
 // InplaceIntersect() only leaves the value pairs that are also present in other.
 func (m CompMap) InplaceIntersect(other CompMap) {
 	for val1, nested := range m {
-		for val2 := range nested {
-			if !other[val1][val2] {
+		for val2, pcs := range nested {
+			for pc := range pcs {
+				if !other[val1][val2][pc] {
+					delete(pcs, pc)
+				}
+			}
+			if len(pcs) == 0 {
 				delete(nested, val2)
 			}
 		}
@@ -373,6 +380,44 @@ func shrinkExpand(v uint64, compMap CompMap, bitsize uint64, image bool) []uint6
 	return res
 }
 
+type HintsLimiter struct {
+	mu       sync.Mutex
+	attempts map[uint64]int // replacement attempts per PC
+}
+
+// Limit restricts hints to at most N replacement attempts per single kernel PC
+// (globally, across all hints mutations for all programs).
+// We are getting too many generated candidates, the fuzzer may not keep up
+// with them at all (hints jobs keep growing infinitely). If a hint indeed came
+// from the input w/o transformation, then we should guess it on the first
+// attempt (or at least after few attempts). If it did not come from the input,
+// or came with a non-trivial transformation, then any number of attempts won't
+// help. So limit the total number of attempts (until the next restart).
+func (limiter *HintsLimiter) Limit(comps CompMap) {
+	const N = 10
+	limiter.mu.Lock()
+	defer limiter.mu.Unlock()
+	if limiter.attempts == nil {
+		limiter.attempts = make(map[uint64]int)
+	}
+	for op1, ops2 := range comps {
+		for op2, pcs := range ops2 {
+			for pc := range pcs {
+				limiter.attempts[pc]++
+				if limiter.attempts[pc] > N {
+					delete(pcs, pc)
+				}
+			}
+			if len(pcs) == 0 {
+				delete(ops2, op2)
+			}
+		}
+		if len(ops2) == 0 {
+			delete(comps, op1)
+		}
+	}
+}
+
 func init() {
 	specialIntsSet = make(map[uint64]bool)
 	for _, v := range specialInts {
diff --git a/prog/hints_test.go b/prog/hints_test.go
index d5a5a1461..44c24bfb8 100644
--- a/prog/hints_test.go
+++ b/prog/hints_test.go
@@ -662,7 +662,7 @@ func TestHintsRandom(t *testing.T) {
 			}
 			comps := make(CompMap)
 			for v := range vals {
-				comps.AddComp(v, r.randInt64())
+				comps.Add(1, v, r.randInt64(), true)
 			}
 			p.MutateWithHints(i, comps, func(p1 *Prog) bool { return true })
 		}
@@ -772,7 +772,7 @@ func BenchmarkHints(b *testing.B) {
 		}
 		comps[i] = make(CompMap)
 		for v := range vals {
-			comps[i].AddComp(v, r.randInt64())
+			comps[i].Add(1, v, r.randInt64(), true)
 		}
 	}
 	b.RunParallel(func(pb *testing.PB) {
@@ -784,10 +784,82 @@ func BenchmarkHints(b *testing.B) {
 	})
 }
 
-func compSet(vals ...uint64) map[uint64]bool {
-	m := make(map[uint64]bool)
+func TestHintsLimiter(t *testing.T) {
+	var limiter HintsLimiter
+
+	// Base case.
+	comps := make(CompMap)
+	comps.Add(1000, 1000, 1100, true)
+	for i := uint64(0); i < 9; i++ {
+		comps.Add(2000, 2000+i, 2100+i, true)
+	}
+	for i := uint64(0); i < 10; i++ {
+		comps.Add(3000, 3000+i, 3100+i, true)
+	}
+	for i := uint64(0); i < 11; i++ {
+		comps.Add(4000, 4000+i, 4100+i, true)
+	}
+	for i := uint64(0); i < 20; i++ {
+		comps.Add(5000, 5000+i, 5100+i, true)
+	}
+	assert.Equal(t, perPCCount(comps), map[uint64]int{
+		1000: 1,
+		2000: 9,
+		3000: 10,
+		4000: 11,
+		5000: 20,
+	})
+	limiter.Limit(comps)
+	assert.Equal(t, perPCCount(comps), map[uint64]int{
+		1000: 1,
+		2000: 9,
+		3000: 10,
+		4000: 10,
+		5000: 10,
+	})
+
+	// Test that counts are accumulated in the limiter.
+	comps = make(CompMap)
+	for i := uint64(0); i < 3; i++ {
+		comps.Add(1000, 1000+i, 1100+i, true)
+	}
+	for i := uint64(0); i < 3; i++ {
+		comps.Add(2000, 2000+i, 2100+i, true)
+	}
+	for i := uint64(0); i < 3; i++ {
+		comps.Add(3000, 3000+i, 3100+i, true)
+	}
+	assert.Equal(t, perPCCount(comps), map[uint64]int{
+		1000: 3,
+		2000: 3,
+		3000: 3,
+	})
+	limiter.Limit(comps)
+	assert.Equal(t, perPCCount(comps), map[uint64]int{
+		1000: 3,
+		2000: 1,
+	})
+}
+
+func perPCCount(comps CompMap) map[uint64]int {
+	res := make(map[uint64]int)
+	for _, ops2 := range comps {
+		for _, pcs := range ops2 {
+			for pc := range pcs {
+				res[pc]++
+			}
+		}
+	}
+	return res
+}
+
+func compSet(vals ...uint64) map[uint64]map[uint64]bool {
+	m := make(map[uint64]map[uint64]bool)
 	for _, v := range vals {
-		m[v] = true
+		if m[v] == nil {
+			m[v] = make(map[uint64]bool)
+		}
+		m[v][1] = true
 	}
 	return m
 }
diff --git a/prog/rand_test.go b/prog/rand_test.go
index 53b8ea338..d1e963595 100644
--- a/prog/rand_test.go
+++ b/prog/rand_test.go
@@ -60,8 +60,8 @@ func generateProg(t *testing.T, target *Target, rs rand.Source, ct *ChoiceTable,
 	for i, c := range p.Calls {
 		comps := make(CompMap)
 		for v := range extractValues(c) {
-			comps.AddComp(v, v+1)
-			comps.AddComp(v, v+10)
+			comps.Add(1, v, v+1, true)
+			comps.Add(1, v, v+10, true)
 		}
 		// If unbounded, this code may take O(N^2) time to complete.
 		// Since large programs are not uncommon, let's limit the number of hint iterations.