petgraph 0.8.3

use alloc::{collections::BinaryHeap, vec, vec::Vec};
use core::hash::Hash;

use hashbrown::hash_map::{
    Entry::{Occupied, Vacant},
    HashMap,
};

use crate::algo::Measure;
use crate::scored::MinScored;
use crate::visit::{EdgeRef, GraphBase, IntoEdges, Visitable};

/// A* shortest path algorithm.
///
/// Computes the shortest path from `start` to `finish`, including the total path cost.
///
/// `finish` is implicitly given via the `is_goal` callback, which should return `true` if the
/// given node is the finish node.
///
/// The function `edge_cost` should return the cost for a particular edge. Edge costs must be
/// non-negative.
///
/// The function `estimate_cost` should return the estimated cost to the finish for a particular
/// node. For the algorithm to find the actual shortest path, it should be admissible, meaning that
/// it should never overestimate the actual cost to get to the nearest goal node. Estimate costs
/// must also be non-negative.
///
/// # Arguments
/// * `graph`: weighted graph.
/// * `start`: the start node.
/// * `is_goal`: the callback defines the goal node.
/// * `edge_cost`: closure that returns cost of a particular edge.
/// * `estimate_cost`: closure that returns the estimated cost to the finish for particular node.
///
/// # Returns
/// * `Some(K, Vec<G::NodeId>)` - the total cost and path from start to finish, if one was found.
/// * `None` - if such a path was not found.
///
/// # Complexity
/// The time complexity largely depends on the heuristic used. Feel free to contribute and provide the exact time complexity :)
///
/// With a trivial heuristic, the algorithm will behave like [`fn@crate::algo::dijkstra`].
///
/// # Example
/// ```
/// use petgraph::Graph;
/// use petgraph::algo::astar;
///
/// let mut g = Graph::new();
/// let a = g.add_node((0., 0.));
/// let b = g.add_node((2., 0.));
/// let c = g.add_node((1., 1.));
/// let d = g.add_node((0., 2.));
/// let e = g.add_node((3., 3.));
/// let f = g.add_node((4., 2.));
/// g.extend_with_edges(&[
///     (a, b, 2),
///     (a, d, 4),
///     (b, c, 1),
///     (b, f, 7),
///     (c, e, 5),
///     (e, f, 1),
///     (d, e, 1),
/// ]);
///
/// // Graph represented with the weight of each edge
/// // Edges with '*' are part of the optimal path.
/// //
/// //     2       1
/// // a ----- b ----- c
/// // | 4*    | 7     |
/// // d       f       | 5
/// // | 1*    | 1*    |
/// // \------ e ------/
///
/// let path = astar(&g, a, |finish| finish == f, |e| *e.weight(), |_| 0);
/// assert_eq!(path, Some((6, vec![a, d, e, f])));
/// ```
pub fn astar<G, F, H, K, IsGoal>(
    graph: G,
    start: G::NodeId,
    mut is_goal: IsGoal,
    mut edge_cost: F,
    mut estimate_cost: H,
) -> Option<(K, Vec<G::NodeId>)>
where
    G: IntoEdges + Visitable,
    IsGoal: FnMut(G::NodeId) -> bool,
    G::NodeId: Eq + Hash,
    F: FnMut(G::EdgeRef) -> K,
    H: FnMut(G::NodeId) -> K,
    K: Measure + Copy,
{
    // The Open set
    let mut visit_next = BinaryHeap::new();
    // A node -> (f, h, g) mapping
    // TODO: Derive `g` from `f` and `h`.
    let mut scores = HashMap::new();
    // The search tree
    let mut path_tracker = PathTracker::<G>::new();

    let zero: K = K::default();
    let g: K = zero;
    let h: K = estimate_cost(start);
    let f: K = g + h;
    scores.insert(start, (f, h, g));
    visit_next.push(MinScored((f, h, g), start));

    while let Some(MinScored((f, h, g), node)) = visit_next.pop() {
        if is_goal(node) {
            let path = path_tracker.reconstruct_path_to(node);
            let (goal_f, goal_h, goal_g) = scores[&node];
            debug_assert_eq!(goal_h, zero);
            debug_assert_eq!(goal_f, goal_g);
            return Some((goal_f, path));
        }

        match scores.entry(node) {
            Occupied(mut entry) => {
                let (_, _, old_g) = *entry.get();
                // The node has already been expanded with a better cost.
                if old_g < g {
                    continue;
                }
                // NOTE: Because there's no closed set, we don't know if we expanded this node.
                // if old_g = g we may be re-expanding this node, but won't insert new neigbours.
                entry.insert((f, h, g));
            }
            Vacant(entry) => {
                entry.insert((f, h, g));
            }
        }

        for edge in graph.edges(node) {
            let neigh = edge.target();
            let neigh_g = g + edge_cost(edge);
            let neigh_h = estimate_cost(neigh);
            let neigh_f = neigh_g + neigh_h;
            let neigh_score = (neigh_f, neigh_h, neigh_g);

            match scores.entry(neigh) {
                Occupied(mut entry) => {
                    let (_, _, old_neigh_g) = *entry.get();
                    if neigh_g >= old_neigh_g {
                        // New cost isn't better
                        continue;
                    }
                    entry.insert(neigh_score);
                }
                Vacant(entry) => {
                    entry.insert(neigh_score);
                }
            }

            path_tracker.set_predecessor(neigh, node);
            visit_next.push(MinScored(neigh_score, neigh));
        }
    }

    None
}

struct PathTracker<G>
where
    G: GraphBase,
    G::NodeId: Eq + Hash,
{
    came_from: HashMap<G::NodeId, G::NodeId>,
}

impl<G> PathTracker<G>
where
    G: GraphBase,
    G::NodeId: Eq + Hash,
{
    fn new() -> PathTracker<G> {
        PathTracker {
            came_from: HashMap::new(),
        }
    }

    fn set_predecessor(&mut self, node: G::NodeId, previous: G::NodeId) {
        self.came_from.insert(node, previous);
    }

    fn reconstruct_path_to(&self, last: G::NodeId) -> Vec<G::NodeId> {
        let mut path = vec![last];

        let mut current = last;
        while let Some(&previous) = self.came_from.get(&current) {
            path.push(previous);
            current = previous;
        }

        path.reverse();

        path
    }
}