use alloc::{
collections::{btree_map::BTreeMap, btree_set::BTreeSet, vec_deque::VecDeque},
vec::Vec,
};
use frame_support::{pallet_prelude::*, traits::DisabledValidators};
use frame_system::pallet_prelude::BlockNumberFor;
use polkadot_primitives::{
vstaging::transpose_claim_queue, CoreIndex, Id, SessionIndex, ValidatorId, ValidatorIndex,
};
use sp_runtime::traits::AtLeast32BitUnsigned;
use rand::{seq::SliceRandom, SeedableRng};
use rand_chacha::ChaCha20Rng;
use crate::configuration::HostConfiguration;
pub use pallet::*;
pub(crate) const SESSION_DELAY: SessionIndex = 2;
#[cfg(test)]
mod tests;
pub mod migration;
#[derive(Encode, Decode, Default, TypeInfo, Debug)]
pub struct RelayParentInfo<Hash> {
pub relay_parent: Hash,
pub state_root: Hash,
pub claim_queue: BTreeMap<Id, BTreeMap<u8, BTreeSet<CoreIndex>>>,
}
#[derive(Encode, Decode, Default, TypeInfo)]
pub struct AllowedRelayParentsTracker<Hash, BlockNumber> {
buffer: VecDeque<RelayParentInfo<Hash>>,
latest_number: BlockNumber,
}
impl<Hash: PartialEq + Copy, BlockNumber: AtLeast32BitUnsigned + Copy>
AllowedRelayParentsTracker<Hash, BlockNumber>
{
pub(crate) fn update(
&mut self,
relay_parent: Hash,
state_root: Hash,
claim_queue: BTreeMap<CoreIndex, VecDeque<Id>>,
number: BlockNumber,
max_ancestry_len: u32,
) {
if self.buffer.iter().any(|info| info.relay_parent == relay_parent) {
return
}
let claim_queue = transpose_claim_queue(claim_queue);
let buffer_size_limit = max_ancestry_len as usize + 1;
self.buffer.push_back(RelayParentInfo { relay_parent, state_root, claim_queue });
self.latest_number = number;
while self.buffer.len() > buffer_size_limit {
let _ = self.buffer.pop_front();
}
}
pub(crate) fn acquire_info(
&self,
relay_parent: Hash,
prev: Option<BlockNumber>,
) -> Option<(&RelayParentInfo<Hash>, BlockNumber)> {
let pos = self.buffer.iter().position(|info| info.relay_parent == relay_parent)?;
let age = (self.buffer.len() - 1) - pos;
let number = self.latest_number - BlockNumber::from(age as u32);
if let Some(prev) = prev {
if prev > number {
return None
}
}
Some((&self.buffer[pos], number))
}
pub(crate) fn hypothetical_earliest_block_number(
&self,
now: BlockNumber,
max_ancestry_len: u32,
) -> BlockNumber {
let allowed_ancestry_len = max_ancestry_len.min(self.buffer.len() as u32);
now - allowed_ancestry_len.into()
}
}
#[frame_support::pallet]
pub mod pallet {
use super::*;
const STORAGE_VERSION: StorageVersion = StorageVersion::new(1);
#[pallet::pallet]
#[pallet::without_storage_info]
#[pallet::storage_version(STORAGE_VERSION)]
pub struct Pallet<T>(_);
#[pallet::config]
pub trait Config: frame_system::Config {
type DisabledValidators: frame_support::traits::DisabledValidators;
}
#[pallet::storage]
pub type CurrentSessionIndex<T: Config> = StorageValue<_, SessionIndex, ValueQuery>;
#[pallet::storage]
pub type ActiveValidatorIndices<T: Config> = StorageValue<_, Vec<ValidatorIndex>, ValueQuery>;
#[pallet::storage]
pub type ActiveValidatorKeys<T: Config> = StorageValue<_, Vec<ValidatorId>, ValueQuery>;
#[pallet::storage]
pub(crate) type AllowedRelayParents<T: Config> =
StorageValue<_, AllowedRelayParentsTracker<T::Hash, BlockNumberFor<T>>, ValueQuery>;
#[pallet::call]
impl<T: Config> Pallet<T> {}
}
impl<T: Config> Pallet<T> {
pub(crate) fn initializer_initialize(_now: BlockNumberFor<T>) -> Weight {
Weight::zero()
}
pub(crate) fn initializer_finalize() {}
pub(crate) fn initializer_on_new_session(
session_index: SessionIndex,
random_seed: [u8; 32],
new_config: &HostConfiguration<BlockNumberFor<T>>,
all_validators: Vec<ValidatorId>,
) -> Vec<ValidatorId> {
AllowedRelayParents::<T>::mutate(|tracker| tracker.buffer.clear());
CurrentSessionIndex::<T>::set(session_index);
let mut rng: ChaCha20Rng = SeedableRng::from_seed(random_seed);
let mut shuffled_indices: Vec<_> = (0..all_validators.len())
.enumerate()
.map(|(i, _)| ValidatorIndex(i as _))
.collect();
shuffled_indices.shuffle(&mut rng);
if let Some(max) = new_config.max_validators {
shuffled_indices.truncate(max as usize);
}
let active_validator_keys =
crate::util::take_active_subset(&shuffled_indices, &all_validators);
ActiveValidatorIndices::<T>::set(shuffled_indices);
ActiveValidatorKeys::<T>::set(active_validator_keys.clone());
active_validator_keys
}
pub fn scheduled_session() -> SessionIndex {
CurrentSessionIndex::<T>::get().saturating_add(SESSION_DELAY)
}
pub fn disabled_validators() -> Vec<ValidatorIndex> {
let shuffled_indices = ActiveValidatorIndices::<T>::get();
let reverse_index = shuffled_indices
.iter()
.enumerate()
.map(|(i, v)| (v.0, ValidatorIndex(i as u32)))
.collect::<BTreeMap<u32, ValidatorIndex>>();
T::DisabledValidators::disabled_validators()
.iter()
.filter_map(|v| reverse_index.get(v).cloned())
.collect()
}
#[cfg(any(feature = "std", feature = "runtime-benchmarks", test))]
pub fn set_session_index(index: SessionIndex) {
CurrentSessionIndex::<T>::set(index);
}
#[cfg(any(feature = "std", feature = "runtime-benchmarks", test))]
pub fn set_active_validators_ascending(active: Vec<ValidatorId>) {
ActiveValidatorIndices::<T>::set(
(0..active.len()).map(|i| ValidatorIndex(i as _)).collect(),
);
ActiveValidatorKeys::<T>::set(active);
}
#[cfg(test)]
pub(crate) fn set_active_validators_with_indices(
indices: Vec<ValidatorIndex>,
keys: Vec<ValidatorId>,
) {
assert_eq!(indices.len(), keys.len());
ActiveValidatorIndices::<T>::set(indices);
ActiveValidatorKeys::<T>::set(keys);
}
#[cfg(test)]
pub(crate) fn add_allowed_relay_parent(
relay_parent: T::Hash,
state_root: T::Hash,
claim_queue: BTreeMap<CoreIndex, VecDeque<Id>>,
number: BlockNumberFor<T>,
max_ancestry_len: u32,
) {
AllowedRelayParents::<T>::mutate(|tracker| {
tracker.update(relay_parent, state_root, claim_queue, number, max_ancestry_len)
})
}
}