github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/internal/sqlsmith/sampler.tmpl (about)

     1  // Copyright 2019 The Cockroach Authors.
     2  //
     3  // Use of this software is governed by the Business Source License
     4  // included in the file licenses/BSL.txt.
     5  //
     6  // As of the Change Date specified in that file, in accordance with
     7  // the Business Source License, use of this software will be governed
     8  // by the Apache License, Version 2.0, included in the file
     9  // licenses/APL.txt.
    10  
    11  package sqlsmith
    12  
    13  import (
    14  	"math/rand"
    15  
    16  	"github.com/cheekybits/genny/generic"
    17  	"github.com/cockroachdb/cockroach/pkg/util/syncutil"
    18  )
    19  
    20  type element generic.Type
    21  
    22  // elementWeight is the generic weight type.
    23  type elementWeight struct {
    24  	weight int
    25  	elem   element
    26  }
    27  
    28  // newWeightedelementSampler creates a elementSampler that produces
    29  // elements. They are returned at the relative frequency of the values of
    30  // weights. All weights must be >= 1.
    31  func newWeightedelementSampler(weights []elementWeight, seed int64) *elementSampler {
    32  	sum := 0
    33  	for _, w := range weights {
    34  		if w.weight < 1 {
    35  			panic("expected weight >= 1")
    36  		}
    37  		sum += w.weight
    38  	}
    39  	if sum == 0 {
    40  		panic("expected weights")
    41  	}
    42  	samples := make([]element, sum)
    43  	pos := 0
    44  	for _, w := range weights {
    45  		for count := 0; count < w.weight; count++ {
    46  			samples[pos] = w.elem
    47  			pos++
    48  		}
    49  	}
    50  	return &elementSampler{
    51  		rnd:     rand.New(rand.NewSource(seed)),
    52  		samples: samples,
    53  	}
    54  }
    55  
    56  type elementSampler struct {
    57  	mu      syncutil.Mutex
    58  	rnd     *rand.Rand
    59  	samples []element
    60  }
    61  
    62  func (w *elementSampler) Next() element {
    63  	w.mu.Lock()
    64  	v := w.samples[w.rnd.Intn(len(w.samples))]
    65  	w.mu.Unlock()
    66  	return v
    67  }