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 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}