mirror of
				https://source.quilibrium.com/quilibrium/ceremonyclient.git
				synced 2025-11-04 03:07:28 +00:00 
			
		
		
		
	
		
			
				
	
	
		
			161 lines
		
	
	
		
			3.5 KiB
		
	
	
	
		
			Go
		
	
	
	
	
	
			
		
		
	
	
			161 lines
		
	
	
		
			3.5 KiB
		
	
	
	
		
			Go
		
	
	
	
	
	
// Copyright 2019 The LevelDB-Go and Pebble Authors. All rights reserved. Use
 | 
						|
// of this source code is governed by a BSD-style license that can be found in
 | 
						|
// the LICENSE file.
 | 
						|
 | 
						|
package randvar
 | 
						|
 | 
						|
import (
 | 
						|
	"encoding/binary"
 | 
						|
	"regexp"
 | 
						|
	"strconv"
 | 
						|
	"strings"
 | 
						|
 | 
						|
	"github.com/cockroachdb/errors"
 | 
						|
	"golang.org/x/exp/rand"
 | 
						|
)
 | 
						|
 | 
						|
var randVarRE = regexp.MustCompile(`^(?:(latest|uniform|zipf):)?(\d+)(?:-(\d+))?$`)
 | 
						|
 | 
						|
// Flag provides a command line flag interface for specifying static random
 | 
						|
// variables.
 | 
						|
type Flag struct {
 | 
						|
	Static
 | 
						|
	spec string
 | 
						|
}
 | 
						|
 | 
						|
// NewFlag creates a new Flag initialized with the specified spec.
 | 
						|
func NewFlag(spec string) *Flag {
 | 
						|
	f := &Flag{}
 | 
						|
	if err := f.Set(spec); err != nil {
 | 
						|
		panic(err)
 | 
						|
	}
 | 
						|
	return f
 | 
						|
}
 | 
						|
 | 
						|
func (f *Flag) String() string {
 | 
						|
	return f.spec
 | 
						|
}
 | 
						|
 | 
						|
// Type implements the Flag.Value interface.
 | 
						|
func (f *Flag) Type() string {
 | 
						|
	return "randvar"
 | 
						|
}
 | 
						|
 | 
						|
// Set implements the Flag.Value interface.
 | 
						|
func (f *Flag) Set(spec string) error {
 | 
						|
	m := randVarRE.FindStringSubmatch(spec)
 | 
						|
	if m == nil {
 | 
						|
		return errors.Errorf("invalid random var spec: %s", errors.Safe(spec))
 | 
						|
	}
 | 
						|
 | 
						|
	min, err := strconv.Atoi(m[2])
 | 
						|
	if err != nil {
 | 
						|
		return err
 | 
						|
	}
 | 
						|
	max := min
 | 
						|
	if m[3] != "" {
 | 
						|
		max, err = strconv.Atoi(m[3])
 | 
						|
		if err != nil {
 | 
						|
			return err
 | 
						|
		}
 | 
						|
	}
 | 
						|
 | 
						|
	switch strings.ToLower(m[1]) {
 | 
						|
	case "", "uniform":
 | 
						|
		f.Static = NewUniform(uint64(min), uint64(max))
 | 
						|
	case "latest":
 | 
						|
		f.Static, err = NewSkewedLatest(uint64(min), uint64(max), 0.99)
 | 
						|
		if err != nil {
 | 
						|
			return err
 | 
						|
		}
 | 
						|
	case "zipf":
 | 
						|
		var err error
 | 
						|
		f.Static, err = NewZipf(uint64(min), uint64(max), 0.99)
 | 
						|
		if err != nil {
 | 
						|
			return err
 | 
						|
		}
 | 
						|
	default:
 | 
						|
		return errors.Errorf("unknown random var distribution: %s", errors.Safe(m[1]))
 | 
						|
	}
 | 
						|
	f.spec = spec
 | 
						|
	return nil
 | 
						|
}
 | 
						|
 | 
						|
// BytesFlag provides a command line flag interface for specifying random
 | 
						|
// bytes. The specification provides for both the length of the random bytes
 | 
						|
// and a target compression ratio.
 | 
						|
type BytesFlag struct {
 | 
						|
	sizeFlag          Flag
 | 
						|
	targetCompression float64
 | 
						|
	spec              string
 | 
						|
}
 | 
						|
 | 
						|
// NewBytesFlag creates a new BytesFlag initialized with the specified spec.
 | 
						|
func NewBytesFlag(spec string) *BytesFlag {
 | 
						|
	f := &BytesFlag{}
 | 
						|
	if err := f.Set(spec); err != nil {
 | 
						|
		panic(err)
 | 
						|
	}
 | 
						|
	return f
 | 
						|
}
 | 
						|
 | 
						|
func (f *BytesFlag) String() string {
 | 
						|
	return f.spec
 | 
						|
}
 | 
						|
 | 
						|
// Type implements the Flag.Value interface.
 | 
						|
func (f *BytesFlag) Type() string {
 | 
						|
	return "randbytes"
 | 
						|
}
 | 
						|
 | 
						|
// Set implements the Flag.Value interface.
 | 
						|
func (f *BytesFlag) Set(spec string) error {
 | 
						|
	parts := strings.Split(spec, "/")
 | 
						|
	if len(parts) == 0 || len(parts) > 2 {
 | 
						|
		return errors.Errorf("invalid randbytes spec: %s", errors.Safe(spec))
 | 
						|
	}
 | 
						|
	if err := f.sizeFlag.Set(parts[0]); err != nil {
 | 
						|
		return err
 | 
						|
	}
 | 
						|
	f.targetCompression = 1.0
 | 
						|
	if len(parts) == 2 {
 | 
						|
		var err error
 | 
						|
		f.targetCompression, err = strconv.ParseFloat(parts[1], 64)
 | 
						|
		if err != nil {
 | 
						|
			return err
 | 
						|
		}
 | 
						|
	}
 | 
						|
	f.spec = spec
 | 
						|
	return nil
 | 
						|
}
 | 
						|
 | 
						|
// Bytes returns random bytes. The length of the random bytes comes from the
 | 
						|
// internal sizeFlag.
 | 
						|
func (f *BytesFlag) Bytes(r *rand.Rand, buf []byte) []byte {
 | 
						|
	size := int(f.sizeFlag.Uint64(r))
 | 
						|
	uniqueSize := int(float64(size) / f.targetCompression)
 | 
						|
	if uniqueSize < 1 {
 | 
						|
		uniqueSize = 1
 | 
						|
	}
 | 
						|
	if cap(buf) < size {
 | 
						|
		buf = make([]byte, size)
 | 
						|
	}
 | 
						|
	data := buf[:size]
 | 
						|
	offset := 0
 | 
						|
	for offset+8 <= uniqueSize {
 | 
						|
		binary.LittleEndian.PutUint64(data[offset:], r.Uint64())
 | 
						|
		offset += 8
 | 
						|
	}
 | 
						|
	word := r.Uint64()
 | 
						|
	for offset < uniqueSize {
 | 
						|
		data[offset] = byte(word)
 | 
						|
		word >>= 8
 | 
						|
		offset++
 | 
						|
	}
 | 
						|
	for offset < size {
 | 
						|
		data[offset] = data[offset-uniqueSize]
 | 
						|
		offset++
 | 
						|
	}
 | 
						|
	return data
 | 
						|
}
 |