use crate::{trace, FxIndexMap};
use indexmap::map::Entry::Vacant;
use std::hash::Hash;
pub struct BFS<STATE, SUCCESSORS, SUCCESS, ITER>
where
STATE: Eq + Hash + Clone,
SUCCESSORS: FnMut(&STATE) -> ITER,
SUCCESS: FnMut(&STATE) -> bool,
ITER: IntoIterator<Item = STATE>,
{
index: usize,
states: FxIndexMap<STATE, usize>,
successors: SUCCESSORS,
success: SUCCESS,
}
impl<STATE, SUCCESSORS, SUCCESS, ITER> BFS<STATE, SUCCESSORS, SUCCESS, ITER>
where
STATE: Eq + Hash + Clone,
SUCCESSORS: FnMut(&STATE) -> ITER,
SUCCESS: FnMut(&STATE) -> bool,
ITER: IntoIterator<Item = STATE>,
{
pub fn new(init: &STATE, successors: SUCCESSORS, success: SUCCESS) -> Self {
let mut states = FxIndexMap::default();
states.insert(init.clone(), 0);
Self {
index: 0,
states,
successors,
success,
}
}
pub fn len(&self) -> usize {
self.states.len()
}
}
impl<STATE, SUCCESSORS, SUCCESS, ITER> Iterator for BFS<STATE, SUCCESSORS, SUCCESS, ITER>
where
STATE: Eq + Hash + Clone,
SUCCESSORS: FnMut(&STATE) -> ITER,
SUCCESS: FnMut(&STATE) -> bool,
ITER: IntoIterator<Item = STATE>,
{
type Item = Option<Vec<STATE>>;
fn next(&mut self) -> Option<Self::Item> {
let (node, _) = self.states.get_index(self.index)?;
for successor in (self.successors)(node) {
if (self.success)(&successor) {
let mut path = trace(&self.states, self.index);
path.push(successor);
self.index += 1;
return Some(Some(path));
}
if let Vacant(e) = self.states.entry(successor) {
e.insert(self.index);
}
}
self.index += 1;
Some(None)
}
}
pub fn solve<STATE, SUCCESSORS, SUCCESS, ITER>(
init: &STATE,
successors: SUCCESSORS,
success: SUCCESS,
) -> Option<Vec<STATE>>
where
STATE: Eq + Hash + Clone,
SUCCESSORS: FnMut(&STATE) -> ITER,
SUCCESS: FnMut(&STATE) -> bool,
ITER: IntoIterator<Item = STATE>,
{
BFS::new(init, successors, success)
.into_iter()
.find(|i| i.is_some())?
}
pub fn paths<STATE, SUCCESSORS, SUCCESS, ITER>(
init: &STATE,
successors: SUCCESSORS,
success: SUCCESS,
) -> impl Iterator<Item = Vec<STATE>>
where
STATE: Eq + Hash + Clone,
SUCCESSORS: FnMut(&STATE) -> ITER,
SUCCESS: FnMut(&STATE) -> bool,
ITER: IntoIterator<Item = STATE>,
{
BFS::new(init, successors, success)
.into_iter()
.filter_map(|i| i)
}