github.com/square/finch@v0.0.0-20240412205204-6530c03e2b96/data/string.go (about)

     1  // Copyright 2024 Block, Inc.
     2  
     3  package data
     4  
     5  import (
     6  	"fmt"
     7  	"math/rand"
     8  	"strings"
     9  	"time"
    10  )
    11  
    12  // StrFillAz implemnts the str-fill-az data generator.
    13  type StrFillAz struct {
    14  	len int64
    15  	src rand.Source
    16  }
    17  
    18  var _ Generator = &StrFillAz{}
    19  
    20  // https://stackoverflow.com/questions/22892120/how-to-generate-a-random-string-of-a-fixed-length-in-go
    21  const letterBytes = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ"
    22  const (
    23  	letterIdxBits = 6                    // 6 bits to represent a letter index
    24  	letterIdxMask = 1<<letterIdxBits - 1 // All 1-bits, as many as letterIdxBits
    25  	letterIdxMax  = 63 / letterIdxBits   // # of letter indices fitting in 63 bits
    26  )
    27  
    28  func NewStrFillAz(params map[string]string) (*StrFillAz, error) {
    29  	g := &StrFillAz{
    30  		len: 100,
    31  		src: rand.NewSource(time.Now().UnixNano()),
    32  	}
    33  	if err := int64From(params, "len", &g.len, false); err != nil {
    34  		return nil, err
    35  	}
    36  	if g.len <= 0 {
    37  		return nil, fmt.Errorf("stra-az param len must be >= 1")
    38  	}
    39  	return g, nil
    40  }
    41  
    42  func (g *StrFillAz) Name() string               { return "str-fill-az" }
    43  func (g *StrFillAz) Format() (uint, string)     { return 1, "'%s'" }
    44  func (g *StrFillAz) Scan(any interface{}) error { return nil }
    45  
    46  func (g *StrFillAz) Copy() Generator {
    47  	return &StrFillAz{
    48  		len: g.len,
    49  		src: rand.NewSource(time.Now().UnixNano()),
    50  	}
    51  }
    52  
    53  func (g *StrFillAz) Values(_ RunCount) []interface{} {
    54  	sb := strings.Builder{}
    55  	sb.Grow(int(g.len))
    56  	// A src.Int63() generates 63 random bits, enough for letterIdxMax characters!
    57  	for i, cache, remain := g.len-1, g.src.Int63(), letterIdxMax; i >= 0; {
    58  		if remain == 0 {
    59  			cache, remain = g.src.Int63(), letterIdxMax
    60  		}
    61  		if idx := int(cache & letterIdxMask); idx < len(letterBytes) {
    62  			sb.WriteByte(letterBytes[idx])
    63  			i--
    64  		}
    65  		cache >>= letterIdxBits
    66  		remain--
    67  	}
    68  	return []interface{}{sb.String()}
    69  }