use crate::{trace, FxIndexMap};
use indexmap::map::Entry::{Occupied, Vacant};
use std::{collections::BinaryHeap, hash::Hash};
#[derive(Debug, Clone, Eq, PartialEq)]
struct Element {
index: usize,
estimate: usize,
}
impl Ord for Element {
fn cmp(&self, other: &Self) -> std::cmp::Ordering {
other.estimate.cmp(&self.estimate)
}
}
impl PartialOrd for Element {
fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
Some(self.cmp(other))
}
}
pub struct GBFS<STATE, SUCCESSORS, SUCCESS, HEURISTIC, ITER>
where
STATE: Eq + Hash + Clone,
SUCCESSORS: FnMut(&STATE) -> ITER,
SUCCESS: FnMut(&STATE) -> bool,
HEURISTIC: FnMut(&STATE) -> usize,
ITER: IntoIterator<Item = STATE>,
{
queue: BinaryHeap<Element>,
states: FxIndexMap<STATE, usize>,
successors: SUCCESSORS,
success: SUCCESS,
heuristic: HEURISTIC,
}
impl<STATE, SUCCESSORS, SUCCESS, HEURISTIC, ITER> GBFS<STATE, SUCCESSORS, SUCCESS, HEURISTIC, ITER>
where
STATE: Eq + Hash + Clone,
SUCCESSORS: FnMut(&STATE) -> ITER,
SUCCESS: FnMut(&STATE) -> bool,
HEURISTIC: FnMut(&STATE) -> usize,
ITER: IntoIterator<Item = STATE>,
{
pub fn new(
init: &STATE,
successors: SUCCESSORS,
success: SUCCESS,
heuristic: HEURISTIC,
) -> Self {
let mut states = FxIndexMap::default();
states.insert(init.clone(), 0);
let queue = BinaryHeap::from(vec![Element {
index: 0,
estimate: 0,
}]);
Self {
queue,
states,
successors,
success,
heuristic,
}
}
}
impl<STATE, SUCCESSORS, SUCCESS, HEURISTIC, ITER> Iterator
for GBFS<STATE, SUCCESSORS, SUCCESS, HEURISTIC, ITER>
where
STATE: Eq + Hash + Clone,
SUCCESSORS: FnMut(&STATE) -> ITER,
SUCCESS: FnMut(&STATE) -> bool,
HEURISTIC: FnMut(&STATE) -> usize,
ITER: IntoIterator<Item = STATE>,
{
type Item = Option<Vec<STATE>>;
fn next(&mut self) -> Option<Self::Item> {
let Element { index, estimate: _ } = self.queue.pop()?;
let (node, _) = self.states.get_index(index).unwrap();
if (self.success)(node) {
return Some(Some(trace(&self.states, index)));
}
for successor in (self.successors)(node) {
let index = match self.states.entry(successor) {
Vacant(e) => {
let i = e.index();
e.insert(index);
i
}
Occupied(_) => continue,
};
let (state, _) = self.states.get_index(index).unwrap();
self.queue.push(Element {
index,
estimate: (self.heuristic)(state),
});
}
Some(None)
}
}
pub fn solve<STATE, SUCCESSORS, SUCCESS, HEURISTIC, ITER>(
init: &STATE,
successors: SUCCESSORS,
success: SUCCESS,
heuristic: HEURISTIC,
) -> Option<Vec<STATE>>
where
STATE: Eq + Hash + Clone,
SUCCESSORS: FnMut(&STATE) -> ITER,
SUCCESS: FnMut(&STATE) -> bool,
HEURISTIC: FnMut(&STATE) -> usize,
ITER: IntoIterator<Item = STATE>,
{
GBFS::new(init, successors, success, heuristic)
.into_iter()
.find(|i| i.is_some())?
}
pub fn paths<STATE, SUCCESSORS, SUCCESS, HEURISTIC, ITER>(
init: &STATE,
successors: SUCCESSORS,
success: SUCCESS,
heuristic: HEURISTIC,
) -> impl Iterator<Item = Vec<STATE>>
where
STATE: Eq + Hash + Clone,
SUCCESSORS: FnMut(&STATE) -> ITER,
SUCCESS: FnMut(&STATE) -> bool,
HEURISTIC: FnMut(&STATE) -> usize,
ITER: IntoIterator<Item = STATE>,
{
GBFS::new(init, successors, success, heuristic)
.into_iter()
.filter_map(|i| i)
}