dfir_lang/graph/
di_mul_graph.rs

Help
1#![warn(missing_docs)]
2
3use std::collections::BTreeSet;
4use std::fmt::Debug;
5use std::iter::FusedIterator;
6
7use serde::{Deserialize, Serialize};
8use slotmap::{Key, SecondaryMap, SlotMap};
9
10/// A directed multigraph where an vertex's inbound and outbound edges are indexed.
11///
12/// `DiMulGraph` does **not** allocate vertices `V`. The user shall use an external
13/// [`SlotMap<V, _>`] for allocating vertices, which also allows the user to associate data with
14/// each vertex.
15///
16/// `DiMulGraph` **does** allocate edges `E` as they are added. Additional data can be associated
17/// with edges via an external [`SecondaryMap<E, _>`].
18#[derive(Clone, Debug, Serialize, Deserialize)]
19#[serde(from = "EdgeList<V, E>", into = "EdgeList<V, E>")]
20pub struct DiMulGraph<V, E>
21where
22    V: Key,
23    E: Key,
24{
25    /// Edge list (src, dst).
26    edges: SlotMap<E, (V, V)>,
27
28    /// Successors for each vert.
29    succs: SecondaryMap<V, Vec<E>>,
30    /// Predecessors for each vert.
31    preds: SecondaryMap<V, Vec<E>>,
32}
33impl<V, E> Default for DiMulGraph<V, E>
34where
35    V: Key,
36    E: Key,
37{
38    fn default() -> Self {
39        let (edges, succs, preds) = Default::default();
40        Self {
41            edges,
42            succs,
43            preds,
44        }
45    }
46}
47impl<V, E> DiMulGraph<V, E>
48where
49    V: Key,
50    E: Key,
51{
52    /// Creates an empty `DiMulGraph`.
53    pub fn new() -> Self {
54        Default::default()
55    }
56
57    /// Creates a `DiMulGraph` with pre-allocated memory for `capacity` vertices and `capacity`
58    /// edges.
59    pub fn with_capacity(capacity: usize) -> Self {
60        Self {
61            // Estimate 1 edge per vertex.
62            edges: SlotMap::with_capacity_and_key(capacity),
63            succs: SecondaryMap::with_capacity(capacity),
64            preds: SecondaryMap::with_capacity(capacity),
65        }
66    }
67
68    /// Assert that `self` is in a consistent state, for debugging. This is computationally
69    /// expensive for large graphs.
70    pub fn assert_valid(&self) {
71        // Ensure each edge exists in the adj lists.
72        for (edge_id, &(src, dst)) in self.edges.iter() {
73            assert!(self.succs[src].contains(&edge_id));
74            assert!(self.preds[dst].contains(&edge_id));
75        }
76
77        // Ensure no duplicate preds or succs
78        for succ_list in self.succs.values() {
79            let set: BTreeSet<&E> = succ_list.iter().collect();
80            assert_eq!(set.len(), succ_list.len());
81        }
82        for pred_list in self.succs.values() {
83            let set: BTreeSet<&E> = pred_list.iter().collect();
84            assert_eq!(set.len(), pred_list.len());
85        }
86
87        // Note: Missing edges and duplicate edges could cancel each other out. But that case is
88        // caught by the above.
89        assert_eq!(
90            self.edges.len(),
91            self.succs.values().map(|vec| vec.len()).sum::<usize>(),
92            "succs broken (contains duplicate or removed edge?)"
93        );
94        assert_eq!(
95            self.edges.len(),
96            self.preds.values().map(|vec| vec.len()).sum::<usize>(),
97            "preds broken (contains duplicate or removed edge?)"
98        );
99    }
100
101    /// HELPER, get the list out of the adj list `adj_list` for a particular vertex `v`.
102    fn get_adj_edges(adj_list: &mut SecondaryMap<V, Vec<E>>, v: V) -> &mut Vec<E> {
103        if !adj_list.contains_key(v) {
104            adj_list.insert(v, Default::default());
105        }
106        &mut adj_list[v]
107    }
108
109    /// Creates an edge going from `src` to `dst` and returns the edge ID.
110    pub fn insert_edge(&mut self, src: V, dst: V) -> E {
111        let e = self.edges.insert((src, dst));
112        Self::get_adj_edges(&mut self.succs, src).push(e);
113        Self::get_adj_edges(&mut self.preds, dst).push(e);
114        e
115    }
116
117    /// For an `edge` from `A --> B`, insert a new vertex `V` along that edge to create
118    /// `A --e0--> V --e1--> B`. Returns the pair of new edge IDs in and out of `V`, i.e.
119    /// `(e0, e1)`.
120    ///
121    /// Returns `None` if the edge doesn't exist.
122    ///
123    /// `edge` is removed from the graph, both returned edge IDs are new.
124    pub fn insert_intermediate_vertex(&mut self, new_vertex: V, edge: E) -> Option<(E, E)> {
125        self.assert_valid();
126
127        // Remove old edge from edges.
128        let (src, dst) = self.edges.remove(edge)?;
129
130        // Insert new edges into edges.
131        let e0 = self.edges.insert((src, new_vertex));
132        let e1 = self.edges.insert((new_vertex, dst));
133
134        // Remove old & add new edges in succs/preds.
135        let succ_vec_idx = self.succs[src].iter().position(|&e| e == edge).unwrap();
136        let pred_vec_idx = self.preds[dst].iter().position(|&e| e == edge).unwrap();
137        assert_eq!(
138            edge,
139            std::mem::replace(&mut self.succs[src][succ_vec_idx], e0)
140        );
141        assert_eq!(
142            edge,
143            std::mem::replace(&mut self.preds[dst][pred_vec_idx], e1)
144        );
145
146        // Insert new vertex succs/preds.
147        assert!(
148            self.preds.insert(new_vertex, vec![e0]).is_none(),
149            "Cannot insert intermediate vertex that already exists"
150        );
151        assert!(
152            self.succs.insert(new_vertex, vec![e1]).is_none(),
153            "Cannot insert intermediate vertex that already exists"
154        );
155
156        self.assert_valid();
157        Some((e0, e1))
158    }
159
160    /// For a vertex with one incoming edge and one outgoing edge, removes the vertex. Inserts a new edge.
161    /// Returns `(new edge, (old edge in, old edge out))`.
162    /// Returns `None` if `vertex` is not in the graph or does not have the right degree in/out.
163    pub fn remove_intermediate_vertex(&mut self, vertex: V) -> Option<(E, (E, E))> {
164        let preds = self.preds.remove(vertex)?;
165        let &[pred_edge] = &*preds else {
166            return None;
167        };
168        let succs = self.succs.remove(vertex).unwrap();
169        let &[succ_edge] = &*succs else {
170            return None;
171        };
172
173        let (src, _v) = self.edges.remove(pred_edge).unwrap();
174        let (_v, dst) = self.edges.remove(succ_edge).unwrap();
175
176        self.succs[src].retain(|&e| e != pred_edge);
177        self.preds[dst].retain(|&e| e != succ_edge);
178
179        let new_edge = self.insert_edge(src, dst);
180        Some((new_edge, (pred_edge, succ_edge)))
181    }
182
183    /// Remove an edge from the graph. If the edgeId is found then the edge is removed from the graph and returned.
184    /// If the edgeId was not found in the graph then nothing is returned and nothing is done.
185    pub fn remove_edge(&mut self, e: E) -> Option<(V, V)> {
186        let (src, dst) = self.edges.remove(e)?;
187
188        self.succs[src].retain(|x| *x != e);
189        self.preds[dst].retain(|x| *x != e);
190
191        Some((src, dst))
192    }
193
194    /// Remove a vertex from the graph, it must have no edges to or from it when doing this.
195    pub fn remove_vertex(&mut self, v: V) {
196        assert!(self.preds[v].is_empty() && self.succs[v].is_empty());
197
198        self.preds.remove(v);
199        self.succs.remove(v);
200    }
201
202    /// Get the source and destination vertex IDs for the given edge ID.
203    pub fn edge(&self, e: E) -> Option<(V, V)> {
204        self.edges.get(e).copied()
205    }
206
207    /// Return an iterator over all edge IDs `E`.
208    pub fn edge_ids(&self) -> slotmap::basic::Keys<E, (V, V)> {
209        self.edges.keys()
210    }
211
212    /// Return an iterator over all edges in form `(E, (V, V))`.
213    pub fn edges(
214        &self,
215    ) -> impl '_ + ExactSizeIterator<Item = (E, (V, V))> + FusedIterator + Clone + Debug {
216        self.edges.iter().map(|(e, &(src, dst))| (e, (src, dst)))
217    }
218
219    /// Return an iterator of all edge IDs coming out of `v`.
220    pub fn successor_edges(&self, v: V) -> std::iter::Copied<std::slice::Iter<'_, E>> {
221        self.succs
222            .get(v)
223            .map(|v| v.iter())
224            .unwrap_or_else(|| [].iter())
225            .copied()
226    }
227
228    /// Return an iterator of all edge IDs going into `v`.
229    pub fn predecessor_edges(&self, v: V) -> std::iter::Copied<std::slice::Iter<'_, E>> {
230        self.preds
231            .get(v)
232            .map(|v| v.iter())
233            .unwrap_or_else(|| [].iter())
234            .copied()
235    }
236
237    /// Return an iterator of all successor vertex IDs of `v`.
238    ///
239    /// If there are multiple edges from `v` to the same vertex, that vertex will appear multiple times in the iterator.
240    pub fn successor_vertices(
241        &self,
242        v: V,
243    ) -> impl '_ + DoubleEndedIterator<Item = V> + ExactSizeIterator + FusedIterator + Clone + Debug
244    {
245        self.successor_edges(v).map(|edge_id| self.edges[edge_id].1)
246    }
247
248    /// Return an iterator of all predecessor vertex IDs of `v`.
249    ///
250    /// If there are multiple edges from a vertex to `v`, that vertex will appear multiple times in the iterator.
251    pub fn predecessor_vertices(
252        &self,
253        v: V,
254    ) -> impl '_ + DoubleEndedIterator<Item = V> + ExactSizeIterator + FusedIterator + Clone + Debug
255    {
256        self.predecessor_edges(v)
257            .map(|edge_id| self.edges[edge_id].0)
258    }
259
260    /// Return an iterator of all successor edge IDs _and_ vertex IDs of `v` in form `(E, V)`.
261    pub fn successors(
262        &self,
263        v: V,
264    ) -> impl '_ + DoubleEndedIterator<Item = (E, V)> + ExactSizeIterator + FusedIterator + Clone + Debug
265    {
266        self.successor_edges(v)
267            .map(|edge_id| (edge_id, self.edges[edge_id].1))
268    }
269
270    /// Return an iterator of all predecessor edge IDs _and_ vertex IDs of `v` in form `(E, V)`.
271    pub fn predecessors(
272        &self,
273        v: V,
274    ) -> impl '_ + DoubleEndedIterator<Item = (E, V)> + ExactSizeIterator + FusedIterator + Clone + Debug
275    {
276        self.predecessor_edges(v)
277            .map(|edge_id| (edge_id, self.edges[edge_id].0))
278    }
279
280    /// The degree (number of edges/vertices) coming out of `v`, i.e. the number of successors.
281    pub fn degree_out(&self, v: V) -> usize {
282        self.succs.get(v).map(Vec::len).unwrap_or_default()
283    }
284
285    /// The degree (number of edges/vertices) going into `v`, i.e. the number of predecessors.
286    pub fn degree_in(&self, v: V) -> usize {
287        self.preds.get(v).map(Vec::len).unwrap_or_default()
288    }
289}
290
291impl<V, E> From<DiMulGraph<V, E>> for EdgeList<V, E>
292where
293    V: Key,
294    E: Key,
295{
296    fn from(value: DiMulGraph<V, E>) -> Self {
297        value.edges
298    }
299}
300
301impl<V, E> From<EdgeList<V, E>> for DiMulGraph<V, E>
302where
303    V: Key,
304    E: Key,
305{
306    fn from(edges: EdgeList<V, E>) -> Self {
307        let mut out = Self {
308            edges,
309            ..Default::default()
310        };
311        for (edge, &(src, dst)) in out.edges.iter() {
312            out.succs.entry(src).unwrap().or_default().push(edge);
313            out.preds.entry(dst).unwrap().or_default().push(edge);
314        }
315        out
316    }
317}
318
319/// A compact edge list representation of a [`DiMulGraph`], used for serialization.
320#[expect(type_alias_bounds, reason = "code readability")]
321pub type EdgeList<V, E>
322where
323    V: Key,
324    E: Key,
325= SlotMap<E, (V, V)>;
dfir_lang/graph/di_mul_graph.rs

dfir_lang/graph/
di_mul_graph.rs
