github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/internal/sqlsmith/random.go (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  func (s *Smither) coin() bool {
    14  	return s.rnd.Intn(2) == 0
    15  }
    16  
    17  func (s *Smither) d6() int {
    18  	return s.rnd.Intn(6) + 1
    19  }
    20  
    21  func (s *Smither) d9() int {
    22  	return s.rnd.Intn(9) + 1
    23  }
    24  
    25  func (s *Smither) d100() int {
    26  	return s.rnd.Intn(100) + 1
    27  }
    28  
    29  // geom returns a sample from a geometric distribution whose mean is 1 / (1 -
    30  // p). Its return value is >= 1. p must satisfy 0 < p < 1. For example, pass
    31  // .5 to this function (whose mean would thus be 2), and 50% of the time this
    32  // function will return 1, 25% will return 2, 12.5% return 3, 6.25% will return
    33  // 4, etc. See: https://en.wikipedia.org/wiki/Geometric_distribution.
    34  func (s *Smither) geom(p float64) int {
    35  	if p <= 0 || p >= 1 {
    36  		panic("bad p")
    37  	}
    38  	count := 1
    39  	for s.rnd.Float64() < p {
    40  		count++
    41  	}
    42  	return count
    43  }
    44  
    45  // sample invokes fn mean number of times (but at most n times) with a
    46  // geometric distribution. The i argument to fn will be unique each time and
    47  // randomly chosen to be between 0 and n-1, inclusive. This can be used to pick
    48  // on average mean samples from a list. If n is <= 0, fn is never invoked.
    49  func (s *Smither) sample(n, mean int, fn func(i int)) {
    50  	if n <= 0 {
    51  		return
    52  	}
    53  	perms := s.rnd.Perm(n)
    54  	m := float64(mean)
    55  	p := (m - 1) / m
    56  	k := s.geom(p)
    57  	if k > n {
    58  		k = n
    59  	}
    60  	for ki := 0; ki < k; ki++ {
    61  		fn(perms[ki])
    62  	}
    63  }
    64  
    65  const letters = "abcdefghijklmnopqrstuvwxyz"
    66  
    67  // randString generates a random string with a target length using characters
    68  // from the input alphabet string.
    69  func (s *Smither) randString(length int, alphabet string) string {
    70  	buf := make([]byte, length)
    71  	for i := range buf {
    72  		buf[i] = alphabet[s.rnd.Intn(len(alphabet))]
    73  	}
    74  	return string(buf)
    75  }