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             :         "sync/atomic"
      20             : 
      21             :         "golang.org/x/exp/rand"
      22             : )
      23             : 
      24             : // Uniform is a random number generator that generates draws from a uniform
      25             : // distribution.
      26             : type Uniform struct {
      27             :         min uint64
      28             :         max atomic.Uint64
      29             : }
      30             : 
      31             : // NewUniform constructs a new Uniform generator with the given
      32             : // parameters. Returns an error if the parameters are outside the accepted
      33             : // range.
      34           0 : func NewUniform(min, max uint64) *Uniform {
      35           0 :         u := &Uniform{min: min}
      36           0 :         u.max.Store(max)
      37           0 :         return u
      38           0 : }
      39             : 
      40             : // IncMax increments max.
      41           0 : func (g *Uniform) IncMax(delta int) {
      42           0 :         g.max.Add(uint64(delta))
      43           0 : }
      44             : 
      45             : // Max returns the max value of the distribution.
      46           0 : func (g *Uniform) Max() uint64 {
      47           0 :         return g.max.Load()
      48           0 : }
      49             : 
      50             : // Uint64 returns a random Uint64 between min and max, drawn from a uniform
      51             : // distribution.
      52           0 : func (g *Uniform) Uint64(rng *rand.Rand) uint64 {
      53           0 :         return rng.Uint64n(g.Max()-g.min+1) + g.min
      54           0 : }