Skip to main content

weavatrix_graph/algo/
astar.rs

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