1use crate::types::{ValidatorId, validator_stake::ValidatorStake};
2use borsh::{BorshDeserialize, BorshSerialize};
3use near_primitives_core::types::Balance;
4use near_schema_checker_lib::ProtocolSchema;
5
6mod compute_price;
7
8#[derive(
11 BorshSerialize,
12 BorshDeserialize,
13 Default,
14 Copy,
15 Clone,
16 Debug,
17 PartialEq,
18 Eq,
19 serde::Serialize,
20 ProtocolSchema,
21)]
22pub struct ValidatorMandatesConfig {
23 target_mandates_per_shard: usize,
25 num_shards: usize,
27}
28
29impl ValidatorMandatesConfig {
30 pub fn new(target_mandates_per_shard: usize, num_shards: usize) -> Self {
39 assert!(num_shards > 0, "there should be at least one shard");
40 Self { target_mandates_per_shard, num_shards }
41 }
42}
43
44#[derive(
53 BorshSerialize,
54 BorshDeserialize,
55 Default,
56 Clone,
57 Debug,
58 PartialEq,
59 Eq,
60 serde::Serialize,
61 ProtocolSchema,
62)]
63pub struct ValidatorMandates {
64 config: ValidatorMandatesConfig,
66 stake_per_mandate: Balance,
68 mandates: Vec<ValidatorId>,
72 partials: Vec<(ValidatorId, Balance)>,
79}
80
81impl ValidatorMandates {
82 pub fn new(config: ValidatorMandatesConfig, validators: &[ValidatorStake]) -> Self {
89 let stakes: Vec<Balance> = validators.iter().map(|v| v.stake()).collect();
90 let stake_per_mandate = compute_price::compute_mandate_price(config, &stakes);
91 let num_mandates_per_validator: Vec<u16> =
92 validators.iter().map(|v| v.num_mandates(stake_per_mandate)).collect();
93 let num_total_mandates =
94 num_mandates_per_validator.iter().map(|&num| usize::from(num)).sum();
95 let mut mandates: Vec<ValidatorId> = Vec::with_capacity(num_total_mandates);
96
97 for i in 0..validators.len() {
98 for _ in 0..num_mandates_per_validator[i] {
99 mandates.push(i as ValidatorId);
101 }
102 }
103
104 let mut partials = Vec::with_capacity(validators.len());
110 for i in 0..validators.len() {
111 let partial_weight = validators[i].partial_mandate_weight(stake_per_mandate);
112 if partial_weight > Balance::ZERO {
113 partials.push((i as ValidatorId, partial_weight));
114 }
115 }
116
117 Self { config, stake_per_mandate, mandates, partials }
118 }
119}
120
121#[cfg(feature = "rand")]
122mod validator_mandates_sample {
123 use super::*;
124 use itertools::Itertools;
125 use rand::{Rng, seq::SliceRandom};
126
127 impl ValidatorMandates {
128 pub fn sample<R>(&self, rng: &mut R) -> ChunkValidatorStakeAssignment
135 where
136 R: Rng + ?Sized,
137 {
138 let shard_ids_for_mandates = ShuffledShardIds::new(rng, self.config.num_shards);
142 let shard_ids_for_partials = ShuffledShardIds::new(rng, self.config.num_shards);
143
144 let shuffled_mandates = self.shuffled_mandates(rng);
145 let shuffled_partials = self.shuffled_partials(rng);
146
147 let stake_per_mandate = self.stake_per_mandate;
154 let mut stake_assignment_per_shard =
155 vec![std::collections::HashMap::new(); self.config.num_shards];
156 for shard_id in 0..self.config.num_shards {
157 let mandates_assignment =
159 &mut stake_assignment_per_shard[shard_ids_for_mandates.get_alias(shard_id)];
160
161 for idx in (shard_id..shuffled_mandates.len()).step_by(self.config.num_shards) {
164 let validator_id = shuffled_mandates[idx];
165 let current_shard_validator_stake: &mut Balance =
166 mandates_assignment.entry(validator_id).or_default();
167 *current_shard_validator_stake =
168 current_shard_validator_stake.checked_add(stake_per_mandate).unwrap();
169 }
170
171 let partials_assignment =
173 &mut stake_assignment_per_shard[shard_ids_for_partials.get_alias(shard_id)];
174
175 for idx in (shard_id..shuffled_partials.len()).step_by(self.config.num_shards) {
178 let (validator_id, partial_weight) = shuffled_partials[idx];
179 let current_shard_validator_stake: &mut Balance =
180 partials_assignment.entry(validator_id).or_default();
181 *current_shard_validator_stake =
182 current_shard_validator_stake.checked_add(partial_weight).unwrap();
183 }
184 }
185
186 let mut ordered_stake_assignment_per_shard = Vec::with_capacity(self.config.num_shards);
188 for shard_id in 0..self.config.num_shards {
189 let stake_assignment = &stake_assignment_per_shard[shard_id];
191 let mut ordered_validator_ids = stake_assignment.keys().sorted().collect_vec();
192 ordered_validator_ids.shuffle(rng);
193 let ordered_mandate_assignment = ordered_validator_ids
194 .into_iter()
195 .map(|validator_id| (*validator_id, stake_assignment[validator_id]))
196 .collect_vec();
197 ordered_stake_assignment_per_shard.push(ordered_mandate_assignment);
198 }
199
200 ordered_stake_assignment_per_shard
201 }
202
203 pub(super) fn shuffled_mandates<R>(&self, rng: &mut R) -> Vec<ValidatorId>
206 where
207 R: Rng + ?Sized,
208 {
209 let mut shuffled_mandates = self.mandates.clone();
210 shuffled_mandates.shuffle(rng);
211 shuffled_mandates
212 }
213
214 pub(super) fn shuffled_partials<R>(&self, rng: &mut R) -> Vec<(ValidatorId, Balance)>
217 where
218 R: Rng + ?Sized,
219 {
220 let mut shuffled_partials = self.partials.clone();
221 shuffled_partials.shuffle(rng);
222 shuffled_partials
223 }
224 }
225}
226
227pub type ChunkValidatorStakeAssignment = Vec<Vec<(ValidatorId, Balance)>>;
236
237#[cfg(feature = "rand")]
253#[derive(Default, Clone, Debug, PartialEq, Eq)]
254struct ShuffledShardIds {
255 shuffled_ids: Vec<usize>,
257}
258
259#[cfg(feature = "rand")]
260impl ShuffledShardIds {
261 fn new<R>(rng: &mut R, num_shards: usize) -> Self
262 where
263 R: rand::Rng + ?Sized,
264 {
265 use rand::seq::SliceRandom;
266
267 let mut shuffled_ids = (0..num_shards).collect::<Vec<_>>();
268 shuffled_ids.shuffle(rng);
269 Self { shuffled_ids }
270 }
271
272 fn get_alias(&self, shard_id: usize) -> usize {
278 self.shuffled_ids[shard_id]
279 }
280}
281
282#[cfg(test)]
283mod tests {
284 use near_crypto::PublicKey;
285 use near_primitives_core::types::Balance;
286 use rand::SeedableRng;
287 use rand_chacha::ChaCha8Rng;
288
289 use crate::{
290 types::ValidatorId, types::validator_stake::ValidatorStake,
291 validator_mandates::ValidatorMandatesConfig,
292 };
293
294 use super::{ChunkValidatorStakeAssignment, ShuffledShardIds, ValidatorMandates};
295
296 fn new_fixed_rng() -> ChaCha8Rng {
299 ChaCha8Rng::seed_from_u64(42)
300 }
301
302 #[test]
303 fn test_validator_mandates_config_new() {
304 let target_mandates_per_shard = 400;
305 let num_shards = 4;
306 assert_eq!(
307 ValidatorMandatesConfig::new(target_mandates_per_shard, num_shards),
308 ValidatorMandatesConfig { target_mandates_per_shard, num_shards },
309 )
310 }
311
312 fn new_validator_stakes() -> Vec<ValidatorStake> {
322 let new_vs = |account_id: &str, balance: Balance| -> ValidatorStake {
323 ValidatorStake::new(
324 account_id.parse().unwrap(),
325 PublicKey::empty(near_crypto::KeyType::ED25519),
326 balance,
327 )
328 };
329
330 vec![
331 new_vs("account_0", Balance::from_yoctonear(30)),
332 new_vs("account_1", Balance::from_yoctonear(27)),
333 new_vs("account_2", Balance::from_yoctonear(9)),
334 new_vs("account_3", Balance::from_yoctonear(12)),
335 new_vs("account_4", Balance::from_yoctonear(35)),
336 new_vs("account_5", Balance::from_yoctonear(4)),
337 new_vs("account_6", Balance::from_yoctonear(6)),
338 ]
339 }
340
341 #[test]
342 fn test_validator_mandates_new() {
343 let validators = new_validator_stakes();
344 let config = ValidatorMandatesConfig::new(3, 4);
345 let mandates = ValidatorMandates::new(config, &validators);
346
347 assert_eq!(mandates.stake_per_mandate, Balance::from_yoctonear(8));
353
354 let expected_mandates: Vec<ValidatorId> = vec![0, 0, 0, 1, 1, 1, 2, 3, 4, 4, 4, 4];
357 assert_eq!(mandates.mandates, expected_mandates);
358
359 assert_eq!(mandates.mandates.len(), config.num_shards * config.target_mandates_per_shard);
361
362 let expected_partials: Vec<(ValidatorId, Balance)> = vec![
365 (0, Balance::from_yoctonear(6)),
366 (1, Balance::from_yoctonear(3)),
367 (2, Balance::from_yoctonear(1)),
368 (3, Balance::from_yoctonear(4)),
369 (4, Balance::from_yoctonear(3)),
370 (5, Balance::from_yoctonear(4)),
371 (6, Balance::from_yoctonear(6)),
372 ];
373 assert_eq!(mandates.partials, expected_partials);
374 }
375
376 #[test]
377 fn test_validator_mandates_shuffled_mandates() {
378 assert_validator_mandates_shuffled_mandates(3, vec![0, 1, 4, 4, 3, 1, 4, 0, 0]);
380 assert_validator_mandates_shuffled_mandates(4, vec![0, 0, 2, 1, 3, 4, 1, 1, 0, 4, 4, 4]);
381 }
382
383 fn assert_validator_mandates_shuffled_mandates(
384 num_shards: usize,
385 expected_assignment: Vec<ValidatorId>,
386 ) {
387 let validators = new_validator_stakes();
388 let config = ValidatorMandatesConfig::new(3, num_shards);
389 let mandates = ValidatorMandates::new(config, &validators);
390 let mut rng = new_fixed_rng();
391
392 let _shard_ids_for_mandates = ShuffledShardIds::new(&mut rng, config.num_shards);
396 let _shard_ids_for_partials = ShuffledShardIds::new(&mut rng, config.num_shards);
397 let assignment = mandates.shuffled_mandates(&mut rng);
398 assert_eq!(
399 assignment, expected_assignment,
400 "Unexpected shuffling for num_shards = {num_shards}"
401 );
402 }
403
404 #[test]
405 fn test_validator_mandates_shuffled_partials() {
406 assert_validator_mandates_shuffled_partials(
408 3,
409 vec![
410 (3, Balance::from_yoctonear(2)),
411 (4, Balance::from_yoctonear(5)),
412 (1, Balance::from_yoctonear(7)),
413 (2, Balance::from_yoctonear(9)),
414 (5, Balance::from_yoctonear(4)),
415 (6, Balance::from_yoctonear(6)),
416 ],
417 );
418 assert_validator_mandates_shuffled_partials(
419 4,
420 vec![
421 (5, Balance::from_yoctonear(4)),
422 (3, Balance::from_yoctonear(4)),
423 (0, Balance::from_yoctonear(6)),
424 (2, Balance::from_yoctonear(1)),
425 (1, Balance::from_yoctonear(3)),
426 (4, Balance::from_yoctonear(3)),
427 (6, Balance::from_yoctonear(6)),
428 ],
429 );
430 }
431
432 fn assert_validator_mandates_shuffled_partials(
433 num_shards: usize,
434 expected_assignment: Vec<(ValidatorId, Balance)>,
435 ) {
436 let validators = new_validator_stakes();
437 let config = ValidatorMandatesConfig::new(3, num_shards);
438 let mandates = ValidatorMandates::new(config, &validators);
439 let mut rng = new_fixed_rng();
440
441 let _shard_ids_for_mandates = ShuffledShardIds::new(&mut rng, config.num_shards);
445 let _shard_ids_for_partials = ShuffledShardIds::new(&mut rng, config.num_shards);
446 let _ = mandates.shuffled_mandates(&mut rng);
447 let assignment = mandates.shuffled_partials(&mut rng);
448 assert_eq!(
449 assignment, expected_assignment,
450 "Unexpected shuffling for num_shards = {num_shards}"
451 );
452 }
453
454 #[test]
457 fn test_validator_mandates_sample_even() {
458 let config = ValidatorMandatesConfig::new(3, 3);
463 let expected_assignment = vec![
464 vec![
465 (1, Balance::from_yoctonear(17)),
466 (4, Balance::from_yoctonear(10)),
467 (6, Balance::from_yoctonear(6)),
468 (0, Balance::from_yoctonear(10)),
469 ],
470 vec![
471 (4, Balance::from_yoctonear(5)),
472 (5, Balance::from_yoctonear(4)),
473 (0, Balance::from_yoctonear(10)),
474 (1, Balance::from_yoctonear(10)),
475 (3, Balance::from_yoctonear(10)),
476 ],
477 vec![
478 (0, Balance::from_yoctonear(10)),
479 (2, Balance::from_yoctonear(9)),
480 (4, Balance::from_yoctonear(20)),
481 (3, Balance::from_yoctonear(2)),
482 ],
483 ];
484 assert_validator_mandates_sample(config, expected_assignment);
485 }
486
487 #[test]
490 fn test_validator_mandates_sample_uneven() {
491 let config = ValidatorMandatesConfig::new(3, 4);
496 let expected_mandates_per_shards = vec![
497 vec![
498 (3, Balance::from_yoctonear(8)),
499 (6, Balance::from_yoctonear(6)),
500 (0, Balance::from_yoctonear(22)),
501 ],
502 vec![
503 (4, Balance::from_yoctonear(8)),
504 (2, Balance::from_yoctonear(9)),
505 (1, Balance::from_yoctonear(8)),
506 ],
507 vec![
508 (0, Balance::from_yoctonear(8)),
509 (3, Balance::from_yoctonear(4)),
510 (4, Balance::from_yoctonear(19)),
511 ],
512 vec![
513 (4, Balance::from_yoctonear(8)),
514 (5, Balance::from_yoctonear(4)),
515 (1, Balance::from_yoctonear(19)),
516 ],
517 ];
518 assert_validator_mandates_sample(config, expected_mandates_per_shards);
519 }
520
521 fn assert_validator_mandates_sample(
523 config: ValidatorMandatesConfig,
524 expected_assignment: ChunkValidatorStakeAssignment,
525 ) {
526 let validators = new_validator_stakes();
527 let mandates = ValidatorMandates::new(config, &validators);
528
529 let mut rng = new_fixed_rng();
530 let assignment = mandates.sample(&mut rng);
531
532 assert_eq!(assignment, expected_assignment);
533 }
534
535 #[test]
536 fn test_shuffled_shard_ids_new() {
537 let mut rng_3_shards = new_fixed_rng();
541 assert_shuffled_shard_ids(&mut rng_3_shards, 3, vec![2, 1, 0], "3 shards, 1st shuffle");
542 assert_shuffled_shard_ids(&mut rng_3_shards, 3, vec![2, 1, 0], "3 shards, 2nd shuffle");
543 let mut rng_4_shards = new_fixed_rng();
544 assert_shuffled_shard_ids(&mut rng_4_shards, 4, vec![0, 2, 1, 3], "4 shards, 1st shuffle");
545 assert_shuffled_shard_ids(&mut rng_4_shards, 4, vec![3, 2, 0, 1], "4 shards, 2nd shuffle");
546 }
547
548 fn assert_shuffled_shard_ids(
549 rng: &mut ChaCha8Rng,
550 num_shards: usize,
551 expected_shuffling: Vec<usize>,
552 test_descriptor: &str,
553 ) {
554 let shuffled_ids_full_mandates = ShuffledShardIds::new(rng, num_shards);
555 assert_eq!(
556 shuffled_ids_full_mandates,
557 ShuffledShardIds { shuffled_ids: expected_shuffling },
558 "Unexpected shuffling for {test_descriptor}",
559 );
560 }
561
562 #[test]
563 fn test_shuffled_shard_ids_get_alias() {
564 let mut rng = new_fixed_rng();
565 let shuffled_ids = ShuffledShardIds::new(&mut rng, 4);
566 assert_eq!(shuffled_ids.get_alias(1), 2);
568 }
569
570 #[test]
571 fn test_deterministic_shuffle() {
572 let config = ValidatorMandatesConfig::new(3, 4);
573 let validators = new_validator_stakes();
574 let mandates = ValidatorMandates::new(config, &validators);
575
576 let mut rng1 = new_fixed_rng();
577 let assignment1 = mandates.sample(&mut rng1);
578
579 let mut rng2 = new_fixed_rng();
580 let assignment2 = mandates.sample(&mut rng2);
581
582 assert_eq!(assignment1, assignment2);
584 }
585}