prog: use type size when generating/mutating ints

Update #1381

prog/hints_test.go

@@ -386,11 +386,11 @@ func TestHintsRandom(t *testing.T) {
 		for i, c := range p.Calls {
 			vals := extractValues(c)
 			for j := 0; j < 5; j++ {
-				vals[r.randInt()] = true
+				vals[r.randInt64()] = true
 			}
 			comps := make(CompMap)
 			for v := range vals {
-				comps.AddComp(v, r.randInt())
+				comps.AddComp(v, r.randInt64())
 			}
 			p.MutateWithHints(i, comps, func(p1 *Prog) {})
 		}
@@ -489,11 +489,11 @@ func BenchmarkHints(b *testing.B) {
 	for i, c := range p.Calls {
 		vals := extractValues(c)
 		for j := 0; j < 5; j++ {
-			vals[r.randInt()] = true
+			vals[r.randInt64()] = true
 		}
 		comps[i] = make(CompMap)
 		for v := range vals {
-			comps[i].AddComp(v, r.randInt())
+			comps[i].AddComp(v, r.randInt64())
 		}
 	}
 	b.RunParallel(func(pb *testing.PB) {

prog/mutation.go

@@ -258,15 +258,17 @@ func mutateInt(r *randGen, s *state, arg Arg) (calls []*Call, retry, preserve bo
 	if r.bin() {
 		return regenerate(r, s, arg)
 	}
+	bits := arg.Type().TypeBitSize()
 	a := arg.(*ConstArg)
 	switch {
 	case r.nOutOf(1, 3):
 		a.Val += uint64(r.Intn(4)) + 1
 	case r.nOutOf(1, 2):
 		a.Val -= uint64(r.Intn(4)) + 1
 	default:
-		a.Val ^= 1 << uint64(r.Intn(64))
+		a.Val ^= 1 << uint64(r.Intn(int(bits)))
 	}
+	a.Val = truncateToBitSize(a.Val, bits)
 	return
 }
 
@@ -706,7 +708,7 @@ var mutateDataFuncs = [...]func(r *randGen, data []byte, minLen, maxLen uint64)
 			return data, false
 		}
 		i := r.Intn(len(data) - width + 1)
-		value := r.randInt()
+		value := r.randInt64()
 		if r.oneOf(10) {
 			value = swap64(value)
 		}

prog/mutation_test.go

@@ -152,6 +152,37 @@ func TestMutateArgument(t *testing.T) {
 	}
 }
 
+func TestSizeMutateArg(t *testing.T) {
+	target, rs, iters := initRandomTargetTest(t, "test", "64")
+	r := newRand(target, rs)
+	ct := target.BuildChoiceTable(nil, nil)
+	for i := 0; i < 100; i++ {
+		p := target.Generate(rs, 10, nil)
+		for it := 0; it < iters; it++ {
+			p1 := p.Clone()
+			ctx := &mutator{
+				p:      p1,
+				r:      r,
+				ncalls: 10,
+				ct:     ct,
+				corpus: nil,
+			}
+			ctx.mutateArg()
+			ForeachArg(p.Calls[0], func(arg Arg, ctx *ArgCtx) {
+				if _, ok := arg.Type().(*IntType); !ok {
+					return
+				}
+				bits := arg.Type().TypeBitSize()
+				limit := uint64(1<<bits - 1)
+				val := arg.(*ConstArg).Val
+				if val > limit {
+					t.Fatalf("Invalid argument value: %d. (arg size: %d; max value: %d)", val, arg.Size(), limit)
+				}
+			})
+		}
+	}
+}
+
 func TestRandomChoice(t *testing.T) {
 	t.Parallel()
 	target, err := GetTarget("test", "64")

prog/rand.go

@@ -55,24 +55,44 @@ func (r *randGen) rand64() uint64 {
 	return v
 }
 
-// Some potentially interesting integers.
-var specialInts = []uint64{
-	0, 1, 31, 32, 63, 64, 127, 128,
-	129, 255, 256, 257, 511, 512,
-	1023, 1024, 1025, 2047, 2048, 4095, 4096,
-	(1 << 15) - 1, (1 << 15), (1 << 15) + 1,
-	(1 << 16) - 1, (1 << 16), (1 << 16) + 1,
-	(1 << 31) - 1, (1 << 31), (1 << 31) + 1,
-	(1 << 32) - 1, (1 << 32), (1 << 32) + 1,
-}
-
-func (r *randGen) randInt() uint64 {
+var (
+	// Some potentially interesting integers.
+	specialInts = []uint64{
+		0, 1, 31, 32, 63, 64, 127, 128,
+		129, 255, 256, 257, 511, 512,
+		1023, 1024, 1025, 2047, 2048, 4095, 4096,
+		(1 << 15) - 1, (1 << 15), (1 << 15) + 1,
+		(1 << 16) - 1, (1 << 16), (1 << 16) + 1,
+		(1 << 31) - 1, (1 << 31), (1 << 31) + 1,
+		(1 << 32) - 1, (1 << 32), (1 << 32) + 1,
+	}
+	// The indexes (exclusive) for the maximum specialInts values that fit in 1, 2, ... 8 bytes.
+	specialIntIndex [9]int
+)
+
+func init() {
+	sort.Slice(specialInts, func(i, j int) bool {
+		return specialInts[i] < specialInts[j]
+	})
+	for i := range specialIntIndex {
+		bitSize := uint64(8 * i)
+		specialIntIndex[i] = sort.Search(len(specialInts), func(i int) bool {
+			return specialInts[i]>>bitSize != 0
+		})
+	}
+}
+
+func (r *randGen) randInt64() uint64 {
+	return r.randInt(64)
+}
+
+func (r *randGen) randInt(bits uint64) uint64 {
 	v := r.rand64()
 	switch {
 	case r.nOutOf(100, 182):
 		v %= 10
-	case r.nOutOf(50, 82):
-		v = specialInts[r.Intn(len(specialInts))]
+	case bits >= 8 && r.nOutOf(50, 82):
+		v = specialInts[r.Intn(specialIntIndex[bits/8])]
 	case r.nOutOf(10, 32):
 		v %= 256
 	case r.nOutOf(10, 22):
@@ -87,14 +107,21 @@ func (r *randGen) randInt() uint64 {
 	case r.nOutOf(5, 7):
 		v = uint64(-int64(v))
 	default:
-		v <<= uint(r.Intn(63))
+		v <<= uint(r.Intn(int(bits)))
 	}
-	return v
+	return truncateToBitSize(v, bits)
+}
+
+func truncateToBitSize(v, bitSize uint64) uint64 {
+	if bitSize == 0 || bitSize > 64 {
+		panic(fmt.Sprintf("invalid bitSize value: %d", bitSize))
+	}
+	return v & uint64(1<<bitSize-1)
 }
 
-func (r *randGen) randRangeInt(begin uint64, end uint64) uint64 {
+func (r *randGen) randRangeInt(begin, end, bitSize uint64) uint64 {
 	if r.oneOf(100) {
-		return r.randInt()
+		return r.randInt(bitSize)
 	}
 	return begin + (r.Uint64() % (end - begin + 1))
 }
@@ -710,7 +737,8 @@ func (a *ConstType) generate(r *randGen, s *state) (arg Arg, calls []*Call) {
 }
 
 func (a *IntType) generate(r *randGen, s *state) (arg Arg, calls []*Call) {
-	v := r.randInt()
+	bits := a.TypeBitSize()
+	v := r.randInt(bits)
 	switch a.Kind {
 	case IntFileoff:
 		switch {
@@ -719,10 +747,10 @@ func (a *IntType) generate(r *randGen, s *state) (arg Arg, calls []*Call) {
 		case r.nOutOf(10, 11):
 			v = r.rand(100)
 		default:
-			v = r.randInt()
+			v = r.randInt(bits)
 		}
 	case IntRange:
-		v = r.randRangeInt(a.RangeBegin, a.RangeEnd)
+		v = r.randRangeInt(a.RangeBegin, a.RangeEnd, bits)
 	}
 	return MakeConstArg(a, v), nil
 }

prog/rand_test.go

@@ -70,3 +70,24 @@ func generateProg(t *testing.T, target *Target, rs rand.Source) *Prog {
 	}
 	return p
 }
+
+func TestSizeGenerateConstArg(t *testing.T) {
+	target, rs, iters := initRandomTargetTest(t, "test", "64")
+	r := newRand(target, rs)
+	for _, c := range target.Syscalls {
+		ForeachType(c, func(typ Type) {
+			if _, ok := typ.(*IntType); !ok {
+				return
+			}
+			bits := typ.TypeBitSize()
+			limit := uint64(1<<bits - 1)
+			for i := 0; i < iters; i++ {
+				newArg, _ := typ.generate(r, nil)
+				newVal := newArg.(*ConstArg).Val
+				if newVal > limit {
+					t.Fatalf("invalid generated value: %d. (arg bitsize: %d; max value: %d)", newVal, bits, limit)
+				}
+			}
+		})
+	}
+}

prog/size.go

@@ -181,9 +181,9 @@ func (r *randGen) mutateSize(arg *ConstArg, parent []Arg) bool {
 	if r.bin() {
 		// Small adjustment to trigger missed size checks.
 		if arg.Val != 0 && r.bin() {
-			arg.Val = r.randRangeInt(0, arg.Val-1)
+			arg.Val = r.randRangeInt(0, arg.Val-1, arg.Type().TypeBitSize())
 		} else {
-			arg.Val = r.randRangeInt(arg.Val+1, arg.Val+1000)
+			arg.Val = r.randRangeInt(arg.Val+1, arg.Val+1000, arg.Type().TypeBitSize())
 		}
 		return true
 	}

prog/types.go

@@ -56,6 +56,7 @@ type Type interface {
 	Optional() bool
 	Varlen() bool
 	Size() uint64
+	TypeBitSize() uint64
 	Format() BinaryFormat
 	BitfieldOffset() uint64
 	BitfieldLength() uint64
@@ -104,6 +105,10 @@ func (t *TypeCommon) Size() uint64 {
 	return t.TypeSize
 }
 
+func (t *TypeCommon) TypeBitSize() uint64 {
+	panic("cannot get the bitsize for a non-integer type")
+}
+
 func (t *TypeCommon) Varlen() bool {
 	return t.IsVarlen
 }
@@ -189,6 +194,19 @@ func (t *IntTypeCommon) Format() BinaryFormat {
 	return t.ArgFormat
 }
 
+// Returns the size in bits for integers in binary format or 64 for string-formatted integers. The return
+// value is used in computing limits and truncating other values.
+func (t *IntTypeCommon) TypeBitSize() uint64 {
+	if t.ArgFormat != FormatNative && t.ArgFormat != FormatBigEndian {
+		// TODO: add special cases for mutation and generation of string-formatted integers.
+		return 64
+	}
+	if t.BitfieldLen != 0 {
+		return t.BitfieldLen
+	}
+	return t.TypeSize * 8
+}
+
 func (t *IntTypeCommon) BitfieldOffset() uint64 {
 	return t.BitfieldOff
 }