1
// Copyright (C) Moondance Labs Ltd.
2
// This file is part of Tanssi.
3

            
4
// Tanssi is free software: you can redistribute it and/or modify
5
// it under the terms of the GNU General Public License as published by
6
// the Free Software Foundation, either version 3 of the License, or
7
// (at your option) any later version.
8

            
9
// Tanssi is distributed in the hope that it will be useful,
10
// but WITHOUT ANY WARRANTY; without even the implied warranty of
11
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
12
// GNU General Public License for more details.
13

            
14
// You should have received a copy of the GNU General Public License
15
// along with Tanssi.  If not, see <http://www.gnu.org/licenses/>
16

            
17
//! # Nimbus Collator Assignment Pallet
18

            
19
#![cfg_attr(not(feature = "std"), no_std)]
20

            
21
pub use pallet::*;
22
use {
23
    dp_collator_assignment::AssignedCollators,
24
    frame_support::pallet_prelude::*,
25
    sp_runtime::{
26
        traits::{AtLeast32BitUnsigned, One, Zero},
27
        Saturating,
28
    },
29
    sp_std::{collections::btree_map::BTreeMap, prelude::*, vec},
30
};
31

            
32
#[cfg(test)]
33
mod mock;
34

            
35
#[cfg(test)]
36
mod tests;
37

            
38
704
#[frame_support::pallet]
39
pub mod pallet {
40
    use super::*;
41

            
42
13620
    #[pallet::pallet]
43
    #[pallet::without_storage_info]
44
    pub struct Pallet<T>(_);
45

            
46
    /// Configure the pallet by specifying the parameters and types on which it depends.
47
    #[pallet::config]
48
    pub trait Config: frame_system::Config {
49
        type SessionIndex: parity_scale_codec::FullCodec + TypeInfo + Copy + AtLeast32BitUnsigned;
50
        type AuthorityId: parity_scale_codec::FullCodec + TypeInfo + Clone;
51
    }
52

            
53
57849
    #[pallet::storage]
54
    #[pallet::getter(fn collator_container_chain)]
55
    pub type CollatorContainerChain<T: Config> = StorageMap<
56
        _,
57
        Twox64Concat,
58
        T::SessionIndex,
59
        AssignedCollators<T::AuthorityId>,
60
        OptionQuery,
61
    >;
62

            
63
176
    #[pallet::call]
64
    impl<T: Config> Pallet<T> {}
65

            
66
    impl<T: Config> Pallet<T> {
67
        /// Assign new collators
68
        /// collators should be queued collators
69
1777
        pub fn assign_collators(
70
1777
            current_session_index: &T::SessionIndex,
71
1777
            queued_id_to_nimbus_map: &BTreeMap<T::AccountId, T::AuthorityId>,
72
1777
            next_collator_assignment: &AssignedCollators<T::AccountId>,
73
1777
        ) {
74
1777
            let next_nimbus_assignment = next_collator_assignment
75
4971
                .map(|account_id| queued_id_to_nimbus_map[account_id].clone());
76
1777

            
77
1777
            // Only applies to session index 0
78
1777
            if current_session_index == &T::SessionIndex::zero() {
79
509
                CollatorContainerChain::<T>::insert(
80
509
                    current_session_index,
81
509
                    next_nimbus_assignment.clone(),
82
509
                );
83
509
                CollatorContainerChain::<T>::insert(
84
509
                    current_session_index.saturating_add(T::SessionIndex::one()),
85
509
                    next_nimbus_assignment,
86
509
                );
87
509

            
88
509
                return;
89
1268
            }
90
1268

            
91
1268
            // Remove value at session - 1, insert new value at session + 1
92
1268
            CollatorContainerChain::<T>::remove(
93
1268
                current_session_index.saturating_sub(T::SessionIndex::one()),
94
1268
            );
95
1268
            CollatorContainerChain::<T>::insert(
96
1268
                current_session_index.saturating_add(T::SessionIndex::one()),
97
1268
                next_nimbus_assignment,
98
1268
            );
99
1777
        }
100

            
101
1777
        pub fn initializer_on_new_session(
102
1777
            current_session_index: &T::SessionIndex,
103
1777
            queued_id_to_nimbus_map: &BTreeMap<T::AccountId, T::AuthorityId>,
104
1777
            next_collator_assignment: &AssignedCollators<T::AccountId>,
105
1777
        ) {
106
1777
            Self::assign_collators(
107
1777
                current_session_index,
108
1777
                queued_id_to_nimbus_map,
109
1777
                next_collator_assignment,
110
1777
            )
111
1777
        }
112
    }
113
}