cyclotron/
lazy_set.rs

1use std::collections::BTreeSet;
2use std::fmt;
3use std::hash::Hash;
4
5// HACK(eddyb) sealed traits/enums to keep them from being used outside of this module.
6mod sealed {
7    pub enum Step<Req, A, R> {
8        Request(Req, A),
9        // FIXME(eddyb) use GATs to move this into a request.
10        Fork(A, A),
11        Return(R),
12    }
13
14    impl<Req, A, R> Step<Req, A, R> {
15        fn map_cont<B>(self, f: impl Fn(A) -> B) -> Step<Req, B, R> {
16            match self {
17                Step::Request(req, a) => Step::Request(req, f(a)),
18                Step::Fork(x, y) => Step::Fork(f(x), f(y)),
19                Step::Return(r) => Step::Return(r),
20            }
21        }
22
23        fn map_ret<S>(self, f: impl FnOnce(R) -> S) -> Step<Req, A, S> {
24            match self {
25                Step::Request(req, a) => Step::Request(req, a),
26                Step::Fork(x, y) => Step::Fork(x, y),
27                Step::Return(r) => Step::Return(f(r)),
28            }
29        }
30    }
31
32    pub trait LazySet<Req, Res>: Clone {
33        type Item;
34        fn step(self, res: Option<Res>) -> Step<Req, Self, Option<Self::Item>>;
35    }
36
37    #[derive(Clone)]
38    pub struct One<T>(pub(super) T);
39
40    impl<Req, Res, T: Clone> LazySet<Req, Res> for One<T> {
41        type Item = T;
42        fn step(self, res: Option<Res>) -> Step<Req, Self, Option<Self::Item>> {
43            assert!(res.is_none());
44            Step::Return(Some(self.0))
45        }
46    }
47
48    impl<Req, Res, T: Clone> LazySet<Req, Res> for Option<T> {
49        type Item = T;
50        fn step(self, res: Option<Res>) -> Step<Req, Self, Option<Self::Item>> {
51            assert!(res.is_none());
52            Step::Return(self)
53        }
54    }
55
56    #[derive(Clone)]
57    pub enum Request<Req> {
58        Start(Req),
59        Complete,
60    }
61
62    impl<Req: Clone, Res> LazySet<Req, Res> for Request<Req> {
63        type Item = Res;
64        fn step(self, res: Option<Res>) -> Step<Req, Self, Option<Self::Item>> {
65            match self {
66                Request::Start(req) => {
67                    assert!(res.is_none());
68                    Step::Request(req, Request::Complete)
69                }
70                Request::Complete => Step::Return(Some(res.unwrap())),
71            }
72        }
73    }
74
75    #[derive(Clone)]
76    pub enum Union<A, B> {
77        Start(A, B),
78        JustA(A),
79        JustB(B),
80    }
81
82    impl<Req, Res, A, B, T> LazySet<Req, Res> for Union<A, B>
83    where
84        A: LazySet<Req, Res, Item = T>,
85        B: LazySet<Req, Res, Item = T>,
86    {
87        type Item = T;
88        fn step(self, res: Option<Res>) -> Step<Req, Self, Option<Self::Item>> {
89            match self {
90                Union::Start(a, b) => {
91                    assert!(res.is_none());
92                    Step::Fork(Union::JustA(a), Union::JustB(b))
93                }
94                Union::JustA(a) => a.step(res).map_cont(Union::JustA),
95                Union::JustB(b) => b.step(res).map_cont(Union::JustB),
96            }
97        }
98    }
99
100    #[derive(Clone)]
101    pub struct Map<A, F>(pub(super) A, pub(super) F);
102
103    impl<Req, Res, A, F, T> LazySet<Req, Res> for Map<A, F>
104    where
105        A: LazySet<Req, Res>,
106        F: FnOnce(A::Item) -> T + Clone,
107    {
108        type Item = T;
109        fn step(self, res: Option<Res>) -> Step<Req, Self, Option<Self::Item>> {
110            let Map(a, f) = self;
111            a.step(res)
112                .map_cont(|a| Map(a, f.clone()))
113                .map_ret(|r| r.map(f))
114        }
115    }
116
117    #[derive(Clone)]
118    pub enum FlatMap<A, B, F> {
119        Start(A, F),
120        Then(B),
121    }
122
123    impl<Req, Res, A, B, F> LazySet<Req, Res> for FlatMap<A, B, F>
124    where
125        A: LazySet<Req, Res>,
126        B: LazySet<Req, Res>,
127        F: FnOnce(A::Item) -> B + Clone,
128    {
129        type Item = B::Item;
130        fn step(self, res: Option<Res>) -> Step<Req, Self, Option<Self::Item>> {
131            match self {
132                FlatMap::Start(a, f) => match a.step(res) {
133                    Step::Request(req, a) => Step::Request(req, FlatMap::Start(a, f)),
134                    Step::Fork(x, y) => {
135                        Step::Fork(FlatMap::Start(x, f.clone()), FlatMap::Start(y, f))
136                    }
137                    Step::Return(None) => Step::Return(None),
138                    Step::Return(Some(r)) => f(r).step(None).map_cont(FlatMap::Then),
139                },
140                FlatMap::Then(b) => b.step(res).map_cont(FlatMap::Then),
141            }
142        }
143    }
144}
145
146// NOTE(eddyb) this would normally be named `LazySetExt`, but we want to allow
147// users to use this trait in bounds, without exposing the sealed one.
148pub trait LazySet<Req, Res>: sealed::LazySet<Req, Res> {
149    fn union<B: LazySet<Req, Res, Item = Self::Item>>(self, b: B) -> sealed::Union<Self, B> {
150        sealed::Union::Start(self, b)
151    }
152
153    fn map<F: FnOnce(Self::Item) -> T, T>(self, f: F) -> sealed::Map<Self, F> {
154        sealed::Map(self, f)
155    }
156
157    fn flat_map<B: LazySet<Req, Res>, F: FnOnce(Self::Item) -> B>(
158        self,
159        f: F,
160    ) -> sealed::FlatMap<Self, B, F> {
161        sealed::FlatMap::Start(self, f)
162    }
163
164    fn run(
165        self,
166        handle: &mut impl FnMut(Req, Self) -> BTreeSet<Res>,
167        res: Option<Res>,
168    ) -> BTreeSet<Self::Item>
169    where
170        Res: Ord,
171        Self::Item: Ord,
172    {
173        let mut output = BTreeSet::new();
174        match self.step(res) {
175            sealed::Step::Request(req, a) => {
176                for res in handle(req, a.clone()) {
177                    output.extend(a.clone().run(handle, Some(res)));
178                }
179            }
180            sealed::Step::Fork(a, b) => {
181                output.extend(a.run(handle, None));
182                output.extend(b.run(handle, None));
183            }
184            sealed::Step::Return(v) => output.extend(v),
185        }
186        output
187    }
188}
189
190impl<Req, Res, A: sealed::LazySet<Req, Res>> LazySet<Req, Res> for A {}
191
192pub fn one<T>(x: T) -> sealed::One<T> {
193    sealed::One(x)
194}
195
196pub fn request<Req>(req: Req) -> sealed::Request<Req> {
197    sealed::Request::Start(req)
198}
199
200#[derive(Clone)]
201pub struct Call<T>(pub T);
202
203pub fn call<T>(x: T) -> sealed::Request<Call<T>> {
204    request(Call(x))
205}
206
207// FIXME(eddyb) this should be in `LazyExt`, but `-> impl Trait`
208// doesn't work in traits yet, move it there whenever that changes.
209pub fn to_eager<K, V, A>(
210    lazy_f: impl FnOnce(K) -> A + Clone,
211) -> impl FnOnce(&mut dyn FnMut(K) -> BTreeSet<V>, K) -> BTreeSet<V> + Clone
212where
213    K: Copy + Eq + Hash + fmt::Debug,
214    V: Clone + Ord + fmt::Debug,
215    A: LazySet<Call<K>, V, Item = V>,
216{
217    move |f, k| lazy_f(k).run(&mut |Call(k), _| f(k), None)
218}
219
220// FIXME(eddyb) figure out module hierarchy here.
221pub mod depth_first {
222    use super::*;
223
224    use std::collections::{hash_map::Entry, HashMap};
225
226    struct CallState<K, V, A> {
227        // FIXME(eddyb) closures can't be Ord or Hash today.
228        // returns: BTreeSet<(K, A)>,
229        returns: Vec<(K, A)>,
230        results: BTreeSet<V>,
231    }
232
233    struct MemoState<K, V, A, F> {
234        calls: HashMap<K, CallState<K, V, A>>,
235        f: F,
236    }
237
238    pub fn memoize<K, V, A>(f: impl FnOnce(K) -> A + Clone) -> impl FnMut(K) -> BTreeSet<V>
239    where
240        K: Copy + Ord + Hash + fmt::Debug,
241        V: Clone + Ord + fmt::Debug,
242        A: LazySet<Call<K>, V, Item = V>,
243    {
244        let mut state = MemoState {
245            calls: HashMap::new(),
246            f,
247        };
248        move |k| match state.calls.entry(k) {
249            // `f(k)` complete (while technically `MemoState` could allow partial
250            // results to be exposed, the only entry-point into `state` is this
251            // closure, which is `FnMut`, disallowing reentrance).
252            Entry::Occupied(entry) => entry.get().results.clone(),
253
254            // `f(k)` never attempted before, we have to run it now.
255            Entry::Vacant(entry) => {
256                entry.insert(CallState {
257                    returns: Default::default(),
258                    results: Default::default(),
259                });
260                let f = state.f.clone();
261                state.run(k, f(k), None);
262                state.calls[&k].results.clone()
263            }
264        }
265    }
266
267    impl<K, V, A, F> MemoState<K, V, A, F>
268    where
269        K: Copy + Eq + Hash + fmt::Debug,
270        V: Clone + Ord + fmt::Debug,
271        A: LazySet<Call<K>, V, Item = V>,
272        F: FnOnce(K) -> A + Clone,
273    {
274        // FIXME(eddyb) try to make `LazySet::run` flexible enough to use it here.
275        fn run(&mut self, current: K, a: A, res: Option<V>) {
276            match a.step(res) {
277                // FIXME(eddyb) deduplicate some of this.
278                sealed::Step::Request(Call(k), a) => match self.calls.entry(k) {
279                    // `f(k)` in progress or completed.
280                    Entry::Occupied(entry) => {
281                        let call = entry.into_mut();
282                        // FIXME(eddyb) closures can't be Ord or Hash today.
283                        // if call.returns.insert((current, a.clone()))
284                        call.returns.push((current, a.clone()));
285                        {
286                            // FIXME(eddyb) try to avoid cloning the set.
287                            for v in call.results.clone() {
288                                self.run(current, a.clone(), Some(v));
289                            }
290                        }
291                    }
292
293                    // `f(k)` never attempted before, we have to start it now.
294                    Entry::Vacant(entry) => {
295                        let call = entry.insert(CallState {
296                            returns: Default::default(),
297                            results: Default::default(),
298                        });
299                        // FIXME(eddyb) closures can't be Ord or Hash today.
300                        // call.returns.insert((current, a.clone()));
301                        call.returns.push((current, a.clone()));
302                        let f = self.f.clone();
303                        self.run(k, f(k), None);
304                    }
305                },
306                sealed::Step::Fork(a, b) => {
307                    self.run(current, a, None);
308                    self.run(current, b, None);
309                }
310                sealed::Step::Return(None) => {}
311                sealed::Step::Return(Some(v)) => {
312                    let call = self.calls.get_mut(&current).unwrap();
313                    if call.results.insert(v.clone()) {
314                        // FIXME(eddyb) try to avoid cloning the set.
315                        for (caller, ret) in call.returns.clone() {
316                            self.run(caller, ret, Some(v.clone()));
317                        }
318                    }
319                }
320            }
321        }
322    }
323}
cyclotron/lazy_set.rs

cyclotron/
lazy_set.rs
