1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
use quote::{ToTokens, quote_spanned};
use super::{
DelayType, OpInstGenerics, OperatorCategory, OperatorConstraints, OperatorInstance,
OperatorWriteOutput, Persistence, RANGE_1, WriteContextArgs,
};
/// > 1 input stream of type `(K, V)`, 1 output stream of type `(K, V)`.
/// > The output will have one tuple for each distinct `K`, with an accumulated (reduced) value of
/// > type `V`.
///
/// If you need the accumulated value to have a different type than the input, use [`fold_keyed`](#fold_keyed).
///
/// > Arguments: one Rust closures. The closure takes two arguments: an `&mut` 'accumulator', and
/// > an element. Accumulator should be updated based on the element.
///
/// A special case of `reduce`, in the spirit of SQL's GROUP BY and aggregation constructs. The input
/// is partitioned into groups by the first field, and for each group the values in the second
/// field are accumulated via the closures in the arguments.
///
/// > Note: The closures have access to the [`context` object](surface_flows.mdx#the-context-object).
///
/// `reduce_keyed` can also be provided with one generic lifetime persistence argument, either
/// `'tick` or `'static`, to specify how data persists. With `'tick`, values will only be collected
/// within the same tick. With `'static`, values will be remembered across ticks and will be
/// aggregated with pairs arriving in later ticks. When not explicitly specified persistence
/// defaults to `'tick`.
///
/// `reduce_keyed` can also be provided with two type arguments, the key and value type. This is
/// required when using `'static` persistence if the compiler cannot infer the types.
///
/// ```dfir
/// source_iter([("toy", 1), ("toy", 2), ("shoe", 11), ("shoe", 35), ("haberdashery", 7)])
/// -> reduce_keyed(|old: &mut u32, val: u32| *old += val)
/// -> assert_eq([("toy", 3), ("shoe", 46), ("haberdashery", 7)]);
/// ```
///
/// Example using `'tick` persistence and type arguments:
/// ```rustbook
/// let (input_send, input_recv) = dfir_rs::util::unbounded_channel::<(&str, &str)>();
/// let mut flow = dfir_rs::dfir_syntax! {
/// source_stream(input_recv)
/// -> reduce_keyed::<'tick, &str>(|old: &mut _, val| *old = std::cmp::max(*old, val))
/// -> for_each(|(k, v)| println!("({:?}, {:?})", k, v));
/// };
///
/// input_send.send(("hello", "oakland")).unwrap();
/// input_send.send(("hello", "berkeley")).unwrap();
/// input_send.send(("hello", "san francisco")).unwrap();
/// flow.run_available();
/// // ("hello", "oakland, berkeley, san francisco, ")
///
/// input_send.send(("hello", "palo alto")).unwrap();
/// flow.run_available();
/// // ("hello", "palo alto, ")
/// ```
pub const REDUCE_KEYED: OperatorConstraints = OperatorConstraints {
name: "reduce_keyed",
categories: &[OperatorCategory::KeyedFold],
hard_range_inn: RANGE_1,
soft_range_inn: RANGE_1,
hard_range_out: RANGE_1,
soft_range_out: RANGE_1,
num_args: 1,
persistence_args: &(0..=1),
type_args: &(0..=2),
is_external_input: false,
has_singleton_output: true,
flo_type: None,
ports_inn: None,
ports_out: None,
input_delaytype_fn: |_| Some(DelayType::Stratum),
write_fn: |wc @ &WriteContextArgs {
df_ident,
context,
op_span,
ident,
inputs,
singleton_output_ident,
is_pull,
work_fn_async,
root,
op_name,
op_inst:
OperatorInstance {
generics: OpInstGenerics { type_args, .. },
..
},
arguments,
..
},
diagnostics| {
assert!(is_pull, "TODO(mingwei): `{}` only supports pull.", op_name);
let [persistence] = wc.persistence_args_disallow_mutable(diagnostics);
let generic_type_args = [
type_args
.first()
.map(ToTokens::to_token_stream)
.unwrap_or(quote_spanned!(op_span=> _)),
type_args
.get(1)
.map(ToTokens::to_token_stream)
.unwrap_or(quote_spanned!(op_span=> _)),
];
let input = &inputs[0];
let aggfn = &arguments[0];
let hashtable_ident = wc.make_ident("hashtable");
let write_prologue = quote_spanned! {op_span=>
let #singleton_output_ident = #df_ident.add_state(::std::cell::RefCell::new(#root::rustc_hash::FxHashMap::<#( #generic_type_args ),*>::default()));
};
let write_prologue_after = wc
.persistence_as_state_lifespan(persistence)
.map(|lifespan| quote_spanned! {op_span=>
#df_ident.set_state_lifespan_hook(#singleton_output_ident, #lifespan, |rcell| { rcell.take(); });
}).unwrap_or_default();
let write_iterator = {
let iter_expr = match persistence {
Persistence::None | Persistence::Tick => quote_spanned! {op_span=>
#hashtable_ident.drain()
},
Persistence::Loop => quote_spanned! {op_span=>
#hashtable_ident.iter().map(
#[allow(suspicious_double_ref_op, clippy::clone_on_copy)]
|(k, v)| (
::std::clone::Clone::clone(k),
::std::clone::Clone::clone(v),
)
)
},
Persistence::Static => quote_spanned! {op_span=>
// Play everything but only on the first run of this tick/stratum.
// (We know we won't have any more inputs, so it is fine to only play once.
// Because of the `DelayType::Stratum` or `DelayType::MonotoneAccum`).
#context.is_first_run_this_tick()
.then_some(#hashtable_ident.iter())
.into_iter()
.flatten()
.map(
#[allow(suspicious_double_ref_op, clippy::clone_on_copy)]
|(k, v)| (
::std::clone::Clone::clone(k),
::std::clone::Clone::clone(v),
)
)
},
Persistence::Mutable => unreachable!(),
};
quote_spanned! {op_span=>
let mut #hashtable_ident = unsafe {
// SAFETY: handle from `#df_ident.add_state(..)`.
#context.state_ref_unchecked(#singleton_output_ident)
}.borrow_mut();
{
#[inline(always)]
fn check_input<Prev, K, V>(prev: Prev)
-> impl #root::dfir_pipes::pull::Pull<Item = (K, V), Meta = Prev::Meta, CanPend = Prev::CanPend, CanEnd = Prev::CanEnd>
where
Prev: #root::dfir_pipes::pull::Pull<Item = (K, V)>,
K: ::std::clone::Clone,
V: ::std::clone::Clone
{
prev
}
/// A: accumulator/item type
#[inline(always)]
fn call_comb_type<A>(acc: &mut A, item: A, f: impl Fn(&mut A, A)) {
let () = (f)(acc, item);
}
let fut = #root::dfir_pipes::pull::Pull::for_each(check_input(#input), |kv| {
match #hashtable_ident.entry(kv.0) {
::std::collections::hash_map::Entry::Vacant(vacant) => {
vacant.insert(kv.1);
}
::std::collections::hash_map::Entry::Occupied(mut occupied) => {
call_comb_type(occupied.get_mut(), kv.1, #aggfn);
}
}
});
let () = #work_fn_async(fut).await;
}
#[allow(clippy::disallowed_methods, reason = "FxHasher is deterministic")]
let #ident = #iter_expr;
let #ident = #root::dfir_pipes::pull::iter(#ident);
}
};
let write_iterator_after = match persistence {
Persistence::None | Persistence::Tick | Persistence::Loop => Default::default(),
Persistence::Static | Persistence::Mutable => quote_spanned! {op_span=>
// Reschedule the subgraph lazily to ensure replay on later ticks.
#context.schedule_subgraph(#context.current_subgraph(), false);
},
};
Ok(OperatorWriteOutput {
write_prologue,
write_prologue_after,
write_iterator,
write_iterator_after,
})
},
};