use indexmap::IndexMap;
use indexmap::map::Entry::{Occupied, Vacant};
use num_traits::Zero;
use std::collections::{BinaryHeap, HashSet};
use std::cmp::Ordering;
use std::hash::Hash;
use std::usize;
use super::reverse_path;
pub fn astar<N, C, FN, IN, FH, FS>(
start: &N,
mut neighbours: FN,
mut heuristic: FH,
mut success: FS,
) -> Option<(Vec<N>, C)>
where
N: Eq + Hash + Clone,
C: Zero + Ord + Copy,
FN: FnMut(&N) -> IN,
IN: IntoIterator<Item = (N, C)>,
FH: FnMut(&N) -> C,
FS: FnMut(&N) -> bool,
{
let mut to_see = BinaryHeap::new();
to_see.push(SmallestCostHolder {
estimated_cost: heuristic(start),
cost: Zero::zero(),
payload: (Zero::zero(), 0),
});
let mut parents: IndexMap<N, (usize, C)> = IndexMap::new();
parents.insert(start.clone(), (usize::MAX, Zero::zero()));
while let Some(SmallestCostHolder {
payload: (cost, i), ..
}) = to_see.pop()
{
let neighbours = {
let (node, &(_, c)) = parents.get_index(i).unwrap();
if success(node) {
let path = reverse_path(&parents, |&(p, _)| p, i);
return Some((path, cost));
}
if cost > c {
continue;
}
neighbours(node)
};
for (neighbour, move_cost) in neighbours {
let new_cost = cost + move_cost;
let h; let n; match parents.entry(neighbour) {
Vacant(e) => {
h = heuristic(e.key());
n = e.index();
e.insert((i, new_cost));
}
Occupied(e) => if e.get().1 > new_cost {
h = heuristic(e.key());
n = e.index();
e.insert((i, new_cost));
} else {
continue;
},
}
to_see.push(SmallestCostHolder {
estimated_cost: new_cost + h,
cost: cost,
payload: (new_cost, n),
});
}
}
None
}
pub fn astar_bag<N, C, FN, IN, FH, FS>(
start: &N,
mut neighbours: FN,
mut heuristic: FH,
mut success: FS,
) -> Option<(AstarSolution<N>, C)>
where
N: Eq + Hash + Clone,
C: Zero + Ord + Copy,
FN: FnMut(&N) -> IN,
IN: IntoIterator<Item = (N, C)>,
FH: FnMut(&N) -> C,
FS: FnMut(&N) -> bool,
{
let mut to_see = BinaryHeap::new();
let mut min_cost = None;
let mut sinks = HashSet::new();
to_see.push(SmallestCostHolder {
estimated_cost: heuristic(start),
cost: Zero::zero(),
payload: (Zero::zero(), 0),
});
let mut parents: IndexMap<N, (HashSet<usize>, C)> = IndexMap::new();
parents.insert(start.clone(), (HashSet::new(), Zero::zero()));
while let Some(SmallestCostHolder {
payload: (cost, i),
estimated_cost,
..
}) = to_see.pop()
{
if let Some(min_cost) = min_cost {
if estimated_cost > min_cost {
break;
}
}
let neighbours = {
let (node, &(_, c)) = parents.get_index(i).unwrap();
if success(node) {
min_cost = Some(cost);
sinks.insert(i);
}
if cost > c {
continue;
}
neighbours(node)
};
for (neighbour, move_cost) in neighbours {
let new_cost = cost + move_cost;
let h; let n; match parents.entry(neighbour) {
Vacant(e) => {
h = heuristic(e.key());
n = e.index();
let mut p = HashSet::new();
p.insert(i);
e.insert((p, new_cost));
}
Occupied(mut e) => if e.get().1 > new_cost {
h = heuristic(e.key());
n = e.index();
let s = e.get_mut();
s.0.clear();
s.0.insert(i);
s.1 = new_cost;
} else {
if e.get().1 == new_cost {
e.get_mut().0.insert(i);
}
continue;
},
}
to_see.push(SmallestCostHolder {
estimated_cost: new_cost + h,
cost: cost,
payload: (new_cost, n),
});
}
}
min_cost.map(|cost| {
let parents = parents
.into_iter()
.map(|(k, (ps, _))| (k, ps.into_iter().collect()))
.collect();
(
AstarSolution {
sinks: sinks.into_iter().collect(),
parents,
current: vec![],
terminated: false,
},
cost,
)
})
}
pub fn astar_bag_collect<N, C, FN, IN, FH, FS>(
start: &N,
neighbours: FN,
heuristic: FH,
success: FS,
) -> Option<(Vec<Vec<N>>, C)>
where
N: Eq + Hash + Clone,
C: Zero + Ord + Copy,
FN: FnMut(&N) -> IN,
IN: IntoIterator<Item = (N, C)>,
FH: FnMut(&N) -> C,
FS: FnMut(&N) -> bool,
{
astar_bag(start, neighbours, heuristic, success)
.map(|(solutions, cost)| (solutions.collect(), cost))
}
struct SmallestCostHolder<K, P> {
estimated_cost: K,
cost: K,
payload: P,
}
impl<K: PartialEq, P> PartialEq for SmallestCostHolder<K, P> {
fn eq(&self, other: &SmallestCostHolder<K, P>) -> bool {
self.estimated_cost.eq(&other.estimated_cost) && self.cost.eq(&other.cost)
}
}
impl<K: PartialEq, P> Eq for SmallestCostHolder<K, P> {}
impl<K: Ord, P> PartialOrd for SmallestCostHolder<K, P> {
fn partial_cmp(&self, other: &SmallestCostHolder<K, P>) -> Option<Ordering> {
Some(self.cmp(other))
}
}
impl<K: Ord, P> Ord for SmallestCostHolder<K, P> {
fn cmp(&self, other: &SmallestCostHolder<K, P>) -> Ordering {
match other.estimated_cost.cmp(&self.estimated_cost) {
Ordering::Equal => self.cost.cmp(&other.cost),
s => s,
}
}
}
#[derive(Clone)]
pub struct AstarSolution<N> {
sinks: Vec<usize>,
parents: Vec<(N, Vec<usize>)>,
current: Vec<Vec<usize>>,
terminated: bool,
}
unsafe impl<N: Send> Send for AstarSolution<N> {}
impl<N: Clone + Eq + Hash> AstarSolution<N> {
fn complete(&mut self) {
loop {
let ps = match self.current.last() {
None => self.sinks.clone(),
Some(last) => {
let &top = last.last().unwrap();
self.parents(top).clone()
}
};
if ps.is_empty() {
break;
}
self.current.push(ps);
}
}
fn next_vec(&mut self) {
while self.current.last().map(|v| v.len()) == Some(1) {
self.current.pop();
}
self.current.last_mut().map(|v| v.pop());
}
fn node(&self, i: usize) -> &N {
&self.parents[i].0
}
fn parents(&self, i: usize) -> &Vec<usize> {
&self.parents[i].1
}
}
impl<N: Clone + Eq + Hash> Iterator for AstarSolution<N> {
type Item = Vec<N>;
fn next(&mut self) -> Option<Self::Item> {
if self.terminated {
return None;
}
self.complete();
let path = self.current
.iter()
.rev()
.map(|v| v.last().cloned().unwrap())
.map(|i| self.node(i).clone())
.collect::<Vec<_>>();
self.next_vec();
self.terminated = self.current.is_empty();
Some(path)
}
}