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
#![cfg_attr(not(feature = "std"), no_std)]
18

            
19
use {
20
    cumulus_primitives_core::ParaId,
21
    parity_scale_codec::{Decode, Encode},
22
    scale_info::prelude::collections::BTreeMap,
23
    sp_std::vec::Vec,
24
};
25

            
26
#[derive(Clone, Encode, Decode, PartialEq, sp_core::RuntimeDebug, scale_info::TypeInfo)]
27
#[cfg_attr(feature = "std", derive(serde::Serialize, serde::Deserialize))]
28
pub struct AssignedCollators<AccountId> {
29
    pub orchestrator_chain: Vec<AccountId>,
30
    pub container_chains: BTreeMap<ParaId, Vec<AccountId>>,
31
}
32

            
33
// Manual default impl that does not require AccountId: Default
34
impl<AccountId> Default for AssignedCollators<AccountId> {
35
271668
    fn default() -> Self {
36
271668
        Self {
37
271668
            orchestrator_chain: Default::default(),
38
271668
            container_chains: Default::default(),
39
271668
        }
40
271668
    }
41
}
42

            
43
impl<AccountId> AssignedCollators<AccountId>
44
where
45
    AccountId: PartialEq,
46
{
47
    pub fn para_id_of(&self, x: &AccountId, orchestrator_chain_para_id: ParaId) -> Option<ParaId> {
48
        for (id, cs) in self.container_chains.iter() {
49
            if cs.contains(x) {
50
                return Some(*id);
51
            }
52
        }
53

            
54
        if self.orchestrator_chain.contains(x) {
55
            return Some(orchestrator_chain_para_id);
56
        }
57

            
58
        None
59
    }
60

            
61
135234
    pub fn map<T, F>(&self, mut f: F) -> AssignedCollators<T>
62
135234
    where
63
135234
        F: FnMut(&AccountId) -> T,
64
135234
    {
65
135234
        let mut a = AssignedCollators {
66
135234
            orchestrator_chain: self.orchestrator_chain.iter().map(&mut f).collect(),
67
135234
            ..Default::default()
68
135234
        };
69

            
70
135234
        for (para_id, collators) in self.container_chains.iter() {
71
            let a_collators = collators.iter().map(&mut f).collect();
72
            a.container_chains.insert(*para_id, a_collators);
73
        }
74

            
75
135234
        a
76
135234
    }
77
}