| use crate::graph::{DirectedGraph, NumEdges, Predecessors, Successors}; |
| use rustc_index::{Idx, IndexVec}; |
| |
| #[cfg(test)] |
| mod tests; |
| |
| /// A directed graph, efficient for cases where node indices are pre-existing. |
| /// |
| /// If `BR` is true, the graph will store back-references, allowing you to get predecessors. |
| pub struct VecGraph<N: Idx, const BR: bool = false> { |
| // This is basically a `HashMap<N, (Vec<N>, If<BR, Vec<N>>)>` -- a map from a node index, to |
| // a list of targets of outgoing edges and (if enabled) a list of sources of incoming edges. |
| // |
| // However, it is condensed into two arrays as an optimization. |
| // |
| // `node_starts[n]` is the start of the list of targets of outgoing edges for node `n`. |
| // So you can get node's successors with `edge_targets[node_starts[n]..node_starts[n + 1]]`. |
| // |
| // If `BR` is true (back references are enabled), then `node_starts[n + edge_count]` is the |
| // start of the list of *sources* of incoming edges. You can get predecessors of a node |
| // similarly to its successors but offsetting by `edge_count`. `edge_count` is |
| // `edge_targets.len()/2` (again, in case BR is true) because half of the vec is back refs. |
| // |
| // All of this might be confusing, so here is an example graph and its representation: |
| // |
| // n3 ----+ |
| // ^ | (if BR = true) |
| // | v outgoing edges incoming edges |
| // n0 -> n1 -> n2 ______________ __________________ |
| // / \ / \ |
| // node indices[1]: n0, n1, n2, n3, n0, n1, n2, n3, n/a |
| // vec indices: n0, n1, n2, n3, n4, n5, n6, n7, n8 |
| // node_starts: [0, 1, 3, 4 4, 4, 5, 7, 8] |
| // | | | | | | | | | |
| // | | +---+ +---+ | +---+ | |
| // | | | | | | | |
| // v v v v v v v |
| // edge_targets: [n1, n2, n3, n2 n0, n1, n3, n1] |
| // / \____/ | | \____/ \ |
| // n0->n1 / | | \ n3<-n1 |
| // / n3->n2 [2] n1<-n0 [2] \ |
| // n1->n2, n1->n3 n2<-n1, n2<-n3 |
| // |
| // The incoming edges are basically stored in the same way as outgoing edges, but offset and |
| // the graph they store is the inverse of the original. Last index in the `node_starts` array |
| // always points to one-past-the-end, so that we don't need to bound check `node_starts[n + 1]` |
| // |
| // [1]: "node indices" are the indices a user of `VecGraph` might use, |
| // note that they are different from "vec indices", |
| // which are the real indices you need to index `node_starts` |
| // |
| // [2]: Note that even though n2 also points to here, |
| // the next index also points here, so n2 has no |
| // successors (`edge_targets[3..3] = []`). |
| // Similarly with n0 and incoming edges |
| // |
| // If this is still confusing... then sorry :( |
| // |
| /// Indices into `edge_targets` that signify a start of list of edges. |
| node_starts: IndexVec<N, usize>, |
| |
| /// Targets (or sources for back refs) of edges |
| edge_targets: Vec<N>, |
| } |
| |
| impl<N: Idx + Ord, const BR: bool> VecGraph<N, BR> { |
| pub fn new(num_nodes: usize, mut edge_pairs: Vec<(N, N)>) -> Self { |
| let num_edges = edge_pairs.len(); |
| |
| let nodes_cap = match BR { |
| // +1 for special entry at the end, pointing one past the end of `edge_targets` |
| false => num_nodes + 1, |
| // *2 for back references |
| true => (num_nodes * 2) + 1, |
| }; |
| |
| let edges_cap = match BR { |
| false => num_edges, |
| // *2 for back references |
| true => num_edges * 2, |
| }; |
| |
| let mut node_starts = IndexVec::with_capacity(nodes_cap); |
| let mut edge_targets = Vec::with_capacity(edges_cap); |
| |
| // Sort the edges by the source -- this is important. |
| edge_pairs.sort(); |
| |
| // Fill forward references |
| create_index( |
| num_nodes, |
| &mut edge_pairs.iter().map(|&(src, _)| src), |
| &mut edge_pairs.iter().map(|&(_, tgt)| tgt), |
| &mut edge_targets, |
| &mut node_starts, |
| ); |
| |
| // Fill back references |
| if BR { |
| // Pop the special "last" entry, it will be replaced by first back ref |
| node_starts.pop(); |
| |
| // Re-sort the edges so that they are sorted by target |
| edge_pairs.sort_by_key(|&(src, tgt)| (tgt, src)); |
| |
| create_index( |
| // Back essentially double the number of nodes |
| num_nodes * 2, |
| // NB: the source/target are switched here too |
| // NB: we double the key index, so that we can later use *2 to get the back references |
| &mut edge_pairs.iter().map(|&(_, tgt)| N::new(tgt.index() + num_nodes)), |
| &mut edge_pairs.iter().map(|&(src, _)| src), |
| &mut edge_targets, |
| &mut node_starts, |
| ); |
| } |
| |
| Self { node_starts, edge_targets } |
| } |
| |
| /// Gets the successors for `source` as a slice. |
| pub fn successors(&self, source: N) -> &[N] { |
| assert!(source.index() < self.num_nodes()); |
| |
| let start_index = self.node_starts[source]; |
| let end_index = self.node_starts[source.plus(1)]; |
| &self.edge_targets[start_index..end_index] |
| } |
| } |
| |
| impl<N: Idx + Ord> VecGraph<N, true> { |
| /// Gets the predecessors for `target` as a slice. |
| pub fn predecessors(&self, target: N) -> &[N] { |
| assert!(target.index() < self.num_nodes()); |
| |
| let target = N::new(target.index() + self.num_nodes()); |
| |
| let start_index = self.node_starts[target]; |
| let end_index = self.node_starts[target.plus(1)]; |
| &self.edge_targets[start_index..end_index] |
| } |
| } |
| |
| /// Creates/initializes the index for the [`VecGraph`]. A helper for [`VecGraph::new`]. |
| /// |
| /// - `num_nodes` is the target number of nodes in the graph |
| /// - `sorted_edge_sources` are the edge sources, sorted |
| /// - `associated_edge_targets` are the edge *targets* in the same order as sources |
| /// - `edge_targets` is the vec of targets to be extended |
| /// - `node_starts` is the index to be filled |
| fn create_index<N: Idx + Ord>( |
| num_nodes: usize, |
| sorted_edge_sources: &mut dyn Iterator<Item = N>, |
| associated_edge_targets: &mut dyn Iterator<Item = N>, |
| edge_targets: &mut Vec<N>, |
| node_starts: &mut IndexVec<N, usize>, |
| ) { |
| let offset = edge_targets.len(); |
| |
| // Store the *target* of each edge into `edge_targets`. |
| edge_targets.extend(associated_edge_targets); |
| |
| // Create the *edge starts* array. We are iterating over the |
| // (sorted) edge pairs. We maintain the invariant that the |
| // length of the `node_starts` array is enough to store the |
| // current source node -- so when we see that the source node |
| // for an edge is greater than the current length, we grow the |
| // edge-starts array by just enough. |
| for (index, source) in sorted_edge_sources.enumerate() { |
| // If we have a list like `[(0, x), (2, y)]`: |
| // |
| // - Start out with `node_starts` of `[]` |
| // - Iterate to `(0, x)` at index 0: |
| // - Push one entry because `node_starts.len()` (0) is <= the source (0) |
| // - Leaving us with `node_starts` of `[0]` |
| // - Iterate to `(2, y)` at index 1: |
| // - Push one entry because `node_starts.len()` (1) is <= the source (2) |
| // - Push one entry because `node_starts.len()` (2) is <= the source (2) |
| // - Leaving us with `node_starts` of `[0, 1, 1]` |
| // - Loop terminates |
| while node_starts.len() <= source.index() { |
| node_starts.push(index + offset); |
| } |
| } |
| |
| // Pad out the `node_starts` array so that it has `num_nodes + |
| // 1` entries. Continuing our example above, if `num_nodes` is |
| // be `3`, we would push one more index: `[0, 1, 1, 2]`. |
| // |
| // Interpretation of that vector: |
| // |
| // [0, 1, 1, 2] |
| // ---- range for N=2 |
| // ---- range for N=1 |
| // ---- range for N=0 |
| while node_starts.len() <= num_nodes { |
| node_starts.push(edge_targets.len()); |
| } |
| |
| assert_eq!(node_starts.len(), num_nodes + 1); |
| } |
| |
| impl<N: Idx, const BR: bool> DirectedGraph for VecGraph<N, BR> { |
| type Node = N; |
| |
| fn num_nodes(&self) -> usize { |
| match BR { |
| false => self.node_starts.len() - 1, |
| // If back refs are enabled, half of the array is said back refs |
| true => (self.node_starts.len() - 1) / 2, |
| } |
| } |
| } |
| |
| impl<N: Idx, const BR: bool> NumEdges for VecGraph<N, BR> { |
| fn num_edges(&self) -> usize { |
| match BR { |
| false => self.edge_targets.len(), |
| // If back refs are enabled, half of the array is reversed edges for them |
| true => self.edge_targets.len() / 2, |
| } |
| } |
| } |
| |
| impl<N: Idx + Ord, const BR: bool> Successors for VecGraph<N, BR> { |
| fn successors(&self, node: N) -> impl Iterator<Item = Self::Node> { |
| self.successors(node).iter().cloned() |
| } |
| } |
| |
| impl<N: Idx + Ord> Predecessors for VecGraph<N, true> { |
| fn predecessors(&self, node: Self::Node) -> impl Iterator<Item = Self::Node> { |
| self.predecessors(node).iter().cloned() |
| } |
| } |
