Skip to main content

proto_vulcan/state/
map_sum.rs

1use crate::engine::Engine;
2use crate::goal::{AnyGoal, Goal};
3use crate::solver::Solver;
4use crate::state::State;
5use crate::stream::{LazyStream, Stream, StreamIterator};
6use crate::user::User;
7use std::marker::PhantomData;
8
9pub fn map_sum<U, E, F, T>(
10    solver: &Solver<U, E>,
11    state: State<U, E>,
12    mut f: F,
13    iter: impl Iterator<Item = T>,
14) -> Stream<U, E>
15where
16    U: User,
17    E: Engine<U>,
18    F: FnMut(T) -> Goal<U, E>,
19{
20    let mut iter = iter.peekable();
21    let mut stream = Stream::empty();
22    loop {
23        match iter.next() {
24            Some(d) => {
25                if iter.peek().is_none() {
26                    // If this is last value in the domain, no need to clone `state`.
27                    let new_stream = f(d).solve(solver, state);
28                    stream = Stream::mplus(new_stream, LazyStream::delay(stream));
29                    break;
30                } else {
31                    let new_stream = f(d).solve(solver, state.clone());
32                    stream = Stream::mplus(new_stream, LazyStream::delay(stream));
33                }
34            }
35            None => {
36                unreachable!();
37            }
38        }
39    }
40    stream
41}
42
43#[derive(Derivative)]
44#[derivative(Clone(bound = "U: User"))]
45pub struct MapSumIterator<U, E, G, F, T, I>
46where
47    U: User,
48    E: Engine<U>,
49    G: AnyGoal<U, E>,
50    F: Fn(T) -> G + Clone + 'static,
51    T: 'static,
52    I: Iterator<Item = T> + Clone,
53{
54    state: State<U, E>,
55    f: F,
56    iter: Box<I>,
57    _phantom: PhantomData<U>,
58    _phantom2: PhantomData<E>,
59}
60
61impl<U, E, G, F, T, I> MapSumIterator<U, E, G, F, T, I>
62where
63    U: User,
64    E: Engine<U>,
65    G: AnyGoal<U, E>,
66    F: Fn(T) -> G + Clone + 'static,
67    T: 'static,
68    I: Iterator<Item = T> + Clone,
69{
70    pub fn new(state: State<U, E>, f: F, iter: Box<I>) -> MapSumIterator<U, E, G, F, T, I> {
71        MapSumIterator {
72            state,
73            f,
74            iter,
75            _phantom: PhantomData,
76            _phantom2: PhantomData,
77        }
78    }
79}
80
81impl<U, E, G, F, T, I> StreamIterator<U, E> for MapSumIterator<U, E, G, F, T, I>
82where
83    U: User,
84    E: Engine<U>,
85    G: AnyGoal<U, E>,
86    F: Fn(T) -> G + Clone + 'static,
87    T: 'static,
88    I: Iterator<Item = T> + Clone + 'static,
89{
90    fn clone_box(&self) -> Box<dyn StreamIterator<U, E>> {
91        Box::new(self.clone())
92    }
93
94    fn next(&mut self, solver: &Solver<U, E>) -> Option<Stream<U, E>> {
95        match self.iter.next() {
96            Some(t) => {
97                let stream = (self.f)(t).solve(solver, self.state.clone());
98                Some(stream)
99            }
100            None => None,
101        }
102    }
103}
104
105pub fn map_sum_iter<U, E, F, T, I>(state: State<U, E>, f: F, iter: Box<I>) -> Stream<U, E>
106where
107    U: User,
108    E: Engine<U>,
109    F: Fn(T) -> Goal<U, E> + Clone + 'static,
110    T: 'static,
111    I: Iterator<Item = T> + Clone + 'static,
112{
113    Stream::iterator(Box::new(MapSumIterator::new(state, f, iter)))
114}