1use std::{collections::HashMap, ops::RangeInclusive};
7
8use candid::Principal;
9use rangemap::RangeInclusiveSet;
10
11use crate::agent::PrincipalStep;
12
13#[derive(Clone)]
19pub struct Subnet {
20 pub(crate) id: Principal,
21 pub(crate) key: Vec<u8>,
24 pub(crate) node_keys: HashMap<Principal, Vec<u8>>,
25 pub(crate) canister_ranges: RangeInclusiveSet<Principal, PrincipalStep>,
26}
27
28impl Subnet {
29 pub fn contains_canister(&self, canister_id: &Principal) -> bool {
31 self.canister_ranges.contains(canister_id)
32 }
33 pub fn iter_canister_ranges(&self) -> CanisterRangesIter<'_> {
35 CanisterRangesIter {
36 inner: self.canister_ranges.iter(),
37 }
38 }
39 pub fn self_reported_key(&self) -> &[u8] {
43 &self.key
44 }
45 pub fn contains_node(&self, node_id: &Principal) -> bool {
47 self.node_keys.contains_key(node_id)
48 }
49 pub fn get_node_key(&self, node_id: &Principal) -> Option<&[u8]> {
51 self.node_keys.get(node_id).map(|k| &k[..])
52 }
53 pub fn iter_nodes(&self) -> SubnetNodeIter<'_> {
55 SubnetNodeIter {
56 inner: self.node_keys.keys(),
57 }
58 }
59 pub fn iter_node_keys(&self) -> SubnetKeysIter<'_> {
61 SubnetKeysIter {
62 inner: self.node_keys.iter(),
63 }
64 }
65 pub fn id(&self) -> Principal {
67 self.id
68 }
69}
70
71pub struct CanisterRangesIter<'a> {
73 inner: rangemap::inclusive_set::Iter<'a, Principal>,
74}
75
76impl Iterator for CanisterRangesIter<'_> {
77 type Item = RangeInclusive<Principal>;
78
79 fn next(&mut self) -> Option<Self::Item> {
80 self.inner.next().cloned()
81 }
82}
83
84pub struct SubnetNodeIter<'a> {
86 inner: std::collections::hash_map::Keys<'a, Principal, Vec<u8>>,
87}
88
89impl<'a> Iterator for SubnetNodeIter<'a> {
90 type Item = Principal;
91
92 fn next(&mut self) -> Option<Self::Item> {
93 self.inner.next().copied()
94 }
95}
96
97pub struct SubnetKeysIter<'a> {
99 inner: std::collections::hash_map::Iter<'a, Principal, Vec<u8>>,
100}
101
102impl<'a> Iterator for SubnetKeysIter<'a> {
103 type Item = (Principal, &'a [u8]);
104
105 fn next(&mut self) -> Option<Self::Item> {
106 self.inner.next().map(|(k, v)| (*k, &v[..]))
107 }
108}