incremental_map/
btree_map.rs

1use std::marker::PhantomData;
2use std::{cell::RefCell, collections::BTreeMap, rc::Rc};
3
4use super::symmetric_fold::{DiffElement, MergeElement, MergeOnceWith, SymmetricDiffMap};
5use super::{FilterMapOperator, MapOperator, Operator};
6use crate::symmetric_fold::SymmetricFoldMap;
7use incremental::expert::{Dependency, Node, WeakNode};
8use incremental::incrsan::NotObserver;
9use incremental::{Cutoff, Incr, Value};
10
11use crate::WithOldIO;
12
13pub trait IncrBTreeMap<K: Value + Ord, V: Value> {
14    fn incr_mapi_<F, V2>(&self, f: F) -> Incr<BTreeMap<K, V2>>
15    where
16        V2: Value,
17        F: FnMut(&K, Incr<V>) -> Incr<V2> + 'static + NotObserver;
18
19    fn incr_mapi_cutoff<F, V2>(&self, f: F, cutoff: Cutoff<V>) -> Incr<BTreeMap<K, V2>>
20    where
21        V2: Value,
22        F: FnMut(&K, Incr<V>) -> Incr<V2> + 'static + NotObserver;
23
24    fn incr_filter_mapi_<F, V2>(&self, f: F) -> Incr<BTreeMap<K, V2>>
25    where
26        V2: Value,
27        F: FnMut(&K, Incr<V>) -> Incr<Option<V2>> + 'static + NotObserver;
28
29    fn incr_filter_mapi_cutoff<F, V2>(&self, f: F, cutoff: Cutoff<V>) -> Incr<BTreeMap<K, V2>>
30    where
31        V2: Value,
32        F: FnMut(&K, Incr<V>) -> Incr<Option<V2>> + 'static + NotObserver;
33
34    fn incr_merge<F, V2, R>(&self, other: &Incr<BTreeMap<K, V2>>, f: F) -> Incr<BTreeMap<K, R>>
35    where
36        V2: Value,
37        R: Value,
38        F: FnMut(&K, MergeElement<&V, &V2>) -> Option<R> + 'static + NotObserver;
39}
40
41impl<K: Value + Ord, V: Value> IncrBTreeMap<K, V> for Incr<BTreeMap<K, V>> {
42    fn incr_mapi_<F, V2>(
43        &self,
44        f: F,
45        // see im_rc
46        // cutoff: Option<Cutoff<V>>
47    ) -> Incr<BTreeMap<K, V2>>
48    where
49        V2: Value,
50        F: FnMut(&K, Incr<V>) -> Incr<V2> + 'static + NotObserver,
51    {
52        incr_filter_mapi_generic_btree_map(self, MapOperator(f, PhantomData), None)
53    }
54
55    fn incr_mapi_cutoff<F, V2>(&self, f: F, cutoff: Cutoff<V>) -> Incr<BTreeMap<K, V2>>
56    where
57        V2: Value,
58        F: FnMut(&K, Incr<V>) -> Incr<V2> + 'static + NotObserver,
59    {
60        incr_filter_mapi_generic_btree_map(self, MapOperator(f, PhantomData), Some(cutoff))
61    }
62
63    fn incr_filter_mapi_<F, V2>(
64        &self,
65        f: F,
66        // see im_rc
67        // cutoff: Option<Cutoff<V>>
68    ) -> Incr<BTreeMap<K, V2>>
69    where
70        V2: Value,
71        F: FnMut(&K, Incr<V>) -> Incr<Option<V2>> + 'static + NotObserver,
72    {
73        incr_filter_mapi_generic_btree_map(self, FilterMapOperator(f, PhantomData), None)
74    }
75
76    fn incr_filter_mapi_cutoff<F, V2>(&self, f: F, cutoff: Cutoff<V>) -> Incr<BTreeMap<K, V2>>
77    where
78        V2: Value,
79        F: FnMut(&K, Incr<V>) -> Incr<Option<V2>> + 'static + NotObserver,
80    {
81        incr_filter_mapi_generic_btree_map(self, FilterMapOperator(f, PhantomData), Some(cutoff))
82    }
83
84    /// Merge two maps incrementally, where
85    ///
86    /// - if a key appears only in self, the predicate runs with [`MergeElement::Left`]
87    /// - if a key appears only in other, the predicate runs with [`MergeElement::Right`]
88    /// - if a key appears in both, the predicate runs with [`MergeElement::Both`]
89    ///
90    /// The predicate is only re-run for added/removed/modified keys in each map, using the
91    /// symmetric diff property.
92    fn incr_merge<F, V2, R>(&self, other: &Incr<BTreeMap<K, V2>>, mut f: F) -> Incr<BTreeMap<K, R>>
93    where
94        V2: Value,
95        R: Value,
96        F: FnMut(&K, MergeElement<&V, &V2>) -> Option<R> + 'static + NotObserver,
97    {
98        let i = self.with_old_input_output2(other, move |old, new_left_map, new_right_map| {
99            let mut did_change = false;
100            let output = merge_shared_impl(
101                old,
102                new_left_map,
103                new_right_map,
104                |mut acc_output, key, merge_elem| {
105                    use MergeElement::*;
106                    did_change = true;
107                    let data = match merge_elem {
108                        Both((_, left_diff), (_, right_diff)) => {
109                            (left_diff.new_data(), right_diff.new_data())
110                        }
111                        Left((_, left_diff)) => (left_diff.new_data(), new_right_map.get(key)),
112                        Right((_, right_diff)) => (new_left_map.get(key), right_diff.new_data()),
113                    };
114                    let output_data_opt = match data {
115                        (None, None) => None,
116                        (Some(x), None) => f(key, MergeElement::Left(x)),
117                        (None, Some(x)) => f(key, MergeElement::Right(x)),
118                        (Some(a), Some(b)) => f(key, MergeElement::Both(a, b)),
119                    };
120                    match output_data_opt {
121                        None => acc_output.remove(key),
122                        Some(r) => acc_output.insert(key.clone(), r),
123                    };
124                    acc_output
125                },
126            );
127            (output, did_change)
128        });
129        #[cfg(debug_assertions)]
130        i.set_graphviz_user_data(Box::new(format!(
131            "incr_merge -> {}",
132            std::any::type_name::<BTreeMap<K, R>>()
133        )));
134        i
135    }
136}
137
138fn incr_filter_mapi_generic_btree_map<K, V, O, V2>(
139    lhs: &Incr<BTreeMap<K, V>>,
140    mut f: O,
141    cutoff: Option<Cutoff<V>>,
142) -> Incr<BTreeMap<K, V2>>
143where
144    K: Value + Ord,
145    V: Value,
146    O: Operator<K, V, V2> + 'static + NotObserver,
147    O::Output: Value,
148    V2: Value,
149{
150    let state = lhs.state();
151    let prev_map: Rc<RefCell<BTreeMap<K, V>>> = Rc::new(RefCell::new(BTreeMap::new()));
152    let acc = Rc::new(RefCell::new(BTreeMap::<K, V2>::new()));
153    let result = Node::<BTreeMap<K, V2>>::new(&state, {
154        let acc_ = acc.clone();
155        move || acc_.borrow().clone()
156    });
157    let on_inner_change = {
158        let acc_ = acc.clone();
159        move |key: &K, opt: &O::Output| {
160            let mut acc = acc_.borrow_mut();
161            let opt = O::as_opt(opt);
162            match opt {
163                None => {
164                    acc.remove(key);
165                }
166                Some(x) => {
167                    acc.insert(key.clone(), x.clone());
168                }
169            }
170            drop(acc);
171        }
172    };
173
174    let mut prev_nodes = BTreeMap::<K, (WeakNode<_>, Dependency<O::Output>)>::new();
175    let result_weak = result.weak();
176
177    let lhs_change = lhs.map_cyclic({
178        move |lhs_change, map| {
179            let mut prev_map_mut = prev_map.borrow_mut();
180            prev_map_mut.symmetric_fold(map, &mut prev_nodes, |nodes, (key, diff)| {
181                match diff {
182                    DiffElement::Unequal(_, _) => {
183                        let (node, _dep) = nodes.get(key).unwrap();
184                        node.make_stale();
185                        nodes
186                    }
187                    DiffElement::Left(_) => {
188                        let (node, dep) = nodes.remove(key).unwrap();
189                        // running remove_dependency will cause node's weak ref to die.
190                        // so we upgrade it first.
191                        let node = node.upgrade().unwrap();
192                        result_weak.remove_dependency(dep);
193                        let mut acc = acc.borrow_mut();
194                        acc.remove(key);
195                        // Invalidate does have to happen after remove_dependency.
196                        node.invalidate();
197                        nodes
198                    }
199                    DiffElement::Right(_) => {
200                        let key = key.clone();
201                        let node = Node::<V>::new(&state, {
202                            let key_ = key.clone();
203                            let prev_map_ = prev_map.clone();
204                            move || {
205                                let prev_map = prev_map_.borrow();
206                                prev_map.get(&key_).unwrap().clone()
207                            }
208                        });
209                        if let Some(cutoff) = cutoff.clone() {
210                            node.watch().set_cutoff(cutoff);
211                        }
212                        let lhs_change = lhs_change.upgrade().unwrap();
213                        node.add_dependency(&lhs_change);
214                        let mapped = f.call_fn(&key, node.watch());
215                        let user_function_dep = result_weak.add_dependency_with(&mapped, {
216                            let key = key.clone();
217                            let on_inner_change = on_inner_change.clone();
218                            move |v| on_inner_change(&key, v)
219                        });
220                        nodes.insert(key, (node.weak(), user_function_dep));
221                        nodes
222                    }
223                }
224            });
225            *prev_map_mut = map.clone();
226        }
227    });
228    result.add_dependency(&lhs_change);
229    result.watch()
230}
231
232pub(crate) fn merge_shared_impl<
233    K: Clone + Ord,
234    V1: Clone + PartialEq,
235    V2: Clone + PartialEq,
236    R: Clone,
237>(
238    old: Option<(BTreeMap<K, V1>, BTreeMap<K, V2>, BTreeMap<K, R>)>,
239    new_left_map: &BTreeMap<K, V1>,
240    new_right_map: &BTreeMap<K, V2>,
241    mut f: impl FnMut(
242        BTreeMap<K, R>,
243        &K,
244        MergeElement<(&K, DiffElement<&V1>), (&K, DiffElement<&V2>)>,
245    ) -> BTreeMap<K, R>,
246) -> BTreeMap<K, R> {
247    let (old_left_map, old_right_map, old_output) = match old {
248        None => (BTreeMap::new(), BTreeMap::new(), BTreeMap::new()),
249        Some(x) => x,
250    };
251    let left_diff = old_left_map.symmetric_diff(new_left_map);
252    let right_diff = old_right_map.symmetric_diff(new_right_map);
253    // relies on the key iteration being sorted, as in BTreeMap.
254    let merge = MergeOnceWith::new(left_diff, right_diff, |(k, _), (k2, _)| k.cmp(k2));
255    merge.fold(old_output, |output, merge_elem| {
256        let key = match merge_elem {
257            MergeElement::Left((key, _)) | MergeElement::Right((key, _)) => key,
258            MergeElement::Both((left_key, _), (_right_key, _)) => {
259                // comparisons can be expensive
260                // assert_eq!(left_key, right_key);
261                left_key
262            }
263        };
264        f(output, key, merge_elem)
265    })
266}