Line data    Source code

       1             : // Copyright 2018 The Cockroach Authors.
       2             : //
       3             : // Licensed under the Apache License, Version 2.0 (the "License");
       4             : // you may not use this file except in compliance with the License.
       5             : // You may obtain a copy of the License at
       6             : //
       7             : //     http://www.apache.org/licenses/LICENSE-2.0
       8             : //
       9             : // Unless required by applicable law or agreed to in writing, software
      10             : // distributed under the License is distributed on an "AS IS" BASIS,
      11             : // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
      12             : // implied. See the License for the specific language governing
      13             : // permissions and limitations under the License. See the AUTHORS file
      14             : // for names of contributors.
      15             : 
      16             : package randvar
      17             : 
      18             : import (
      19             :         "math/rand/v2"
      20             :         "sync/atomic"
      21             : )
      22             : 
      23             : // Uniform is a random number generator that generates draws from a uniform
      24             : // distribution.
      25             : type Uniform struct {
      26             :         min uint64
      27             :         max atomic.Uint64
      28             : }
      29             : 
      30             : // NewUniform constructs a new Uniform generator with the given
      31             : // parameters. Returns an error if the parameters are outside the accepted
      32             : // range.
      33           1 : func NewUniform(min, max uint64) *Uniform {
      34           1 :         u := &Uniform{min: min}
      35           1 :         u.max.Store(max)
      36           1 :         return u
      37           1 : }
      38             : 
      39             : // IncMax increments max.
      40           0 : func (g *Uniform) IncMax(delta int) {
      41           0 :         g.max.Add(uint64(delta))
      42           0 : }
      43             : 
      44             : // Max returns the max value of the distribution.
      45           1 : func (g *Uniform) Max() uint64 {
      46           1 :         return g.max.Load()
      47           1 : }
      48             : 
      49             : // Uint64 returns a random Uint64 between min and max, drawn from a uniform
      50             : // distribution.
      51           1 : func (g *Uniform) Uint64(rng *rand.Rand) uint64 {
      52           1 :         return rng.Uint64N(g.Max()-g.min+1) + g.min
      53           1 : }