use crate::{trace, FxIndexMap};
use indexmap::map::Entry::Vacant;
use std::hash::Hash;
pub struct DFS<STATE, SUCCESSORS, SUCCESS, ITER>
where
STATE: Eq + Hash + Clone,
SUCCESSORS: FnMut(&STATE) -> ITER,
SUCCESS: FnMut(&STATE) -> bool,
ITER: IntoIterator<Item = STATE>,
{
stack: Vec<usize>,
states: FxIndexMap<STATE, usize>,
successors: SUCCESSORS,
success: SUCCESS,
}
impl<STATE, SUCCESSORS, SUCCESS, ITER> DFS<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 {
stack: vec![0],
states,
successors,
success,
}
}
}
impl<STATE, SUCCESSORS, SUCCESS, ITER> Iterator for DFS<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 index = self.stack.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) {
if let Vacant(e) = self.states.entry(successor) {
let index = {
let i = e.index();
e.insert(index);
i
};
self.stack.push(index);
}
}
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>,
{
DFS::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>,
{
DFS::new(init, successors, success)
.into_iter()
.filter_map(|i| i)
}