prog: implemented argument and call priorities

3 files changed, 377 insertions, 76 deletions

diff --git a/prog/mutation.go b/prog/mutation.go
index 6c356c4b4..1031217b4 100644
--- a/prog/mutation.go
+++ b/prog/mutation.go
@@ -5,7 +5,9 @@ package prog
 
 import (
 	"fmt"
+	"math"
 	"math/rand"
+	"sort"
 	"unsafe"
 )
 
@@ -155,8 +157,9 @@ func (ctx *mutator) mutateArg() bool {
 	if len(p.Calls) == 0 {
 		return false
 	}
-	c := p.Calls[r.Intn(len(p.Calls))]
-	if len(c.Args) == 0 {
+
+	c, ok := chooseCall(p, r)
+	if !ok {
 		return false
 	}
 	s := analyze(ctx.ct, ctx.corpus, p, c)
@@ -168,8 +171,8 @@ func (ctx *mutator) mutateArg() bool {
 		if len(ma.args) == 0 {
 			return false
 		}
-		idx := r.Intn(len(ma.args))
-		arg, ctx := ma.args[idx], ma.ctxes[idx]
+		chosenIdx := randomChoice(ma.priorities, r)
+		arg, ctx := ma.args[chosenIdx], ma.ctxes[chosenIdx]
 		calls, ok1 := p.Target.mutateArg(r, s, arg, ctx, &updateSizes)
 		if !ok1 {
 			ok = false
@@ -184,6 +187,43 @@ func (ctx *mutator) mutateArg() bool {
 	return true
 }
 
+// Select a call based on the complexity of the arguments.
+func chooseCall(p *Prog, r *randGen) (*Call, bool) {
+	var callPriorities []float64
+	noArgs := true
+
+	for _, c := range p.Calls {
+		totalPrio := float64(0)
+		ForeachArg(c, func(arg Arg, ctx *ArgCtx) {
+			prio, stopRecursion := arg.Type().getMutationPrio(p.Target, arg, false)
+			totalPrio += prio
+			ctx.Stop = stopRecursion
+		})
+		callPriorities = append(callPriorities, totalPrio)
+		if len(c.Args) > 0 {
+			noArgs = false
+		}
+	}
+
+	// Calls without arguments.
+	if noArgs {
+		return nil, false
+	}
+
+	return p.Calls[randomChoice(callPriorities, r)], true
+}
+
+// Generate a random index from a given 1-D array of priorities.
+func randomChoice(priorities []float64, r *randGen) int {
+	sum := float64(0)
+	probs := make([]float64, len(priorities))
+	for i, prio := range priorities {
+		sum += prio
+		probs[i] = sum
+	}
+	return sort.SearchFloat64s(probs, sum*r.Float64())
+}
+
 func (target *Target) mutateArg(r *randGen, s *state, arg Arg, ctx ArgCtx, updateSizes *bool) ([]*Call, bool) {
 	var baseSize uint64
 	if ctx.Base != nil {
@@ -400,49 +440,149 @@ type mutationArgs struct {
 	target        *Target
 	args          []Arg
 	ctxes         []ArgCtx
+	priorities    []float64
 	ignoreSpecial bool
 }
 
+const (
+	maxPriority = float64(10)
+	minPriority = float64(1)
+	dontMutate  = float64(0)
+)
+
 func (ma *mutationArgs) collectArg(arg Arg, ctx *ArgCtx) {
 	ignoreSpecial := ma.ignoreSpecial
 	ma.ignoreSpecial = false
-	switch typ := arg.Type().(type) {
-	case *StructType:
-		if ma.target.SpecialTypes[typ.Name()] == nil || ignoreSpecial {
-			return // For structs only individual fields are updated.
-		}
-		// These special structs are mutated as a whole.
-		ctx.Stop = true
-	case *UnionType:
-		if ma.target.SpecialTypes[typ.Name()] == nil && len(typ.Fields) == 1 || ignoreSpecial {
-			return
-		}
-		ctx.Stop = true
-	case *ArrayType:
-		// Don't mutate fixed-size arrays.
-		if typ.Kind == ArrayRangeLen && typ.RangeBegin == typ.RangeEnd {
-			return
-		}
-	case *CsumType:
-		return // Checksum is updated when the checksummed data changes.
-	case *ConstType:
-		return // Well, this is const.
-	case *BufferType:
-		if typ.Kind == BufferString && len(typ.Values) == 1 {
-			return // string const
-		}
-	case *PtrType:
-		if arg.(*PointerArg).IsSpecial() {
-			// TODO: we ought to mutate this, but we don't have code for this yet.
-			return
-		}
-	}
+
 	typ := arg.Type()
-	if typ == nil || typ.Dir() == DirOut || !typ.Varlen() && typ.Size() == 0 {
+	prio, stopRecursion := typ.getMutationPrio(ma.target, arg, ignoreSpecial)
+	ctx.Stop = stopRecursion
+
+	if prio == dontMutate {
 		return
 	}
+
+	if typ.Dir() == DirOut || !typ.Varlen() && typ.Size() == 0 {
+		return
+	}
+
 	ma.args = append(ma.args, arg)
 	ma.ctxes = append(ma.ctxes, *ctx)
+	ma.priorities = append(ma.priorities, prio)
+}
+
+// TODO: find a way to estimate optimal priority values.
+// Assign a priority for each type. The boolean is the reference type and it has
+// the minimum priority, since it has only two possible values.
+func (t *IntType) getMutationPrio(target *Target, arg Arg, ignoreSpecial bool) (prio float64, stopRecursion bool) {
+	// For a integer without a range of values, the priority is based on
+	// the number of bits occupied by the underlying type.
+	plainPrio := math.Log2((float64(t.Size() * 8))) + 0.1*maxPriority
+	if t.Kind != IntRange {
+		return plainPrio, false
+	}
+
+	switch size := t.RangeEnd - t.RangeBegin + 1; {
+	case size <= 15:
+		// For a small range, we assume that it is effectively
+		// similar with FlagsType and we need to try all possible values.
+		prio = rangeSizePrio(size)
+	case size <= 256:
+		// We consider that a relevant range has at most 256
+		// values (the number of values that can be represented on a byte).
+		prio = maxPriority
+	default:
+		// Ranges larger than 256 are equivalent with a plain integer.
+		prio = plainPrio
+	}
+	return prio, false
+}
+
+func (t *StructType) getMutationPrio(target *Target, arg Arg, ignoreSpecial bool) (prio float64, stopRecursion bool) {
+	if target.SpecialTypes[t.Name()] == nil || ignoreSpecial {
+		return dontMutate, false
+	}
+	return maxPriority, true
+}
+
+func (t *UnionType) getMutationPrio(target *Target, arg Arg, ignoreSpecial bool) (prio float64, stopRecursion bool) {
+	if target.SpecialTypes[t.Name()] == nil && len(t.Fields) == 1 || ignoreSpecial {
+		return dontMutate, false
+	}
+	// For a non-special type union with more than one option
+	// we mutate the union itself and also the value of the current option.
+	if target.SpecialTypes[t.Name()] == nil {
+		return maxPriority, false
+	}
+	return maxPriority, true
+}
+
+func (t *FlagsType) getMutationPrio(target *Target, arg Arg, ignoreSpecial bool) (prio float64, stopRecursion bool) {
+	prio = rangeSizePrio(uint64(len(t.Vals)))
+	if t.BitMask {
+		// We want a higher priority because the mutation will include
+		// more possible operations (bitwise operations).
+		prio += 0.1 * maxPriority
+	}
+	return prio, false
+}
+
+// Assigns a priority based on the range size.
+func rangeSizePrio(size uint64) (prio float64) {
+	switch size {
+	case 0:
+		prio = dontMutate
+	case 1:
+		prio = minPriority
+	default:
+		// Priority proportional with the number of values. After a threshold, the priority is constant.
+		// The threshold is 15 because most of the calls have <= 15 possible values for a flag.
+		prio = math.Min(float64(size)/3+0.4*maxPriority, 0.9*maxPriority)
+	}
+	return prio
+}
+
+func (t *PtrType) getMutationPrio(target *Target, arg Arg, ignoreSpecial bool) (prio float64, stopRecursion bool) {
+	if arg.(*PointerArg).IsSpecial() {
+		// TODO: we ought to mutate this, but we don't have code for this yet.
+		return dontMutate, false
+	}
+	return 0.3 * maxPriority, false
+}
+
+func (t *ConstType) getMutationPrio(target *Target, arg Arg, ignoreSpecial bool) (prio float64, stopRecursion bool) {
+	return dontMutate, false
+}
+
+func (t *CsumType) getMutationPrio(target *Target, arg Arg, ignoreSpecial bool) (prio float64, stopRecursion bool) {
+	return dontMutate, false
+}
+
+func (t *ProcType) getMutationPrio(target *Target, arg Arg, ignoreSpecial bool) (prio float64, stopRecursion bool) {
+	return 0.5 * maxPriority, false
+}
+
+func (t *ResourceType) getMutationPrio(target *Target, arg Arg, ignoreSpecial bool) (prio float64, stopRecursion bool) {
+	return 0.5 * maxPriority, false
+}
+
+func (t *VmaType) getMutationPrio(target *Target, arg Arg, ignoreSpecial bool) (prio float64, stopRecursion bool) {
+	return 0.5 * maxPriority, false
+}
+
+func (t *LenType) getMutationPrio(target *Target, arg Arg, ignoreSpecial bool) (prio float64, stopRecursion bool) {
+	return 0.6 * maxPriority, false
+}
+
+func (t *BufferType) getMutationPrio(target *Target, arg Arg, ignoreSpecial bool) (prio float64, stopRecursion bool) {
+	return 0.8 * maxPriority, false
+}
+
+func (t *ArrayType) getMutationPrio(target *Target, arg Arg, ignoreSpecial bool) (prio float64, stopRecursion bool) {
+	if t.Kind == ArrayRangeLen && t.RangeBegin == t.RangeEnd {
+		return dontMutate, false
+	}
+	return maxPriority, false
 }
 
 func mutateData(r *randGen, data []byte, minLen, maxLen uint64) []byte {
diff --git a/prog/mutation_test.go b/prog/mutation_test.go
index b041f02d2..9ae807e0a 100644
--- a/prog/mutation_test.go
+++ b/prog/mutation_test.go
@@ -6,6 +6,7 @@ package prog
 import (
 	"bytes"
 	"fmt"
+	"math"
 	"math/rand"
 	"sync"
 	"testing"
@@ -22,7 +23,6 @@ func TestMutationFlags(t *testing.T) {
 			`r0 = mutate$flags2(&(0x7f0000000000)="2e2f66696c653000", 0x0)`,
 			`r0 = mutate$flags2(&(0x7f0000000000)="2e2f66696c653000", 0xd9)`,
 		},
-
 		// Mutate flags (bitmask = false).
 		{
 			`r0 = mutate$flags3(&(0x7f0000000000)="2e2f66696c653000", 0x0)`,
@@ -33,10 +33,157 @@ func TestMutationFlags(t *testing.T) {
 			`r0 = mutate$flags3(&(0x7f0000000000)="2e2f66696c653000", 0xaaaaaaaaaaaaaaaa)`,
 		},
 	}
+	runMutationTests(t, tests)
+}
 
+func TestChooseCall(t *testing.T) {
+	tests := [][2]string{
+		// The call with many arguments has a higher mutation probability.
+		{
+			`mutate0()
+mutate$integer(0x0, 0x1, 0x1, 0x1, 0x0, 0x1, 0x0, 0x0, 0x1)
+mutate$integer2(0x00, 0x00, 0x20, 0x00, 0x01)`,
+			`mutate0()
+mutate$integer(0x0, 0x1, 0x1, 0x1, 0x0, 0x1, 0x0, 0x0, 0xffffffff)
+mutate$integer2(0x00, 0x00, 0x20, 0x00, 0x01)`,
+		},
+		// Calls with the same probability.
+		{
+			`mutate$integer(0x0, 0x1, 0x1, 0x1, 0x0, 0x1, 0x0, 0x0, 0x1)
+mutate$integer(0x0, 0x1, 0x1, 0x1, 0x0, 0x1, 0x0, 0x0, 0x1)
+mutate$integer(0x0, 0x1, 0x1, 0x1, 0x0, 0x1, 0x0, 0x0, 0x1)
+mutate$integer(0x0, 0x1, 0x1, 0x1, 0x0, 0x1, 0x0, 0x0, 0x1)
+mutate$integer(0x0, 0x1, 0x1, 0x1, 0x0, 0x1, 0x0, 0x0, 0x1)`,
+			`mutate$integer(0x0, 0x1, 0x1, 0x1, 0x0, 0x1, 0x0, 0x0, 0x1)
+mutate$integer(0x0, 0x1, 0x1, 0x1, 0x0, 0x1, 0x0, 0x0, 0x1)
+mutate$integer(0x0, 0x1, 0x1, 0x1, 0x0, 0x1, 0x0, 0x0, 0xff)
+mutate$integer(0x0, 0x1, 0x1, 0x1, 0x0, 0x1, 0x0, 0x0, 0x1)
+mutate$integer(0x0, 0x1, 0x1, 0x1, 0x0, 0x1, 0x0, 0x0, 0x1)`,
+		},
+		// The call with a lower probability can be mutated.
+		{
+			`mutate7(&(0x7f0000000000)='123', 0x3)
+mutate$integer(0x0, 0x1, 0x1, 0x1, 0x0, 0x1, 0x0, 0x0, 0x1)
+r0 = mutate$flags(&(0x7f0000000000)="2e2f66696c653000", 0x0, 0x1, 0x1)`,
+			`mutate7(&(0x7f0000000000)='123', 0x2)
+mutate$integer(0x0, 0x1, 0x1, 0x1, 0x0, 0x1, 0x0, 0x0, 0x1)
+r0 = mutate$flags(&(0x7f0000000000)="2e2f66696c653000", 0x0, 0x1, 0x1)`,
+		},
+		// Complex arguments.
+		{
+			`test$struct(&(0x7f0000000000)={0x0, {0x0}})
+test$array0(&(0x7f0000001000)={0x1, [@f0=0x2, @f1=0x3], 0x4})
+mutate$integer(0x0, 0x1, 0x1, 0x1, 0x0, 0x1, 0x0, 0x0, 0x1)`,
+			`test$struct(&(0x7f0000000000)={0xff, {0x0}})
+test$array0(&(0x7f0000001000)={0x1, [@f0=0x2, @f1=0x3], 0x4})
+mutate$integer(0x0, 0x1, 0x1, 0x1, 0x0, 0x1, 0x0, 0x0, 0x1)`,
+		},
+	}
 	runMutationTests(t, tests)
 }
 
+func TestMutateArgument(t *testing.T) {
+	tests := [][2]string{
+		// Mutate an integer with a higher priority than the boolean arguments.
+		{
+			`mutate$integer(0x0, 0x1, 0x1, 0x1, 0x0, 0x1, 0x0, 0x0, 0x1)`,
+			`mutate$integer(0x0, 0x1, 0x1, 0x1, 0x0, 0x1, 0x0, 0x0, 0xffffffff)`,
+		},
+		// Mutate a boolean.
+		{
+			`mutate$integer(0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0)`,
+			`mutate$integer(0x0, 0x0, 0x0, 0x1, 0x0, 0x0, 0x0, 0x0, 0x0)`,
+		},
+		// Mutate flags (bitmask = true).
+		{
+			`r0 = mutate$flags(&(0x7f0000000000)="2e2f66696c653000", 0x0, 0x1, 0x1)`,
+			`r0 = mutate$flags(&(0x7f0000000000)="2e2f66696c653000", 0x20, 0x1, 0x9)`,
+		},
+		// Mutate an int8 from a set of other arguments with higher priority.
+		{
+			`mutate$integer2(0x00, 0x00, 0x20, 0x00, 0x01)`,
+			`mutate$integer2(0x00, 0x00, 0x20, 0x00, 0x07)`,
+		},
+		// Mutate an array of structs
+		{
+			`mutate$array2(&(0x7f0000000000)=[{0x0}, {0x0}, {0x0}, {0x0}, {0x0}])`,
+			`mutate$array2(&(0x7f0000000000)=[{0x0}, {0x0}, {0x3}, {0x0}, {0x0}])`,
+		},
+		// Mutate a non-special union that have more than 1 option
+		{
+			`mutate$union(&(0x7f0000000000)=@f1=[0x0, 0x1, 0x2, 0x3, 0x0, 0x1, 0x2, 0x3, 0x0, 0x0])`,
+			`mutate$union(&(0x7f0000000000)=@f0=0x2)`,
+		},
+		// Mutate the value of the current option in union
+		{
+			`mutate$union(&(0x7f0000000000)=@f1=[0x0, 0x1, 0x2, 0x3, 0x0, 0x1, 0x2, 0x3, 0x0, 0x0])`,
+			`mutate$union(&(0x7f0000000000)=@f1=[0x0, 0x1, 0xff, 0x3, 0x0, 0x1, 0x2, 0x3, 0x0, 0x0])`,
+		},
+	}
+
+	target := initTargetTest(t, "test", "64")
+	for ti, test := range tests {
+		test := test
+		t.Run(fmt.Sprint(ti), func(t *testing.T) {
+			t.Parallel()
+			rs, ct, p, goal, err := buildTestContext(test, target)
+			if err != nil {
+				t.Fatalf("failed to deserialize the program: %v", err)
+			}
+			want := goal.Serialize()
+			for i := 0; i < 1e5; i++ {
+				p1 := p.Clone()
+				ctx := &mutator{
+					p:      p1,
+					r:      newRand(p1.Target, rs),
+					ncalls: 0,
+					ct:     ct,
+					corpus: nil,
+				}
+				ctx.mutateArg()
+				data1 := p1.Serialize()
+				if bytes.Equal(want, data1) {
+					t.Logf("success on iter %v", i)
+					return
+				}
+			}
+			t.Fatalf("failed to achieve goal, original:%s\ngoal:%s", test[0], test[1])
+		})
+	}
+}
+
+func TestRandomChoice(t *testing.T) {
+	t.Parallel()
+	target, err := GetTarget("test", "64")
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	r := newRand(target, randSource(t))
+	priorities := []float64{1, 1, 1, 1, 1, 1, 1, 1, 2}
+
+	const (
+		maxIters    = 100000
+		searchedIdx = 8
+		prob        = 0.2
+		eps         = 0.01
+	)
+
+	var index, count int
+	for i := 0; i < maxIters; i++ {
+		index = randomChoice(priorities, r)
+
+		if index == searchedIdx {
+			count++
+		}
+	}
+
+	diff := math.Abs(prob*maxIters - float64(count))
+	if diff > eps*maxIters {
+		t.Fatalf("The difference (%f) is higher than %f%%", diff, eps*100)
+	}
+}
+
 func TestClone(t *testing.T) {
 	target, rs, iters := initTest(t)
 	for i := 0; i < iters; i++ {
@@ -157,6 +304,12 @@ mutate5(&(0x7f0000001000)="2e2f66696c653000", 0x22c0)
 mutate5(&(0x7f0000001000)="2e2f66696c653000", 0x22c0)
 mutate5(&(0x7f0000001000)="2e2f66696c653100", 0x22c0)
 `},
+		// Mutate the array.
+		{`
+mutate$array(0x1, 0x30, &(0x7f0000000000)=[0x1, 0x1, 0x1, 0x1, 0x1])
+`, `
+mutate$array(0x1, 0x30, &(0x7f0000000000)=[0x1, 0x1, 0x1, 0x1])
+`},
 		// Extend an array.
 		{`
 mutate3(&(0x7f0000000000)=[0x1, 0x1], 0x2)
@@ -180,45 +333,6 @@ mutate8(0xffffffffffffffff)
 	runMutationTests(t, tests)
 }
 
-func runMutationTests(t *testing.T, tests [][2]string) {
-	target := initTargetTest(t, "test", "64")
-
-	for ti, test := range tests {
-		test := test
-		t.Run(fmt.Sprint(ti), func(t *testing.T) {
-			t.Parallel()
-			p, err := target.Deserialize([]byte(test[0]), Strict)
-			if err != nil {
-				t.Fatalf("failed to deserialize original program: %v", err)
-			}
-			goal, err := target.Deserialize([]byte(test[1]), Strict)
-			if err != nil {
-				t.Fatalf("failed to deserialize goal program: %v", err)
-			}
-			want := goal.Serialize()
-			enabled := make(map[*Syscall]bool)
-			for _, c := range p.Calls {
-				enabled[c.Meta] = true
-			}
-			for _, c := range goal.Calls {
-				enabled[c.Meta] = true
-			}
-			ct := target.BuildChoiceTable(nil, enabled)
-			rs := rand.NewSource(0)
-			for i := 0; i < 1e5; i++ {
-				p1 := p.Clone()
-				p1.Mutate(rs, len(goal.Calls), ct, nil)
-				data1 := p1.Serialize()
-				if bytes.Equal(want, data1) {
-					t.Logf("success on iter %v", i)
-					return
-				}
-			}
-			t.Fatalf("failed to achieve goal, original:%s\ngoal:%s", test[0], test[1])
-		})
-	}
-}
-
 func BenchmarkMutate(b *testing.B) {
 	target, cleanup := initBench(b)
 	defer cleanup()
@@ -259,3 +373,49 @@ func linuxAmd64ChoiceTable(target *Target) *ChoiceTable {
 	})
 	return linuxCT
 }
+
+func runMutationTests(t *testing.T, tests [][2]string) {
+	target := initTargetTest(t, "test", "64")
+	for ti, test := range tests {
+		test := test
+		t.Run(fmt.Sprint(ti), func(t *testing.T) {
+			t.Parallel()
+			rs, ct, p, goal, err := buildTestContext(test, target)
+			if err != nil {
+				t.Fatalf("failed to deserialize the program: %v", err)
+			}
+			want := goal.Serialize()
+			for i := 0; i < 1e5; i++ {
+				p1 := p.Clone()
+				p1.Mutate(rs, len(goal.Calls), ct, nil)
+				data1 := p1.Serialize()
+				if bytes.Equal(want, data1) {
+					t.Logf("success on iter %v", i)
+					return
+				}
+			}
+			t.Fatalf("failed to achieve goal, original:%s\ngoal:%s", test[0], test[1])
+		})
+	}
+}
+
+func buildTestContext(test [2]string, target *Target) (rs rand.Source, ct *ChoiceTable, p, goal *Prog, err error) {
+	p, err = target.Deserialize([]byte(test[0]), Strict)
+	if err != nil {
+		return
+	}
+	goal, err = target.Deserialize([]byte(test[1]), Strict)
+	if err != nil {
+		return
+	}
+	enabled := make(map[*Syscall]bool)
+	for _, c := range p.Calls {
+		enabled[c.Meta] = true
+	}
+	for _, c := range goal.Calls {
+		enabled[c.Meta] = true
+	}
+	ct = target.BuildChoiceTable(nil, enabled)
+	rs = rand.NewSource(0)
+	return
+}
diff --git a/prog/types.go b/prog/types.go
index c36dc076c..9a58a9eee 100644
--- a/prog/types.go
+++ b/prog/types.go
@@ -65,6 +65,7 @@ type Type interface {
 	isDefaultArg(arg Arg) bool
 	generate(r *randGen, s *state) (arg Arg, calls []*Call)
 	mutate(r *randGen, s *state, arg Arg, ctx ArgCtx) (calls []*Call, retry, preserve bool)
+	getMutationPrio(target *Target, arg Arg, ignoreSpecial bool) (prio float64, stopRecursion bool)
 	minimize(ctx *minimizeArgsCtx, arg Arg, path string) bool
 }