Skip to main content

weavatrix_graph/algo/
shortest.rs

1use super::traversal::{Direction, for_each_adjacent};
2use crate::IndexGraphView;
3use std::cmp::Reverse;
4use std::collections::BinaryHeap;
5
6#[derive(Debug, Clone, PartialEq, Eq)]
7pub struct WeightedPath<Node> {
8    nodes: Vec<Node>,
9    total_cost: u64,
10}
11
12impl<Node> WeightedPath<Node> {
13    pub(super) fn from_parts(nodes: Vec<Node>, total_cost: u64) -> Self {
14        Self { nodes, total_cost }
15    }
16
17    #[must_use]
18    pub fn nodes(&self) -> &[Node] {
19        &self.nodes
20    }
21
22    #[must_use]
23    pub const fn total_cost(&self) -> u64 {
24        self.total_cost
25    }
26
27    #[must_use]
28    pub fn into_nodes(self) -> Vec<Node> {
29        self.nodes
30    }
31}
32
33pub fn dijkstra<G, F>(
34    graph: &G,
35    source: G::Node,
36    target: G::Node,
37    mut edge_cost: F,
38) -> Option<WeightedPath<G::Node>>
39where
40    G: IndexGraphView,
41    F: FnMut(G::Edge) -> u64,
42{
43    dijkstra_filtered(graph, source, target, Direction::Outgoing, |edge| {
44        Some(edge_cost(edge))
45    })
46}
47
48pub fn dijkstra_filtered<G, F>(
49    graph: &G,
50    source: G::Node,
51    target: G::Node,
52    direction: Direction,
53    mut edge_cost: F,
54) -> Option<WeightedPath<G::Node>>
55where
56    G: IndexGraphView,
57    F: FnMut(G::Edge) -> Option<u64>,
58{
59    if !graph.contains_node(source) || !graph.contains_node(target) {
60        return None;
61    }
62    let bound = graph.node_bound();
63    let mut nodes = vec![None; bound];
64    for node in graph.node_indices() {
65        nodes[G::node_slot(node)] = Some(node);
66    }
67    let mut costs = vec![u64::MAX; bound];
68    let mut predecessor = vec![None; bound];
69    let source_slot = G::node_slot(source);
70    costs[source_slot] = 0;
71    let mut queue = BinaryHeap::new();
72    queue.push(Reverse((0_u64, source_slot)));
73
74    while let Some(Reverse((cost, slot))) = queue.pop() {
75        if cost != costs[slot] {
76            continue;
77        }
78        let Some(node) = nodes[slot] else {
79            continue;
80        };
81        if node == target {
82            return Some(WeightedPath {
83                nodes: reconstruct::<G>(source, target, &predecessor)?,
84                total_cost: cost,
85            });
86        }
87        relax_neighbors(
88            graph,
89            node,
90            direction,
91            cost,
92            &mut edge_cost,
93            &mut costs,
94            &mut predecessor,
95            &mut queue,
96        );
97    }
98    None
99}
100
101#[allow(clippy::too_many_arguments)]
102fn relax_neighbors<G, F>(
103    graph: &G,
104    node: G::Node,
105    direction: Direction,
106    cost: u64,
107    edge_cost: &mut F,
108    costs: &mut [u64],
109    predecessor: &mut [Option<G::Node>],
110    queue: &mut BinaryHeap<Reverse<(u64, usize)>>,
111) where
112    G: IndexGraphView,
113    F: FnMut(G::Edge) -> Option<u64>,
114{
115    for_each_adjacent(graph, node, direction, &mut |_| true, |edge, neighbor| {
116        let Some(weight) = edge_cost(edge) else {
117            return;
118        };
119        let Some(candidate) = cost.checked_add(weight) else {
120            return;
121        };
122        let slot = G::node_slot(neighbor);
123        if candidate < costs[slot] {
124            costs[slot] = candidate;
125            predecessor[slot] = Some(node);
126            queue.push(Reverse((candidate, slot)));
127        }
128    });
129}
130
131pub(super) fn reconstruct<G: IndexGraphView>(
132    source: G::Node,
133    target: G::Node,
134    predecessor: &[Option<G::Node>],
135) -> Option<Vec<G::Node>> {
136    let mut path = vec![target];
137    let mut cursor = target;
138    while cursor != source {
139        cursor = predecessor[G::node_slot(cursor)]?;
140        path.push(cursor);
141    }
142    path.reverse();
143    Some(path)
144}