1use std::collections::{BTreeMap, HashMap};
2
3use borsh::{BorshDeserialize, BorshSerialize};
4use near_primitives_core::types::{Balance, EpochHeight, ProtocolVersion, ValidatorId};
5use near_primitives_core::{
6 hash::hash,
7 types::{BlockHeight, ShardId},
8};
9use smart_default::SmartDefault;
10
11use crate::rand::StakeWeightedIndex;
12use crate::shard_layout::ShardLayout;
13use crate::types::validator_stake::{ValidatorStake, ValidatorStakeIter};
14use crate::types::{AccountId, ValidatorKickoutReason, ValidatorStakeV1};
15use crate::validator_mandates::ValidatorMandates;
16use crate::version::PROTOCOL_VERSION;
17use near_schema_checker_lib::ProtocolSchema;
18
19#[derive(
21 BorshSerialize, BorshDeserialize, Clone, Debug, PartialEq, Eq, serde::Serialize, ProtocolSchema,
22)]
23pub enum EpochInfo {
24 V1(EpochInfoV1),
25 V2(EpochInfoV2),
26 V3(EpochInfoV3),
27 V4(EpochInfoV4),
28}
29
30pub type RngSeed = [u8; 32];
31
32#[derive(
33 Default,
34 BorshSerialize,
35 BorshDeserialize,
36 Clone,
37 Debug,
38 PartialEq,
39 Eq,
40 serde::Serialize,
41 ProtocolSchema,
42)]
43pub struct ValidatorWeight(ValidatorId, u64);
44
45#[derive(
47 SmartDefault,
48 BorshSerialize,
49 BorshDeserialize,
50 Clone,
51 Debug,
52 PartialEq,
53 Eq,
54 serde::Serialize,
55 ProtocolSchema,
56)]
57pub struct EpochInfoV4 {
58 pub epoch_height: EpochHeight,
59 pub validators: Vec<ValidatorStake>,
60 pub validator_to_index: HashMap<AccountId, ValidatorId>,
61 pub block_producers_settlement: Vec<ValidatorId>,
62 pub chunk_producers_settlement: Vec<Vec<ValidatorId>>,
63 pub _hidden_validators_settlement: Vec<ValidatorWeight>,
65 pub _fishermen: Vec<crate::types::validator_stake::ValidatorStake>,
67 pub _fishermen_to_index: HashMap<AccountId, ValidatorId>,
69 pub stake_change: BTreeMap<AccountId, Balance>,
70 pub validator_reward: HashMap<AccountId, Balance>,
71 pub validator_kickout: HashMap<AccountId, ValidatorKickoutReason>,
72 pub minted_amount: Balance,
73 pub seat_price: Balance,
74 #[default(PROTOCOL_VERSION)]
75 pub protocol_version: ProtocolVersion,
76 rng_seed: RngSeed,
78 block_producers_sampler: crate::rand::StakeWeightedIndex,
79 chunk_producers_sampler: Vec<crate::rand::StakeWeightedIndex>,
80 validator_mandates: crate::validator_mandates::ValidatorMandates,
82}
83
84impl Default for EpochInfo {
85 fn default() -> Self {
86 Self::V2(EpochInfoV2::default())
87 }
88}
89
90#[derive(
92 SmartDefault,
93 BorshSerialize,
94 BorshDeserialize,
95 Clone,
96 Debug,
97 PartialEq,
98 Eq,
99 serde::Serialize,
100 ProtocolSchema,
101)]
102pub struct EpochInfoV2 {
103 pub epoch_height: EpochHeight,
106 pub validators: Vec<ValidatorStake>,
108 pub validator_to_index: HashMap<AccountId, ValidatorId>,
110 pub block_producers_settlement: Vec<ValidatorId>,
112 pub chunk_producers_settlement: Vec<Vec<ValidatorId>>,
114 pub hidden_validators_settlement: Vec<ValidatorWeight>,
116 pub fishermen: Vec<ValidatorStake>,
118 pub fishermen_to_index: HashMap<AccountId, ValidatorId>,
120 pub stake_change: BTreeMap<AccountId, Balance>,
122 pub validator_reward: HashMap<AccountId, Balance>,
124 pub validator_kickout: HashMap<AccountId, ValidatorKickoutReason>,
126 pub minted_amount: Balance,
128 pub seat_price: Balance,
130 #[default(PROTOCOL_VERSION)]
132 pub protocol_version: ProtocolVersion,
133}
134
135#[derive(
138 SmartDefault,
139 BorshSerialize,
140 BorshDeserialize,
141 Clone,
142 Debug,
143 PartialEq,
144 Eq,
145 serde::Serialize,
146 ProtocolSchema,
147)]
148pub struct EpochInfoV3 {
149 pub epoch_height: EpochHeight,
150 pub validators: Vec<ValidatorStake>,
151 pub validator_to_index: HashMap<AccountId, ValidatorId>,
152 pub block_producers_settlement: Vec<ValidatorId>,
153 pub chunk_producers_settlement: Vec<Vec<ValidatorId>>,
154 pub hidden_validators_settlement: Vec<ValidatorWeight>,
155 pub fishermen: Vec<ValidatorStake>,
156 pub fishermen_to_index: HashMap<AccountId, ValidatorId>,
157 pub stake_change: BTreeMap<AccountId, Balance>,
158 pub validator_reward: HashMap<AccountId, Balance>,
159 pub validator_kickout: HashMap<AccountId, ValidatorKickoutReason>,
160 pub minted_amount: Balance,
161 pub seat_price: Balance,
162 #[default(PROTOCOL_VERSION)]
163 pub protocol_version: ProtocolVersion,
164 rng_seed: RngSeed,
166 block_producers_sampler: StakeWeightedIndex,
167 chunk_producers_sampler: Vec<StakeWeightedIndex>,
168}
169
170impl EpochInfo {
171 pub fn new(
172 epoch_height: EpochHeight,
173 validators: Vec<ValidatorStake>,
174 validator_to_index: HashMap<AccountId, ValidatorId>,
175 block_producers_settlement: Vec<ValidatorId>,
176 chunk_producers_settlement: Vec<Vec<ValidatorId>>,
177 stake_change: BTreeMap<AccountId, Balance>,
178 validator_reward: HashMap<AccountId, Balance>,
179 validator_kickout: HashMap<AccountId, ValidatorKickoutReason>,
180 minted_amount: Balance,
181 seat_price: Balance,
182 protocol_version: ProtocolVersion,
183 rng_seed: RngSeed,
184 validator_mandates: ValidatorMandates,
185 ) -> Self {
186 let stake_weights = |ids: &[ValidatorId]| -> StakeWeightedIndex {
187 StakeWeightedIndex::new(
188 ids.iter()
189 .copied()
190 .map(|validator_id| validators[validator_id as usize].stake())
191 .collect(),
192 )
193 };
194 let block_producers_sampler = stake_weights(&block_producers_settlement);
195 let chunk_producers_sampler =
196 chunk_producers_settlement.iter().map(|vs| stake_weights(vs)).collect();
197 Self::V4(EpochInfoV4 {
198 epoch_height,
199 validators,
200 _fishermen: Default::default(),
201 validator_to_index,
202 block_producers_settlement,
203 chunk_producers_settlement,
204 _hidden_validators_settlement: Default::default(),
205 stake_change,
206 validator_reward,
207 validator_kickout,
208 _fishermen_to_index: Default::default(),
209 minted_amount,
210 seat_price,
211 protocol_version,
212 rng_seed,
213 block_producers_sampler,
214 chunk_producers_sampler,
215 validator_mandates,
216 })
217 }
218
219 #[inline]
220 pub fn epoch_height_mut(&mut self) -> &mut EpochHeight {
221 match self {
222 Self::V1(v1) => &mut v1.epoch_height,
223 Self::V2(v2) => &mut v2.epoch_height,
224 Self::V3(v3) => &mut v3.epoch_height,
225 Self::V4(v4) => &mut v4.epoch_height,
226 }
227 }
228
229 #[inline]
230 pub fn epoch_height(&self) -> EpochHeight {
231 match self {
232 Self::V1(v1) => v1.epoch_height,
233 Self::V2(v2) => v2.epoch_height,
234 Self::V3(v3) => v3.epoch_height,
235 Self::V4(v4) => v4.epoch_height,
236 }
237 }
238
239 #[inline]
240 pub fn seat_price(&self) -> Balance {
241 match self {
242 Self::V1(v1) => v1.seat_price,
243 Self::V2(v2) => v2.seat_price,
244 Self::V3(v3) => v3.seat_price,
245 Self::V4(v4) => v4.seat_price,
246 }
247 }
248
249 #[inline]
250 pub fn minted_amount(&self) -> Balance {
251 match self {
252 Self::V1(v1) => v1.minted_amount,
253 Self::V2(v2) => v2.minted_amount,
254 Self::V3(v3) => v3.minted_amount,
255 Self::V4(v4) => v4.minted_amount,
256 }
257 }
258
259 #[inline]
260 pub fn block_producers_settlement(&self) -> &[ValidatorId] {
261 match self {
262 Self::V1(v1) => &v1.block_producers_settlement,
263 Self::V2(v2) => &v2.block_producers_settlement,
264 Self::V3(v3) => &v3.block_producers_settlement,
265 Self::V4(v4) => &v4.block_producers_settlement,
266 }
267 }
268
269 #[inline]
270 pub fn chunk_producers_settlement(&self) -> &[Vec<ValidatorId>] {
271 match self {
272 Self::V1(v1) => &v1.chunk_producers_settlement,
273 Self::V2(v2) => &v2.chunk_producers_settlement,
274 Self::V3(v3) => &v3.chunk_producers_settlement,
275 Self::V4(v4) => &v4.chunk_producers_settlement,
276 }
277 }
278
279 #[inline]
280 pub fn chunk_producers_settlement_mut(&mut self) -> &mut Vec<Vec<ValidatorId>> {
281 match self {
282 Self::V1(v1) => &mut v1.chunk_producers_settlement,
283 Self::V2(v2) => &mut v2.chunk_producers_settlement,
284 Self::V3(v3) => &mut v3.chunk_producers_settlement,
285 Self::V4(v4) => &mut v4.chunk_producers_settlement,
286 }
287 }
288
289 #[inline]
290 pub fn validator_kickout(&self) -> &HashMap<AccountId, ValidatorKickoutReason> {
291 match self {
292 Self::V1(v1) => &v1.validator_kickout,
293 Self::V2(v2) => &v2.validator_kickout,
294 Self::V3(v3) => &v3.validator_kickout,
295 Self::V4(v4) => &v4.validator_kickout,
296 }
297 }
298
299 #[inline]
300 pub fn protocol_version(&self) -> ProtocolVersion {
301 match self {
302 Self::V1(v1) => v1.protocol_version,
303 Self::V2(v2) => v2.protocol_version,
304 Self::V3(v3) => v3.protocol_version,
305 Self::V4(v4) => v4.protocol_version,
306 }
307 }
308
309 #[inline]
310 pub fn stake_change(&self) -> &BTreeMap<AccountId, Balance> {
311 match self {
312 Self::V1(v1) => &v1.stake_change,
313 Self::V2(v2) => &v2.stake_change,
314 Self::V3(v3) => &v3.stake_change,
315 Self::V4(v4) => &v4.stake_change,
316 }
317 }
318
319 #[inline]
320 pub fn validator_reward(&self) -> &HashMap<AccountId, Balance> {
321 match self {
322 Self::V1(v1) => &v1.validator_reward,
323 Self::V2(v2) => &v2.validator_reward,
324 Self::V3(v3) => &v3.validator_reward,
325 Self::V4(v4) => &v4.validator_reward,
326 }
327 }
328
329 #[inline]
330 pub fn validators_iter(&self) -> ValidatorStakeIter {
331 match self {
332 Self::V1(v1) => ValidatorStakeIter::v1(&v1.validators),
333 Self::V2(v2) => ValidatorStakeIter::new(&v2.validators),
334 Self::V3(v3) => ValidatorStakeIter::new(&v3.validators),
335 Self::V4(v4) => ValidatorStakeIter::new(&v4.validators),
336 }
337 }
338
339 #[inline]
340 pub fn fishermen_iter(&self) -> ValidatorStakeIter {
341 match self {
342 Self::V1(v1) => ValidatorStakeIter::v1(&v1.fishermen),
343 Self::V2(v2) => ValidatorStakeIter::new(&v2.fishermen),
344 Self::V3(v3) => ValidatorStakeIter::new(&v3.fishermen),
345 Self::V4(v4) => ValidatorStakeIter::new(&v4._fishermen),
346 }
347 }
348
349 #[inline]
350 pub fn validator_stake(&self, validator_id: u64) -> Balance {
351 match self {
352 Self::V1(v1) => v1.validators[validator_id as usize].stake,
353 Self::V2(v2) => v2.validators[validator_id as usize].stake(),
354 Self::V3(v3) => v3.validators[validator_id as usize].stake(),
355 Self::V4(v4) => v4.validators[validator_id as usize].stake(),
356 }
357 }
358
359 #[inline]
360 pub fn validator_account_id(&self, validator_id: u64) -> &AccountId {
361 match self {
362 Self::V1(v1) => &v1.validators[validator_id as usize].account_id,
363 Self::V2(v2) => v2.validators[validator_id as usize].account_id(),
364 Self::V3(v3) => v3.validators[validator_id as usize].account_id(),
365 Self::V4(v4) => v4.validators[validator_id as usize].account_id(),
366 }
367 }
368
369 #[inline]
370 pub fn account_is_validator(&self, account_id: &AccountId) -> bool {
371 match self {
372 Self::V1(v1) => v1.validator_to_index.contains_key(account_id),
373 Self::V2(v2) => v2.validator_to_index.contains_key(account_id),
374 Self::V3(v3) => v3.validator_to_index.contains_key(account_id),
375 Self::V4(v4) => v4.validator_to_index.contains_key(account_id),
376 }
377 }
378
379 pub fn get_validator_id(&self, account_id: &AccountId) -> Option<&ValidatorId> {
380 match self {
381 Self::V1(v1) => v1.validator_to_index.get(account_id),
382 Self::V2(v2) => v2.validator_to_index.get(account_id),
383 Self::V3(v3) => v3.validator_to_index.get(account_id),
384 Self::V4(v4) => v4.validator_to_index.get(account_id),
385 }
386 }
387
388 pub fn get_validator_by_account(&self, account_id: &AccountId) -> Option<ValidatorStake> {
389 match self {
390 Self::V1(v1) => v1.validator_to_index.get(account_id).map(|validator_id| {
391 ValidatorStake::V1(v1.validators[*validator_id as usize].clone())
392 }),
393 Self::V2(v2) => v2
394 .validator_to_index
395 .get(account_id)
396 .map(|validator_id| v2.validators[*validator_id as usize].clone()),
397 Self::V3(v3) => v3
398 .validator_to_index
399 .get(account_id)
400 .map(|validator_id| v3.validators[*validator_id as usize].clone()),
401 Self::V4(v4) => v4
402 .validator_to_index
403 .get(account_id)
404 .map(|validator_id| v4.validators[*validator_id as usize].clone()),
405 }
406 }
407
408 pub fn get_validator_stake(&self, account_id: &AccountId) -> Option<Balance> {
409 match self {
410 Self::V1(v1) => v1
411 .validator_to_index
412 .get(account_id)
413 .map(|validator_id| v1.validators[*validator_id as usize].stake),
414 Self::V2(v2) => v2
415 .validator_to_index
416 .get(account_id)
417 .map(|validator_id| v2.validators[*validator_id as usize].stake()),
418 Self::V3(v3) => v3
419 .validator_to_index
420 .get(account_id)
421 .map(|validator_id| v3.validators[*validator_id as usize].stake()),
422 Self::V4(v4) => v4
423 .validator_to_index
424 .get(account_id)
425 .map(|validator_id| v4.validators[*validator_id as usize].stake()),
426 }
427 }
428
429 #[inline]
430 pub fn get_validator(&self, validator_id: u64) -> ValidatorStake {
431 match self {
432 Self::V1(v1) => ValidatorStake::V1(v1.validators[validator_id as usize].clone()),
433 Self::V2(v2) => v2.validators[validator_id as usize].clone(),
434 Self::V3(v3) => v3.validators[validator_id as usize].clone(),
435 Self::V4(v4) => v4.validators[validator_id as usize].clone(),
436 }
437 }
438
439 #[inline]
440 pub fn account_is_fisherman(&self, account_id: &AccountId) -> bool {
441 match self {
442 Self::V1(v1) => v1.fishermen_to_index.contains_key(account_id),
443 Self::V2(v2) => v2.fishermen_to_index.contains_key(account_id),
444 Self::V3(v3) => v3.fishermen_to_index.contains_key(account_id),
445 Self::V4(v4) => v4._fishermen_to_index.contains_key(account_id),
446 }
447 }
448
449 pub fn get_fisherman_by_account(&self, account_id: &AccountId) -> Option<ValidatorStake> {
450 match self {
451 Self::V1(v1) => v1.fishermen_to_index.get(account_id).map(|validator_id| {
452 ValidatorStake::V1(v1.fishermen[*validator_id as usize].clone())
453 }),
454 Self::V2(v2) => v2
455 .fishermen_to_index
456 .get(account_id)
457 .map(|validator_id| v2.fishermen[*validator_id as usize].clone()),
458 Self::V3(v3) => v3
459 .fishermen_to_index
460 .get(account_id)
461 .map(|validator_id| v3.fishermen[*validator_id as usize].clone()),
462 Self::V4(v4) => v4
463 ._fishermen_to_index
464 .get(account_id)
465 .map(|validator_id| v4._fishermen[*validator_id as usize].clone()),
466 }
467 }
468
469 #[inline]
470 pub fn get_fisherman(&self, fisherman_id: u64) -> ValidatorStake {
471 match self {
472 Self::V1(v1) => ValidatorStake::V1(v1.fishermen[fisherman_id as usize].clone()),
473 Self::V2(v2) => v2.fishermen[fisherman_id as usize].clone(),
474 Self::V3(v3) => v3.fishermen[fisherman_id as usize].clone(),
475 Self::V4(v4) => v4._fishermen[fisherman_id as usize].clone(),
476 }
477 }
478
479 #[inline]
480 pub fn validators_len(&self) -> usize {
481 match self {
482 Self::V1(v1) => v1.validators.len(),
483 Self::V2(v2) => v2.validators.len(),
484 Self::V3(v3) => v3.validators.len(),
485 Self::V4(v4) => v4.validators.len(),
486 }
487 }
488
489 #[inline]
490 pub fn rng_seed(&self) -> RngSeed {
491 match self {
492 Self::V1(_) | Self::V2(_) => Default::default(),
493 Self::V3(v3) => v3.rng_seed,
494 Self::V4(v4) => v4.rng_seed,
495 }
496 }
497
498 #[inline]
499 pub fn validator_mandates(&self) -> ValidatorMandates {
500 match self {
501 Self::V1(_) | Self::V2(_) | Self::V3(_) => Default::default(),
502 Self::V4(v4) => v4.validator_mandates.clone(),
503 }
504 }
505
506 pub fn sample_block_producer(&self, height: BlockHeight) -> ValidatorId {
507 match &self {
508 Self::V1(v1) => {
509 let bp_settlement = &v1.block_producers_settlement;
510 bp_settlement[(height % (bp_settlement.len() as u64)) as usize]
511 }
512 Self::V2(v2) => {
513 let bp_settlement = &v2.block_producers_settlement;
514 bp_settlement[(height % (bp_settlement.len() as u64)) as usize]
515 }
516 Self::V3(v3) => {
517 let seed = Self::block_produce_seed(height, &v3.rng_seed);
518 v3.block_producers_settlement[v3.block_producers_sampler.sample(seed)]
519 }
520 Self::V4(v4) => {
521 let seed = Self::block_produce_seed(height, &v4.rng_seed);
522 v4.block_producers_settlement[v4.block_producers_sampler.sample(seed)]
523 }
524 }
525 }
526
527 pub fn sample_chunk_producer(
528 &self,
529 shard_layout: &ShardLayout,
530 shard_id: ShardId,
531 height: BlockHeight,
532 ) -> Option<ValidatorId> {
533 let shard_index = shard_layout.get_shard_index(shard_id).ok()?;
534 match &self {
535 Self::V1(v1) => {
536 let cp_settlement = &v1.chunk_producers_settlement;
537 let shard_cps = cp_settlement.get(shard_index)?;
538 shard_cps.get((height as u64 % (shard_cps.len() as u64)) as usize).copied()
539 }
540 Self::V2(v2) => {
541 let cp_settlement = &v2.chunk_producers_settlement;
542 let shard_cps = cp_settlement.get(shard_index)?;
543 shard_cps.get((height as u64 % (shard_cps.len() as u64)) as usize).copied()
544 }
545 Self::V3(v3) => {
546 let seed = Self::chunk_produce_seed(&v3.rng_seed, height, shard_id);
547 let sample = v3.chunk_producers_sampler.get(shard_index)?.sample(seed);
548 v3.chunk_producers_settlement.get(shard_index)?.get(sample).copied()
549 }
550 Self::V4(v4) => {
551 let seed = Self::chunk_produce_seed(&v4.rng_seed, height, shard_id);
552 let sample = v4.chunk_producers_sampler.get(shard_index)?.sample(seed);
553 v4.chunk_producers_settlement.get(shard_index)?.get(sample).copied()
554 }
555 }
556 }
557
558 #[cfg(feature = "rand")]
559 pub fn sample_chunk_validators(
560 &self,
561 height: BlockHeight,
562 ) -> crate::validator_mandates::ChunkValidatorStakeAssignment {
563 match &self {
565 Self::V1(_) | Self::V2(_) | Self::V3(_) => Default::default(),
566 Self::V4(v4) => {
567 let mut rng = Self::chunk_validate_rng(&v4.rng_seed, height);
568 v4.validator_mandates.sample(&mut rng)
569 }
570 }
571 }
572
573 fn block_produce_seed(height: BlockHeight, seed: &RngSeed) -> [u8; 32] {
575 let mut buffer = [0u8; 40];
576 buffer[0..32].copy_from_slice(seed);
577 buffer[32..40].copy_from_slice(&height.to_le_bytes());
578 hash(&buffer).0
579 }
580
581 fn chunk_produce_seed(seed: &RngSeed, height: BlockHeight, shard_id: ShardId) -> [u8; 32] {
582 let mut buffer = [0u8; 48];
584 buffer[0..32].copy_from_slice(seed);
585 buffer[32..40].copy_from_slice(&height.to_le_bytes());
586 buffer[40..48].copy_from_slice(&shard_id.to_le_bytes());
587 hash(&buffer).0
588 }
589}
590
591#[cfg(feature = "rand")]
592impl EpochInfo {
593 fn chunk_validate_rng(seed: &RngSeed, height: BlockHeight) -> rand_chacha::ChaCha20Rng {
597 let mut buffer = [0u8; 40];
600 buffer[0..32].copy_from_slice(seed);
601 buffer[32..40].copy_from_slice(&height.to_le_bytes());
602
603 let seed = hash(&buffer);
608 rand::SeedableRng::from_seed(seed.0)
609 }
610
611 pub fn shard_assignment_rng(seed: &RngSeed) -> rand_chacha::ChaCha20Rng {
614 let mut buffer = [0u8; 62];
615 buffer[0..32].copy_from_slice(seed);
616 buffer[32..62].copy_from_slice(b"shard_assignment_shuffling_rng");
618 let seed = hash(&buffer);
619 rand::SeedableRng::from_seed(seed.0)
620 }
621}
622
623#[derive(
625 SmartDefault,
626 BorshSerialize,
627 BorshDeserialize,
628 Clone,
629 Debug,
630 PartialEq,
631 Eq,
632 serde::Serialize,
633 ProtocolSchema,
634)]
635pub struct EpochInfoV1 {
636 pub epoch_height: EpochHeight,
639 pub validators: Vec<ValidatorStakeV1>,
641 pub validator_to_index: HashMap<AccountId, ValidatorId>,
643 pub block_producers_settlement: Vec<ValidatorId>,
645 pub chunk_producers_settlement: Vec<Vec<ValidatorId>>,
647 pub hidden_validators_settlement: Vec<ValidatorWeight>,
649 pub fishermen: Vec<ValidatorStakeV1>,
651 pub fishermen_to_index: HashMap<AccountId, ValidatorId>,
653 pub stake_change: BTreeMap<AccountId, Balance>,
655 pub validator_reward: HashMap<AccountId, Balance>,
657 pub validator_kickout: HashMap<AccountId, ValidatorKickoutReason>,
659 pub minted_amount: Balance,
661 pub seat_price: Balance,
663 #[default(PROTOCOL_VERSION)]
665 pub protocol_version: ProtocolVersion,
666}