Skip to main content

hyphae/traits/collections/
select_cell.rs

1//! Select-cell plan node implementing [`MapQuery`].
2//!
3//! `select_cell` is the reactive variant of `select`: each row's inclusion is
4//! gated by a [`Watchable`] producing `bool`. Returns an uncompiled plan node; call
5//! [`MapQuery::materialize`] to compile a plan into a subscribable
6//! [`CellMap`](crate::CellMap).
7
8use std::{hash::Hash, marker::PhantomData, sync::Arc};
9
10use super::ProjectCellExt;
11use crate::{
12    map_query::{
13        BuildQueryRuntime, MapQuery,
14        properties::{ByMapKey, PlanProperties, ZeroOrOne},
15    },
16    pipeline::{Materialize, Pipeline},
17    subscription::SubscriptionGuard,
18    traits::{CellValue, Gettable, MapExt},
19};
20
21/// Plan node for [`SelectCellExt::select_cell`].
22///
23/// Reactive filter: each row's inclusion is decided by a per-row
24/// [`Watchable`] producing `bool` from `predicate(&key, &value)`. The row is
25/// included while the gate is `true` and excluded when `false`.
26///
27/// Not [`Clone`]: cloning a plan would silently duplicate per-row
28/// subscription work; share by materializing once.
29#[allow(private_bounds)]
30pub struct SelectCellPlan<S, K, V, W, F>
31where
32    S: MapQuery<Key = K, Value = V>,
33    K: Hash + Eq + CellValue,
34    V: CellValue,
35    W: Pipeline<bool>
36        + crate::pipeline::PipelineSeed<bool>
37        + Gettable<bool>
38        + Clone
39        + Send
40        + Sync
41        + 'static,
42    F: Fn(&K, &V) -> W + Send + Sync + 'static,
43{
44    pub(crate) source: S,
45    pub(crate) predicate: Arc<F>,
46    pub(crate) _types: PhantomData<fn() -> (K, V, W)>,
47}
48
49#[allow(private_bounds)]
50impl<S, K, V, W, F> PlanProperties for SelectCellPlan<S, K, V, W, F>
51where
52    S: MapQuery<Key = K, Value = V>,
53    K: Hash + Eq + CellValue,
54    V: CellValue,
55    W: Pipeline<bool>
56        + crate::pipeline::PipelineSeed<bool>
57        + Gettable<bool>
58        + Clone
59        + Send
60        + Sync
61        + 'static,
62    F: Fn(&K, &V) -> W + Send + Sync + 'static,
63{
64    type Cardinality = ZeroOrOne;
65    type InputPartition = S::OutputPartition;
66    type OutputPartition = ByMapKey<K>;
67}
68
69impl<S, K, V, W, F> BuildQueryRuntime<K, V> for SelectCellPlan<S, K, V, W, F>
70where
71    S: MapQuery<Key = K, Value = V>,
72    K: Hash + Eq + CellValue,
73    V: CellValue,
74    W: Pipeline<bool>
75        + crate::pipeline::PipelineSeed<bool>
76        + Gettable<bool>
77        + Clone
78        + Send
79        + Sync
80        + 'static,
81    F: Fn(&K, &V) -> W + Send + Sync + 'static,
82{
83    fn build_into(
84        self,
85        cx: &mut crate::map_query::compiler::CompileContext,
86        sink: crate::map_query::BoxedMapDiffSink<K, V>,
87    ) -> Vec<SubscriptionGuard> {
88        // Implement select_cell as a project_cell whose inner Watchable maps
89        // the boolean gate into Option<(K, V)>.
90        let predicate = self.predicate;
91        let inner_plan = self.source.project_cell(move |k: &K, v: &V| {
92            let k = k.clone();
93            let v = v.clone();
94            predicate(&k, &v)
95                .map(move |include| {
96                    if *include {
97                        Some((k.clone(), v.clone()))
98                    } else {
99                        None
100                    }
101                })
102                .materialize()
103        });
104        crate::map_query::compile_runtime_into(inner_plan, cx, sink)
105    }
106}
107
108#[allow(private_bounds)]
109impl<S, K, V, W, F> MapQuery for SelectCellPlan<S, K, V, W, F>
110where
111    S: MapQuery<Key = K, Value = V>,
112    K: Hash + Eq + CellValue,
113    V: CellValue,
114    W: Pipeline<bool>
115        + crate::pipeline::PipelineSeed<bool>
116        + Gettable<bool>
117        + Clone
118        + Send
119        + Sync
120        + 'static,
121    F: Fn(&K, &V) -> W + Send + Sync + 'static,
122{
123    type Key = K;
124    type Value = V;
125}
126
127/// Select-cell operator returning a [`MapQuery`] plan node.
128///
129/// `select_cell` consumes `self` and returns an uncompiled plan node; call
130/// [`MapQuery::materialize`] on the result to obtain a subscribable
131/// [`CellMap`](crate::CellMap).
132pub trait SelectCellExt<K, V>: MapQuery<Key = K, Value = V>
133where
134    K: Hash + Eq + CellValue,
135    V: CellValue,
136{
137    /// Reactive filter: each row has a watchable boolean gate.
138    ///
139    /// `predicate(&key, &value)` returns a watchable bool; the row is
140    /// included when true and excluded when false.
141    #[track_caller]
142    #[allow(private_bounds)]
143    fn select_cell<W, F>(self, predicate: F) -> impl MapQuery<Key = K, Value = V>
144    where
145        W: Pipeline<bool>
146            + crate::pipeline::PipelineSeed<bool>
147            + Gettable<bool>
148            + Clone
149            + Send
150            + Sync
151            + 'static,
152        F: Fn(&K, &V) -> W + Send + Sync + 'static,
153    {
154        SelectCellPlan {
155            source: self,
156            predicate: Arc::new(predicate),
157            _types: PhantomData,
158        }
159    }
160}
161
162impl<K, V, M> SelectCellExt<K, V> for M
163where
164    K: Hash + Eq + CellValue,
165    V: CellValue,
166    M: MapQuery<Key = K, Value = V>,
167{
168}
169
170#[cfg(test)]
171mod tests {
172    use std::sync::mpsc;
173
174    use super::*;
175    use crate::{Cell, CellMap, MapExt, Materialize, cell_map::MapDiff};
176
177    #[test]
178    fn select_cell_reacts_to_predicate_changes() {
179        let values = CellMap::<String, i32>::new();
180        let gates = CellMap::<String, bool>::new();
181
182        values.insert("a".to_string(), 10);
183        values.insert("b".to_string(), 20);
184        gates.insert("a".to_string(), false);
185        gates.insert("b".to_string(), true);
186
187        let filtered = values
188            .select_cell({
189                let gates = gates.clone();
190                move |key, _value| gates.get(key).map(|v| v.unwrap_or(false)).materialize()
191            })
192            .materialize();
193
194        assert_eq!(filtered.entries().materialize().get().len(), 1);
195        assert!(!filtered.contains_key(&"a".to_string()));
196        assert!(filtered.contains_key(&"b".to_string()));
197
198        gates.insert("a".to_string(), true);
199        assert_eq!(filtered.entries().materialize().get().len(), 2);
200        gates.insert("b".to_string(), false);
201        assert_eq!(filtered.entries().materialize().get().len(), 1);
202    }
203
204    #[test]
205    fn select_cell_preserves_upstream_batch_without_extra_emissions() {
206        let source = CellMap::<String, i32>::new();
207        let out = source
208            .clone()
209            .select_cell(|_, _| Cell::new(true).lock())
210            .materialize();
211
212        let (tx, rx) = mpsc::channel::<MapDiff<String, i32>>();
213        let _guard = out.subscribe_diffs(move |diff| {
214            let _ = tx.send(diff.clone());
215        });
216
217        source.insert_many(vec![("a".to_string(), 1), ("b".to_string(), 2)]);
218        let seen: Vec<_> = rx.try_iter().collect();
219        assert_eq!(seen.len(), 2);
220        assert!(matches!(
221            seen.last(),
222            Some(MapDiff::Batch { changes }) if changes.len() == 2
223        ));
224    }
225}