smallgraph/
lib.rs

1#![no_std]
2use smallvec::SmallVec;
3
4pub type NodeIndex = usize;
5pub type Generation = usize;
6pub type NodeHandle = (NodeIndex, Generation);
7
8pub struct SmallGraph<T> {
9    pub(crate) free: SmallVec<[NodeHandle; 128]>,
10    pub(crate) nodes: SmallVec<[(Generation, Option<T>); 128]>,
11    pub(crate) connections: SmallVec<[(NodeIndex, NodeIndex); 256]>,
12}
13
14impl<T> SmallGraph<T> {
15    /// Create a new SmallGraph
16    pub fn new() -> SmallGraph<T> {
17        SmallGraph {
18            free: SmallVec::new(),
19            nodes: SmallVec::new(),
20            connections: SmallVec::new(),
21        }
22    }
23
24    /// Insert a value into graph
25    pub fn insert(&mut self, value: T) -> NodeHandle {
26        if self.free.len() == 0 {
27            let index = self.nodes.len();
28            let gen = 0;
29            self.nodes.push((gen, Some(value)));
30            (index, gen)
31        } else {
32            let n = self.free.remove(0);
33            let (index, gen) = n;
34            self.nodes[index] = (gen + 1, Some(value));
35            (n.0, gen + 1)
36        }
37    }
38
39    /// Create a directed connection between parent and child
40    pub fn connect_to(&mut self, parent: NodeHandle, child: NodeHandle) {
41        self.connections.push((parent.0, child.0));
42    }
43
44    /// Get nodes this node has an edge to
45    pub fn neighbors(&self, node: NodeHandle) -> SmallVec<[NodeHandle; 8]> {
46        let mut children = SmallVec::new();
47        for i in 0..self.connections.len() {
48            if self.connections[i].0 == node.0 {
49                let child_handle = (
50                    self.connections[i].1,
51                    self.nodes[self.connections[i].1].0,
52                );
53                children.push(child_handle);
54            }
55        }
56        children
57    }
58
59    /// Get nodes that have an edge that can reach node
60    pub fn nodes_with_neighbor(&self, node: NodeHandle) -> SmallVec<[NodeHandle; 8]> {
61        let mut children = SmallVec::new();
62        for i in 0..self.connections.len() {
63            if self.connections[i].1 == node.0 {
64                let child_handle = (
65                    self.connections[i].1,
66                    self.nodes[self.connections[i].1].0,
67                );
68                children.push(child_handle);
69            }
70        }
71        children
72    }
73
74    /// Create a two way connection between two nodes
75    pub fn connect(&mut self, a: NodeHandle, b: NodeHandle) {
76        self.connections.push((a.0, b.0));
77        self.connections.push((b.0, a.0));
78    }
79
80    /// Disconnect all connections a node has
81    pub fn disconnect_all(&mut self, n: NodeHandle) {
82        self.connections
83            .retain(|&mut connection| (connection).0 != n.0 && (connection).1 != n.0);
84    }
85
86    /// Disconnect all connections between two nodes
87    pub fn disconnect(&mut self, a: NodeHandle, b: NodeHandle) {
88        self.connections.retain(|&mut connection| {
89            !((connection.0 == a.0 && connection.1 == b.1)
90                && (connection.0 == b.0 && connection.1 == a.1))
91        });
92    }
93
94    /// Disconnect edge connection between source and destination
95    pub fn disconnect_from(&mut self, source: NodeHandle, destination: NodeHandle) {
96        self.connections
97            .retain(|&mut connection| !(connection.0 == source.0 && connection.1 == destination.1));
98    }
99
100    /// Determine if there is a connection connection between a source and destination node
101    pub fn is_connected_to(&mut self, source: NodeHandle, destination: NodeHandle) -> bool {
102        self.connections
103            .iter()
104            .find(|&connection| connection.0 == source.0 && connection.1 == destination.0)
105            .is_some()
106    }
107
108    /// Remove a node and it's connections from graph
109    pub fn remove(&mut self, n: NodeHandle) -> Option<T> {
110        let (index, gen) = n;
111        if self.nodes[index].0 == gen {
112            self.disconnect_all(n);
113            let mut r = (gen + 1, None);
114            core::mem::swap(&mut self.nodes[index], &mut r);
115            self.free.push(n);
116            return Some(r.1.unwrap());
117        }
118        None
119    }
120
121    /// Returns the count of nodes
122    pub fn node_count(&self) -> usize {
123        let mut c = 0;
124        for i in 0..self.nodes.len() {
125            if self.nodes[i].1.is_some() {
126                c += 1;
127            }
128        }
129        c
130    }
131
132    /// Get the value of a node
133    pub fn get(&self, n: NodeHandle) -> Option<&T> {
134        let (index, gen) = n;
135        if self.nodes[index].0 == gen {
136            if let Some(value) = &self.nodes[index].1 {
137                return Some(value);
138            }
139        }
140        None
141    }
142
143    /// Get a mutable value of a node
144    pub fn get_mut(&mut self, n: NodeHandle) -> Option<&mut T> {
145        let (index, gen) = n;
146        if self.nodes[index].0 == gen {
147            if let Some(value) = &mut self.nodes[index].1 {
148                return Some(value);
149            }
150        }
151        None
152    }
153}
154
155#[cfg(test)]
156mod tests {
157    // Note this useful idiom: importing names from outer (for mod tests) scope.
158    use super::*;
159
160    #[derive(Debug, PartialEq)]
161    struct Foo {
162        v: u8,
163    }
164
165    #[test]
166    fn test_basic_0() {
167        let g = SmallGraph::<Foo>::new();
168        assert_eq!(0, g.nodes.len());
169        assert_eq!(0, g.free.len());
170        assert_eq!(0, g.connections.len());
171    }
172
173    #[test]
174    fn test_basic_1() {
175        let mut g = SmallGraph::<Foo>::new();
176        let f1 = g.insert(Foo { v: 24 });
177        let f2 = g.insert(Foo { v: 42 });
178        assert_eq!(2, g.nodes.len());
179        assert_eq!(0, f1.0);
180        assert_eq!(0, f1.1);
181        assert_eq!(1, f2.0);
182        assert_eq!(0, f2.1);
183    }
184
185    #[test]
186    fn test_basic_2() {
187        let mut g = SmallGraph::<Foo>::new();
188        let f1 = g.insert(Foo { v: 42 });
189        let r = g.remove(f1).expect("could not remove");
190        assert_eq!(42, r.v);
191        let f2 = g.insert(Foo { v: 55 });
192        assert_eq!(0, f2.0);
193        assert_eq!(1, f2.1);
194        assert_eq!(None, g.get(f1));
195        assert_eq!(55, g.get(f2).unwrap().v);
196    }
197
198    #[test]
199    fn test_basic_3() {
200        let mut g = SmallGraph::<Foo>::new();
201        let f1 = g.insert(Foo { v: 24 });
202        let f2 = g.insert(Foo { v: 42 });
203        g.connect(f1, f2);
204        assert_eq!(2, g.connections.len());
205        assert_eq!((0, 1), g.connections[0]);
206        assert_eq!((1, 0), g.connections[1]);
207    }
208
209    #[test]
210    fn test_basic_4() {
211        let mut g = SmallGraph::<Foo>::new();
212        let f1 = g.insert(Foo { v: 24 });
213        let f2 = g.insert(Foo { v: 42 });
214        let f3 = g.insert(Foo { v: 33 });
215        g.connect_to(f1, f2);
216        g.connect_to(f2, f3);
217        g.connect_to(f1, f3);
218        assert_eq!(3, g.connections.len());
219        assert_eq!((0, 1), g.connections[0]);
220        assert_eq!((1, 2), g.connections[1]);
221        g.remove(f2);
222        assert_eq!(1, g.connections.len());
223        assert_eq!(true, g.is_connected_to(f1, f3));
224        assert_eq!(false, g.is_connected_to(f1, f2));
225        assert_eq!(false, g.is_connected_to(f2, f3));
226    }
227
228    #[test]
229    fn test_basic_5() {
230        let mut g = SmallGraph::<Foo>::new();
231        g.insert(Foo { v: 24 });
232        g.insert(Foo { v: 42 });
233        let f3 = g.insert(Foo { v: 33 });
234        assert_eq!(3, g.node_count());
235        g.remove(f3);
236        assert_eq!(2, g.node_count());
237    }
238}