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