diff options
Diffstat (limited to 'pkg/corpus/prio.go')
| -rw-r--r-- | pkg/corpus/prio.go | 29 |
1 files changed, 20 insertions, 9 deletions
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 { |