dusk_consensus/user/
cluster.rs

1// This Source Code Form is subject to the terms of the Mozilla Public
2// License, v. 2.0. If a copy of the MPL was not distributed with this
3// file, You can obtain one at http://mozilla.org/MPL/2.0/.
4//
5// Copyright (c) DUSK NETWORK. All rights reserved.
6
7use std::collections::btree_map::Iter;
8use std::collections::BTreeMap;
9
10#[derive(Debug, Default)]
11pub struct Cluster<T>(BTreeMap<T, usize>);
12
13impl<T> Cluster<T>
14where
15    T: Default + std::cmp::Ord + Clone + std::fmt::Debug,
16{
17    pub(crate) fn new() -> Self {
18        Self(Default::default())
19    }
20
21    /// total_occurrences return the sum of all the values.
22    pub fn total_occurrences(&self) -> usize {
23        self.0.values().sum()
24    }
25
26    /// Adds key with specified weight. Weight per key can be set only once.
27    ///
28    /// Return None if `weight` is 0 or key is already set.
29    pub fn add(&mut self, key: &T, weight: usize) -> Option<usize> {
30        if weight == 0 {
31            return None;
32        }
33        if self.0.contains_key(key) {
34            // already updated
35            return None;
36        }
37
38        self.0.insert(key.clone(), weight);
39        Some(weight)
40    }
41
42    pub fn iter(&self) -> Iter<T, usize> {
43        self.0.iter()
44    }
45
46    pub fn into_vec(self) -> Vec<(T, usize)> {
47        self.0.into_iter().collect()
48    }
49
50    pub fn contains_key(&self, key: &T) -> bool {
51        self.0.contains_key(key)
52    }
53}
54
55#[cfg(test)]
56mod tests {
57    use crate::user::cluster::Cluster;
58
59    #[test]
60    pub fn test_set_weight() {
61        let mut a = Cluster::new();
62
63        a.add(&'a', 3);
64        a.add(&'b', 0);
65        a.add(&'b', 11);
66        assert_eq!(a.total_occurrences(), 14);
67
68        let res = a.add(&'b', 1);
69        assert!(res.is_none());
70        assert_eq!(a.total_occurrences(), 14);
71    }
72}
dusk_consensus/user/cluster.rs

dusk_consensus/user/
cluster.rs
