github.com/google/syzkaller@v0.0.0-20251211124644-a066d2bc4b02/pkg/corpus/prio.go (about)

     1  // Copyright 2024 syzkaller project authors. All rights reserved.
     2  // Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.
     3  
     4  package corpus
     5  
     6  import (
     7  	"math/rand"
     8  	"sort"
     9  
    10  	"github.com/google/syzkaller/pkg/signal"
    11  	"github.com/google/syzkaller/prog"
    12  )
    13  
    14  type ProgramsList struct {
    15  	progs    []*prog.Prog
    16  	sumPrios int64
    17  	accPrios []int64
    18  }
    19  
    20  func (pl *ProgramsList) chooseProgram(r *rand.Rand) *prog.Prog {
    21  	if len(pl.progs) == 0 {
    22  		return nil
    23  	}
    24  	randVal := r.Int63n(pl.sumPrios + 1)
    25  	idx := sort.Search(len(pl.accPrios), func(i int) bool {
    26  		return pl.accPrios[i] >= randVal
    27  	})
    28  	return pl.progs[idx]
    29  }
    30  
    31  func (pl *ProgramsList) saveProgram(p *prog.Prog, signal signal.Signal) {
    32  	prio := int64(len(signal))
    33  	if prio == 0 {
    34  		prio = 1
    35  	}
    36  	pl.sumPrios += prio
    37  	pl.accPrios = append(pl.accPrios, pl.sumPrios)
    38  	pl.progs = append(pl.progs, p)
    39  }
    40  
    41  func (corpus *Corpus) ChooseProgram(r *rand.Rand) *prog.Prog {
    42  	corpus.mu.RLock()
    43  	defer corpus.mu.RUnlock()
    44  	if len(corpus.progsMap) == 0 {
    45  		return nil
    46  	}
    47  	// We could have used an approach similar to chooseProgram(), but for small number
    48  	// of focus areas that is an overkill.
    49  	var randArea *focusAreaState
    50  	if len(corpus.focusAreas) > 0 {
    51  		sum := 0.0
    52  		nonEmpty := make([]*focusAreaState, 0, len(corpus.focusAreas))
    53  		for _, area := range corpus.focusAreas {
    54  			if len(area.progs) == 0 {
    55  				continue
    56  			}
    57  			sum += area.Weight
    58  			nonEmpty = append(nonEmpty, area)
    59  		}
    60  		val := r.Float64() * sum
    61  		currSum := 0.0
    62  		for _, area := range nonEmpty {
    63  			if val >= currSum {
    64  				randArea = area
    65  			}
    66  			currSum += area.Weight
    67  		}
    68  	}
    69  	if randArea != nil {
    70  		return randArea.chooseProgram(r)
    71  	}
    72  	return corpus.chooseProgram(r)
    73  }
    74  
    75  func (corpus *Corpus) Programs() []*prog.Prog {
    76  	corpus.mu.RLock()
    77  	defer corpus.mu.RUnlock()
    78  	return corpus.progs
    79  }