1use crate::ecies;
5use crate::types::ShareIndex;
6use fastcrypto::error::{FastCryptoError, FastCryptoResult};
7use fastcrypto::groups::GroupElement;
8use fastcrypto::hash::{Blake2b256, Digest, HashFunction};
9use serde::{Deserialize, Serialize};
10use tracing::debug;
11
12pub type PartyId = u16;
13
14#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
16pub struct Node<G: GroupElement> {
17 pub id: PartyId,
18 pub pk: ecies::PublicKey<G>,
19 pub weight: u16, }
21
22#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
24pub struct Nodes<G: GroupElement> {
25 nodes: Vec<Node<G>>, total_weight: u16, accumulated_weights: Vec<u16>, nodes_with_nonzero_weight: Vec<u16>, }
31
32impl<G: GroupElement + Serialize> Nodes<G> {
33 const MAX_NODES: usize = 1000;
35
36 pub fn new(nodes: Vec<Node<G>>) -> FastCryptoResult<Self> {
38 let mut nodes = nodes;
39 nodes.sort_by_key(|n| n.id);
40 if (0..nodes.len()).any(|i| (nodes[i].id as usize) != i) {
42 return Err(FastCryptoError::InvalidInput);
43 }
44 if nodes.is_empty() || nodes.len() > Self::MAX_NODES {
46 return Err(FastCryptoError::InvalidInput);
47 }
48 let total_weight = nodes.iter().map(|n| n.weight as u32).sum::<u32>();
50 if total_weight > u16::MAX as u32 || total_weight == 0 {
51 return Err(FastCryptoError::InvalidInput);
52 }
53 let total_weight = total_weight as u16;
54
55 let accumulated_weights = Self::get_accumulated_weights(&nodes);
57 let nodes_with_nonzero_weight = Self::filter_nonzero_weights(&nodes);
58
59 Ok(Self {
60 nodes,
61 total_weight,
62 accumulated_weights,
63 nodes_with_nonzero_weight,
64 })
65 }
66
67 fn filter_nonzero_weights(nodes: &[Node<G>]) -> Vec<u16> {
68 nodes
69 .iter()
70 .enumerate()
71 .filter_map(|(i, n)| if n.weight > 0 { Some(i as u16) } else { None })
72 .collect::<Vec<_>>()
73 }
74
75 fn get_accumulated_weights(nodes: &[Node<G>]) -> Vec<u16> {
76 nodes
77 .iter()
78 .filter_map(|n| if n.weight > 0 { Some(n.weight) } else { None })
79 .scan(0, |accumulated_weight, weight| {
80 *accumulated_weight += weight;
81 Some(*accumulated_weight)
82 })
83 .collect::<Vec<_>>()
84 }
85
86 pub fn total_weight(&self) -> u16 {
88 self.total_weight
89 }
90
91 pub fn num_nodes(&self) -> usize {
93 self.nodes.len()
94 }
95
96 pub fn share_ids_iter(&self) -> impl Iterator<Item = ShareIndex> {
98 (1..=self.total_weight).map(|i| ShareIndex::new(i).expect("nonzero"))
99 }
100
101 pub fn share_id_to_node(&self, share_id: &ShareIndex) -> FastCryptoResult<&Node<G>> {
103 let nonzero_node_id = self
104 .accumulated_weights
105 .binary_search(&share_id.get())
106 .unwrap_or_else(|i| i);
107 match self.nodes_with_nonzero_weight.get(nonzero_node_id) {
108 Some(node_id) => self.node_id_to_node(*node_id),
109 None => Err(FastCryptoError::InvalidInput),
110 }
111 }
112
113 pub fn node_id_to_node(&self, party_id: PartyId) -> FastCryptoResult<&Node<G>> {
114 self.nodes
115 .get(party_id as usize)
116 .ok_or(FastCryptoError::InvalidInput)
117 }
118
119 pub fn share_ids_of(&self, id: PartyId) -> FastCryptoResult<Vec<ShareIndex>> {
121 self.node_id_to_node(id)?;
123
124 Ok(self
126 .share_ids_iter()
127 .filter(|share_id| self.share_id_to_node(share_id).expect("valid share id").id == id)
128 .collect::<Vec<_>>())
129 }
130
131 pub fn iter(&self) -> impl Iterator<Item = &Node<G>> {
133 self.nodes.iter()
134 }
135
136 pub fn hash(&self) -> Digest<32> {
138 let mut hash = Blake2b256::default();
139 hash.update(bcs::to_bytes(&self.nodes).expect("should serialize"));
140 hash.finalize()
141 }
142
143 pub fn new_reduced(
155 nodes_vec: Vec<Node<G>>,
156 t: u16,
157 allowed_delta: u16,
158 total_weight_lower_bound: u16,
159 ) -> FastCryptoResult<(Self, u16)> {
160 let n = Self::new(nodes_vec)?; assert!(total_weight_lower_bound <= n.total_weight && total_weight_lower_bound > 0);
162 let mut max_d = 1;
163 for d in 2..=40 {
164 let new_total_weight = n.nodes.iter().map(|n| n.weight / d).sum::<u16>();
167 if new_total_weight < total_weight_lower_bound {
168 break;
169 }
170 let delta = n.nodes.iter().map(|n| n.weight % d).sum::<u16>();
173 if delta <= allowed_delta {
174 max_d = d;
175 }
176 }
177 debug!(
178 "Nodes::reduce reducing from {} with max_d {}, allowed_delta {}, total_weight_lower_bound {}",
179 n.total_weight, max_d, allowed_delta, total_weight_lower_bound
180 );
181
182 let nodes = n
183 .nodes
184 .iter()
185 .map(|n| Node {
186 id: n.id,
187 pk: n.pk.clone(),
188 weight: n.weight / max_d,
189 })
190 .collect::<Vec<_>>();
191 let accumulated_weights = Self::get_accumulated_weights(&nodes);
192 let nodes_with_nonzero_weight = Self::filter_nonzero_weights(&nodes);
193 let total_weight = nodes.iter().map(|n| n.weight).sum::<u16>();
195 let new_t = t / max_d + (t % max_d != 0) as u16;
196 Ok((
197 Self {
198 nodes,
199 total_weight,
200 accumulated_weights,
201 nodes_with_nonzero_weight,
202 },
203 new_t,
204 ))
205 }
206}