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
use quote::{ToTokens, quote_spanned};
use super::{
OpInstGenerics, OperatorCategory, OperatorConstraints, OperatorInstance, OperatorWriteOutput,
Persistence, RANGE_1, WriteContextArgs,
};
use crate::diagnostic::{Diagnostic, Level};
/// List state operator, but with a closure to map the input to the state lattice and a factory
/// function to initialize the internal data structure.
///
/// The emitted outputs (both the referencable singleton and the optional pass-through stream) are
/// of the same type as the inputs to the state_by operator and are not required to be a lattice
/// type. This is useful receiving pass-through context information on the output side.
///
/// ```dfir
/// use std::collections::HashSet;
///
///
/// use lattices::set_union::{CartesianProductBimorphism, SetUnionHashSet, SetUnionSingletonSet};
///
/// my_state = source_iter(0..3)
/// -> state_by::<SetUnionHashSet<usize>>(SetUnionSingletonSet::new_from, std::default::Default::default);
/// ```
/// The 2nd argument into `state_by` is a factory function that can be used to supply a custom
/// initial value for the backing state. The initial value is still expected to be bottom (and will
/// be checked). This is useful for doing things like pre-allocating buffers, etc. In the above
/// example, it is just using `Default::default()`
///
/// An example of preallocating the capacity in a hashmap:
///
/// ```dfir
/// use std::collections::HashSet;
/// use lattices::set_union::{SetUnion, CartesianProductBimorphism, SetUnionHashSet, SetUnionSingletonSet};
///
/// my_state = source_iter(0..3)
/// -> state_by::<SetUnionHashSet<usize>>(SetUnionSingletonSet::new_from, {|| SetUnion::new(HashSet::<usize>::with_capacity(1_000)) });
/// ```
///
/// The `state` operator is equivalent to `state_by` used with an identity mapping operator with
/// `Default::default` providing the factory function.
pub const STATE_BY: OperatorConstraints = OperatorConstraints {
name: "state_by",
categories: &[OperatorCategory::Persistence],
hard_range_inn: RANGE_1,
soft_range_inn: RANGE_1,
hard_range_out: &(0..=1),
soft_range_out: &(0..=1),
num_args: 2,
persistence_args: &(0..=1),
type_args: &(0..=1),
is_external_input: false,
has_singleton_output: true,
flo_type: None,
ports_inn: None,
ports_out: None,
input_delaytype_fn: |_| None,
write_fn: |wc @ &WriteContextArgs {
root,
context,
df_ident,
op_span,
ident,
inputs,
outputs,
is_pull,
singleton_output_ident,
op_name,
op_inst:
OperatorInstance {
generics:
OpInstGenerics {
type_args,
persistence_args,
..
},
..
},
arguments,
..
},
diagnostics| {
let lattice_type = type_args
.first()
.map(ToTokens::to_token_stream)
.unwrap_or(quote_spanned!(op_span=> _));
let persistence = match persistence_args[..] {
[] => Persistence::Tick,
[Persistence::Mutable] => {
diagnostics.push(Diagnostic::spanned(
op_span,
Level::Error,
format!("{} does not support `'mut`.", op_name),
));
Persistence::Tick
}
[a] => a,
_ => unreachable!(),
};
let state_ident = singleton_output_ident;
let factory_fn = &arguments[1];
let write_prologue = quote_spanned! {op_span=>
let #state_ident = {
let data_struct: #lattice_type = (#factory_fn)();
::std::debug_assert!(#root::lattices::IsBot::is_bot(&data_struct));
#df_ident.add_state(::std::cell::RefCell::new(data_struct))
};
};
let write_prologue_after = wc
.persistence_as_state_lifespan(persistence)
.map(|lifespan| quote_spanned! {op_span=>
#df_ident.set_state_lifespan_hook(#state_ident, #lifespan, |rcell| { rcell.take(); });
}).unwrap_or_default();
let by_fn = &arguments[0];
// TODO(mingwei): deduplicate codegen
let write_iterator = if is_pull {
let input = &inputs[0];
quote_spanned! {op_span=>
let #ident = {
let state_ref = unsafe {
// SAFETY: handle from `#df_ident.add_state(..)`.
#context.state_ref_unchecked(#state_ident)
};
fn check_input<'a, Item, MappingFn, MappedItem, Prev, Lat>(
prev: Prev,
mapfn: MappingFn,
state_ref: &'a ::std::cell::RefCell<Lat>,
) -> impl 'a + #root::dfir_pipes::pull::Pull<Item = Item, Meta = Prev::Meta>
where
Item: ::std::clone::Clone,
MappingFn: 'a + Fn(Item) -> MappedItem,
Prev: 'a + #root::dfir_pipes::pull::Pull<Item = Item>,
Lat: 'static + #root::lattices::Merge<MappedItem>,
{
#root::dfir_pipes::pull::Pull::filter(
prev,
move |item| {
let mut state = state_ref.borrow_mut();
#root::lattices::Merge::merge(&mut *state, (mapfn)(::std::clone::Clone::clone(item)))
},
)
}
check_input::<_, _, _, _, #lattice_type>(#input, #by_fn, state_ref)
};
}
} else if let Some(output) = outputs.first() {
quote_spanned! {op_span=>
let #ident = {
let state_ref = unsafe {
// SAFETY: handle from `#df_ident.add_state(..)`.
#context.state_ref_unchecked(#state_ident)
};
fn check_output<'a, Item, MappingFn, MappedItem, Psh, Lat>(
push: Psh,
mapfn: MappingFn,
state_ref: &'a ::std::cell::RefCell<Lat>,
) -> impl 'a + #root::dfir_pipes::push::Push<Item, ()>
where
Item: 'a + ::std::clone::Clone,
MappingFn: 'a + Fn(Item) -> MappedItem,
Psh: 'a + #root::dfir_pipes::push::Push<Item, ()>,
Lat: 'static + #root::lattices::Merge<MappedItem>,
{
#root::dfir_pipes::push::filter(move |item| {
let mut state = state_ref.borrow_mut();
#root::lattices::Merge::merge(&mut *state, (mapfn)(::std::clone::Clone::clone(item)))
}, push)
}
check_output::<_, _, _, _, #lattice_type>(#output, #by_fn, state_ref)
};
}
} else {
quote_spanned! {op_span=>
let #ident = {
let state_ref = unsafe {
// SAFETY: handle from `#df_ident.add_state(..)`.
#context.state_ref_unchecked(#state_ident)
};
fn check_output<'a, Item, MappingFn, MappedItem, Lat>(
state_ref: &'a ::std::cell::RefCell<Lat>,
mapfn: MappingFn,
) -> impl 'a + #root::dfir_pipes::push::Push<Item, ()>
where
Item: 'a,
MappedItem: 'a,
MappingFn: 'a + Fn(Item) -> MappedItem,
Lat: 'static + #root::lattices::Merge<MappedItem>,
{
#root::dfir_pipes::push::for_each(move |item| {
let mut state = state_ref.borrow_mut();
#root::lattices::Merge::merge(&mut *state, (mapfn)(item));
})
}
check_output::<_, _, _, #lattice_type>(state_ref, #by_fn)
};
}
};
Ok(OperatorWriteOutput {
write_prologue,
write_prologue_after,
write_iterator,
..Default::default()
})
},
};