prog: enable MutateWithHints() abortion

The call may potentially generate a very large number of possible
mutations. Add a way to abort the process.

3 files changed, 39 insertions, 20 deletions

diff --git a/prog/hints.go b/prog/hints.go
index 98aeccce2..cbfe9c859 100644
--- a/prog/hints.go
+++ b/prog/hints.go
@@ -65,32 +65,40 @@ func (m CompMap) String() string {
 
 // Mutates the program using the comparison operands stored in compMaps.
 // For each of the mutants executes the exec callback.
-func (p *Prog) MutateWithHints(callIndex int, comps CompMap, exec func(p *Prog)) {
+// The callback must return whether we should continue substitution (true)
+// or abort the process (false).
+func (p *Prog) MutateWithHints(callIndex int, comps CompMap, exec func(p *Prog) bool) {
 	p = p.Clone()
 	c := p.Calls[callIndex]
-	execValidate := func() {
+	doMore := true
+	execValidate := func() bool {
 		// Don't try to fix the candidate program.
 		// Assuming the original call was sanitized, we've got a bad call
 		// as the result of hint substitution, so just throw it away.
 		if p.Target.sanitize(c, false) != nil {
-			return
+			return true
 		}
 		if p.checkConditions() != nil {
 			// Patching unions that no longer satisfy conditions would
 			// require much deeped changes to prog arguments than
 			// generateHints() expects.
 			// Let's just ignore such mutations.
-			return
+			return true
 		}
 		p.debugValidate()
-		exec(p)
+		doMore = exec(p)
+		return doMore
 	}
-	ForeachArg(c, func(arg Arg, _ *ArgCtx) {
+	ForeachArg(c, func(arg Arg, ctx *ArgCtx) {
+		if !doMore {
+			ctx.Stop = true
+			return
+		}
 		generateHints(comps, arg, execValidate)
 	})
 }
 
-func generateHints(compMap CompMap, arg Arg, exec func()) {
+func generateHints(compMap CompMap, arg Arg, exec func() bool) {
 	typ := arg.Type()
 	if typ == nil || arg.Dir() == DirOut {
 		return
@@ -134,18 +142,20 @@ func generateHints(compMap CompMap, arg Arg, exec func()) {
 	}
 }
 
-func checkConstArg(arg *ConstArg, compMap CompMap, exec func()) {
+func checkConstArg(arg *ConstArg, compMap CompMap, exec func() bool) {
 	original := arg.Val
 	// Note: because shrinkExpand returns a map, order of programs is non-deterministic.
 	// This can affect test coverage reports.
 	for _, replacer := range shrinkExpand(original, compMap, arg.Type().TypeBitSize(), false) {
 		arg.Val = replacer
-		exec()
+		if !exec() {
+			break
+		}
 	}
 	arg.Val = original
 }
 
-func checkDataArg(arg *DataArg, compMap CompMap, exec func()) {
+func checkDataArg(arg *DataArg, compMap CompMap, exec func() bool) {
 	bytes := make([]byte, 8)
 	data := arg.Data()
 	size := len(data)
@@ -159,13 +169,15 @@ func checkDataArg(arg *DataArg, compMap CompMap, exec func()) {
 		for _, replacer := range shrinkExpand(val, compMap, 64, false) {
 			binary.LittleEndian.PutUint64(bytes, replacer)
 			copy(data[i:], bytes)
-			exec()
+			if !exec() {
+				break
+			}
 		}
 		copy(data[i:], original)
 	}
 }
 
-func checkCompressedArg(arg *DataArg, compMap CompMap, exec func()) {
+func checkCompressedArg(arg *DataArg, compMap CompMap, exec func() bool) {
 	data0 := arg.Data()
 	data, dtor := image.MustDecompress(data0)
 	defer dtor()
@@ -183,7 +195,9 @@ func checkCompressedArg(arg *DataArg, compMap CompMap, exec func()) {
 			binary.LittleEndian.PutUint64(bytes, replacer)
 			copy(data[i:], bytes)
 			arg.SetData(image.Compress(data))
-			exec()
+			if !exec() {
+				break
+			}
 		}
 		copy(data[i:], original)
 	}
diff --git a/prog/hints_test.go b/prog/hints_test.go
index b01ac0204..a4b8da7bf 100644
--- a/prog/hints_test.go
+++ b/prog/hints_test.go
@@ -164,8 +164,9 @@ func TestHintsCheckConstArg(t *testing.T) {
 			var res []uint64
 			typ := types[fmt.Sprintf("int%v_%v", test.size, test.bitsize)]
 			constArg := MakeConstArg(typ, DirIn, test.in)
-			checkConstArg(constArg, test.comps, func() {
+			checkConstArg(constArg, test.comps, func() bool {
 				res = append(res, constArg.Val)
+				return true
 			})
 			if !reflect.DeepEqual(res, test.res) {
 				t.Fatalf("\ngot : %v\nwant: %v", res, test.res)
@@ -302,8 +303,9 @@ func TestHintsCheckDataArg(t *testing.T) {
 		t.Run(fmt.Sprintf("%v", test.name), func(t *testing.T) {
 			res := make(map[string]bool)
 			dataArg := MakeDataArg(typ, DirIn, []byte(test.in))
-			checkDataArg(dataArg, test.comps, func() {
+			checkDataArg(dataArg, test.comps, func() bool {
 				res[string(dataArg.Data())] = true
+				return true
 			})
 			if !reflect.DeepEqual(res, test.res) {
 				s := "\ngot:  ["
@@ -390,8 +392,9 @@ func TestHintsCompressedImage(t *testing.T) {
 		t.Run(fmt.Sprintf("%v", i), func(t *testing.T) {
 			var res []string
 			arg := MakeDataArg(typ, DirIn, image.Compress([]byte(test.input)))
-			generateHints(test.comps, arg, func() {
+			generateHints(test.comps, arg, func() bool {
 				res = append(res, string(arg.Data()))
+				return true
 			})
 			for i, compressed := range res {
 				data, dtor := image.MustDecompress([]byte(compressed))
@@ -588,7 +591,7 @@ func TestHintsRandom(t *testing.T) {
 			for v := range vals {
 				comps.AddComp(v, r.randInt64())
 			}
-			p.MutateWithHints(i, comps, func(p1 *Prog) {})
+			p.MutateWithHints(i, comps, func(p1 *Prog) bool { return true })
 		}
 	}
 }
@@ -649,9 +652,10 @@ func TestHintsData(t *testing.T) {
 			t.Fatal(err)
 		}
 		var got []string
-		p.MutateWithHints(0, test.comps, func(newP *Prog) {
+		p.MutateWithHints(0, test.comps, func(newP *Prog) bool {
 			got = append(got, hex.EncodeToString(
 				newP.Calls[0].Args[0].(*PointerArg).Res.(*DataArg).Data()))
+			return true
 		})
 		sort.Strings(test.out)
 		sort.Strings(got)
@@ -683,7 +687,7 @@ func BenchmarkHints(b *testing.B) {
 	b.RunParallel(func(pb *testing.PB) {
 		for pb.Next() {
 			for i := range p.Calls {
-				p.MutateWithHints(i, comps[i], func(p1 *Prog) {})
+				p.MutateWithHints(i, comps[i], func(p1 *Prog) bool { return true })
 			}
 		}
 	})
diff --git a/prog/rand_test.go b/prog/rand_test.go
index ffe9bd822..9973134d1 100644
--- a/prog/rand_test.go
+++ b/prog/rand_test.go
@@ -63,8 +63,9 @@ func generateProg(t *testing.T, target *Target, rs rand.Source, ct *ChoiceTable,
 			comps.AddComp(v, v+1)
 			comps.AddComp(v, v+10)
 		}
-		p.MutateWithHints(i, comps, func(p1 *Prog) {
+		p.MutateWithHints(i, comps, func(p1 *Prog) bool {
 			p = p1.Clone()
+			return true
 		})
 	}
 	for _, crash := range []bool{false, true} {