prog: refactor HeatMap

Provide NumMutations method instead of Size.
It allows HeatMap to choose number of mutations better
(e.g. for completely empty/flat images w/o interesting data).

3 files changed, 20 insertions, 18 deletions

diff --git a/prog/heatmap.go b/prog/heatmap.go
index 51daaf251..d1c5ea3b9 100644
--- a/prog/heatmap.go
+++ b/prog/heatmap.go
@@ -16,8 +16,8 @@ import (
 //  2. Select random indices according to the probability distribution:
 //     `idx := hm.ChooseLocation(r)`.
 type Heatmap interface {
-	ChooseLocation(r *rand.Rand) int
-	Size() int
+	NumMutations() int
+	ChooseLocation() int
 }
 
 // Generic heatmaps model a probability distribution based on sparse data,
@@ -25,27 +25,31 @@ type Heatmap interface {
 // views data as a series of chunks of length `granularity`, ignoring chunks
 // which are a single repeated byte. Indices are chosen uniformly amongst the
 // remaining "interesting" segments.
-func MakeGenericHeatmap(data []byte) Heatmap {
+func MakeGenericHeatmap(data []byte, r *rand.Rand) Heatmap {
 	if len(data) == 0 {
 		panic("cannot create a GenericHeatmap with no data")
 	}
-	var hm GenericHeatmap
+	hm := &GenericHeatmap{
+		r: r,
+	}
 	hm.length, hm.segments = calculateLengthAndSegments(data, granularity)
-	return &hm
+	return hm
+}
+
+func (hm *GenericHeatmap) NumMutations() int {
+	// At least two mutations, up to about one mutation every 128 KB of heatmap size.
+	return hm.r.Intn(hm.length/(1<<17)+1) + 2
 }
 
-func (hm *GenericHeatmap) ChooseLocation(r *rand.Rand) int {
+func (hm *GenericHeatmap) ChooseLocation() int {
 	// Uniformly choose an index within one of the segments.
-	heatmapIdx := r.Intn(hm.length)
+	heatmapIdx := hm.r.Intn(hm.length)
 	rawIdx := translateIdx(heatmapIdx, hm.segments)
 	return rawIdx
 }
 
-func (hm *GenericHeatmap) Size() int {
-	return hm.length
-}
-
 type GenericHeatmap struct {
+	r        *rand.Rand
 	segments []segment // "Interesting" parts of the data.
 	length   int       // Sum of all segment lengths.
 }
diff --git a/prog/heatmap_test.go b/prog/heatmap_test.go
index 4103a4f97..9c891f258 100644
--- a/prog/heatmap_test.go
+++ b/prog/heatmap_test.go
@@ -57,10 +57,10 @@ func TestGenericHeatmap(t *testing.T) {
 			t.Fatalf("bad decode: %v", err)
 		}
 		for i := 0; i < iters; i++ {
-			hm := MakeGenericHeatmap(data).(*GenericHeatmap)
+			hm := MakeGenericHeatmap(data, r).(*GenericHeatmap)
 
 			for j := 0; j < tries; j++ {
-				index := hm.ChooseLocation(r)
+				index := hm.ChooseLocation()
 				if !checkIndex(index, len(data), test.regions) {
 					hm.debugPrint(t, data, test.regions)
 					t.Fatalf("selected index %d does not fall in a region\n", index)
diff --git a/prog/mutation.go b/prog/mutation.go
index 881403247..96e11516c 100644
--- a/prog/mutation.go
+++ b/prog/mutation.go
@@ -401,11 +401,9 @@ func (r *randGen) mutateImage(image []byte) (data []byte, retry bool) {
 	if len(data) == 0 {
 		return image, true // Do not mutate empty data.
 	}
-	hm := MakeGenericHeatmap(data)
-	// At least two mutations, up to about one mutation every 128 KB of heatmap size.
-	numMutations := r.Intn(hm.Size()/(1<<17)+1) + 2
-	for i := 0; i < numMutations; i++ {
-		index := hm.ChooseLocation(r.Rand)
+	hm := MakeGenericHeatmap(data, r.Rand)
+	for i := hm.NumMutations(); i > 0; i-- {
+		index := hm.ChooseLocation()
 		width := 1 << uint(r.Intn(4))
 		if index+width > len(data) {
 			width = 1