flmodules/random_connections/
nodes.rs

1use core::cmp::min;
2use flarch::nodeids::{NodeID, NodeIDs};
3use serde::{Deserialize, Serialize};
4
5#[derive(PartialEq, Clone, Debug, Serialize, Deserialize)]
6pub struct NodeTime {
7    pub id: NodeID,
8    pub ticks: u32,
9}
10
11#[derive(PartialEq, Clone, Debug, Serialize, Deserialize, Default)]
12pub struct Nodes(pub Vec<NodeTime>);
13
14impl Nodes {
15    pub fn new() -> Nodes {
16        Self { 0: vec![] }
17    }
18
19    pub fn get_nodes(&self) -> NodeIDs {
20        NodeIDs {
21            0: self.0.iter().map(|n| n.id).collect(),
22        }
23    }
24
25    pub fn remove_missing(&mut self, nodes: &NodeIDs) -> NodeIDs {
26        let removed = self.get_nodes().remove_missing(nodes);
27        self.0.retain(|nt| !removed.0.contains(&nt.id));
28        removed
29    }
30
31    pub fn remove(&mut self, nodes: &NodeIDs) {
32        self.0.retain(|nt| !nodes.0.contains(&nt.id));
33    }
34
35    pub fn add_new(&mut self, nodes: Vec<NodeID>) {
36        let nodes_new: Vec<NodeID> = nodes.into_iter().filter(|n| !self.contains(n)).collect();
37        let mut nts = nodes_new
38            .iter()
39            .map(|n| NodeTime { id: *n, ticks: 0 })
40            .collect();
41        self.0.append(&mut nts);
42    }
43
44    pub fn contains(&self, node: &NodeID) -> bool {
45        self.get_nodes().0.contains(node)
46    }
47
48    // Removes the oldest n nodes. If less than n nodes are stored, only remove
49    // these nodes.
50    // It returns the removed nodes.
51    pub fn remove_oldest_n(&mut self, mut n: usize) -> NodeIDs {
52        self.0.sort_by(|a, b| b.ticks.cmp(&a.ticks));
53        n = min(n, self.0.len());
54        NodeIDs {
55            0: self.0.splice(..n, vec![]).map(|n| n.id).collect(),
56        }
57    }
58
59    // Removes the newest n nodes. If less than n nodes are stored, only remove
60    // these nodes.
61    // It returns the removed nodes.
62    pub fn remove_newest_n(&mut self, mut n: usize) -> NodeIDs {
63        self.0.sort_by(|a, b| a.ticks.cmp(&b.ticks));
64        n = min(n, self.0.len());
65        NodeIDs {
66            0: self.0.splice(..n, vec![]).map(|n| n.id).collect(),
67        }
68    }
69
70    // Only keeps the newest n nodes. If there are more than n nodes, the
71    // oldest nodes are discarded.
72    pub fn keep_newest_n(&mut self, n: usize) -> NodeIDs {
73        if self.0.len() > n {
74            self.remove_oldest_n(self.0.len() - n)
75        } else {
76            NodeIDs::empty()
77        }
78    }
79
80    // Only keeps the oldest n nodes. If there are more than n nodes, the
81    // newest nodes are discarded.
82    pub fn keep_oldest_n(&mut self, n: usize) -> NodeIDs {
83        if self.0.len() > n {
84            self.remove_newest_n(self.0.len() - n)
85        } else {
86            NodeIDs::empty()
87        }
88    }
89
90    // Returns nodes that are as old or older than age.
91    pub fn oldest_ticks(&self, ticks: u32) -> NodeIDs {
92        let nodes: Vec<NodeID> = self
93            .sorted(false)
94            .0
95            .iter()
96            .take_while(|n| n.ticks >= ticks)
97            .map(|n| n.id)
98            .collect();
99        NodeIDs { 0: nodes }
100    }
101
102    pub fn sorted(&self, rising: bool) -> Nodes {
103        let mut sorted = self.0.clone();
104        sorted.sort_by(|a, b| a.ticks.cmp(&b.ticks));
105        if !rising {
106            sorted.reverse();
107        }
108        Nodes { 0: sorted }
109    }
110
111    // Counts number of nodes that are as old or older than ticks.
112    pub fn count_expired(&self, ticks: u32) -> usize {
113        self.0
114            .iter()
115            .fold(0, |a, n| (n.ticks >= ticks).then(|| a + 1).unwrap_or(a))
116    }
117
118    // Removes nodes that are as old or older than age and returns the removed nodes.
119    pub fn remove_oldest_ticks(&mut self, ticks: u32) -> NodeIDs {
120        let nodes = self.oldest_ticks(ticks);
121        self.0.splice(..nodes.0.len(), vec![]);
122        nodes
123    }
124
125    // Increases the tick of all nodes by 1
126    pub fn tick(&mut self) {
127        for node in self.0.iter_mut() {
128            node.ticks += 1;
129        }
130    }
131}
132
133#[cfg(test)]
134mod tests {
135    use flarch::start_logging;
136
137    use super::*;
138
139    #[test]
140    fn expired() {
141        let mut n = Nodes::new();
142        n.add_new(NodeIDs::new(3).0);
143        let mut ticks = 1;
144        for node in n.0.iter_mut() {
145            node.ticks = ticks;
146            ticks += 1;
147        }
148
149        assert_eq!(0, n.count_expired(4));
150        assert_eq!(1, n.count_expired(3));
151        assert_eq!(2, n.count_expired(2));
152        assert_eq!(3, n.count_expired(1));
153    }
154
155    fn make_nodes(n: usize) -> Nodes {
156        let mut nodes = Nodes::new();
157        nodes.add_new(NodeIDs::new(n as u32).0);
158        for node in 0..n {
159            nodes.0[node].ticks = node as u32 + 1;
160        }
161        nodes
162    }
163
164    // Tests the nodes and the remove methods
165    #[test]
166    fn test_nodes_remove() -> anyhow::Result<()> {
167        start_logging();
168
169        let mut nodes = make_nodes(5);
170        let mut removed = nodes.remove_oldest_n(2);
171        assert_eq!(removed.0.len(), 2);
172        assert_eq!(nodes.0.len(), 3);
173
174        removed = nodes.remove_oldest_n(5);
175        assert_eq!(removed.0.len(), 3);
176        assert_eq!(nodes.0.len(), 0);
177
178        nodes = make_nodes(5);
179        removed = nodes.remove_oldest_ticks(6);
180        assert_eq!(nodes.0.len(), 5);
181        assert_eq!(removed.0.len(), 0);
182
183        removed = nodes.remove_oldest_ticks(4);
184        assert_eq!(nodes.0.len(), 3);
185        assert_eq!(removed.0.len(), 2);
186
187        Ok(())
188    }
189}