Line data Source code
1 : #ifndef HEADER_fd_src_ballet_wsample_fd_wsample_h 2 : #define HEADER_fd_src_ballet_wsample_fd_wsample_h 3 : 4 : /* This header defines methods for computing weighted "random" samples 5 : of the specific type used by Solana for computing the leader 6 : schedule and the Turbine tree. 7 : 8 : In particular, those random samples are generated by: 9 : 10 : 1. Seeding ChaCha20 with a 32B seed (different for leader schedule vs 11 : Turbine). 12 : 2. Using the ChaCha20 random stream to generate a uniformly 13 : distributed random integer x in [0, total_stake_weight). 14 : 3. Returning the first index such that the cumulative sum of stake is 15 : at least x. */ 16 : 17 : #include "../fd_ballet_base.h" 18 : #include "../chacha/fd_chacha_rng.h" 19 : 20 : struct fd_wsample_private; 21 : typedef struct fd_wsample_private fd_wsample_t; 22 : 23 : #define FD_WSAMPLE_RNG_CHACHA20 (0U) 24 : #define FD_WSAMPLE_RNG_CHACHA8 (1U) 25 : 26 10746 : #define FD_WSAMPLE_ALIGN (64UL) 27 : /* fd_leaders really wants a compile time-compatible footprint... The 28 : internal count is 1/8 * (9^ceil(log_9(ele_cnt)) - 1) */ 29 : #define FD_WSAMPLE_FOOTPRINT( ele_cnt, restore_enabled ) \ 30 : (192UL + 64UL*((restore_enabled)?2UL:1UL)*( \ 31 : ((ele_cnt)<= 1UL)? 0UL : \ 32 : ((ele_cnt)<= 9UL)? 1UL : \ 33 : ((ele_cnt)<= 81UL)? 10UL : \ 34 : ((ele_cnt)<= 729UL)? 91UL : \ 35 : ((ele_cnt)<= 6561UL)? 820UL : \ 36 : ((ele_cnt)<= 59049UL)? 7381UL : \ 37 : ((ele_cnt)<= 531441UL)? 66430UL : \ 38 : ((ele_cnt)<= 4782969UL)? 597871UL : \ 39 : ((ele_cnt)<= 43046721UL)? 5380840UL : \ 40 : ((ele_cnt)<= 387420489UL)? 48427561UL : 435848050UL )) 41 : 42 : /* fd_wsample_{align, footprint} give the alignment and footprint 43 : respectively required to create a weighted sampler with at most 44 : ele_cnt stake weights. If restore_enabled is zero, calls to 45 : wsample_restore_all will be no-ops, but the footprint required will 46 : be smaller. ele_cnt in [0, INT_MAX) (note, not ULONG MAX). 47 : 48 : fd_wsample_{join,leave} join and leave a memory region formatted as a 49 : weighted sampler, respectively. They both are simple casts. 50 : 51 : fd_wsample_delete unformats a memory region used as a weighted 52 : sampler. Releases all interest in rng. */ 53 : 54 : FD_FN_CONST ulong fd_wsample_align ( void ); 55 : FD_FN_CONST ulong fd_wsample_footprint( ulong ele_cnt, int restore_enabled ); 56 : fd_wsample_t * fd_wsample_join ( void * shmem ); 57 : void * fd_wsample_leave ( fd_wsample_t * sampler ); 58 : void * fd_wsample_delete ( void * shmem ); 59 : 60 : 61 : /* FD_WSAMPLE_HINT_*: Hints that can be passed to fd_wsample_new in the 62 : opt_hint field. The hint specifies the distribution of weights: 63 : FLAT: The all weights are approximately constant. 64 : POWERLAW: The weights are sorted largest to smallest and decay with a 65 : power law distribution. At the moment, this assumes it's 66 : proportional to 1/x. 67 : 68 : The hint can also specify whether sampling will be done with deleting 69 : or without deleting: 70 : REMOVE: Sampling will be done without replacement, i.e. mostly with 71 : fd_wsample_sample_and_remove and/or 72 : fd_wsample_sample_and_remove_many. 73 : NOREMOVE: Sampling will be done with replacement, i.e. mostly with 74 : fd_wsample_sample and/or 75 : fd_wsample_sample_many. */ 76 1554 : #define FD_WSAMPLE_HINT_FLAT 0 77 10788 : #define FD_WSAMPLE_HINT_POWERLAW_NOREMOVE 2 78 283431 : #define FD_WSAMPLE_HINT_POWERLAW_REMOVE 3 79 : 80 : /* fd_wsample_new_init, fd_wsample_new_add_weight, and 81 : fd_wsample_new_fini format a memory region with the appropriate 82 : alignment and footprint to be usable as a weighted sampler. This 83 : multi-function initialization process prevents needing to construct a 84 : flat array of weights, which is often inconvenient. 85 : 86 : The caller must first call fd_wsample_new_init, then new_add_weight 87 : ele_cnt times, and finally new_fini. Only at that point will the 88 : region of memory be ready to be joined. 89 : 90 : fd_wsample_new_init begins the formatting a memory region. shmem is 91 : a pointer to the first byte of the memory region to use. rng must be 92 : a local join of a ChaCha20 RNG struct. The weighted sampler will use 93 : rng to generate random numbers. It may seem more natural for the 94 : weighted sampler to own its own rng, but this is done to facilitate 95 : sharing of rngs between weighted samplers, which is useful for 96 : Turbine. ele_cnt specifies the number of elements that can be 97 : sampled from and must be less than INT_MAX. If restore_enabled is 98 : set to 0, fd_wsample_restore_all will not work but the required 99 : footprint is smaller. opt_hint gives a hint of the shape of the 100 : weights and the style of queries that will be most common; this hint 101 : impacts query performance but not correctness. opt_hint must be one 102 : of FD_WSAMPLE_HINT_*. 103 : 104 : fd_wsample_new_add adds a weight to a partially formatted memory 105 : region. shmem must be a partially constructed region of memory, as 106 : returned by fd_wsample_new_init or fd_wsample_new_add_weight, weight 107 : must be strictly positive, and the cumulative sum of this weight and 108 : all other weights must be no more than ULONG_MAX. 109 : 110 : fd_wsample_new_fini finalizes the formatting of a partially formatted 111 : memory region. shmem must be a partially constructed region of 112 : memory, as returned by fd_wsample_new_add_weight (or 113 : fd_wsample_new_init if ele_cnt==0). If poisoned_weight is non-zero, 114 : the weighted sampler will end with a poisoned region representing an 115 : indeterminate number of unknown elements with total weight equal to 116 : poisoned_weight. This is useful for chopping off a long tail so that 117 : the number of elements can be bounded easily. The sum of all weights 118 : and poisoned_weight must be no more than ULONG_MAX. 119 : 120 : Retains read/write interest in rng. 121 : 122 : Each function returns shmem on success and NULL on failure. It's 123 : safe to pass NULL as shmem, in which case NULL will be returned, so 124 : you only need to check the final result. 125 : 126 : On successful completion of the formatting process, the weighted 127 : sampler will contain an element corresponding to each provided 128 : weight. Caller is not joined on return. */ 129 : void * fd_wsample_new_init( void * shmem, 130 : fd_chacha_rng_t * rng, 131 : ulong ele_cnt, 132 : int restore_enabled, 133 : int opt_hint ); 134 : void * fd_wsample_new_add ( void * shmem, ulong weight ); 135 : void * fd_wsample_new_fini( void * shmem, ulong poisoned_weight ); 136 : 137 : /* fd_wsample_get_rng returns the value provided for rng in new. */ 138 : fd_chacha_rng_t * fd_wsample_get_rng( fd_wsample_t * sampler ); 139 : 140 : /* fd_wsample_seed_rng seeds the ChaCha8 rng with the provided seed in 141 : preparation for sampling. This function is compatible with Solana's 142 : ChaCha8Rng::from_seed. */ 143 : void fd_wsample_seed_rng( fd_wsample_t * sampler, 144 : uchar seed[ 32 ] ); 145 : 146 : /* fd_wsample_sample{_and_remove}{,_many} produces one or cnt (in the 147 : _many case) weighted random samples from the sampler. If the 148 : _and_remove variant of the function is called, the returned node will 149 : be temporarily removed from the sampler, i.e. for sampling without 150 : replacement. Random samples are produced using the Solana-required 151 : method. Sampler's RNG must be seeded appropriately 152 : prior to using these functions. 153 : 154 : The _many variants of the function store the ith index they sampled 155 : in idxs[i] for i in [0, cnt). The other variants of the function 156 : simply return the sampled index. 157 : 158 : If the sampler has no unremoved elements, these functions will 159 : return/store FD_WSAMPLE_EMPTY. 160 : 161 : If the RNG sample lands in the poisoned region, these functions will 162 : return/store FD_WSAMPLE_INDETERMINATE. If such a sample is produced 163 : in the without replacement mode, all subsequent calls (with and 164 : without replacement) will also return FD_WSAMPLE_INDETERMINATE until 165 : the next call to fd_wsample_restore_all. This is necessary because 166 : the poisoned region represents an indeterminate number of elements, 167 : so it's not possible to know how removing one of them will affect the 168 : total weight, and thus all subsequent random samples. 169 : 170 : For each index i, fd_wsample_sample returns i with 171 : probability weights[i]/sum(weights), only considering weights of 172 : elements that have not been removed. */ 173 331032 : #define FD_WSAMPLE_EMPTY UINT_MAX 174 22916934 : #define FD_WSAMPLE_INDETERMINATE (UINT_MAX-1UL) 175 : ulong fd_wsample_sample ( fd_wsample_t * sampler ); 176 : ulong fd_wsample_sample_and_remove ( fd_wsample_t * sampler ); 177 : void fd_wsample_sample_many ( fd_wsample_t * sampler, ulong * idxs, ulong cnt ); 178 : void fd_wsample_sample_and_remove_many( fd_wsample_t * sampler, ulong * idxs, ulong cnt ); 179 : 180 : /* fd_wsample_remove_idx removes an element by index as if it had been selected 181 : for sampling without replacement. Unless restore_all is called, this 182 : index will no longer be returned by any of the sample methods, and 183 : its weight will be excluded from the remaining sampling operations. 184 : Removing an element that has already been removed is a no-op. idx 185 : must be in [0, ele_cnt). */ 186 : void fd_wsample_remove_idx( fd_wsample_t * sampler, ulong idx ); 187 : 188 : /* fd_wsample_restore_all restores all elements removed with one of the 189 : sample_and_remove functions or with fd_wsample_remove_idx. The 190 : elements with have their original weight. This is faster than 191 : recreating the weighted sampler, even if all elements have been 192 : removed. Returns sampler on success and NULL on failure. The only 193 : error case is if sampler was constructed with restore_enabled set to 0, 194 : in which case no elements are restored. */ 195 : fd_wsample_t * fd_wsample_restore_all( fd_wsample_t * sampler ); 196 : 197 : #endif /* HEADER_fd_src_ballet_wsample_fd_wsample_h */
