pkg/corpus: track the total coverage

2 files changed, 35 insertions, 2 deletions

diff --git a/pkg/corpus/corpus.go b/pkg/corpus/corpus.go
index 1d14ba026..66ac4d4bb 100644
--- a/pkg/corpus/corpus.go
+++ b/pkg/corpus/corpus.go
@@ -20,7 +20,8 @@ type Corpus struct {
 	ctx     context.Context
 	mu      sync.RWMutex
 	progs   map[string]*Item
-	signal  signal.Signal // signal of inputs in corpus
+	signal  signal.Signal // total signal of all items
+	cover   cover.Cover   // total coverage of all items
 	updates chan<- NewItemEvent
 	ProgramsList
 }
@@ -104,6 +105,7 @@ type NewItemEvent struct {
 	Sig      string
 	Exists   bool
 	ProgData []byte
+	NewCover []uint32
 }
 
 func (corpus *Corpus) Save(inp NewInput) {
@@ -154,6 +156,7 @@ func (corpus *Corpus) Save(inp NewInput) {
 		corpus.saveProgram(inp.Prog, inp.Signal)
 	}
 	corpus.signal.Merge(inp.Signal)
+	newCover := corpus.cover.MergeDiff(inp.Cover)
 	if corpus.updates != nil {
 		select {
 		case <-corpus.ctx.Done():
@@ -161,6 +164,7 @@ func (corpus *Corpus) Save(inp NewInput) {
 			Sig:      sig,
 			Exists:   exists,
 			ProgData: progData,
+			NewCover: newCover,
 		}:
 		}
 	}
@@ -198,6 +202,7 @@ func (corpus *Corpus) Item(sig string) *Item {
 type Stat struct {
 	Progs  int
 	Signal int
+	Cover  int
 }
 
 func (corpus *Corpus) Stat() Stat {
@@ -206,6 +211,7 @@ func (corpus *Corpus) Stat() Stat {
 	return Stat{
 		Progs:  len(corpus.progs),
 		Signal: len(corpus.signal),
+		Cover:  len(corpus.cover),
 	}
 }
 
diff --git a/pkg/corpus/corpus_test.go b/pkg/corpus/corpus_test.go
index cce537087..04768838a 100644
--- a/pkg/corpus/corpus_test.go
+++ b/pkg/corpus/corpus_test.go
@@ -48,11 +48,38 @@ func TestCorpusOperation(t *testing.T) {
 	}
 
 	// Verify the total signal.
-	assert.Len(t, corpus.Signal(), 5)
+	stat := corpus.Stat()
+	assert.Equal(t, Stat{
+		Signal: 5,
+		Cover:  0,
+		Progs:  2,
+	}, stat)
 
 	corpus.Minimize(true)
 }
 
+func TestCorpusCoverage(t *testing.T) {
+	target := getTarget(t, targets.TestOS, targets.TestArch64)
+	ch := make(chan NewItemEvent)
+	corpus := NewMonitoredCorpus(context.Background(), ch)
+	rs := rand.NewSource(0)
+
+	inp := generateInput(target, rs, 5, 5)
+	inp.Cover = []uint32{10, 11}
+	go corpus.Save(inp)
+	event := <-ch
+	assert.Equal(t, []uint32{10, 11}, event.NewCover)
+
+	inp.Call = 1
+	inp.Cover = []uint32{11, 12}
+	go corpus.Save(inp)
+	event = <-ch
+	assert.Equal(t, []uint32{12}, event.NewCover)
+
+	// Check the total corpus size.
+	assert.Equal(t, corpus.Stat().Cover, 3)
+}
+
 func TestCorpusSaveConcurrency(t *testing.T) {
 	target := getTarget(t, targets.TestOS, targets.TestArch64)
 	corpus := NewCorpus(context.Background())