pkg/corpus: update Programs List priority storage

Static prefix sums have been replaced with a Fenwick tree.

In the current syzkaller, program priority was set based on a Signal
received by a single system call. This commit allows priority to be
changed dynamically, making it possible to maintain priority based on
Signals from all system calls.

Signed-off-by: Grigory Bazilevich <g.bazilevich@ispras.ru>

5 files changed, 171 insertions, 17 deletions

diff --git a/pkg/corpus/corpus.go b/pkg/corpus/corpus.go
index 594521fcc..54c8d96cd 100644
--- a/pkg/corpus/corpus.go
+++ b/pkg/corpus/corpus.go
@@ -131,6 +131,7 @@ func (corpus *Corpus) Save(inp NewInput) {
 	corpus.mu.Lock()
 	defer corpus.mu.Unlock()
 
+	var newItem *Item = nil
 	update := ItemUpdate{
 		Call:     inp.Call,
 		RawCover: inp.RawCover,
@@ -143,7 +144,7 @@ func (corpus *Corpus) Save(inp NewInput) {
 		var newCover cover.Cover
 		newCover.Merge(old.Cover)
 		newCover.Merge(inp.Cover)
-		newItem := &Item{
+		newItem = &Item{
 			Sig:     sig,
 			Prog:    old.Prog,
 			Call:    old.Call,
@@ -157,10 +158,8 @@ func (corpus *Corpus) Save(inp NewInput) {
 		if len(newItem.Updates) < maxUpdates {
 			newItem.Updates = append(newItem.Updates, update)
 		}
-		corpus.progsMap[sig] = newItem
-		corpus.applyFocusAreas(newItem, inp.Cover)
 	} else {
-		item := &Item{
+		newItem = &Item{
 			Sig:     sig,
 			Call:    inp.Call,
 			Prog:    inp.Prog,
@@ -169,10 +168,11 @@ func (corpus *Corpus) Save(inp NewInput) {
 			Cover:   inp.Cover,
 			Updates: []ItemUpdate{update},
 		}
-		corpus.progsMap[sig] = item
-		corpus.applyFocusAreas(item, inp.Cover)
-		corpus.saveProgram(inp.Prog, inp.Signal)
 	}
+	corpus.progsMap[sig] = newItem
+	corpus.applyFocusAreas(newItem, newItem.Cover)
+	corpus.saveProgram(inp.Prog, newItem.Signal)
+
 	corpus.signal.Merge(inp.Signal)
 	newCover := corpus.cover.MergeDiff(inp.Cover)
 	if corpus.updates != nil {
diff --git a/pkg/corpus/fenwik.go b/pkg/corpus/fenwik.go
new file mode 100644
index 000000000..3a49219ac
--- /dev/null
+++ b/pkg/corpus/fenwik.go
@@ -0,0 +1,61 @@
+package corpus
+
+type fenwikTree struct {
+	array   []int64
+	data    []int64
+	fullSum int64
+
+	size     int
+	capacity int
+}
+
+func (t *fenwikTree) addNew(val int64) {
+	if t.size+1 >= t.capacity {
+		t.extendAndRebuild()
+	}
+
+	t.size++
+	t.add(t.size, val)
+}
+
+func (t *fenwikTree) add(ind int, val int64) {
+	t.fullSum += val
+	t.array[ind] += val
+	for ind < t.capacity {
+		t.data[ind] += val
+		ind += ind & -ind
+	}
+}
+
+func (t *fenwikTree) get(ind int) int64 {
+	return t.array[ind]
+}
+
+func (t *fenwikTree) set(ind int, val int64) {
+	t.add(ind, val-t.get(ind))
+}
+
+func (t *fenwikTree) prefixSumTo(ind int) (sum int64) {
+	for ind > 0 {
+		sum += t.data[ind]
+		ind -= ind & -ind
+	}
+	return
+}
+
+func (t *fenwikTree) extendAndRebuild() {
+	if t.capacity == 0 {
+		t.capacity = 16
+	} else {
+		t.capacity *= 2
+	}
+
+	oldArray := t.array
+	t.array = make([]int64, t.capacity)
+	t.data = make([]int64, t.capacity)
+	t.fullSum = 0
+
+	for i := 1; i < len(oldArray); i++ {
+		t.add(i, oldArray[i])
+	}
+}
diff --git a/pkg/corpus/fenwik_test.go b/pkg/corpus/fenwik_test.go
new file mode 100644
index 000000000..177e21aa2
--- /dev/null
+++ b/pkg/corpus/fenwik_test.go
@@ -0,0 +1,82 @@
+package corpus
+
+import (
+	"math/rand"
+	"testing"
+)
+
+func naivePrefixSum(a []int64, r int) int64 {
+	var s int64
+	for i := 1; i <= r; i++ {
+		s += a[i]
+	}
+	return s
+}
+
+func TestFenwickBasic(t *testing.T) {
+	var ft fenwikTree
+
+	values := []int64{0, 5, -3, 7, 10}
+	for _, v := range values[1:] {
+		ft.addNew(v)
+	}
+
+	if ft.size != len(values)-1 {
+		t.Fatalf("size mismatch: got %d want %d", ft.size, len(values)-1)
+	}
+
+	for i := 1; i < ft.size; i++ {
+		got := ft.prefixSumTo(i)
+		want := naivePrefixSum(values, i)
+		if got != want {
+			t.Fatalf("prefixSumTo(%d) = %d, want %d", i, got, want)
+		}
+	}
+}
+
+func TestFenwickExtend(t *testing.T) {
+	var ft fenwikTree
+
+	n := 1000
+	ref := make([]int64, 1)
+
+	for i := 1; i <= n; i++ {
+		ft.addNew(int64(i))
+		ref = append(ref, int64(i))
+
+		if ft.prefixSumTo(i) != naivePrefixSum(ref, i) {
+			t.Fatalf("extend failed at i=%d, fw=%d and naive=%d", i, ft.prefixSumTo(i), naivePrefixSum(ref, i))
+		}
+	}
+}
+
+func TestFenwickRandom(t *testing.T) {
+	var ft fenwikTree
+	ft.extendAndRebuild()
+
+	ref := make([]int64, 1)
+
+	rng := rand.New(rand.NewSource(1))
+
+	for step := range 10_000 {
+		if rng.Intn(2) == 0 || ft.size <= 2 {
+			v := int64(rng.Intn(1000) - 500)
+			ft.addNew(v)
+			ref = append(ref, v)
+		} else {
+			i := rng.Intn(ft.size-1) + 1
+			v := int64(rng.Intn(1000) - 500)
+			ft.add(i, v)
+			ref[i] += v
+		}
+
+		i := rng.Intn(ft.size)
+		if ft.prefixSumTo(i) != naivePrefixSum(ref, i) {
+			t.Fatalf("mismatch at step %d, i=%d", step, i)
+		}
+
+		if ft.prefixSumTo(ft.size) != ft.fullSum {
+			t.Fatalf("fullSum mismatch %d, sm=%d, fullSum=%d, naive=%d", step, ft.prefixSumTo(ft.size), ft.fullSum, naivePrefixSum(ref, ft.size))
+		}
+	}
+}
diff --git a/pkg/corpus/prio.go b/pkg/corpus/prio.go
index f23466df6..4d0053850 100644
--- a/pkg/corpus/prio.go
+++ b/pkg/corpus/prio.go
@@ -12,30 +12,41 @@ import (
 )
 
 type ProgramsList struct {
-	progs    []*prog.Prog
-	sumPrios int64
-	accPrios []int64
+	progs     []*prog.Prog
+	progToInd map[*prog.Prog]int
+
+	prios fenwikTree
 }
 
 func (pl *ProgramsList) chooseProgram(r *rand.Rand) *prog.Prog {
 	if len(pl.progs) == 0 {
 		return nil
 	}
-	randVal := r.Int63n(pl.sumPrios + 1)
-	idx := sort.Search(len(pl.accPrios), func(i int) bool {
-		return pl.accPrios[i] >= randVal
+	randVal := r.Int63n(pl.prios.fullSum + 1)
+	idx := sort.Search(len(pl.progs), func(i int) bool {
+		return pl.prios.prefixSumTo(i+1) >= randVal
 	})
 	return pl.progs[idx]
 }
 
 func (pl *ProgramsList) saveProgram(p *prog.Prog, signal signal.Signal) {
+	if len(pl.progs) == 0 {
+		pl.progToInd = make(map[*prog.Prog]int)
+	}
+
 	prio := int64(len(signal))
 	if prio == 0 {
 		prio = 1
 	}
-	pl.sumPrios += prio
-	pl.accPrios = append(pl.accPrios, pl.sumPrios)
-	pl.progs = append(pl.progs, p)
+
+	if ind, ok := pl.progToInd[p]; ok {
+		pl.prios.set(ind, prio)
+	} else {
+		pl.progs = append(pl.progs, p)
+		pl.progToInd[p] = len(pl.progs)
+
+		pl.prios.addNew(prio)
+	}
 }
 
 func (corpus *Corpus) ChooseProgram(r *rand.Rand) *prog.Prog {
diff --git a/pkg/corpus/prio_test.go b/pkg/corpus/prio_test.go
index 71507c2b1..5f327896e 100644
--- a/pkg/corpus/prio_test.go
+++ b/pkg/corpus/prio_test.go
@@ -41,7 +41,7 @@ func TestChooseProgram(t *testing.T) {
 		counters[corpus.chooseProgram(r)]++
 	}
 	for p, prio := range priorities {
-		prob := float64(prio) / float64(corpus.sumPrios)
+		prob := float64(prio) / float64(corpus.prios.fullSum)
 		diff := math.Abs(prob*maxIters - float64(counters[p]))
 		if diff > eps*maxIters {
 			t.Fatalf("the difference (%f) is higher than %f%%", diff, eps*100)