github.com/vipernet-xyz/tm@v0.34.24/test/e2e/generator/random.go (about)

     1  package main
     2  
     3  import (
     4  	"math/rand"
     5  	"sort"
     6  )
     7  
     8  // combinations takes input in the form of a map of item lists, and returns a
     9  // list of all combinations of each item for each key. E.g.:
    10  //
    11  // {"foo": [1, 2, 3], "bar": [4, 5, 6]}
    12  //
    13  // Will return the following maps:
    14  //
    15  // {"foo": 1, "bar": 4}
    16  // {"foo": 1, "bar": 5}
    17  // {"foo": 1, "bar": 6}
    18  // {"foo": 2, "bar": 4}
    19  // {"foo": 2, "bar": 5}
    20  // {"foo": 2, "bar": 6}
    21  // {"foo": 3, "bar": 4}
    22  // {"foo": 3, "bar": 5}
    23  // {"foo": 3, "bar": 6}
    24  func combinations(items map[string][]interface{}) []map[string]interface{} {
    25  	keys := []string{}
    26  	for key := range items {
    27  		keys = append(keys, key)
    28  	}
    29  	sort.Strings(keys)
    30  	return combiner(map[string]interface{}{}, keys, items)
    31  }
    32  
    33  // combiner is a utility function for combinations.
    34  func combiner(head map[string]interface{}, pending []string, items map[string][]interface{}) []map[string]interface{} {
    35  	if len(pending) == 0 {
    36  		return []map[string]interface{}{head}
    37  	}
    38  	key, pending := pending[0], pending[1:]
    39  
    40  	result := []map[string]interface{}{}
    41  	for _, value := range items[key] {
    42  		path := map[string]interface{}{}
    43  		for k, v := range head {
    44  			path[k] = v
    45  		}
    46  		path[key] = value
    47  		result = append(result, combiner(path, pending, items)...)
    48  	}
    49  	return result
    50  }
    51  
    52  // uniformChoice chooses a single random item from the argument list, uniformly weighted.
    53  type uniformChoice []interface{}
    54  
    55  func (uc uniformChoice) Choose(r *rand.Rand) interface{} {
    56  	return uc[r.Intn(len(uc))]
    57  }
    58  
    59  // weightedChoice chooses a single random key from a map of keys and weights.
    60  type weightedChoice map[interface{}]uint
    61  
    62  func (wc weightedChoice) Choose(r *rand.Rand) interface{} {
    63  	total := 0
    64  	choices := make([]interface{}, 0, len(wc))
    65  	for choice, weight := range wc {
    66  		total += int(weight)
    67  		choices = append(choices, choice)
    68  	}
    69  
    70  	rem := r.Intn(total)
    71  	for _, choice := range choices {
    72  		rem -= int(wc[choice])
    73  		if rem <= 0 {
    74  			return choice
    75  		}
    76  	}
    77  
    78  	return nil
    79  }
    80  
    81  // probSetChoice picks a set of strings based on each string's probability (0-1).
    82  type probSetChoice map[string]float64
    83  
    84  func (pc probSetChoice) Choose(r *rand.Rand) []string {
    85  	choices := []string{}
    86  	for item, prob := range pc {
    87  		if r.Float64() <= prob {
    88  			choices = append(choices, item)
    89  		}
    90  	}
    91  	return choices
    92  }
    93  
    94  // uniformSetChoice picks a set of strings with uniform probability, picking at least one.
    95  type uniformSetChoice []string
    96  
    97  func (usc uniformSetChoice) Choose(r *rand.Rand) []string {
    98  	choices := []string{}
    99  	indexes := r.Perm(len(usc))
   100  	if len(indexes) > 1 {
   101  		indexes = indexes[:1+r.Intn(len(indexes)-1)]
   102  	}
   103  	for _, i := range indexes {
   104  		choices = append(choices, usc[i])
   105  	}
   106  	return choices
   107  }