1use std::collections::HashMap;
4use std::fmt::{Debug, Display};
5use std::ops::{Bound, RangeBounds};
6use std::sync::OnceLock;
7
8use documented::DocumentedVariants;
9use proc_macro2::{Ident, Literal, Span, TokenStream};
10use quote::quote_spanned;
11use serde::{Deserialize, Serialize};
12use slotmap::Key;
13use syn::punctuated::Punctuated;
14use syn::{Expr, Token, parse_quote_spanned};
15
16use super::{
17 GraphLoopId, GraphNode, GraphNodeId, GraphSubgraphId, OpInstGenerics, OperatorInstance,
18 PortIndexValue,
19};
20use crate::diagnostic::{Diagnostic, Diagnostics, Level};
21use crate::parse::{Operator, PortIndex};
22
23#[derive(Clone, Copy, PartialOrd, Ord, PartialEq, Eq, Debug, Serialize, Deserialize)]
25pub enum DelayType {
26 Tick,
28 TickLazy,
30 Loop,
32 LoopLazy,
34}
35
36pub enum PortListSpec {
38 Variadic,
40 Fixed(Punctuated<PortIndex, Token![,]>),
42}
43
44pub struct OperatorConstraints {
46 pub name: &'static str,
48 pub categories: &'static [OperatorCategory],
50
51 pub hard_range_inn: &'static dyn RangeTrait<usize>,
54 pub soft_range_inn: &'static dyn RangeTrait<usize>,
56 pub hard_range_out: &'static dyn RangeTrait<usize>,
58 pub soft_range_out: &'static dyn RangeTrait<usize>,
60 pub num_args: usize,
62 pub persistence_args: &'static dyn RangeTrait<usize>,
64 pub type_args: &'static dyn RangeTrait<usize>,
68 pub is_external_input: bool,
71 pub flo_type: Option<FloType>,
73
74 pub ports_inn: Option<fn() -> PortListSpec>,
76 pub ports_out: Option<fn() -> PortListSpec>,
78
79 pub input_delaytype_fn: fn(&PortIndexValue) -> Option<DelayType>,
81 pub write_fn: WriteFn,
83}
84
85pub type WriteFn = fn(&WriteContextArgs<'_>, &mut Diagnostics) -> Result<OperatorWriteOutput, ()>;
87
88impl Debug for OperatorConstraints {
89 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
90 f.debug_struct("OperatorConstraints")
91 .field("name", &self.name)
92 .field("hard_range_inn", &self.hard_range_inn)
93 .field("soft_range_inn", &self.soft_range_inn)
94 .field("hard_range_out", &self.hard_range_out)
95 .field("soft_range_out", &self.soft_range_out)
96 .field("num_args", &self.num_args)
97 .field("persistence_args", &self.persistence_args)
98 .field("type_args", &self.type_args)
99 .field("is_external_input", &self.is_external_input)
100 .field("ports_inn", &self.ports_inn)
101 .field("ports_out", &self.ports_out)
102 .finish()
106 }
107}
108
109#[derive(Default)]
113pub struct OperatorWriteOutput {
114 pub write_prologue: TokenStream,
117 pub write_iterator: TokenStream,
124 pub write_iterator_after: TokenStream,
126 pub write_tick_end: TokenStream,
129}
130
131pub const RANGE_ANY: &'static dyn RangeTrait<usize> = &(0..);
133pub const RANGE_0: &'static dyn RangeTrait<usize> = &(0..=0);
135pub const RANGE_1: &'static dyn RangeTrait<usize> = &(1..=1);
137
138pub fn identity_write_iterator_fn(
141 &WriteContextArgs {
142 root,
143 op_span,
144 ident,
145 inputs,
146 outputs,
147 is_pull,
148 op_inst:
149 OperatorInstance {
150 generics: OpInstGenerics { type_args, .. },
151 ..
152 },
153 ..
154 }: &WriteContextArgs,
155) -> TokenStream {
156 let generic_type = type_args
157 .first()
158 .map(quote::ToTokens::to_token_stream)
159 .unwrap_or(quote_spanned!(op_span=> _));
160
161 if is_pull {
162 let input = &inputs[0];
163 quote_spanned! {op_span=>
164 let #ident = {
165 fn check_input<Pull, Item>(pull: Pull) -> impl #root::dfir_pipes::pull::Pull<Item = Item, Meta = Pull::Meta, CanPend = Pull::CanPend, CanEnd = Pull::CanEnd>
166 where
167 Pull: #root::dfir_pipes::pull::Pull<Item = Item>,
168 {
169 pull
170 }
171 check_input::<_, #generic_type>(#input)
172 };
173 }
174 } else {
175 let output = &outputs[0];
176 quote_spanned! {op_span=>
177 let #ident = {
178 fn check_output<Psh, Item>(push: Psh) -> impl #root::dfir_pipes::push::Push<Item, (), CanPend = Psh::CanPend>
179 where
180 Psh: #root::dfir_pipes::push::Push<Item, ()>,
181 {
182 push
183 }
184 check_output::<_, #generic_type>(#output)
185 };
186 }
187 }
188}
189
190pub const IDENTITY_WRITE_FN: WriteFn = |write_context_args, _| {
192 let write_iterator = identity_write_iterator_fn(write_context_args);
193 Ok(OperatorWriteOutput {
194 write_iterator,
195 ..Default::default()
196 })
197};
198
199pub fn null_write_iterator_fn(
202 &WriteContextArgs {
203 root,
204 op_span,
205 ident,
206 inputs,
207 outputs,
208 is_pull,
209 op_inst:
210 OperatorInstance {
211 generics: OpInstGenerics { type_args, .. },
212 ..
213 },
214 ..
215 }: &WriteContextArgs,
216) -> TokenStream {
217 let default_type = parse_quote_spanned! {op_span=> _};
218 let iter_type = type_args.first().unwrap_or(&default_type);
219
220 if is_pull {
221 quote_spanned! {op_span=>
222 let #ident = #root::dfir_pipes::pull::poll_fn({
223 #(
224 let mut #inputs = ::std::boxed::Box::pin(#inputs);
225 )*
226 move |_cx| {
227 #(
231 let #inputs = #root::dfir_pipes::pull::Pull::pull(
232 ::std::pin::Pin::as_mut(&mut #inputs),
233 <_ as #root::dfir_pipes::Context>::from_task(_cx),
234 );
235 )*
236 #(
237 if let #root::dfir_pipes::pull::PullStep::Pending(_) = #inputs {
238 return #root::dfir_pipes::pull::PullStep::Pending(#root::dfir_pipes::Yes);
239 }
240 )*
241 #root::dfir_pipes::pull::PullStep::<_, _, #root::dfir_pipes::Yes, _>::Ended(#root::dfir_pipes::Yes)
242 }
243 });
244 }
245 } else {
246 quote_spanned! {op_span=>
247 #[allow(clippy::let_unit_value)]
248 let _ = (#(#outputs),*);
249 let #ident = #root::dfir_pipes::push::for_each::<_, #iter_type>(::std::mem::drop::<#iter_type>);
250 }
251 }
252}
253
254pub const NULL_WRITE_FN: WriteFn = |write_context_args, _| {
257 let write_iterator = null_write_iterator_fn(write_context_args);
258 Ok(OperatorWriteOutput {
259 write_iterator,
260 ..Default::default()
261 })
262};
263
264macro_rules! declare_ops {
265 ( $( $mod:ident :: $op:ident, )* ) => {
266 $( pub(crate) mod $mod; )*
267 pub const OPERATORS: &[OperatorConstraints] = &[
269 $( $mod :: $op, )*
270 ];
271 };
272}
273declare_ops![
274 all_iterations::ALL_ITERATIONS,
275 anti_join::ANTI_JOIN,
276 assert::ASSERT,
277 assert_eq::ASSERT_EQ,
278 batch::BATCH,
279 batch_lazy::BATCH_LAZY,
280 chain::CHAIN,
281 chain_first_n::CHAIN_FIRST_N,
282 _counter::_COUNTER,
283 cross_join::CROSS_JOIN,
284 cross_join_multiset::CROSS_JOIN_MULTISET,
285 cross_singleton::CROSS_SINGLETON,
286 demux_enum::DEMUX_ENUM,
287 dest_file::DEST_FILE,
288 dest_sink::DEST_SINK,
289 dest_sink_serde::DEST_SINK_SERDE,
290 difference::DIFFERENCE,
291 enumerate::ENUMERATE,
292 filter::FILTER,
293 filter_map::FILTER_MAP,
294 flat_map::FLAT_MAP,
295 flat_map_stream_blocking::FLAT_MAP_STREAM_BLOCKING,
296 flatten::FLATTEN,
297 flatten_stream_blocking::FLATTEN_STREAM_BLOCKING,
298 fold::FOLD,
299 fold_no_replay::FOLD_NO_REPLAY,
300 for_each::FOR_EACH,
301 identity::IDENTITY,
302 initialize::INITIALIZE,
303 inspect::INSPECT,
304 iter_ref::ITER_REF,
305 join::JOIN,
306 join_fused::JOIN_FUSED,
307 join_fused_lhs::JOIN_FUSED_LHS,
308 join_fused_rhs::JOIN_FUSED_RHS,
309 join_multiset::JOIN_MULTISET,
310 join_multiset_half::JOIN_MULTISET_HALF,
311 fold_keyed::FOLD_KEYED,
312 reduce_keyed::REDUCE_KEYED,
313 lattice_bimorphism::LATTICE_BIMORPHISM,
314 _lattice_fold_batch::_LATTICE_FOLD_BATCH,
315 lattice_fold::LATTICE_FOLD,
316 _lattice_join_fused_join::_LATTICE_JOIN_FUSED_JOIN,
317 lattice_reduce::LATTICE_REDUCE,
318 map::MAP,
319 union::UNION,
320 multiset_delta::MULTISET_DELTA,
321 defer_signal::DEFER_SIGNAL,
322 defer_tick::DEFER_TICK,
323 defer_tick_lazy::DEFER_TICK_LAZY,
324 null::NULL,
325 partition::PARTITION,
326 persist::PERSIST,
327 resolve_futures::RESOLVE_FUTURES,
328 resolve_futures_blocking::RESOLVE_FUTURES_BLOCKING,
329 resolve_futures_blocking_ordered::RESOLVE_FUTURES_BLOCKING_ORDERED,
330 resolve_futures_ordered::RESOLVE_FUTURES_ORDERED,
331 reduce::REDUCE,
332 reduce_no_replay::REDUCE_NO_REPLAY,
333 scan::SCAN,
334 scan_async_blocking::SCAN_ASYNC_BLOCKING,
335 spin::SPIN,
336 sort::SORT,
337 sort_by_key::SORT_BY_KEY,
338 source_file::SOURCE_FILE,
339 source_interval::SOURCE_INTERVAL,
340 source_iter::SOURCE_ITER,
341 source_json::SOURCE_JSON,
342 source_stdin::SOURCE_STDIN,
343 source_stream::SOURCE_STREAM,
344 source_stream_serde::SOURCE_STREAM_SERDE,
345 state::STATE,
346 state_by::STATE_BY,
347 tee::TEE,
348 unique::UNIQUE,
349 unzip::UNZIP,
350 zip::ZIP,
351 zip_longest::ZIP_LONGEST,
352];
353
354pub fn operator_lookup() -> &'static HashMap<&'static str, &'static OperatorConstraints> {
356 pub static OPERATOR_LOOKUP: OnceLock<HashMap<&'static str, &'static OperatorConstraints>> =
357 OnceLock::new();
358 OPERATOR_LOOKUP.get_or_init(|| OPERATORS.iter().map(|op| (op.name, op)).collect())
359}
360pub fn find_node_op_constraints(node: &GraphNode) -> Option<&'static OperatorConstraints> {
362 if let GraphNode::Operator(operator) = node {
363 find_op_op_constraints(operator)
364 } else {
365 None
366 }
367}
368pub fn find_op_op_constraints(operator: &Operator) -> Option<&'static OperatorConstraints> {
370 let name = &*operator.name_string();
371 operator_lookup().get(name).copied()
372}
373
374#[derive(Clone)]
376pub struct WriteContextArgs<'a> {
377 pub root: &'a TokenStream,
379 pub context: &'a Ident,
382 pub df_ident: &'a Ident,
386 pub subgraph_id: GraphSubgraphId,
388 pub node_id: GraphNodeId,
390 pub loop_id: Option<GraphLoopId>,
392 pub op_span: Span,
394 pub op_tag: Option<String>,
396 pub work_fn: &'a Ident,
398 pub work_fn_async: &'a Ident,
400
401 pub ident: &'a Ident,
403 pub is_pull: bool,
405 pub inputs: &'a [Ident],
407 pub outputs: &'a [Ident],
409
410 pub op_name: &'static str,
412 pub op_inst: &'a OperatorInstance,
414 pub arguments: &'a Punctuated<Expr, Token![,]>,
420}
421impl WriteContextArgs<'_> {
422 pub fn make_ident(&self, suffix: impl AsRef<str>) -> Ident {
428 Ident::new(
429 &format!(
430 "sg_{:?}_node_{:?}_{}",
431 self.subgraph_id.data(),
432 self.node_id.data(),
433 suffix.as_ref(),
434 ),
435 self.op_span,
436 )
437 }
438
439 pub fn persistence_args<const N: usize>(
442 &self,
443 diagnostics: &mut Diagnostics,
444 ) -> [Persistence; N] {
445 let len = self.op_inst.generics.persistence_args.len();
446 if 0 != len && 1 != len && N != len {
447 diagnostics.push(Diagnostic::spanned(
448 self.op_span,
449 Level::Error,
450 format!(
451 "The operator `{}` only accepts 0, 1, or {} persistence arguments",
452 self.op_name, N
453 ),
454 ));
455 }
456
457 let default_persistence = if self.loop_id.is_some() {
458 Persistence::None
459 } else {
460 Persistence::Tick
461 };
462 let mut out = [default_persistence; N];
463 self.op_inst
464 .generics
465 .persistence_args
466 .iter()
467 .copied()
468 .cycle() .take(N)
470 .enumerate()
471 .for_each(|(i, p)| {
472 out[i] = p;
473 });
474 out
475 }
476}
477
478pub trait RangeTrait<T>: Send + Sync + Debug
480where
481 T: ?Sized,
482{
483 fn start_bound(&self) -> Bound<&T>;
485 fn end_bound(&self) -> Bound<&T>;
487 fn contains(&self, item: &T) -> bool
489 where
490 T: PartialOrd<T>;
491
492 fn human_string(&self) -> String
494 where
495 T: Display + PartialEq,
496 {
497 match (self.start_bound(), self.end_bound()) {
498 (Bound::Unbounded, Bound::Unbounded) => "any number of".to_owned(),
499
500 (Bound::Included(n), Bound::Included(x)) if n == x => {
501 format!("exactly {}", n)
502 }
503 (Bound::Included(n), Bound::Included(x)) => {
504 format!("at least {} and at most {}", n, x)
505 }
506 (Bound::Included(n), Bound::Excluded(x)) => {
507 format!("at least {} and less than {}", n, x)
508 }
509 (Bound::Included(n), Bound::Unbounded) => format!("at least {}", n),
510 (Bound::Excluded(n), Bound::Included(x)) => {
511 format!("more than {} and at most {}", n, x)
512 }
513 (Bound::Excluded(n), Bound::Excluded(x)) => {
514 format!("more than {} and less than {}", n, x)
515 }
516 (Bound::Excluded(n), Bound::Unbounded) => format!("more than {}", n),
517 (Bound::Unbounded, Bound::Included(x)) => format!("at most {}", x),
518 (Bound::Unbounded, Bound::Excluded(x)) => format!("less than {}", x),
519 }
520 }
521}
522
523impl<R, T> RangeTrait<T> for R
524where
525 R: RangeBounds<T> + Send + Sync + Debug,
526{
527 fn start_bound(&self) -> Bound<&T> {
528 self.start_bound()
529 }
530
531 fn end_bound(&self) -> Bound<&T> {
532 self.end_bound()
533 }
534
535 fn contains(&self, item: &T) -> bool
536 where
537 T: PartialOrd<T>,
538 {
539 self.contains(item)
540 }
541}
542
543#[derive(Clone, Copy, PartialOrd, Ord, PartialEq, Eq, Debug, Serialize, Deserialize)]
545pub enum Persistence {
546 None,
548 Loop,
550 Tick,
552 Static,
554}
555impl Persistence {
556 pub fn to_str_lowercase(self) -> &'static str {
558 match self {
559 Persistence::None => "none",
560 Persistence::Tick => "tick",
561 Persistence::Loop => "loop",
562 Persistence::Static => "static",
563 }
564 }
565}
566
567fn make_missing_runtime_msg(op_name: &str) -> Literal {
569 Literal::string(&format!(
570 "`{}()` must be used within a Tokio runtime. For example, use `#[dfir_rs::main]` on your main method.",
571 op_name
572 ))
573}
574
575#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash, DocumentedVariants)]
577pub enum OperatorCategory {
578 Map,
580 Filter,
582 Flatten,
584 Fold,
586 KeyedFold,
588 LatticeFold,
590 Persistence,
592 MultiIn,
594 MultiOut,
596 Source,
598 Sink,
600 Control,
602 CompilerFusionOperator,
604 Windowing,
606 Unwindowing,
608}
609impl OperatorCategory {
610 pub fn name(self) -> &'static str {
612 self.get_variant_docs().split_once(":").unwrap().0
613 }
614 pub fn description(self) -> &'static str {
616 self.get_variant_docs().split_once(":").unwrap().1
617 }
618}
619
620#[derive(Clone, Copy, PartialOrd, Ord, PartialEq, Eq, Debug)]
622pub enum FloType {
623 Source,
625 Windowing,
627 WindowingLazy,
630 Unwindowing,
632}