github.com/celestiaorg/celestia-node@v0.15.0-beta.1/share/availability/light/sample.go (about)

     1  // TODO(@Wondertan): Instead of doing sampling over the coordinates do a random walk over NMT trees.
     2  package light
     3  
     4  import (
     5  	crand "crypto/rand"
     6  	"math/big"
     7  )
     8  
     9  // Sample is a point in 2D space over square.
    10  type Sample struct {
    11  	Row, Col int
    12  }
    13  
    14  // SampleSquare randomly picks *num* unique points from the given *width* square
    15  // and returns them as samples.
    16  func SampleSquare(squareWidth int, num int) ([]Sample, error) {
    17  	ss := newSquareSampler(squareWidth, num)
    18  	err := ss.generateSample(num)
    19  	if err != nil {
    20  		return nil, err
    21  	}
    22  	return ss.samples(), nil
    23  }
    24  
    25  type squareSampler struct {
    26  	squareWidth int
    27  	smpls       map[Sample]struct{}
    28  }
    29  
    30  func newSquareSampler(squareWidth int, expectedSamples int) *squareSampler {
    31  	return &squareSampler{
    32  		squareWidth: squareWidth,
    33  		smpls:       make(map[Sample]struct{}, expectedSamples),
    34  	}
    35  }
    36  
    37  // generateSample randomly picks unique point on a 2D spaces.
    38  func (ss *squareSampler) generateSample(num int) error {
    39  	if num > ss.squareWidth*ss.squareWidth {
    40  		num = ss.squareWidth
    41  	}
    42  
    43  	done := 0
    44  	for done < num {
    45  		s := Sample{
    46  			Row: randInt(ss.squareWidth),
    47  			Col: randInt(ss.squareWidth),
    48  		}
    49  
    50  		if _, ok := ss.smpls[s]; ok {
    51  			continue
    52  		}
    53  
    54  		done++
    55  		ss.smpls[s] = struct{}{}
    56  	}
    57  
    58  	return nil
    59  }
    60  
    61  func (ss *squareSampler) samples() []Sample {
    62  	samples := make([]Sample, 0, len(ss.smpls))
    63  	for s := range ss.smpls {
    64  		samples = append(samples, s)
    65  	}
    66  	return samples
    67  }
    68  
    69  func randInt(max int) int {
    70  	n, err := crand.Int(crand.Reader, big.NewInt(int64(max)))
    71  	if err != nil {
    72  		panic(err) // won't panic as rand.Reader is endless
    73  	}
    74  
    75  	return int(n.Int64())
    76  }