Skip to main content

datum/stream/
flow.rs

1use super::*;
2use crate::Attributes;
3use crate::context::FlowWithContext;
4use crate::stream::error::{decide_supervision, panic_stream_error};
5use futures::{FutureExt, task::noop_waker};
6use std::any::Any;
7use std::task::Context;
8use std::{
9    collections::{HashMap, HashSet},
10    thread,
11};
12
13pub(super) enum FlowTransform<In, Out> {
14    Pure(PureTransform<In, Out>),
15    Runtime(RuntimeTransform<In, Out>),
16}
17
18pub struct Flow<In, Out, Mat = NotUsed> {
19    pub(super) transform: FlowTransform<In, Out>,
20    pub(super) materialize: Arc<dyn Fn() -> StreamResult<Mat> + Send + Sync>,
21    pub(super) hints: FlowHints,
22    pub(super) attributes: Attributes,
23}
24
25#[derive(Clone, Copy, Debug, PartialEq, Eq)]
26pub(super) enum GroupByBatchMode {
27    Immediate,
28    FiniteEagerNoRecreate,
29}
30
31#[derive(Clone)]
32pub struct BidiFlow<I1, O1, I2, O2> {
33    top: Flow<I1, O1, NotUsed>,
34    bottom: Flow<I2, O2, NotUsed>,
35    attributes: Attributes,
36}
37
38impl<In, Out> Clone for FlowTransform<In, Out> {
39    fn clone(&self) -> Self {
40        match self {
41            Self::Pure(transform) => Self::Pure(Arc::clone(transform)),
42            Self::Runtime(transform) => Self::Runtime(Arc::clone(transform)),
43        }
44    }
45}
46
47impl<In, Out, Mat> Clone for Flow<In, Out, Mat> {
48    fn clone(&self) -> Self {
49        Self {
50            transform: self.transform.clone(),
51            materialize: Arc::clone(&self.materialize),
52            hints: self.hints,
53            attributes: self.attributes.clone(),
54        }
55    }
56}
57
58fn call_supervised<T, F>(context: &str, f: F) -> StreamResult<T>
59where
60    F: FnOnce() -> StreamResult<T>,
61{
62    catch_unwind(AssertUnwindSafe(f)).unwrap_or_else(|_| Err(panic_stream_error(context)))
63}
64
65impl<T: Send + 'static> Flow<T, T, NotUsed> {
66    #[must_use]
67    pub fn identity() -> Self {
68        Self::from_preserving_transform(|input| input)
69    }
70}
71
72impl<In: Send + 'static, Out: Send + 'static> Flow<In, Out, NotUsed> {
73    pub(crate) fn from_transform<F>(transform: F) -> Self
74    where
75        F: Fn(BoxStream<In>) -> BoxStream<Out> + Send + Sync + 'static,
76    {
77        Self::from_parts_with_hints(transform, || Ok(NotUsed), FlowHints::default())
78    }
79
80    pub(crate) fn from_preserving_transform<F>(transform: F) -> Self
81    where
82        F: Fn(BoxStream<In>) -> BoxStream<Out> + Send + Sync + 'static,
83    {
84        Self::from_parts_with_hints(
85            transform,
86            || Ok(NotUsed),
87            FlowHints::PRESERVES_INLINE_HEAD_TERMINAL,
88        )
89    }
90
91    pub(crate) fn from_runtime_transform<F>(transform: F) -> Self
92    where
93        F: Fn(BoxStream<In>, &Materializer) -> StreamResult<BoxStream<Out>> + Send + Sync + 'static,
94    {
95        Self::from_runtime_transform_with_hints(transform, FlowHints::default())
96    }
97
98    fn from_runtime_transform_with_hints<F>(transform: F, hints: FlowHints) -> Self
99    where
100        F: Fn(BoxStream<In>, &Materializer) -> StreamResult<BoxStream<Out>> + Send + Sync + 'static,
101    {
102        Self {
103            transform: FlowTransform::Runtime(Arc::new(transform)),
104            materialize: Arc::new(|| Ok(NotUsed)),
105            hints,
106            attributes: Attributes::default(),
107        }
108    }
109
110    #[must_use]
111    pub fn from_sink_and_source<InMat, OutMat>(
112        sink: Sink<In, InMat>,
113        source: Source<Out, OutMat>,
114    ) -> Self
115    where
116        InMat: Send + 'static,
117        OutMat: Send + 'static,
118    {
119        Self::from_runtime_transform(move |input, materializer| {
120            let sink_keepalive = Arc::new(MaterializedKeepalive::default());
121            let sink_input = Box::new(InputKeepaliveStream {
122                inner: input,
123                keepalive: Arc::clone(&sink_keepalive),
124                peer_keepalive: None,
125            });
126            let sink_mat = sink.clone().run(sink_input, materializer)?;
127            sink_keepalive.store(Box::new(sink_mat));
128            let (output, source_mat) = Arc::clone(&source.factory).create(materializer)?;
129            let source_keepalive = Arc::new(MaterializedKeepalive::default());
130            source_keepalive.store(Box::new(source_mat));
131            Ok(Box::new(CoupledStream {
132                inner: output,
133                source_keepalive,
134                sink_keepalive: None,
135                coupled: false,
136            }))
137        })
138    }
139
140    #[must_use]
141    pub fn from_sink_and_source_coupled<InMat, OutMat>(
142        sink: Sink<In, InMat>,
143        source: Source<Out, OutMat>,
144    ) -> Self
145    where
146        InMat: Send + 'static,
147        OutMat: Send + 'static,
148    {
149        Self::from_runtime_transform(move |input, materializer| {
150            let source_keepalive = Arc::new(MaterializedKeepalive::default());
151            let sink_keepalive = Arc::new(MaterializedKeepalive::default());
152            let sink_input = Box::new(InputKeepaliveStream {
153                inner: input,
154                keepalive: Arc::clone(&sink_keepalive),
155                peer_keepalive: Some(Arc::clone(&source_keepalive)),
156            });
157            let sink_mat = sink.clone().run(sink_input, materializer)?;
158            sink_keepalive.store(Box::new(sink_mat));
159            let (output, source_mat) = Arc::clone(&source.factory).create(materializer)?;
160            source_keepalive.store(Box::new(source_mat));
161            Ok(Box::new(CoupledStream {
162                inner: output,
163                source_keepalive,
164                sink_keepalive: Some(sink_keepalive),
165                coupled: true,
166            }))
167        })
168    }
169    #[must_use]
170    pub fn future_flow<InnerMat, F, Fut>(future: F) -> Flow<In, Out, StreamCompletion<InnerMat>>
171    where
172        InnerMat: Send + 'static,
173        F: Fn() -> Fut + Send + Sync + 'static,
174        Fut: Future<Output = StreamResult<Flow<In, Out, InnerMat>>> + Send + 'static,
175    {
176        let future = Arc::new(future);
177        Flow::from_runtime_materialized_factory(move || {
178            let (sender, receiver) = oneshot::channel();
179            let sender = Arc::new(Mutex::new(Some(sender)));
180            let future = Arc::clone(&future);
181            let transform: RuntimeTransform<In, Out> =
182                Arc::new(move |input, materializer: &Materializer| {
183                    let mat_sender = sender
184                        .lock()
185                        .expect("future_flow materialized sender poisoned")
186                        .take()
187                        .ok_or_else(|| {
188                            StreamError::Failed("future_flow transform already materialized".into())
189                        })?;
190                    Ok(Box::new(FutureFlowStream {
191                        future: Arc::clone(&future),
192                        materializer: materializer
193                            .with_name_prefix(materializer.name_prefix().to_owned()),
194                        input: Some(input),
195                        current: None,
196                        mat_sender: Some(mat_sender),
197                        initialized: false,
198                        terminated: false,
199                        _marker: PhantomData,
200                    }) as BoxStream<Out>)
201                });
202            (transform, StreamCompletion::from_receiver(receiver, None))
203        })
204    }
205
206    #[must_use]
207    pub fn lazy_flow<InnerMat, F>(create: F) -> Flow<In, Out, StreamCompletion<InnerMat>>
208    where
209        InnerMat: Send + 'static,
210        F: Fn() -> Flow<In, Out, InnerMat> + Send + Sync + 'static,
211    {
212        let create = Arc::new(create);
213        Self::lazy_future_flow(move || {
214            let create = Arc::clone(&create);
215            async move { catch_unwind_failed("lazy_flow factory", || create()) }
216        })
217    }
218
219    #[must_use]
220    pub fn lazy_future_flow<InnerMat, F, Fut>(
221        create: F,
222    ) -> Flow<In, Out, StreamCompletion<InnerMat>>
223    where
224        InnerMat: Send + 'static,
225        F: Fn() -> Fut + Send + Sync + 'static,
226        Fut: Future<Output = StreamResult<Flow<In, Out, InnerMat>>> + Send + 'static,
227    {
228        let create = Arc::new(create);
229        Flow::from_runtime_materialized_factory(move || {
230            let (sender, receiver) = oneshot::channel();
231            let sender = Arc::new(Mutex::new(Some(sender)));
232            let create = Arc::clone(&create);
233            let transform: RuntimeTransform<In, Out> =
234                Arc::new(move |input, materializer: &Materializer| {
235                    let mat_sender = sender
236                        .lock()
237                        .expect("lazy_future_flow materialized sender poisoned")
238                        .take()
239                        .ok_or_else(|| {
240                            StreamError::Failed(
241                                "lazy_future_flow transform already materialized".into(),
242                            )
243                        })?;
244                    Ok(Box::new(LazyFutureFlowStream {
245                        create: Arc::clone(&create),
246                        materializer: materializer
247                            .with_name_prefix(materializer.name_prefix().to_owned()),
248                        input: Some(input),
249                        current: None,
250                        mat_sender: Some(mat_sender),
251                        initialized: false,
252                        terminated: false,
253                        _marker: PhantomData,
254                    }) as BoxStream<Out>)
255                });
256            (transform, StreamCompletion::from_receiver(receiver, None))
257        })
258    }
259}
260
261#[derive(Default)]
262struct MaterializedKeepalive {
263    released: AtomicBool,
264    value: Mutex<Option<Box<dyn Any + Send>>>,
265}
266
267impl MaterializedKeepalive {
268    fn store(&self, value: Box<dyn Any + Send>) {
269        let mut slot = self.value.lock().expect("materialized keepalive poisoned");
270        if !self.released.load(Ordering::SeqCst) {
271            *slot = Some(value);
272            return;
273        }
274        // Already released before the mat arrived (e.g. the sink finished on
275        // an empty input before the source side finished materializing):
276        // late stores must get the same ACTIVE release, not a plain drop —
277        // dropping a Cancellable handle does not cancel it.
278        drop(slot);
279        release_materialized_value(value);
280    }
281
282    fn release(&self) {
283        self.released.store(true, Ordering::SeqCst);
284        if let Some(value) = self
285            .value
286            .lock()
287            .expect("materialized keepalive poisoned")
288            .take()
289        {
290            release_materialized_value(value);
291        }
292    }
293
294    fn is_released(&self) -> bool {
295        self.released.load(Ordering::SeqCst)
296    }
297}
298
299fn release_materialized_value(value: Box<dyn Any + Send>) {
300    match value.downcast::<Cancellable>() {
301        Ok(cancellable) => {
302            cancellable.cancel();
303        }
304        Err(value) => drop(value),
305    }
306}
307
308struct InputKeepaliveStream<In> {
309    inner: BoxStream<In>,
310    keepalive: Arc<MaterializedKeepalive>,
311    peer_keepalive: Option<Arc<MaterializedKeepalive>>,
312}
313
314impl<In> Iterator for InputKeepaliveStream<In> {
315    type Item = StreamResult<In>;
316
317    fn next(&mut self) -> Option<Self::Item> {
318        self.inner.next()
319    }
320}
321
322impl<In> Drop for InputKeepaliveStream<In> {
323    fn drop(&mut self) {
324        self.keepalive.release();
325        if let Some(peer_keepalive) = &self.peer_keepalive {
326            peer_keepalive.release();
327        }
328    }
329}
330
331struct CoupledStream<Out> {
332    inner: BoxStream<Out>,
333    source_keepalive: Arc<MaterializedKeepalive>,
334    sink_keepalive: Option<Arc<MaterializedKeepalive>>,
335    coupled: bool,
336}
337
338impl<Out> Iterator for CoupledStream<Out> {
339    type Item = StreamResult<Out>;
340
341    fn next(&mut self) -> Option<Self::Item> {
342        if self.coupled && self.source_keepalive.is_released() {
343            return None;
344        }
345        let next = self.inner.next();
346        if next.is_none() || next.as_ref().is_some_and(|item| item.is_err()) {
347            self.source_keepalive.release();
348            if self.coupled
349                && let Some(sink_keepalive) = &self.sink_keepalive
350            {
351                sink_keepalive.release();
352            }
353        }
354        next
355    }
356}
357
358impl<Out> Drop for CoupledStream<Out> {
359    fn drop(&mut self) {
360        self.source_keepalive.release();
361        if self.coupled
362            && let Some(sink_keepalive) = &self.sink_keepalive
363        {
364            sink_keepalive.release();
365        }
366    }
367}
368
369impl<In: Send + 'static, Out: Send + 'static, Mat: Send + 'static> Flow<In, Out, Mat> {
370    pub fn as_flow_with_context<U, CtxIn, CtxOut, Collapse, Extract>(
371        self,
372        collapse_context: Collapse,
373        extract_context: Extract,
374    ) -> FlowWithContext<U, CtxIn, Out, CtxOut, Mat>
375    where
376        U: Send + 'static,
377        CtxIn: Send + 'static,
378        CtxOut: Send + 'static,
379        Collapse: Fn(U, CtxIn) -> In + Send + Sync + 'static,
380        Extract: Fn(&Out) -> CtxOut + Send + Sync + 'static,
381    {
382        FlowWithContext::from_flow(
383            Flow::identity()
384                .map(move |(value, context)| collapse_context(value, context))
385                .via_mat(self, |_, flow_mat| flow_mat)
386                .map(move |value| {
387                    let context = extract_context(&value);
388                    (value, context)
389                }),
390        )
391    }
392
393    pub(crate) fn from_parts<F, M>(transform: F, materialize: M) -> Self
394    where
395        F: Fn(BoxStream<In>) -> BoxStream<Out> + Send + Sync + 'static,
396        M: Fn() -> StreamResult<Mat> + Send + Sync + 'static,
397    {
398        Self::from_parts_with_hints(transform, materialize, FlowHints::default())
399    }
400
401    pub(crate) fn from_materialized_factory<F>(factory: F) -> Self
402    where
403        F: Fn() -> (PureTransform<In, Out>, Mat) + Send + Sync + 'static,
404    {
405        struct PendingSlot<In, Out, Mat> {
406            transform: Option<PureTransform<In, Out>>,
407            mat: Option<Mat>,
408            transform_taken: bool,
409            mat_taken: bool,
410        }
411
412        let pending = Arc::new(Mutex::new(HashMap::<
413            thread::ThreadId,
414            Vec<PendingSlot<In, Out, Mat>>,
415        >::new()));
416        let factory = Arc::new(factory);
417
418        let transform = {
419            let pending = Arc::clone(&pending);
420            let factory = Arc::clone(&factory);
421            move |input| {
422                let transform = {
423                    let mut pending = pending.lock().expect("flow materialized factory poisoned");
424                    let thread_id = thread::current().id();
425                    let slots = pending.entry(thread_id).or_default();
426                    let index = slots
427                        .iter()
428                        .position(|slot| !slot.transform_taken && slot.mat_taken)
429                        .unwrap_or_else(|| {
430                            let (transform, mat) = factory();
431                            slots.push(PendingSlot {
432                                transform: Some(transform),
433                                mat: Some(mat),
434                                transform_taken: false,
435                                mat_taken: false,
436                            });
437                            slots.len() - 1
438                        });
439                    let slot = slots
440                        .get_mut(index)
441                        .expect("pending flow materialization slot exists");
442                    slot.transform_taken = true;
443                    let transform = slot
444                        .transform
445                        .take()
446                        .expect("pending flow transform present");
447                    if slot.transform_taken && slot.mat_taken {
448                        slots.remove(index);
449                    }
450                    if slots.is_empty() {
451                        pending.remove(&thread_id);
452                    }
453                    transform
454                };
455                transform(input)
456            }
457        };
458
459        let materialize = {
460            let pending = Arc::clone(&pending);
461            let factory = Arc::clone(&factory);
462            move || {
463                let mat = {
464                    let mut pending = pending.lock().expect("flow materialized factory poisoned");
465                    let thread_id = thread::current().id();
466                    let slots = pending.entry(thread_id).or_default();
467                    let index = slots
468                        .iter()
469                        .position(|slot| !slot.mat_taken && slot.transform_taken)
470                        .unwrap_or_else(|| {
471                            let (transform, mat) = factory();
472                            slots.push(PendingSlot {
473                                transform: Some(transform),
474                                mat: Some(mat),
475                                transform_taken: false,
476                                mat_taken: false,
477                            });
478                            slots.len() - 1
479                        });
480                    let slot = slots
481                        .get_mut(index)
482                        .expect("pending flow materialization slot exists");
483                    slot.mat_taken = true;
484                    let mat = slot
485                        .mat
486                        .take()
487                        .expect("pending flow materialized value present");
488                    if slot.transform_taken && slot.mat_taken {
489                        slots.remove(index);
490                    }
491                    if slots.is_empty() {
492                        pending.remove(&thread_id);
493                    }
494                    mat
495                };
496                Ok(mat)
497            }
498        };
499
500        Self::from_parts_with_hints(transform, materialize, FlowHints::default())
501    }
502
503    pub(crate) fn from_runtime_materialized_factory<F>(factory: F) -> Self
504    where
505        F: Fn() -> (RuntimeTransform<In, Out>, Mat) + Send + Sync + 'static,
506    {
507        struct PendingSlot<In, Out, Mat> {
508            transform: Option<RuntimeTransform<In, Out>>,
509            mat: Option<Mat>,
510            transform_taken: bool,
511            mat_taken: bool,
512        }
513
514        let pending = Arc::new(Mutex::new(HashMap::<
515            thread::ThreadId,
516            Vec<PendingSlot<In, Out, Mat>>,
517        >::new()));
518        let factory = Arc::new(factory);
519
520        let transform = {
521            let pending = Arc::clone(&pending);
522            let factory = Arc::clone(&factory);
523            move |input, materializer: &Materializer| {
524                let thread_id = thread::current().id();
525                let transform = {
526                    let mut pending = pending.lock().expect("flow materialized factory poisoned");
527                    let slots = pending.entry(thread_id).or_default();
528                    if let Some(index) = slots
529                        .iter()
530                        .position(|slot| !slot.transform_taken && slot.mat_taken)
531                    {
532                        let slot = slots
533                            .get_mut(index)
534                            .expect("pending flow materialization slot exists");
535                        slot.transform_taken = true;
536                        let transform = slot
537                            .transform
538                            .take()
539                            .expect("pending flow transform present");
540                        if slot.transform_taken && slot.mat_taken {
541                            slots.remove(index);
542                        }
543                        if slots.is_empty() {
544                            pending.remove(&thread_id);
545                        }
546                        Some(transform)
547                    } else {
548                        None
549                    }
550                };
551
552                let transform = match transform {
553                    Some(transform) => transform,
554                    None => {
555                        let (transform, mat) = factory();
556                        let mut pending =
557                            pending.lock().expect("flow materialized factory poisoned");
558                        let slots = pending.entry(thread_id).or_default();
559                        slots.push(PendingSlot {
560                            transform: None,
561                            mat: Some(mat),
562                            transform_taken: true,
563                            mat_taken: false,
564                        });
565                        transform
566                    }
567                };
568
569                transform(input, materializer)
570            }
571        };
572
573        let materialize = {
574            let pending = Arc::clone(&pending);
575            let factory = Arc::clone(&factory);
576            move || {
577                let thread_id = thread::current().id();
578                let mat = {
579                    let mut pending = pending.lock().expect("flow materialized factory poisoned");
580                    let slots = pending.entry(thread_id).or_default();
581                    if let Some(index) = slots
582                        .iter()
583                        .position(|slot| !slot.mat_taken && slot.transform_taken)
584                    {
585                        let slot = slots
586                            .get_mut(index)
587                            .expect("pending flow materialization slot exists");
588                        slot.mat_taken = true;
589                        let mat = slot
590                            .mat
591                            .take()
592                            .expect("pending flow materialized value present");
593                        if slot.transform_taken && slot.mat_taken {
594                            slots.remove(index);
595                        }
596                        if slots.is_empty() {
597                            pending.remove(&thread_id);
598                        }
599                        Some(mat)
600                    } else {
601                        None
602                    }
603                };
604
605                match mat {
606                    Some(mat) => Ok(mat),
607                    None => {
608                        let (transform, mat) = factory();
609                        let mut pending =
610                            pending.lock().expect("flow materialized factory poisoned");
611                        let slots = pending.entry(thread_id).or_default();
612                        slots.push(PendingSlot {
613                            transform: Some(transform),
614                            mat: None,
615                            transform_taken: false,
616                            mat_taken: true,
617                        });
618                        Ok(mat)
619                    }
620                }
621            }
622        };
623
624        Flow {
625            transform: FlowTransform::Runtime(Arc::new(transform)),
626            materialize: Arc::new(materialize),
627            hints: FlowHints::default(),
628            attributes: Attributes::default(),
629        }
630    }
631
632    fn from_parts_with_hints<F, M>(transform: F, materialize: M, hints: FlowHints) -> Self
633    where
634        F: Fn(BoxStream<In>) -> BoxStream<Out> + Send + Sync + 'static,
635        M: Fn() -> StreamResult<Mat> + Send + Sync + 'static,
636    {
637        Self {
638            transform: FlowTransform::Pure(Arc::new(transform)),
639            materialize: Arc::new(materialize),
640            hints,
641            attributes: Attributes::default(),
642        }
643    }
644
645    #[must_use]
646    pub fn attributes(&self) -> &Attributes {
647        &self.attributes
648    }
649
650    #[must_use]
651    pub fn with_attributes(mut self, attributes: Attributes) -> Self {
652        self.attributes = attributes;
653        self
654    }
655
656    #[must_use]
657    pub fn add_attributes(mut self, attributes: Attributes) -> Self {
658        self.attributes = self.attributes.and(attributes);
659        self
660    }
661
662    #[must_use]
663    pub fn named(self, name: impl Into<String>) -> Self {
664        self.add_attributes(Attributes::named(name))
665    }
666
667    /// Insert an async boundary after this flow.
668    ///
669    /// The boundary uses the same [`crate::AsyncBoundaryExecutionConfig`] defaults
670    /// as the GraphDSL `AsyncBoundary` runner and hands elements across a bounded
671    /// Ractor-backed queue. `async_boundary` is the Rust-friendly primary name;
672    /// `Flow::r#async` is provided as the Akka-mirroring alias.
673    #[must_use]
674    pub fn async_boundary(self) -> Self {
675        self.async_boundary_with_config(crate::graph::AsyncBoundaryExecutionConfig::default())
676    }
677
678    /// Akka-mirroring alias for [`Flow::async_boundary`].
679    #[must_use]
680    pub fn r#async(self) -> Self {
681        self.async_boundary()
682    }
683
684    /// Insert an async boundary with an explicit bounded handoff configuration.
685    ///
686    /// `config.buffer_size` controls the number of elements that may be queued
687    /// between the upstream and downstream fused regions. A zero buffer is
688    /// rejected when the stream is materialized.
689    #[must_use]
690    pub fn async_boundary_with_config(
691        self,
692        config: crate::graph::AsyncBoundaryExecutionConfig,
693    ) -> Self {
694        self.via(Flow::from_runtime_transform(move |input, materializer| {
695            super::async_boundary::linear_async_boundary_stream(input, materializer, config)
696        }))
697    }
698
699    /// Insert an async boundary with a custom bounded handoff size.
700    #[must_use]
701    pub fn async_boundary_with_buffer(self, buffer_size: usize) -> Self {
702        self.async_boundary_with_config(crate::graph::AsyncBoundaryExecutionConfig {
703            buffer_size,
704            ..crate::graph::AsyncBoundaryExecutionConfig::default()
705        })
706    }
707
708    #[must_use]
709    pub fn via<Next, NextMat>(self, next: Flow<Out, Next, NextMat>) -> Flow<In, Next, Mat>
710    where
711        Next: Send + 'static,
712        NextMat: Send + 'static,
713    {
714        self.via_mat(next, Keep::left)
715    }
716
717    #[must_use]
718    pub fn via_mat<Next, NextMat, Combined, F>(
719        self,
720        next: Flow<Out, Next, NextMat>,
721        combine: F,
722    ) -> Flow<In, Next, Combined>
723    where
724        Next: Send + 'static,
725        NextMat: Send + 'static,
726        Combined: Send + 'static,
727        F: Fn(Mat, NextMat) -> Combined + Send + Sync + 'static,
728    {
729        let Flow {
730            transform: first,
731            materialize: materialize_first,
732            hints: first_hints,
733            attributes: first_attributes,
734        } = self;
735        let Flow {
736            transform: second,
737            materialize: materialize_second,
738            hints: second_hints,
739            attributes: second_attributes,
740        } = next;
741        let combine = Arc::new(combine);
742        match (first, second) {
743            (FlowTransform::Pure(first), FlowTransform::Pure(second)) => {
744                let hints = first_hints.then(second_hints);
745                Flow::from_parts_with_hints(
746                    move |input| second(first(input)),
747                    move || {
748                        let left = materialize_first()?;
749                        let right = materialize_second()?;
750                        Ok(combine(left, right))
751                    },
752                    hints,
753                )
754                .with_attributes(first_attributes.and(second_attributes))
755            }
756            (first, second) => {
757                let hints = first_hints.then(second_hints);
758                Flow {
759                    transform: FlowTransform::Runtime(Arc::new(move |input, materializer| {
760                        let stream = match &first {
761                            FlowTransform::Pure(first) => first(input),
762                            FlowTransform::Runtime(first) => first(input, materializer)?,
763                        };
764                        match &second {
765                            FlowTransform::Pure(second) => Ok(second(stream)),
766                            FlowTransform::Runtime(second) => second(stream, materializer),
767                        }
768                    })),
769                    materialize: Arc::new(move || {
770                        let left = materialize_first()?;
771                        let right = materialize_second()?;
772                        Ok(combine(left, right))
773                    }),
774                    hints,
775                    attributes: first_attributes.and(second_attributes),
776                }
777            }
778        }
779    }
780
781    #[must_use]
782    pub fn via_mat_with<Next, NextMat, Combined, F>(
783        self,
784        next: Flow<Out, Next, NextMat>,
785        combine: F,
786    ) -> Flow<In, Next, Combined>
787    where
788        Next: Send + 'static,
789        NextMat: Send + 'static,
790        Combined: Send + 'static,
791        F: Fn(Mat, NextMat) -> Combined + Send + Sync + 'static,
792    {
793        self.via_mat(next, combine)
794    }
795
796    #[must_use]
797    pub fn map<Next, F>(self, f: F) -> Flow<In, Next, Mat>
798    where
799        Next: Send + 'static,
800        F: Fn(Out) -> Next + Send + Sync + 'static,
801    {
802        let stage = Arc::new(f);
803        match &self.transform {
804            FlowTransform::Pure(_) => {
805                let Flow {
806                    transform,
807                    materialize,
808                    hints,
809                    attributes,
810                } = self;
811                let FlowTransform::Pure(transform) = transform else {
812                    unreachable!("pure transform checked above");
813                };
814                Flow::from_parts_with_hints(
815                    move |input| {
816                        let stage = Arc::clone(&stage);
817                        Box::new(transform(input).map(move |item| item.map(|item| stage(item))))
818                    },
819                    move || materialize(),
820                    hints,
821                )
822                .with_attributes(attributes)
823            }
824            FlowTransform::Runtime(_) => self.via(Flow::from_transform(move |input| {
825                let stage = Arc::clone(&stage);
826                Box::new(input.map(move |item| item.map(|item| stage(item))))
827            })),
828        }
829    }
830
831    /// Fallible `map`. An error stops the stream, matching Datum's existing
832    /// default supervision behavior.
833    #[must_use]
834    pub fn map_result<Next, F>(self, f: F) -> Flow<In, Next, Mat>
835    where
836        Next: Send + 'static,
837        F: Fn(Out) -> StreamResult<Next> + Send + Sync + 'static,
838    {
839        let stage = Arc::new(f);
840        self.via(Flow::from_transform(move |input| {
841            let stage = Arc::clone(&stage);
842            Box::new(input.map(move |item| item.and_then(|item| stage(item))))
843        }))
844    }
845
846    /// Fallible `map` with an Akka-style supervision decider.
847    ///
848    /// `Resume` drops the failing element. `Restart` is equivalent to `Resume`
849    /// for stateless `map`.
850    #[must_use]
851    pub fn map_result_with_supervision<Next, F>(
852        self,
853        f: F,
854        decider: SupervisionDecider,
855    ) -> Flow<In, Next, Mat>
856    where
857        Next: Send + 'static,
858        F: Fn(Out) -> StreamResult<Next> + Send + Sync + 'static,
859    {
860        let stage = Arc::new(f);
861        self.via(Flow::from_transform(move |input| {
862            let stage = Arc::clone(&stage);
863            let decider = Arc::clone(&decider);
864            Box::new(input.filter_map(move |item| match item {
865                Ok(item) => match call_supervised("map_result callback", || stage(item)) {
866                    Ok(next) => Some(Ok(next)),
867                    Err(error) => match decide_supervision(&decider, &error) {
868                        SupervisionDirective::Stop => Some(Err(error)),
869                        SupervisionDirective::Resume | SupervisionDirective::Restart => None,
870                    },
871                },
872                Err(error) => Some(Err(error)),
873            }))
874        }))
875    }
876
877    #[must_use]
878    pub fn filter<F>(self, predicate: F) -> Flow<In, Out, Mat>
879    where
880        F: Fn(&Out) -> bool + Send + Sync + 'static,
881    {
882        let predicate = Arc::new(predicate);
883        self.via(Flow::from_preserving_transform(move |input| {
884            let predicate = Arc::clone(&predicate);
885            Box::new(input.filter_map(move |item| match item {
886                Ok(item) if predicate(&item) => Some(Ok(item)),
887                Ok(_) => None,
888                Err(error) => Some(Err(error)),
889            }))
890        }))
891    }
892
893    #[must_use]
894    pub fn filter_result<F>(self, predicate: F) -> Flow<In, Out, Mat>
895    where
896        F: Fn(&Out) -> StreamResult<bool> + Send + Sync + 'static,
897    {
898        let predicate = Arc::new(predicate);
899        self.via(Flow::from_transform(move |input| {
900            let predicate = Arc::clone(&predicate);
901            Box::new(input.filter_map(move |item| match item {
902                Ok(item) => match predicate(&item) {
903                    Ok(true) => Some(Ok(item)),
904                    Ok(false) => None,
905                    Err(error) => Some(Err(error)),
906                },
907                Err(error) => Some(Err(error)),
908            }))
909        }))
910    }
911
912    #[must_use]
913    pub fn filter_result_with_supervision<F>(
914        self,
915        predicate: F,
916        decider: SupervisionDecider,
917    ) -> Flow<In, Out, Mat>
918    where
919        F: Fn(&Out) -> StreamResult<bool> + Send + Sync + 'static,
920    {
921        let predicate = Arc::new(predicate);
922        self.via(Flow::from_transform(move |input| {
923            let predicate = Arc::clone(&predicate);
924            let decider = Arc::clone(&decider);
925            Box::new(input.filter_map(move |item| match item {
926                Ok(item) => match call_supervised("filter_result callback", || predicate(&item)) {
927                    Ok(true) => Some(Ok(item)),
928                    Ok(false) => None,
929                    Err(error) => match decide_supervision(&decider, &error) {
930                        SupervisionDirective::Stop => Some(Err(error)),
931                        SupervisionDirective::Resume | SupervisionDirective::Restart => None,
932                    },
933                },
934                Err(error) => Some(Err(error)),
935            }))
936        }))
937    }
938
939    #[must_use]
940    pub fn filter_not<F>(self, predicate: F) -> Flow<In, Out, Mat>
941    where
942        F: Fn(&Out) -> bool + Send + Sync + 'static,
943    {
944        let predicate = Arc::new(predicate);
945        self.filter(move |item| !predicate(item))
946    }
947
948    #[must_use]
949    pub fn filter_map<Next, F>(self, f: F) -> Flow<In, Next, Mat>
950    where
951        Next: Send + 'static,
952        F: Fn(Out) -> Option<Next> + Send + Sync + 'static,
953    {
954        let stage = Arc::new(f);
955        self.via(Flow::from_transform(move |input| {
956            let stage = Arc::clone(&stage);
957            Box::new(input.filter_map(move |item| match item {
958                Ok(item) => stage(item).map(Ok),
959                Err(error) => Some(Err(error)),
960            }))
961        }))
962    }
963
964    #[must_use]
965    pub fn filter_map_result<Next, F>(self, f: F) -> Flow<In, Next, Mat>
966    where
967        Next: Send + 'static,
968        F: Fn(Out) -> StreamResult<Option<Next>> + Send + Sync + 'static,
969    {
970        let stage = Arc::new(f);
971        self.via(Flow::from_transform(move |input| {
972            let stage = Arc::clone(&stage);
973            Box::new(input.filter_map(move |item| match item {
974                Ok(item) => match stage(item) {
975                    Ok(Some(next)) => Some(Ok(next)),
976                    Ok(None) => None,
977                    Err(error) => Some(Err(error)),
978                },
979                Err(error) => Some(Err(error)),
980            }))
981        }))
982    }
983
984    #[must_use]
985    pub fn filter_map_result_with_supervision<Next, F>(
986        self,
987        f: F,
988        decider: SupervisionDecider,
989    ) -> Flow<In, Next, Mat>
990    where
991        Next: Send + 'static,
992        F: Fn(Out) -> StreamResult<Option<Next>> + Send + Sync + 'static,
993    {
994        let stage = Arc::new(f);
995        self.via(Flow::from_transform(move |input| {
996            let stage = Arc::clone(&stage);
997            let decider = Arc::clone(&decider);
998            Box::new(input.filter_map(move |item| match item {
999                Ok(item) => match call_supervised("filter_map_result callback", || stage(item)) {
1000                    Ok(Some(next)) => Some(Ok(next)),
1001                    Ok(None) => None,
1002                    Err(error) => match decide_supervision(&decider, &error) {
1003                        SupervisionDirective::Stop => Some(Err(error)),
1004                        SupervisionDirective::Resume | SupervisionDirective::Restart => None,
1005                    },
1006                },
1007                Err(error) => Some(Err(error)),
1008            }))
1009        }))
1010    }
1011
1012    #[must_use]
1013    pub fn map_concat<Next, F, I>(self, f: F) -> Flow<In, Next, Mat>
1014    where
1015        Next: Send + 'static,
1016        F: Fn(Out) -> I + Send + Sync + 'static,
1017        I: IntoIterator<Item = Next>,
1018        I::IntoIter: Send + 'static,
1019    {
1020        let stage = Arc::new(f);
1021        self.via(Flow::from_transform(move |mut input| {
1022            let stage = Arc::clone(&stage);
1023            let mut current = None::<I::IntoIter>;
1024            let mut done = false;
1025            Box::new(std::iter::from_fn(move || {
1026                loop {
1027                    if let Some(iter) = &mut current {
1028                        if let Some(item) = iter.next() {
1029                            return Some(Ok(item));
1030                        }
1031                        current = None;
1032                    }
1033
1034                    if done {
1035                        return None;
1036                    }
1037
1038                    match input.next()? {
1039                        Ok(item) => current = Some(stage(item).into_iter()),
1040                        Err(error) => {
1041                            done = true;
1042                            return Some(Err(error));
1043                        }
1044                    }
1045                }
1046            }))
1047        }))
1048    }
1049
1050    #[must_use]
1051    pub fn map_concat_result<Next, F, I>(self, f: F) -> Flow<In, Next, Mat>
1052    where
1053        Next: Send + 'static,
1054        F: Fn(Out) -> StreamResult<I> + Send + Sync + 'static,
1055        I: IntoIterator<Item = Next>,
1056        I::IntoIter: Send + 'static,
1057    {
1058        let stage = Arc::new(f);
1059        self.via(Flow::from_transform(move |mut input| {
1060            let stage = Arc::clone(&stage);
1061            let mut current = None::<I::IntoIter>;
1062            let mut done = false;
1063            Box::new(std::iter::from_fn(move || {
1064                loop {
1065                    if let Some(iter) = &mut current {
1066                        if let Some(item) = iter.next() {
1067                            return Some(Ok(item));
1068                        }
1069                        current = None;
1070                    }
1071
1072                    if done {
1073                        return None;
1074                    }
1075
1076                    match input.next()? {
1077                        Ok(item) => match stage(item) {
1078                            Ok(items) => current = Some(items.into_iter()),
1079                            Err(error) => {
1080                                done = true;
1081                                return Some(Err(error));
1082                            }
1083                        },
1084                        Err(error) => {
1085                            done = true;
1086                            return Some(Err(error));
1087                        }
1088                    }
1089                }
1090            }))
1091        }))
1092    }
1093
1094    #[must_use]
1095    pub fn map_concat_result_with_supervision<Next, F, I>(
1096        self,
1097        f: F,
1098        decider: SupervisionDecider,
1099    ) -> Flow<In, Next, Mat>
1100    where
1101        Next: Send + 'static,
1102        F: Fn(Out) -> StreamResult<I> + Send + Sync + 'static,
1103        I: IntoIterator<Item = Next>,
1104        I::IntoIter: Send + 'static,
1105    {
1106        let stage = Arc::new(f);
1107        self.via(Flow::from_transform(move |mut input| {
1108            let stage = Arc::clone(&stage);
1109            let decider = Arc::clone(&decider);
1110            let mut current = None::<I::IntoIter>;
1111            let mut done = false;
1112            Box::new(std::iter::from_fn(move || {
1113                loop {
1114                    if let Some(iter) = &mut current {
1115                        if let Some(item) = iter.next() {
1116                            return Some(Ok(item));
1117                        }
1118                        current = None;
1119                    }
1120
1121                    if done {
1122                        return None;
1123                    }
1124
1125                    match input.next()? {
1126                        Ok(item) => {
1127                            match call_supervised("map_concat_result callback", || stage(item)) {
1128                                Ok(items) => current = Some(items.into_iter()),
1129                                Err(error) => match decide_supervision(&decider, &error) {
1130                                    SupervisionDirective::Stop => {
1131                                        done = true;
1132                                        return Some(Err(error));
1133                                    }
1134                                    SupervisionDirective::Resume
1135                                    | SupervisionDirective::Restart => {}
1136                                },
1137                            }
1138                        }
1139                        Err(error) => {
1140                            done = true;
1141                            return Some(Err(error));
1142                        }
1143                    }
1144                }
1145            }))
1146        }))
1147    }
1148
1149    #[must_use]
1150    pub fn stateful_map<State, Next, F>(self, seed: State, f: F) -> Flow<In, Next, Mat>
1151    where
1152        State: Clone + Send + Sync + 'static,
1153        Next: Send + 'static,
1154        F: Fn(&mut State, Out) -> Next + Send + Sync + 'static,
1155    {
1156        let stage = Arc::new(f);
1157        self.via(Flow::from_transform(move |input| {
1158            let stage = Arc::clone(&stage);
1159            let mut state = seed.clone();
1160            Box::new(input.map(move |item| item.map(|item| stage(&mut state, item))))
1161        }))
1162    }
1163
1164    #[must_use]
1165    pub fn stateful_map_result<State, Next, F>(self, seed: State, f: F) -> Flow<In, Next, Mat>
1166    where
1167        State: Clone + Send + Sync + 'static,
1168        Next: Send + 'static,
1169        F: Fn(&mut State, Out) -> StreamResult<Next> + Send + Sync + 'static,
1170    {
1171        let stage = Arc::new(f);
1172        self.via(Flow::from_transform(move |input| {
1173            let stage = Arc::clone(&stage);
1174            let mut state = seed.clone();
1175            Box::new(input.map(move |item| match item {
1176                Ok(item) => stage(&mut state, item),
1177                Err(error) => Err(error),
1178            }))
1179        }))
1180    }
1181
1182    #[must_use]
1183    pub fn stateful_map_result_with_supervision<State, Next, F>(
1184        self,
1185        seed: State,
1186        f: F,
1187        decider: SupervisionDecider,
1188    ) -> Flow<In, Next, Mat>
1189    where
1190        State: Clone + Send + Sync + 'static,
1191        Next: Send + 'static,
1192        F: Fn(&mut State, Out) -> StreamResult<Next> + Send + Sync + 'static,
1193    {
1194        let stage = Arc::new(f);
1195        self.via(Flow::from_transform(move |input| {
1196            let stage = Arc::clone(&stage);
1197            let decider = Arc::clone(&decider);
1198            let seed = seed.clone();
1199            let mut state = seed.clone();
1200            Box::new(input.filter_map(move |item| match item {
1201                Ok(item) => match call_supervised("stateful_map_result callback", || {
1202                    stage(&mut state, item)
1203                }) {
1204                    Ok(next) => Some(Ok(next)),
1205                    Err(error) => match decide_supervision(&decider, &error) {
1206                        SupervisionDirective::Stop => Some(Err(error)),
1207                        SupervisionDirective::Resume => None,
1208                        SupervisionDirective::Restart => {
1209                            state = seed.clone();
1210                            None
1211                        }
1212                    },
1213                },
1214                Err(error) => Some(Err(error)),
1215            }))
1216        }))
1217    }
1218
1219    #[must_use]
1220    pub fn stateful_map_concat<State, Next, F, I>(self, seed: State, f: F) -> Flow<In, Next, Mat>
1221    where
1222        State: Clone + Send + Sync + 'static,
1223        Next: Send + 'static,
1224        F: Fn(&mut State, Out) -> I + Send + Sync + 'static,
1225        I: IntoIterator<Item = Next>,
1226        I::IntoIter: Send + 'static,
1227    {
1228        let stage = Arc::new(f);
1229        self.via(Flow::from_transform(move |mut input| {
1230            let stage = Arc::clone(&stage);
1231            let mut state = seed.clone();
1232            let mut current = None::<I::IntoIter>;
1233            let mut done = false;
1234            Box::new(std::iter::from_fn(move || {
1235                loop {
1236                    if let Some(iter) = &mut current {
1237                        if let Some(item) = iter.next() {
1238                            return Some(Ok(item));
1239                        }
1240                        current = None;
1241                    }
1242
1243                    if done {
1244                        return None;
1245                    }
1246
1247                    match input.next()? {
1248                        Ok(item) => current = Some(stage(&mut state, item).into_iter()),
1249                        Err(error) => {
1250                            done = true;
1251                            return Some(Err(error));
1252                        }
1253                    }
1254                }
1255            }))
1256        }))
1257    }
1258
1259    #[must_use]
1260    pub fn stateful_map_concat_result<State, Next, F, I>(
1261        self,
1262        seed: State,
1263        f: F,
1264    ) -> Flow<In, Next, Mat>
1265    where
1266        State: Clone + Send + Sync + 'static,
1267        Next: Send + 'static,
1268        F: Fn(&mut State, Out) -> StreamResult<I> + Send + Sync + 'static,
1269        I: IntoIterator<Item = Next>,
1270        I::IntoIter: Send + 'static,
1271    {
1272        let stage = Arc::new(f);
1273        self.via(Flow::from_transform(move |mut input| {
1274            let stage = Arc::clone(&stage);
1275            let mut state = seed.clone();
1276            let mut current = None::<I::IntoIter>;
1277            let mut done = false;
1278            Box::new(std::iter::from_fn(move || {
1279                loop {
1280                    if let Some(iter) = &mut current {
1281                        if let Some(item) = iter.next() {
1282                            return Some(Ok(item));
1283                        }
1284                        current = None;
1285                    }
1286
1287                    if done {
1288                        return None;
1289                    }
1290
1291                    match input.next()? {
1292                        Ok(item) => match stage(&mut state, item) {
1293                            Ok(items) => current = Some(items.into_iter()),
1294                            Err(error) => {
1295                                done = true;
1296                                return Some(Err(error));
1297                            }
1298                        },
1299                        Err(error) => {
1300                            done = true;
1301                            return Some(Err(error));
1302                        }
1303                    }
1304                }
1305            }))
1306        }))
1307    }
1308
1309    #[must_use]
1310    pub fn stateful_map_concat_result_with_supervision<State, Next, F, I>(
1311        self,
1312        seed: State,
1313        f: F,
1314        decider: SupervisionDecider,
1315    ) -> Flow<In, Next, Mat>
1316    where
1317        State: Clone + Send + Sync + 'static,
1318        Next: Send + 'static,
1319        F: Fn(&mut State, Out) -> StreamResult<I> + Send + Sync + 'static,
1320        I: IntoIterator<Item = Next>,
1321        I::IntoIter: Send + 'static,
1322    {
1323        let stage = Arc::new(f);
1324        self.via(Flow::from_transform(move |mut input| {
1325            let stage = Arc::clone(&stage);
1326            let decider = Arc::clone(&decider);
1327            let seed = seed.clone();
1328            let mut state = seed.clone();
1329            let mut current = None::<I::IntoIter>;
1330            let mut done = false;
1331            Box::new(std::iter::from_fn(move || {
1332                loop {
1333                    if let Some(iter) = &mut current {
1334                        if let Some(item) = iter.next() {
1335                            return Some(Ok(item));
1336                        }
1337                        current = None;
1338                    }
1339
1340                    if done {
1341                        return None;
1342                    }
1343
1344                    match input.next()? {
1345                        Ok(item) => {
1346                            match call_supervised("stateful_map_concat_result callback", || {
1347                                stage(&mut state, item)
1348                            }) {
1349                                Ok(items) => current = Some(items.into_iter()),
1350                                Err(error) => match decide_supervision(&decider, &error) {
1351                                    SupervisionDirective::Stop => {
1352                                        done = true;
1353                                        return Some(Err(error));
1354                                    }
1355                                    SupervisionDirective::Resume => {}
1356                                    SupervisionDirective::Restart => state = seed.clone(),
1357                                },
1358                            }
1359                        }
1360                        Err(error) => {
1361                            done = true;
1362                            return Some(Err(error));
1363                        }
1364                    }
1365                }
1366            }))
1367        }))
1368    }
1369
1370    #[must_use]
1371    /// Polls each future once on the drain thread and moves only pending
1372    /// futures onto the Tokio runtime. Contract-conforming futures behave
1373    /// identically across that handoff; futures must not block inside `poll`.
1374    pub fn map_async<Next, F, Fut>(self, parallelism: usize, f: F) -> Flow<In, Next, Mat>
1375    where
1376        Next: Send + 'static,
1377        F: Fn(Out) -> Fut + Send + Sync + 'static,
1378        Fut: Future<Output = StreamResult<Next>> + Send + 'static,
1379    {
1380        assert!(
1381            parallelism > 0,
1382            "map_async parallelism must be greater than zero"
1383        );
1384        let stage = Arc::new(f);
1385        self.via(Flow::from_runtime_transform_with_hints(
1386            move |input, _materializer| {
1387                let stage = Arc::clone(&stage);
1388                Ok(map_async_ordered(input, parallelism, stage))
1389            },
1390            FlowHints::PRESERVES_TERMINAL_CONSUMER_BATCH,
1391        ))
1392    }
1393
1394    #[must_use]
1395    /// Fallible `map_async` with an Akka-style supervision decider.
1396    ///
1397    /// `Resume` and `Restart` drop the failed future result and keep the
1398    /// ordered output sequence moving. Upstream errors are not supervised.
1399    pub fn map_async_with_supervision<Next, F, Fut>(
1400        self,
1401        parallelism: usize,
1402        f: F,
1403        decider: SupervisionDecider,
1404    ) -> Flow<In, Next, Mat>
1405    where
1406        Next: Send + 'static,
1407        F: Fn(Out) -> Fut + Send + Sync + 'static,
1408        Fut: Future<Output = StreamResult<Next>> + Send + 'static,
1409    {
1410        assert!(
1411            parallelism > 0,
1412            "map_async parallelism must be greater than zero"
1413        );
1414        let stage = Arc::new(f);
1415        self.via(Flow::from_runtime_transform(move |input, _materializer| {
1416            let stage = Arc::clone(&stage);
1417            let decider = Arc::clone(&decider);
1418            Ok(map_async_ordered_supervised(
1419                input,
1420                parallelism,
1421                stage,
1422                decider,
1423            ))
1424        }))
1425    }
1426
1427    #[must_use]
1428    /// Polls each future once on the drain thread and moves only pending
1429    /// futures onto the Tokio runtime. Contract-conforming futures behave
1430    /// identically across that handoff; futures must not block inside `poll`.
1431    pub fn map_async_unordered<Next, F, Fut>(self, parallelism: usize, f: F) -> Flow<In, Next, Mat>
1432    where
1433        Next: Send + 'static,
1434        F: Fn(Out) -> Fut + Send + Sync + 'static,
1435        Fut: Future<Output = StreamResult<Next>> + Send + 'static,
1436    {
1437        assert!(
1438            parallelism > 0,
1439            "map_async_unordered parallelism must be greater than zero"
1440        );
1441        let stage = Arc::new(f);
1442        self.via(Flow::from_runtime_transform_with_hints(
1443            move |input, _materializer| {
1444                let stage = Arc::clone(&stage);
1445                Ok(map_async_unordered(input, parallelism, stage))
1446            },
1447            FlowHints::PRESERVES_TERMINAL_CONSUMER_BATCH,
1448        ))
1449    }
1450
1451    #[must_use]
1452    pub fn map_async_unordered_with_supervision<Next, F, Fut>(
1453        self,
1454        parallelism: usize,
1455        f: F,
1456        decider: SupervisionDecider,
1457    ) -> Flow<In, Next, Mat>
1458    where
1459        Next: Send + 'static,
1460        F: Fn(Out) -> Fut + Send + Sync + 'static,
1461        Fut: Future<Output = StreamResult<Next>> + Send + 'static,
1462    {
1463        assert!(
1464            parallelism > 0,
1465            "map_async_unordered parallelism must be greater than zero"
1466        );
1467        let stage = Arc::new(f);
1468        self.via(Flow::from_runtime_transform(move |input, _materializer| {
1469            let stage = Arc::clone(&stage);
1470            let decider = Arc::clone(&decider);
1471            Ok(map_async_unordered_supervised(
1472                input,
1473                parallelism,
1474                stage,
1475                decider,
1476            ))
1477        }))
1478    }
1479
1480    #[must_use]
1481    /// Polls each future once on the drain thread and moves only pending
1482    /// futures onto the Tokio runtime. Contract-conforming futures behave
1483    /// identically across that handoff; futures must not block inside `poll`.
1484    pub fn map_async_partitioned<Key, Next, Partition, F, Fut>(
1485        self,
1486        parallelism: usize,
1487        per_partition: usize,
1488        partition: Partition,
1489        f: F,
1490    ) -> Flow<In, Next, Mat>
1491    where
1492        Key: Clone + Eq + Hash + Send + 'static,
1493        Next: Send + 'static,
1494        Partition: Fn(&Out) -> Key + Send + Sync + 'static,
1495        F: Fn(Out) -> Fut + Send + Sync + 'static,
1496        Fut: Future<Output = StreamResult<Next>> + Send + 'static,
1497    {
1498        assert!(
1499            parallelism > 0,
1500            "map_async_partitioned parallelism must be greater than zero"
1501        );
1502        assert!(
1503            per_partition > 0,
1504            "map_async_partitioned per_partition must be greater than zero"
1505        );
1506        let partition = Arc::new(partition);
1507        let stage = Arc::new(f);
1508        self.via(Flow::from_runtime_transform_with_hints(
1509            move |input, _materializer| {
1510                let partition = Arc::clone(&partition);
1511                let stage = Arc::clone(&stage);
1512                Ok(map_async_partitioned(
1513                    input,
1514                    parallelism,
1515                    per_partition,
1516                    partition,
1517                    stage,
1518                ))
1519            },
1520            FlowHints::PRESERVES_TERMINAL_CONSUMER_BATCH,
1521        ))
1522    }
1523
1524    #[must_use]
1525    pub fn take(self, n: usize) -> Flow<In, Out, Mat> {
1526        self.via(Flow::from_transform(move |input| Box::new(input.take(n))))
1527    }
1528
1529    #[must_use]
1530    pub fn drop(self, n: usize) -> Flow<In, Out, Mat> {
1531        self.via(Flow::from_transform(move |input| {
1532            let mut remaining = n;
1533            Box::new(input.filter_map(move |item| match item {
1534                Ok(_) if remaining > 0 => {
1535                    remaining -= 1;
1536                    None
1537                }
1538                other => Some(other),
1539            }))
1540        }))
1541    }
1542
1543    #[must_use]
1544    pub fn take_while<F>(self, predicate: F) -> Flow<In, Out, Mat>
1545    where
1546        F: Fn(&Out) -> bool + Send + Sync + 'static,
1547    {
1548        let predicate = Arc::new(predicate);
1549        self.via(Flow::from_transform(move |mut input| {
1550            let predicate = Arc::clone(&predicate);
1551            let mut open = true;
1552            Box::new(std::iter::from_fn(move || {
1553                if !open {
1554                    return None;
1555                }
1556
1557                match input.next() {
1558                    Some(Ok(item)) if predicate(&item) => Some(Ok(item)),
1559                    Some(Ok(_)) | None => {
1560                        open = false;
1561                        None
1562                    }
1563                    Some(Err(error)) => Some(Err(error)),
1564                }
1565            }))
1566        }))
1567    }
1568
1569    #[must_use]
1570    pub fn drop_while<F>(self, predicate: F) -> Flow<In, Out, Mat>
1571    where
1572        F: Fn(&Out) -> bool + Send + Sync + 'static,
1573    {
1574        let predicate = Arc::new(predicate);
1575        self.via(Flow::from_transform(move |mut input| {
1576            let predicate = Arc::clone(&predicate);
1577            let mut dropping = true;
1578            Box::new(std::iter::from_fn(move || {
1579                loop {
1580                    let next = input.next()?;
1581                    match next {
1582                        Ok(item) if dropping && predicate(&item) => continue,
1583                        Ok(item) => {
1584                            dropping = false;
1585                            return Some(Ok(item));
1586                        }
1587                        Err(error) => return Some(Err(error)),
1588                    }
1589                }
1590            }))
1591        }))
1592    }
1593
1594    #[must_use]
1595    pub fn limit(self, max: u64) -> Flow<In, Out, Mat> {
1596        self.via(Flow::from_transform(move |input| {
1597            let mut seen = 0_u64;
1598            Box::new(input.map(move |item| match item {
1599                Ok(item) if seen < max => {
1600                    seen += 1;
1601                    Ok(item)
1602                }
1603                Ok(_) => Err(StreamError::LimitExceeded { max }),
1604                Err(error) => Err(error),
1605            }))
1606        }))
1607    }
1608
1609    #[must_use]
1610    pub fn grouped(self, size: usize) -> Flow<In, Vec<Out>, Mat> {
1611        assert!(size > 0, "grouped size must be greater than zero");
1612        self.via(Flow::from_transform(move |mut input| {
1613            Box::new(std::iter::from_fn(move || {
1614                let mut group = Vec::with_capacity(size);
1615                while group.len() < size {
1616                    match input.next() {
1617                        Some(Ok(item)) => group.push(item),
1618                        Some(Err(error)) => return Some(Err(error)),
1619                        None => break,
1620                    }
1621                }
1622
1623                if group.is_empty() {
1624                    None
1625                } else {
1626                    Some(Ok(group))
1627                }
1628            }))
1629        }))
1630    }
1631
1632    #[must_use]
1633    pub fn scan<State, F>(self, seed: State, f: F) -> Flow<In, State, Mat>
1634    where
1635        State: Clone + Send + Sync + 'static,
1636        F: Fn(State, Out) -> State + Send + Sync + 'static,
1637    {
1638        let stage = Arc::new(f);
1639        self.via(Flow::from_transform(move |mut input| {
1640            let stage = Arc::clone(&stage);
1641            let mut state = Some(seed.clone());
1642            let mut emit_seed = true;
1643            Box::new(std::iter::from_fn(move || {
1644                if emit_seed {
1645                    emit_seed = false;
1646                    return Some(Ok(state.as_ref().expect("scan state present").clone()));
1647                }
1648
1649                match input.next()? {
1650                    Ok(item) => {
1651                        // Move the previous state into the stage (one clone for
1652                        // the emitted value) instead of cloning it twice.
1653                        let prev = state.take().expect("scan state present");
1654                        let next = stage(prev, item);
1655                        state = Some(next.clone());
1656                        Some(Ok(next))
1657                    }
1658                    Err(error) => Some(Err(error)),
1659                }
1660            }))
1661        }))
1662    }
1663
1664    #[must_use]
1665    pub fn scan_async<State, F, Fut>(self, seed: State, f: F) -> Flow<In, State, Mat>
1666    where
1667        State: Clone + Send + Sync + 'static,
1668        F: Fn(State, Out) -> Fut + Send + Sync + 'static,
1669        Fut: Future<Output = StreamResult<State>> + Send + 'static,
1670    {
1671        let stage = Arc::new(f);
1672        self.via(Flow::from_transform(move |mut input| {
1673            let stage = Arc::clone(&stage);
1674            let mut state = Some(seed.clone());
1675            let mut emit_seed = true;
1676            let mut terminated = false;
1677            Box::new(std::iter::from_fn(move || {
1678                if terminated {
1679                    return None;
1680                }
1681                if emit_seed {
1682                    emit_seed = false;
1683                    return Some(Ok(state
1684                        .as_ref()
1685                        .expect("scan_async state present")
1686                        .clone()));
1687                }
1688
1689                match input.next()? {
1690                    Ok(item) => {
1691                        let prev = state.take().expect("scan_async state present");
1692                        match catch_unwind_failed("scan_async factory", || stage(prev, item))
1693                            .and_then(run_future_inline_or_spawn)
1694                        {
1695                            Ok(next) => {
1696                                state = Some(next.clone());
1697                                Some(Ok(next))
1698                            }
1699                            Err(error) => {
1700                                terminated = true;
1701                                Some(Err(error))
1702                            }
1703                        }
1704                    }
1705                    Err(error) => {
1706                        terminated = true;
1707                        Some(Err(error))
1708                    }
1709                }
1710            }))
1711        }))
1712    }
1713
1714    #[must_use]
1715    pub fn scan_result<State, F>(self, seed: State, f: F) -> Flow<In, State, Mat>
1716    where
1717        State: Clone + Send + Sync + 'static,
1718        F: Fn(State, Out) -> StreamResult<State> + Send + Sync + 'static,
1719    {
1720        let stage = Arc::new(f);
1721        self.via(Flow::from_transform(move |mut input| {
1722            let stage = Arc::clone(&stage);
1723            let mut state = Some(seed.clone());
1724            let mut emit_seed = true;
1725            Box::new(std::iter::from_fn(move || {
1726                if emit_seed {
1727                    emit_seed = false;
1728                    return Some(Ok(state.as_ref().expect("scan state present").clone()));
1729                }
1730
1731                match input.next()? {
1732                    Ok(item) => {
1733                        let prev = state.take().expect("scan state present");
1734                        match stage(prev, item) {
1735                            Ok(next) => {
1736                                state = Some(next.clone());
1737                                Some(Ok(next))
1738                            }
1739                            Err(error) => Some(Err(error)),
1740                        }
1741                    }
1742                    Err(error) => Some(Err(error)),
1743                }
1744            }))
1745        }))
1746    }
1747
1748    #[must_use]
1749    pub fn scan_result_with_supervision<State, F>(
1750        self,
1751        seed: State,
1752        f: F,
1753        decider: SupervisionDecider,
1754    ) -> Flow<In, State, Mat>
1755    where
1756        State: Clone + Send + Sync + 'static,
1757        F: Fn(State, Out) -> StreamResult<State> + Send + Sync + 'static,
1758    {
1759        let stage = Arc::new(f);
1760        self.via(Flow::from_transform(move |mut input| {
1761            let stage = Arc::clone(&stage);
1762            let decider = Arc::clone(&decider);
1763            let seed = seed.clone();
1764            let mut state = Some(seed.clone());
1765            let mut emit_seed = true;
1766            Box::new(std::iter::from_fn(move || {
1767                loop {
1768                    if emit_seed {
1769                        emit_seed = false;
1770                        return Some(Ok(state.as_ref().expect("scan state present").clone()));
1771                    }
1772
1773                    match input.next()? {
1774                        Ok(item) => {
1775                            let prev = state.take().expect("scan state present");
1776                            match call_supervised("scan_result callback", || {
1777                                stage(prev.clone(), item)
1778                            }) {
1779                                Ok(next) => {
1780                                    state = Some(next.clone());
1781                                    return Some(Ok(next));
1782                                }
1783                                Err(error) => match decide_supervision(&decider, &error) {
1784                                    SupervisionDirective::Stop => return Some(Err(error)),
1785                                    SupervisionDirective::Resume => {
1786                                        state = Some(prev);
1787                                    }
1788                                    SupervisionDirective::Restart => {
1789                                        state = Some(seed.clone());
1790                                        emit_seed = true;
1791                                    }
1792                                },
1793                            }
1794                        }
1795                        Err(error) => return Some(Err(error)),
1796                    }
1797                }
1798            }))
1799        }))
1800    }
1801
1802    #[must_use]
1803    pub fn sliding(self, size: usize, step: usize) -> Flow<In, Vec<Out>, Mat>
1804    where
1805        Out: Clone,
1806    {
1807        assert!(size > 0, "sliding size must be greater than zero");
1808        assert!(step > 0, "sliding step must be greater than zero");
1809        self.via(Flow::from_transform(move |mut input| {
1810            // Mirrors Akka's `Sliding` stage (akka-stream Ops.scala): a full
1811            // window is emitted once the buffer reaches `size`, the emitted
1812            // window is retained and `step` elements are dropped lazily on the
1813            // following pull, and a trailing partial window is emitted at
1814            // upstream finish only when `0 < len < size`.
1815            let mut buffer = VecDeque::with_capacity(size.max(step));
1816            let mut terminated = false;
1817
1818            Box::new(std::iter::from_fn(move || {
1819                if terminated {
1820                    return None;
1821                }
1822
1823                loop {
1824                    match input.next() {
1825                        Some(Ok(item)) => {
1826                            buffer.push_back(item);
1827                            if buffer.len() < size {
1828                                continue;
1829                            } else if buffer.len() == size {
1830                                return Some(Ok(buffer.iter().cloned().collect()));
1831                            } else if step <= size {
1832                                for _ in 0..step {
1833                                    buffer.pop_front();
1834                                }
1835                                if buffer.len() == size {
1836                                    return Some(Ok(buffer.iter().cloned().collect()));
1837                                }
1838                            } else if buffer.len() == step {
1839                                // step > size: discard the gap between windows.
1840                                buffer.clear();
1841                            }
1842                        }
1843                        Some(Err(error)) => {
1844                            terminated = true;
1845                            return Some(Err(error));
1846                        }
1847                        None => {
1848                            terminated = true;
1849                            if !buffer.is_empty() && buffer.len() < size {
1850                                return Some(Ok(buffer.iter().cloned().collect()));
1851                            }
1852                            return None;
1853                        }
1854                    }
1855                }
1856            }))
1857        }))
1858    }
1859
1860    #[must_use]
1861    pub fn fold<Acc, F>(self, zero: Acc, f: F) -> Flow<In, Acc, Mat>
1862    where
1863        Acc: Clone + Send + Sync + 'static,
1864        F: Fn(Acc, Out) -> Acc + Send + Sync + 'static,
1865    {
1866        let stage = Arc::new(f);
1867        self.via(Flow::from_transform(move |input| {
1868            let stage = Arc::clone(&stage);
1869            let mut acc = zero.clone();
1870            for item in input {
1871                match item {
1872                    Ok(item) => acc = stage(acc, item),
1873                    Err(error) => return Box::new(std::iter::once(Err(error))),
1874                }
1875            }
1876            Box::new(std::iter::once(Ok(acc)))
1877        }))
1878    }
1879
1880    #[must_use]
1881    pub fn fold_async<Acc, F, Fut>(self, zero: Acc, f: F) -> Flow<In, Acc, Mat>
1882    where
1883        Acc: Clone + Send + Sync + 'static,
1884        F: Fn(Acc, Out) -> Fut + Send + Sync + 'static,
1885        Fut: Future<Output = StreamResult<Acc>> + Send + 'static,
1886    {
1887        let stage = Arc::new(f);
1888        self.via(Flow::from_transform(move |mut input| {
1889            let stage = Arc::clone(&stage);
1890            let mut acc = Some(zero.clone());
1891            let mut done = false;
1892            Box::new(std::iter::from_fn(move || {
1893                if done {
1894                    return None;
1895                }
1896                done = true;
1897
1898                for item in input.by_ref() {
1899                    let item = match item {
1900                        Ok(item) => item,
1901                        Err(error) => return Some(Err(error)),
1902                    };
1903                    let current = acc.take().expect("fold_async accumulator present");
1904                    match catch_unwind_failed("fold_async factory", || stage(current, item))
1905                        .and_then(run_future_inline_or_spawn)
1906                    {
1907                        Ok(next) => acc = Some(next),
1908                        Err(error) => return Some(Err(error)),
1909                    }
1910                }
1911
1912                Some(Ok(acc.take().expect("fold_async accumulator present")))
1913            }))
1914        }))
1915    }
1916
1917    #[must_use]
1918    pub fn map_with_resource<Resource, Next, Create, F, Close>(
1919        self,
1920        create: Create,
1921        f: F,
1922        close: Close,
1923    ) -> Flow<In, Next, Mat>
1924    where
1925        Resource: Send + 'static,
1926        Next: Send + 'static,
1927        Create: Fn() -> StreamResult<Resource> + Send + Sync + 'static,
1928        F: Fn(&mut Resource, Out) -> StreamResult<Next> + Send + Sync + 'static,
1929        Close: Fn(Resource) -> StreamResult<Option<Next>> + Send + Sync + 'static,
1930    {
1931        let create = Arc::new(create);
1932        let stage = Arc::new(f);
1933        let close = Arc::new(close);
1934        self.via(Flow::from_transform(move |input| {
1935            Box::new(MapWithResourceStream {
1936                input,
1937                create: Arc::clone(&create),
1938                stage: Arc::clone(&stage),
1939                close: Arc::clone(&close),
1940                resource: None,
1941                created: false,
1942                pending_terminal: None,
1943                terminated: false,
1944                _marker: PhantomData,
1945            }) as BoxStream<Next>
1946        }))
1947    }
1948
1949    #[must_use]
1950    pub fn fold_result<Acc, F>(self, zero: Acc, f: F) -> Flow<In, Acc, Mat>
1951    where
1952        Acc: Clone + Send + Sync + 'static,
1953        F: Fn(Acc, Out) -> StreamResult<Acc> + Send + Sync + 'static,
1954    {
1955        let stage = Arc::new(f);
1956        self.via(Flow::from_transform(move |input| {
1957            let stage = Arc::clone(&stage);
1958            let mut acc = zero.clone();
1959            for item in input {
1960                match item {
1961                    Ok(item) => match stage(acc, item) {
1962                        Ok(next) => acc = next,
1963                        Err(error) => return Box::new(std::iter::once(Err(error))),
1964                    },
1965                    Err(error) => return Box::new(std::iter::once(Err(error))),
1966                }
1967            }
1968            Box::new(std::iter::once(Ok(acc)))
1969        }))
1970    }
1971
1972    #[must_use]
1973    pub fn fold_result_with_supervision<Acc, F>(
1974        self,
1975        zero: Acc,
1976        f: F,
1977        decider: SupervisionDecider,
1978    ) -> Flow<In, Acc, Mat>
1979    where
1980        Acc: Clone + Send + Sync + 'static,
1981        F: Fn(Acc, Out) -> StreamResult<Acc> + Send + Sync + 'static,
1982    {
1983        let stage = Arc::new(f);
1984        self.via(Flow::from_transform(move |input| {
1985            let stage = Arc::clone(&stage);
1986            let decider = Arc::clone(&decider);
1987            let mut acc = zero.clone();
1988            for item in input {
1989                match item {
1990                    Ok(item) => {
1991                        let previous = acc;
1992                        match call_supervised("fold_result callback", || {
1993                            stage(previous.clone(), item)
1994                        }) {
1995                            Ok(next) => acc = next,
1996                            Err(error) => match decide_supervision(&decider, &error) {
1997                                SupervisionDirective::Stop => {
1998                                    return Box::new(std::iter::once(Err(error)));
1999                                }
2000                                SupervisionDirective::Resume => acc = previous,
2001                                SupervisionDirective::Restart => acc = zero.clone(),
2002                            },
2003                        }
2004                    }
2005                    Err(error) => return Box::new(std::iter::once(Err(error))),
2006                }
2007            }
2008            Box::new(std::iter::once(Ok(acc)))
2009        }))
2010    }
2011
2012    #[must_use]
2013    pub fn reduce<F>(self, f: F) -> Flow<In, Out, Mat>
2014    where
2015        F: Fn(Out, Out) -> Out + Send + Sync + 'static,
2016    {
2017        let stage = Arc::new(f);
2018        self.via(Flow::from_transform(move |mut input| {
2019            let Some(first) = input.next() else {
2020                return Box::new(std::iter::once(Err(StreamError::EmptyStream)));
2021            };
2022            let mut acc = match first {
2023                Ok(item) => item,
2024                Err(error) => return Box::new(std::iter::once(Err(error))),
2025            };
2026            for item in input {
2027                match item {
2028                    Ok(item) => acc = stage(acc, item),
2029                    Err(error) => return Box::new(std::iter::once(Err(error))),
2030                }
2031            }
2032            Box::new(std::iter::once(Ok(acc)))
2033        }))
2034    }
2035
2036    #[must_use]
2037    pub fn reduce_result<F>(self, f: F) -> Flow<In, Out, Mat>
2038    where
2039        Out: Clone,
2040        F: Fn(Out, Out) -> StreamResult<Out> + Send + Sync + 'static,
2041    {
2042        let stage = Arc::new(f);
2043        self.via(Flow::from_transform(move |mut input| {
2044            let Some(first) = input.next() else {
2045                return Box::new(std::iter::once(Err(StreamError::EmptyStream)));
2046            };
2047            let mut acc = match first {
2048                Ok(item) => item,
2049                Err(error) => return Box::new(std::iter::once(Err(error))),
2050            };
2051            for item in input {
2052                match item {
2053                    Ok(item) => match stage(acc, item) {
2054                        Ok(next) => acc = next,
2055                        Err(error) => return Box::new(std::iter::once(Err(error))),
2056                    },
2057                    Err(error) => return Box::new(std::iter::once(Err(error))),
2058                }
2059            }
2060            Box::new(std::iter::once(Ok(acc)))
2061        }))
2062    }
2063
2064    #[must_use]
2065    pub fn reduce_result_with_supervision<F>(
2066        self,
2067        f: F,
2068        decider: SupervisionDecider,
2069    ) -> Flow<In, Out, Mat>
2070    where
2071        Out: Clone,
2072        F: Fn(Out, Out) -> StreamResult<Out> + Send + Sync + 'static,
2073    {
2074        let stage = Arc::new(f);
2075        self.via(Flow::from_transform(move |input| {
2076            let stage = Arc::clone(&stage);
2077            let decider = Arc::clone(&decider);
2078            let mut acc = None::<Out>;
2079            for item in input {
2080                match item {
2081                    Ok(item) => {
2082                        let Some(previous) = acc.take() else {
2083                            acc = Some(item);
2084                            continue;
2085                        };
2086                        match call_supervised("reduce_result callback", || {
2087                            stage(previous.clone(), item)
2088                        }) {
2089                            Ok(next) => acc = Some(next),
2090                            Err(error) => match decide_supervision(&decider, &error) {
2091                                SupervisionDirective::Stop => {
2092                                    return Box::new(std::iter::once(Err(error)));
2093                                }
2094                                SupervisionDirective::Resume => acc = Some(previous),
2095                                SupervisionDirective::Restart => acc = None,
2096                            },
2097                        }
2098                    }
2099                    Err(error) => return Box::new(std::iter::once(Err(error))),
2100                }
2101            }
2102            match acc {
2103                Some(acc) => Box::new(std::iter::once(Ok(acc))),
2104                None => Box::new(std::iter::once(Err(StreamError::EmptyStream))),
2105            }
2106        }))
2107    }
2108
2109    #[must_use]
2110    pub fn map_error<F>(self, f: F) -> Flow<In, Out, Mat>
2111    where
2112        F: Fn(StreamError) -> StreamError + Send + Sync + 'static,
2113    {
2114        let stage = Arc::new(f);
2115        self.via(Flow::from_transform(move |input| {
2116            let stage = Arc::clone(&stage);
2117            Box::new(input.map(move |item| item.map_err(|error| stage(error))))
2118        }))
2119    }
2120
2121    #[must_use]
2122    pub fn recover<F>(self, f: F) -> Flow<In, Out, Mat>
2123    where
2124        F: Fn(StreamError) -> Option<Out> + Send + Sync + 'static,
2125    {
2126        let stage = Arc::new(f);
2127        self.via(Flow::from_transform(move |mut input| {
2128            let stage = Arc::clone(&stage);
2129            let mut done = false;
2130            Box::new(std::iter::from_fn(move || {
2131                if done {
2132                    return None;
2133                }
2134                match input.next()? {
2135                    Ok(item) => Some(Ok(item)),
2136                    Err(error) => {
2137                        done = true;
2138                        match stage(error.clone()) {
2139                            Some(item) => Some(Ok(item)),
2140                            None => Some(Err(error)),
2141                        }
2142                    }
2143                }
2144            }))
2145        }))
2146    }
2147
2148    #[must_use]
2149    pub fn recover_with<F>(self, f: F) -> Flow<In, Out, Mat>
2150    where
2151        F: Fn(StreamError) -> Option<Source<Out>> + Send + Sync + 'static,
2152    {
2153        // Akka's `recoverWith` retries indefinitely (it is `recoverWithRetries(-1, _)`).
2154        self.recover_with_attempts(None, f)
2155    }
2156
2157    #[must_use]
2158    pub fn recover_with_retries<F>(self, retries: usize, f: F) -> Flow<In, Out, Mat>
2159    where
2160        F: Fn(StreamError) -> Option<Source<Out>> + Send + Sync + 'static,
2161    {
2162        self.recover_with_attempts(Some(retries), f)
2163    }
2164
2165    fn recover_with_attempts<F>(self, attempts: Option<usize>, f: F) -> Flow<In, Out, Mat>
2166    where
2167        F: Fn(StreamError) -> Option<Source<Out>> + Send + Sync + 'static,
2168    {
2169        let stage = Arc::new(f);
2170        self.via(Flow::from_runtime_transform(move |input, materializer| {
2171            let replacement_materializer =
2172                materializer.with_name_prefix(materializer.name_prefix().to_owned());
2173            let stage = Arc::clone(&stage);
2174            // `None` means unbounded retries (Akka's `recoverWith`).
2175            let mut remaining_retries = attempts;
2176            let mut current = input;
2177            let mut terminated = false;
2178
2179            Ok(Box::new(std::iter::from_fn(move || {
2180                if terminated {
2181                    return None;
2182                }
2183
2184                loop {
2185                    match current.next() {
2186                        Some(Ok(item)) => return Some(Ok(item)),
2187                        Some(Err(error)) if remaining_retries != Some(0) => {
2188                            if let Some(remaining) = remaining_retries.as_mut() {
2189                                *remaining -= 1;
2190                            }
2191                            match stage(error.clone()) {
2192                                Some(source) => match Arc::clone(&source.factory)
2193                                    .create(&replacement_materializer)
2194                                {
2195                                    Ok((stream, _)) => current = stream,
2196                                    Err(error) => {
2197                                        terminated = true;
2198                                        return Some(Err(error));
2199                                    }
2200                                },
2201                                None => {
2202                                    terminated = true;
2203                                    return Some(Err(error));
2204                                }
2205                            }
2206                        }
2207                        Some(Err(error)) => {
2208                            terminated = true;
2209                            return Some(Err(error));
2210                        }
2211                        None => {
2212                            terminated = true;
2213                            return None;
2214                        }
2215                    }
2216                }
2217            })) as BoxStream<Out>)
2218        }))
2219    }
2220
2221    #[must_use]
2222    pub fn on_error_complete(self) -> Flow<In, Out, Mat> {
2223        self.via(Flow::from_transform(move |mut input| {
2224            let mut done = false;
2225            Box::new(std::iter::from_fn(move || {
2226                if done {
2227                    return None;
2228                }
2229                match input.next()? {
2230                    Ok(item) => Some(Ok(item)),
2231                    Err(_) => {
2232                        done = true;
2233                        None
2234                    }
2235                }
2236            }))
2237        }))
2238    }
2239
2240    #[must_use]
2241    pub fn prefix_and_tail(self, n: usize) -> Flow<In, (Vec<Out>, Source<Out>), Mat> {
2242        self.via(Flow::from_runtime_transform(move |input, _materializer| {
2243            Ok(prefix_and_tail_stream(input, n))
2244        }))
2245    }
2246
2247    #[must_use]
2248    pub fn flat_map_prefix<Next, NextMat, F>(self, n: usize, f: F) -> Flow<In, Next, Mat>
2249    where
2250        Next: Send + 'static,
2251        NextMat: Send + 'static,
2252        F: Fn(Vec<Out>) -> Flow<Out, Next, NextMat> + Send + Sync + 'static,
2253        Out: Clone,
2254    {
2255        let stage = Arc::new(f);
2256        self.via(Flow::from_runtime_transform(
2257            move |mut input, materializer| {
2258                let mut prefix = Vec::with_capacity(n);
2259                while prefix.len() < n {
2260                    match input.next() {
2261                        Some(Ok(item)) => prefix.push(item),
2262                        Some(Err(error)) => {
2263                            return Ok(Box::new(std::iter::once(Err(error))) as BoxStream<Next>);
2264                        }
2265                        None => break,
2266                    }
2267                }
2268
2269                let flow = stage(prefix);
2270                let transform = flow.transform;
2271                let _ = (flow.materialize)()?;
2272                match transform {
2273                    FlowTransform::Pure(transform) => Ok(transform(input)),
2274                    FlowTransform::Runtime(transform) => transform(input, materializer),
2275                }
2276            },
2277        ))
2278    }
2279
2280    #[must_use]
2281    pub fn group_by<Key, F>(
2282        self,
2283        max_substreams: usize,
2284        f: F,
2285        allow_closed_substream_recreation: bool,
2286    ) -> Flow<In, Source<Out>, Mat>
2287    where
2288        Key: Clone + Eq + Hash + Send + 'static,
2289        F: Fn(&Out) -> Key + Send + Sync + 'static,
2290        Out: Clone,
2291    {
2292        self.group_by_with_batching(
2293            max_substreams,
2294            f,
2295            allow_closed_substream_recreation,
2296            GroupByBatchMode::Immediate,
2297        )
2298    }
2299
2300    pub(super) fn group_by_with_batching<Key, F>(
2301        self,
2302        max_substreams: usize,
2303        f: F,
2304        allow_closed_substream_recreation: bool,
2305        batch_mode: GroupByBatchMode,
2306    ) -> Flow<In, Source<Out>, Mat>
2307    where
2308        Key: Clone + Eq + Hash + Send + 'static,
2309        F: Fn(&Out) -> Key + Send + Sync + 'static,
2310        Out: Clone,
2311    {
2312        assert!(
2313            max_substreams > 0,
2314            "group_by max_substreams must be greater than zero"
2315        );
2316        let key_fn = Arc::new(f);
2317        self.via(Flow::from_runtime_transform(move |input, materializer| {
2318            Ok(group_by_stream(
2319                input,
2320                max_substreams,
2321                allow_closed_substream_recreation,
2322                Arc::clone(&key_fn),
2323                batch_mode,
2324                materializer,
2325            ))
2326        }))
2327    }
2328
2329    #[must_use]
2330    pub fn split_when<F>(self, predicate: F) -> Flow<In, Source<Out>, Mat>
2331    where
2332        F: Fn(&Out) -> bool + Send + Sync + 'static,
2333        Out: Clone,
2334    {
2335        let predicate = Arc::new(predicate);
2336        self.via(Flow::from_runtime_transform(move |input, materializer| {
2337            Ok(split_streams(
2338                input,
2339                SplitMode::When,
2340                Arc::clone(&predicate),
2341                materializer,
2342            ))
2343        }))
2344    }
2345
2346    #[must_use]
2347    pub fn split_after<F>(self, predicate: F) -> Flow<In, Source<Out>, Mat>
2348    where
2349        F: Fn(&Out) -> bool + Send + Sync + 'static,
2350        Out: Clone,
2351    {
2352        let predicate = Arc::new(predicate);
2353        self.via(Flow::from_runtime_transform(move |input, materializer| {
2354            Ok(split_streams(
2355                input,
2356                SplitMode::After,
2357                Arc::clone(&predicate),
2358                materializer,
2359            ))
2360        }))
2361    }
2362
2363    #[must_use]
2364    pub fn flat_map_concat<Next, NextMat, F>(self, f: F) -> Flow<In, Next, Mat>
2365    where
2366        Next: Send + 'static,
2367        NextMat: Send + 'static,
2368        F: Fn(Out) -> Source<Next, NextMat> + Send + Sync + 'static,
2369    {
2370        let stage = Arc::new(f);
2371        self.via(Flow::from_runtime_transform(move |input, materializer| {
2372            Ok(flat_map_concat_stream(
2373                input,
2374                Arc::clone(&stage),
2375                materializer,
2376            ))
2377        }))
2378    }
2379
2380    #[must_use]
2381    pub fn flat_map_merge<Next, NextMat, F>(self, breadth: usize, f: F) -> Flow<In, Next, Mat>
2382    where
2383        Next: Send + 'static,
2384        NextMat: Send + 'static,
2385        F: Fn(Out) -> Source<Next, NextMat> + Send + Sync + 'static,
2386    {
2387        assert!(
2388            breadth > 0,
2389            "flat_map_merge breadth must be greater than zero"
2390        );
2391        let stage = Arc::new(f);
2392        self.via(Flow::from_runtime_transform_with_hints(
2393            move |input, materializer| {
2394                Ok(flat_map_merge_stream(
2395                    input,
2396                    breadth,
2397                    Arc::clone(&stage),
2398                    materializer,
2399                ))
2400            },
2401            FlowHints::PRESERVES_TERMINAL_CONSUMER_BATCH,
2402        ))
2403    }
2404
2405    #[must_use]
2406    pub fn concat<Mat2>(self, that: Source<Out, Mat2>) -> Flow<In, Out, Mat>
2407    where
2408        Mat2: Send + 'static,
2409    {
2410        self.concat_sources([that])
2411    }
2412
2413    #[must_use]
2414    pub fn concat_lazy<Mat2>(self, that: Source<Out, Mat2>) -> Flow<In, Out, Mat>
2415    where
2416        Mat2: Send + 'static,
2417    {
2418        let that_factory = that.factory;
2419        self.via(Flow::from_runtime_transform(move |input, materializer| {
2420            let primary = input;
2421            Ok(concat_streams_lazy(
2422                primary,
2423                vec![Arc::clone(&that_factory)],
2424                materializer,
2425            ))
2426        }))
2427    }
2428
2429    #[must_use]
2430    pub fn concat_all_lazy<Mat2, I>(self, those: I) -> Flow<In, Out, Mat>
2431    where
2432        Mat2: Send + 'static,
2433        I: IntoIterator<Item = Source<Out, Mat2>>,
2434    {
2435        let source_factories: Vec<_> = those.into_iter().map(|source| source.factory).collect();
2436        self.via(Flow::from_runtime_transform(move |input, materializer| {
2437            Ok(concat_streams_lazy(
2438                input,
2439                source_factories.clone(),
2440                materializer,
2441            ))
2442        }))
2443    }
2444
2445    #[must_use]
2446    pub fn prepend<Mat2>(self, that: Source<Out, Mat2>) -> Flow<In, Out, Mat>
2447    where
2448        Mat2: Send + 'static,
2449    {
2450        self.prepend_sources([that])
2451    }
2452
2453    #[must_use]
2454    pub fn prepend_lazy<Mat2>(self, that: Source<Out, Mat2>) -> Flow<In, Out, Mat>
2455    where
2456        Mat2: Send + 'static,
2457    {
2458        self.prepend(that)
2459    }
2460
2461    #[must_use]
2462    pub fn or_else<Mat2>(self, secondary: Source<Out, Mat2>) -> Flow<In, Out, Mat>
2463    where
2464        Mat2: Send + 'static,
2465    {
2466        let secondary_factory = secondary.factory;
2467        self.via(Flow::from_runtime_transform(move |input, materializer| {
2468            let secondary = match Arc::clone(&secondary_factory).create(materializer) {
2469                Ok((stream, _)) => stream,
2470                Err(error) => Box::new(std::iter::once(Err(error))) as BoxStream<Out>,
2471            };
2472            Ok(or_else_stream(input, secondary))
2473        }))
2474    }
2475
2476    #[must_use]
2477    pub fn interleave<Mat2>(
2478        self,
2479        that: Source<Out, Mat2>,
2480        segment_size: usize,
2481    ) -> Flow<In, Out, Mat>
2482    where
2483        Mat2: Send + 'static,
2484    {
2485        self.interleave_all([that], segment_size, false)
2486    }
2487
2488    #[must_use]
2489    pub fn interleave_all<Mat2, I>(
2490        self,
2491        those: I,
2492        segment_size: usize,
2493        eager_close: bool,
2494    ) -> Flow<In, Out, Mat>
2495    where
2496        Mat2: Send + 'static,
2497        I: IntoIterator<Item = Source<Out, Mat2>>,
2498    {
2499        let source_factories: Vec<_> = those.into_iter().map(|source| source.factory).collect();
2500        self.via(Flow::from_runtime_transform(move |input, materializer| {
2501            let mut streams = Vec::with_capacity(source_factories.len() + 1);
2502            streams.push(input);
2503            for factory in &source_factories {
2504                let stream = match Arc::clone(factory).create(materializer) {
2505                    Ok((stream, _)) => stream,
2506                    Err(error) => {
2507                        return Ok(Box::new(std::iter::once(Err(error))) as BoxStream<Out>);
2508                    }
2509                };
2510                streams.push(stream);
2511            }
2512            Ok(interleave_streams(streams, segment_size, eager_close))
2513        }))
2514    }
2515
2516    #[must_use]
2517    pub fn merge_sorted<Mat2>(self, that: Source<Out, Mat2>) -> Flow<In, Out, Mat>
2518    where
2519        Out: Ord,
2520        Mat2: Send + 'static,
2521    {
2522        let source_factory = that.factory;
2523        self.via(Flow::from_runtime_transform(move |input, materializer| {
2524            let other = match Arc::clone(&source_factory).create(materializer) {
2525                Ok((stream, _)) => stream,
2526                Err(error) => return Ok(Box::new(std::iter::once(Err(error))) as BoxStream<Out>),
2527            };
2528            Ok(merge_sorted_stream(input, other))
2529        }))
2530    }
2531
2532    #[must_use]
2533    pub fn merge_latest<Mat2>(
2534        self,
2535        that: Source<Out, Mat2>,
2536        eager_complete: bool,
2537    ) -> Flow<In, Vec<Out>, Mat>
2538    where
2539        Out: Clone,
2540        Mat2: Send + 'static,
2541    {
2542        let source_factory = that.factory;
2543        self.via(Flow::from_runtime_transform(move |input, materializer| {
2544            let other = match Arc::clone(&source_factory).create(materializer) {
2545                Ok((stream, _)) => stream,
2546                Err(error) => {
2547                    return Ok(Box::new(std::iter::once(Err(error))) as BoxStream<Vec<Out>>);
2548                }
2549            };
2550            Ok(merge_latest_streams(vec![input, other], eager_complete))
2551        }))
2552    }
2553
2554    #[must_use]
2555    pub fn merge_all<Mat2, I>(self, those: I, eager_complete: bool) -> Flow<In, Out, Mat>
2556    where
2557        Mat2: Send + 'static,
2558        I: IntoIterator<Item = Source<Out, Mat2>>,
2559    {
2560        let source_factories: Vec<_> = those.into_iter().map(|source| source.factory).collect();
2561        self.via(Flow::from_runtime_transform(move |input, materializer| {
2562            let mut streams = Vec::with_capacity(source_factories.len() + 1);
2563            streams.push(input);
2564            for factory in &source_factories {
2565                let stream = match Arc::clone(factory).create(materializer) {
2566                    Ok((stream, _)) => stream,
2567                    Err(error) => {
2568                        return Ok(Box::new(std::iter::once(Err(error))) as BoxStream<Out>);
2569                    }
2570                };
2571                streams.push(stream);
2572            }
2573            Ok(merge_streams(streams, eager_complete))
2574        }))
2575    }
2576
2577    #[must_use]
2578    pub fn zip_with<Mat2, Out2, Next, F>(
2579        self,
2580        that: Source<Out2, Mat2>,
2581        combine: F,
2582    ) -> Flow<In, Next, Mat>
2583    where
2584        Out2: Send + 'static,
2585        Next: Send + 'static,
2586        Mat2: Send + 'static,
2587        F: Fn(Out, Out2) -> Next + Send + Sync + 'static,
2588    {
2589        let source_factory = that.factory;
2590        let combine = Arc::new(combine);
2591        self.via(Flow::from_runtime_transform(move |input, materializer| {
2592            let other = match Arc::clone(&source_factory).create(materializer) {
2593                Ok((stream, _)) => stream,
2594                Err(error) => return Ok(Box::new(std::iter::once(Err(error))) as BoxStream<Next>),
2595            };
2596            let combine = Arc::clone(&combine);
2597            Ok(Box::new(
2598                zip_streams(input, other)
2599                    .map(move |item| item.map(|(left, right)| combine(left, right))),
2600            ) as BoxStream<Next>)
2601        }))
2602    }
2603
2604    #[must_use]
2605    pub fn zip_latest<Mat2, Out2>(self, that: Source<Out2, Mat2>) -> Flow<In, (Out, Out2), Mat>
2606    where
2607        Out: Clone,
2608        Out2: Clone + Send + 'static,
2609        Mat2: Send + 'static,
2610    {
2611        self.zip_latest_with(that, true, |left, right| (left, right))
2612    }
2613
2614    #[must_use]
2615    pub fn zip_latest_with<Mat2, Out2, Next, F>(
2616        self,
2617        that: Source<Out2, Mat2>,
2618        eager_complete: bool,
2619        combine: F,
2620    ) -> Flow<In, Next, Mat>
2621    where
2622        Out: Clone,
2623        Out2: Clone + Send + 'static,
2624        Next: Send + 'static,
2625        Mat2: Send + 'static,
2626        F: Fn(Out, Out2) -> Next + Send + Sync + 'static,
2627    {
2628        let source_factory = that.factory;
2629        let combine = Arc::new(combine);
2630        self.via(Flow::from_runtime_transform(move |input, materializer| {
2631            let other = match Arc::clone(&source_factory).create(materializer) {
2632                Ok((stream, _)) => stream,
2633                Err(error) => return Ok(Box::new(std::iter::once(Err(error))) as BoxStream<Next>),
2634            };
2635            Ok(zip_latest_with_stream(
2636                input,
2637                other,
2638                eager_complete,
2639                Arc::clone(&combine),
2640            ))
2641        }))
2642    }
2643
2644    #[must_use]
2645    pub fn zip_with_index(self) -> Flow<In, (Out, u64), Mat> {
2646        self.via(Flow::from_runtime_transform(
2647            move |mut input, _materializer| {
2648                let mut index = 0_u64;
2649                Ok(Box::new(std::iter::from_fn(move || {
2650                    input.next().map(|item| {
2651                        item.map(|value| {
2652                            let pair = (value, index);
2653                            index = index.wrapping_add(1);
2654                            pair
2655                        })
2656                    })
2657                })) as BoxStream<(Out, u64)>)
2658            },
2659        ))
2660    }
2661
2662    #[must_use]
2663    pub fn zip_all<Mat2, Out2>(
2664        self,
2665        that: Source<Out2, Mat2>,
2666        this_elem: Out,
2667        that_elem: Out2,
2668    ) -> Flow<In, (Out, Out2), Mat>
2669    where
2670        Out: Clone + Sync,
2671        Out2: Clone + Send + Sync + 'static,
2672        Mat2: Send + 'static,
2673    {
2674        let source_factory = that.factory;
2675        self.via(Flow::from_runtime_transform(move |input, materializer| {
2676            let other = match Arc::clone(&source_factory).create(materializer) {
2677                Ok((stream, _)) => stream,
2678                Err(error) => {
2679                    return Ok(Box::new(std::iter::once(Err(error))) as BoxStream<(Out, Out2)>);
2680                }
2681            };
2682            Ok(zip_all_stream(
2683                input,
2684                other,
2685                this_elem.clone(),
2686                that_elem.clone(),
2687            ))
2688        }))
2689    }
2690
2691    #[must_use]
2692    pub fn also_to<SideMat>(self, sink: Sink<Out, SideMat>) -> Flow<In, Out, Mat>
2693    where
2694        Out: Clone,
2695        SideMat: Send + 'static,
2696    {
2697        self.via(Flow::from_runtime_transform(
2698            move |mut input: BoxStream<Out>, materializer| {
2699                let (side_sender, side_mat) = materialize_side_sink(&sink, materializer, 0)?;
2700                let mut sender = Some(side_sender);
2701                let side_mat = side_mat;
2702                Ok(Box::new(std::iter::from_fn(move || match input.next() {
2703                    Some(Ok(item)) => {
2704                        let _ = &side_mat;
2705                        if sender
2706                            .as_ref()
2707                            .is_some_and(|sender| sender.send(Ok(item.clone())).is_err())
2708                        {
2709                            sender = None;
2710                            return None;
2711                        }
2712                        Some(Ok(item))
2713                    }
2714                    Some(Err(error)) => {
2715                        let _ = &side_mat;
2716                        let _ = sender
2717                            .as_ref()
2718                            .and_then(|sender| sender.send(Err(error.clone())).ok());
2719                        sender = None;
2720                        Some(Err(error))
2721                    }
2722                    None => {
2723                        let _ = &side_mat;
2724                        sender = None;
2725                        None
2726                    }
2727                })) as BoxStream<Out>)
2728            },
2729        ))
2730    }
2731
2732    #[must_use]
2733    pub fn also_to_all<SideMat, I>(self, sinks: I) -> Flow<In, Out, Mat>
2734    where
2735        Out: Clone,
2736        SideMat: Send + 'static,
2737        I: IntoIterator<Item = Sink<Out, SideMat>>,
2738    {
2739        let sinks: Vec<_> = sinks.into_iter().collect();
2740        if sinks.is_empty() {
2741            return self;
2742        }
2743
2744        self.via(Flow::from_runtime_transform(
2745            move |mut input: BoxStream<Out>, materializer| {
2746                let mut sides = sinks
2747                    .iter()
2748                    .map(|sink| materialize_side_sink(sink, materializer, 0))
2749                    .collect::<StreamResult<Vec<_>>>()?;
2750                Ok(Box::new(std::iter::from_fn(move || match input.next() {
2751                    Some(Ok(item)) => {
2752                        for (sender, _) in &sides {
2753                            if sender.send(Ok(item.clone())).is_err() {
2754                                sides.clear();
2755                                return None;
2756                            }
2757                        }
2758                        Some(Ok(item))
2759                    }
2760                    Some(Err(error)) => {
2761                        for (sender, _) in &sides {
2762                            let _ = sender.send(Err(error.clone())).ok();
2763                        }
2764                        sides.clear();
2765                        Some(Err(error))
2766                    }
2767                    None => {
2768                        sides.clear();
2769                        None
2770                    }
2771                })) as BoxStream<Out>)
2772            },
2773        ))
2774    }
2775
2776    #[must_use]
2777    pub fn divert_to<SideMat, F>(self, sink: Sink<Out, SideMat>, predicate: F) -> Flow<In, Out, Mat>
2778    where
2779        SideMat: Send + 'static,
2780        F: Fn(&Out) -> bool + Send + Sync + 'static,
2781    {
2782        let predicate = Arc::new(predicate);
2783        self.via(Flow::from_runtime_transform(
2784            move |mut input: BoxStream<Out>, materializer| {
2785                let predicate = Arc::clone(&predicate);
2786                let (side_sender, side_mat) = materialize_side_sink(&sink, materializer, 0)?;
2787                let mut sender = Some(side_sender);
2788                let side_mat = side_mat;
2789                Ok(Box::new(std::iter::from_fn(move || {
2790                    loop {
2791                        let _ = &side_mat;
2792                        match input.next() {
2793                            Some(Ok(item)) if predicate(&item) => {
2794                                if sender
2795                                    .as_ref()
2796                                    .is_some_and(|sender| sender.send(Ok(item)).is_err())
2797                                {
2798                                    sender = None;
2799                                    return None;
2800                                }
2801                            }
2802                            Some(Ok(item)) => return Some(Ok(item)),
2803                            Some(Err(error)) => {
2804                                let _ = sender
2805                                    .as_ref()
2806                                    .and_then(|sender| sender.send(Err(error.clone())).ok());
2807                                sender = None;
2808                                return Some(Err(error));
2809                            }
2810                            None => {
2811                                sender = None;
2812                                return None;
2813                            }
2814                        }
2815                    }
2816                })) as BoxStream<Out>)
2817            },
2818        ))
2819    }
2820
2821    #[must_use]
2822    pub fn wire_tap<SideMat>(self, sink: Sink<Out, SideMat>) -> Flow<In, Out, Mat>
2823    where
2824        Out: Clone,
2825        SideMat: Send + 'static,
2826    {
2827        self.via(Flow::from_runtime_transform(
2828            move |mut input: BoxStream<Out>, materializer| {
2829                let (side_sender, side_mat) = materialize_side_sink(&sink, materializer, 1)?;
2830                let mut sender = Some(side_sender);
2831                let side_mat = side_mat;
2832                Ok(Box::new(std::iter::from_fn(move || match input.next() {
2833                    Some(Ok(item)) => {
2834                        let _ = &side_mat;
2835                        if let Some(side) = sender.as_ref() {
2836                            match side.try_send(Ok(item.clone())) {
2837                                Ok(()) | Err(std::sync::mpsc::TrySendError::Full(_)) => {}
2838                                Err(std::sync::mpsc::TrySendError::Disconnected(_)) => {
2839                                    sender = None
2840                                }
2841                            }
2842                        }
2843                        Some(Ok(item))
2844                    }
2845                    Some(Err(error)) => {
2846                        let _ = &side_mat;
2847                        if let Some(side) = sender.as_ref() {
2848                            match side.try_send(Err(error.clone())) {
2849                                Ok(())
2850                                | Err(std::sync::mpsc::TrySendError::Full(_))
2851                                | Err(std::sync::mpsc::TrySendError::Disconnected(_)) => {}
2852                            }
2853                        }
2854                        sender = None;
2855                        Some(Err(error))
2856                    }
2857                    None => {
2858                        let _ = &side_mat;
2859                        sender = None;
2860                        None
2861                    }
2862                })) as BoxStream<Out>)
2863            },
2864        ))
2865    }
2866
2867    fn concat_sources<Mat2, I>(self, those: I) -> Flow<In, Out, Mat>
2868    where
2869        Mat2: Send + 'static,
2870        I: IntoIterator<Item = Source<Out, Mat2>>,
2871    {
2872        let source_factories: Vec<_> = those.into_iter().map(|source| source.factory).collect();
2873        self.via(Flow::from_runtime_transform(move |input, materializer| {
2874            let mut streams = Vec::with_capacity(source_factories.len() + 1);
2875            streams.push(input);
2876            for factory in &source_factories {
2877                let stream = match Arc::clone(factory).create(materializer) {
2878                    Ok((stream, _)) => stream,
2879                    Err(error) => {
2880                        return Ok(Box::new(std::iter::once(Err(error))) as BoxStream<Out>);
2881                    }
2882                };
2883                streams.push(stream);
2884            }
2885            Ok(concat_streams(streams))
2886        }))
2887    }
2888
2889    fn prepend_sources<Mat2, I>(self, those: I) -> Flow<In, Out, Mat>
2890    where
2891        Mat2: Send + 'static,
2892        I: IntoIterator<Item = Source<Out, Mat2>>,
2893    {
2894        let source_factories: Vec<_> = those.into_iter().map(|source| source.factory).collect();
2895        self.via(Flow::from_runtime_transform(move |input, materializer| {
2896            let mut streams = Vec::with_capacity(source_factories.len() + 1);
2897            for factory in &source_factories {
2898                let stream = match Arc::clone(factory).create(materializer) {
2899                    Ok((stream, _)) => stream,
2900                    Err(error) => {
2901                        return Ok(Box::new(std::iter::once(Err(error))) as BoxStream<Out>);
2902                    }
2903                };
2904                streams.push(stream);
2905            }
2906            streams.push(input);
2907            Ok(concat_streams(streams))
2908        }))
2909    }
2910
2911    // ── small operators (WP-5) ───────────────────────────────────────
2912
2913    #[must_use]
2914    pub fn intersperse(self, inject: Out) -> Flow<In, Out, Mat>
2915    where
2916        Out: Clone + Sync,
2917    {
2918        let inject = Arc::new(inject);
2919        self.via(Flow::from_transform(move |mut input| {
2920            let inject = Arc::clone(&inject);
2921            let mut first = true;
2922            Box::new(std::iter::from_fn(move || {
2923                if first {
2924                    first = false;
2925                    match input.next() {
2926                        None => None,
2927                        Some(item) => {
2928                            if item.is_err() {
2929                                first = true;
2930                            }
2931                            Some(item)
2932                        }
2933                    }
2934                } else {
2935                    match input.next() {
2936                        None => None,
2937                        Some(item) => {
2938                            if item.is_err() {
2939                                first = true;
2940                                Some(item)
2941                            } else {
2942                                Some(Ok((*inject).clone()))
2943                            }
2944                        }
2945                    }
2946                }
2947            }))
2948        }))
2949    }
2950
2951    #[must_use]
2952    pub fn flatten_optional<Inner>(self) -> Flow<In, Inner, Mat>
2953    where
2954        Out: Into<Option<Inner>>,
2955        Inner: Send + 'static,
2956    {
2957        self.filter_map(|item| item.into())
2958    }
2959
2960    #[must_use]
2961    pub fn grouped_weighted<F>(self, max_weight: usize, cost_fn: F) -> Flow<In, Vec<Out>, Mat>
2962    where
2963        F: Fn(&Out) -> usize + Send + Sync + 'static,
2964    {
2965        let cost_fn = Arc::new(cost_fn);
2966        self.via(Flow::from_transform(move |mut input| {
2967            let cost_fn = Arc::clone(&cost_fn);
2968            Box::new(std::iter::from_fn(move || {
2969                let mut group = Vec::new();
2970                let mut weight = 0usize;
2971                while weight < max_weight {
2972                    match input.next() {
2973                        Some(Ok(item)) => {
2974                            let item_weight = cost_fn(&item);
2975                            if weight > 0 && weight + item_weight > max_weight {
2976                                group.push(item);
2977                                break;
2978                            }
2979                            weight += item_weight;
2980                            group.push(item);
2981                        }
2982                        Some(Err(error)) => return Some(Err(error)),
2983                        None => break,
2984                    }
2985                }
2986                if group.is_empty() {
2987                    None
2988                } else {
2989                    Some(Ok(group))
2990                }
2991            }))
2992        }))
2993    }
2994
2995    #[must_use]
2996    pub fn limit_weighted<F>(self, max_weight: usize, cost_fn: F) -> Flow<In, Out, Mat>
2997    where
2998        F: Fn(&Out) -> usize + Send + Sync + 'static,
2999    {
3000        let cost_fn = Arc::new(cost_fn);
3001        self.via(Flow::from_transform(move |mut input| {
3002            let cost_fn = Arc::clone(&cost_fn);
3003            let mut weight = 0usize;
3004            Box::new(std::iter::from_fn(move || match input.next()? {
3005                Ok(item) => {
3006                    let item_weight = cost_fn(&item);
3007                    if weight + item_weight > max_weight {
3008                        Some(Err(StreamError::LimitExceeded {
3009                            max: max_weight as u64,
3010                        }))
3011                    } else {
3012                        weight += item_weight;
3013                        Some(Ok(item))
3014                    }
3015                }
3016                Err(error) => Some(Err(error)),
3017            }))
3018        }))
3019    }
3020
3021    #[must_use]
3022    pub fn contramap<NewIn, F>(self, f: F) -> Flow<NewIn, Out, Mat>
3023    where
3024        NewIn: Send + 'static,
3025        F: Fn(NewIn) -> In + Send + Sync + 'static,
3026    {
3027        let stage = Arc::new(f);
3028        Flow::from_transform(move |input| {
3029            let stage = Arc::clone(&stage);
3030            Box::new(input.map(move |item| item.map(|item| stage(item))))
3031        })
3032        .via_mat(self, Keep::right)
3033    }
3034
3035    #[must_use]
3036    pub fn monitor<F>(self, f: F) -> Flow<In, Out, Mat>
3037    where
3038        Out: Clone,
3039        F: Fn(&Out) + Send + Sync + 'static,
3040    {
3041        let stage = Arc::new(f);
3042        self.via(Flow::from_transform(move |input| {
3043            let stage = Arc::clone(&stage);
3044            Box::new(input.map(move |item| match item {
3045                Ok(item) => {
3046                    stage(&item);
3047                    Ok(item)
3048                }
3049                Err(error) => Err(error),
3050            }))
3051        }))
3052    }
3053
3054    #[must_use]
3055    pub fn watch_termination<CallbackMat, F>(self, materialize_callback: F) -> Flow<In, Out, Mat>
3056    where
3057        Mat: Clone,
3058        CallbackMat: Send + 'static,
3059        F: Fn(Mat) -> CallbackMat + Send + Sync + 'static,
3060    {
3061        let cb = Arc::new(materialize_callback);
3062        self.map_materialized_value(move |mat| {
3063            let _ = cb(mat.clone());
3064            mat
3065        })
3066    }
3067
3068    #[must_use]
3069    pub fn to<SinkMat>(self, sink: Sink<Out, SinkMat>) -> Sink<In, Mat>
3070    where
3071        SinkMat: Send + 'static,
3072    {
3073        self.to_mat(sink, Keep::left)
3074    }
3075
3076    #[must_use]
3077    pub fn to_mat<SinkMat, Combined, F>(
3078        self,
3079        sink: Sink<Out, SinkMat>,
3080        combine: F,
3081    ) -> Sink<In, Combined>
3082    where
3083        SinkMat: Send + 'static,
3084        Combined: Send + 'static,
3085        F: Fn(Mat, SinkMat) -> Combined + Send + Sync + 'static,
3086    {
3087        let transform = self.transform;
3088        let materialize = self.materialize;
3089        let combine = Arc::new(combine);
3090        Sink::from_runner(move |input, materializer| {
3091            let flow_mat = materialize()?;
3092            let input = match &transform {
3093                FlowTransform::Pure(transform) => transform(input),
3094                FlowTransform::Runtime(transform) => transform(input, materializer)?,
3095            };
3096            let sink_mat = sink.run(input, materializer)?;
3097            Ok(combine(flow_mat, sink_mat))
3098        })
3099    }
3100
3101    #[must_use]
3102    pub fn map_materialized_value<NextMat, F>(self, f: F) -> Flow<In, Out, NextMat>
3103    where
3104        NextMat: Send + 'static,
3105        F: Fn(Mat) -> NextMat + Send + Sync + 'static,
3106    {
3107        let transform = self.transform;
3108        let materialize = self.materialize;
3109        let hints = self.hints;
3110        let attributes = self.attributes;
3111        let f = Arc::new(f);
3112        match transform {
3113            FlowTransform::Pure(transform) => Flow::from_parts_with_hints(
3114                move |input| transform(input),
3115                move || {
3116                    let mat = materialize()?;
3117                    Ok(f(mat))
3118                },
3119                hints,
3120            )
3121            .with_attributes(attributes),
3122            FlowTransform::Runtime(transform) => Flow {
3123                transform: FlowTransform::Runtime(transform),
3124                materialize: Arc::new(move || {
3125                    let mat = materialize()?;
3126                    Ok(f(mat))
3127                }),
3128                hints,
3129                attributes,
3130            },
3131        }
3132    }
3133}
3134
3135struct MapWithResourceStream<In, Out, Resource, Create, F, Close>
3136where
3137    Create: Fn() -> StreamResult<Resource>,
3138    F: Fn(&mut Resource, In) -> StreamResult<Out>,
3139    Close: Fn(Resource) -> StreamResult<Option<Out>>,
3140{
3141    input: BoxStream<In>,
3142    create: Arc<Create>,
3143    stage: Arc<F>,
3144    close: Arc<Close>,
3145    resource: Option<Resource>,
3146    created: bool,
3147    pending_terminal: Option<StreamError>,
3148    terminated: bool,
3149    _marker: PhantomData<fn() -> Out>,
3150}
3151
3152impl<In, Out, Resource, Create, F, Close> MapWithResourceStream<In, Out, Resource, Create, F, Close>
3153where
3154    Create: Fn() -> StreamResult<Resource>,
3155    F: Fn(&mut Resource, In) -> StreamResult<Out>,
3156    Close: Fn(Resource) -> StreamResult<Option<Out>>,
3157{
3158    fn ensure_created(&mut self) -> StreamResult<()> {
3159        if self.created {
3160            return Ok(());
3161        }
3162        self.created = true;
3163        let resource = catch_unwind_failed("map_with_resource create", || (self.create)())
3164            .and_then(|result| result)?;
3165        self.resource = Some(resource);
3166        Ok(())
3167    }
3168
3169    fn close_resource(&mut self) -> StreamResult<Option<Out>> {
3170        match self.resource.take() {
3171            Some(resource) => {
3172                catch_unwind_failed("map_with_resource close", || (self.close)(resource))
3173                    .and_then(|result| result)
3174            }
3175            None => Ok(None),
3176        }
3177    }
3178
3179    fn close_with_terminal(&mut self, terminal: Option<StreamError>) -> Option<StreamResult<Out>> {
3180        self.terminated = terminal.is_none();
3181        match self.close_resource() {
3182            Ok(Some(item)) => {
3183                self.pending_terminal = terminal;
3184                Some(Ok(item))
3185            }
3186            Ok(None) => match terminal {
3187                Some(error) => {
3188                    self.terminated = true;
3189                    Some(Err(error))
3190                }
3191                None => {
3192                    self.terminated = true;
3193                    None
3194                }
3195            },
3196            Err(error) => {
3197                self.terminated = true;
3198                Some(Err(terminal.unwrap_or(error)))
3199            }
3200        }
3201    }
3202}
3203
3204impl<In, Out, Resource, Create, F, Close> Iterator
3205    for MapWithResourceStream<In, Out, Resource, Create, F, Close>
3206where
3207    Create: Fn() -> StreamResult<Resource>,
3208    F: Fn(&mut Resource, In) -> StreamResult<Out>,
3209    Close: Fn(Resource) -> StreamResult<Option<Out>>,
3210{
3211    type Item = StreamResult<Out>;
3212
3213    fn next(&mut self) -> Option<Self::Item> {
3214        if let Some(error) = self.pending_terminal.take() {
3215            self.terminated = true;
3216            return Some(Err(error));
3217        }
3218        if self.terminated {
3219            return None;
3220        }
3221        if let Err(error) = self.ensure_created() {
3222            self.terminated = true;
3223            return Some(Err(error));
3224        }
3225
3226        match self.input.next() {
3227            Some(Ok(item)) => {
3228                let result = {
3229                    let resource = self
3230                        .resource
3231                        .as_mut()
3232                        .expect("map_with_resource resource is open");
3233                    catch_unwind_failed("map_with_resource function", || {
3234                        (self.stage)(resource, item)
3235                    })
3236                    .and_then(|result| result)
3237                };
3238                match result {
3239                    Ok(item) => Some(Ok(item)),
3240                    Err(error) => self.close_with_terminal(Some(error)),
3241                }
3242            }
3243            Some(Err(error)) => self.close_with_terminal(Some(error)),
3244            None => self.close_with_terminal(None),
3245        }
3246    }
3247}
3248
3249impl<I1: Send + 'static, O1: Send + 'static, I2: Send + 'static, O2: Send + 'static>
3250    BidiFlow<I1, O1, I2, O2>
3251{
3252    #[must_use]
3253    pub fn from_flows<Mat1, Mat2>(top: Flow<I1, O1, Mat1>, bottom: Flow<I2, O2, Mat2>) -> Self
3254    where
3255        Mat1: Send + 'static,
3256        Mat2: Send + 'static,
3257    {
3258        Self {
3259            top: top.map_materialized_value(|_| NotUsed),
3260            bottom: bottom.map_materialized_value(|_| NotUsed),
3261            attributes: Attributes::default(),
3262        }
3263    }
3264}
3265
3266impl<I1: Send + 'static, O1: Send + 'static, I2: Send + 'static, O2: Send + 'static>
3267    BidiFlow<I1, O1, I2, O2>
3268{
3269    #[must_use]
3270    pub fn attributes(&self) -> &Attributes {
3271        &self.attributes
3272    }
3273
3274    #[must_use]
3275    pub fn with_attributes(mut self, attributes: Attributes) -> Self {
3276        self.attributes = attributes;
3277        self
3278    }
3279
3280    #[must_use]
3281    pub fn add_attributes(mut self, attributes: Attributes) -> Self {
3282        self.attributes = self.attributes.and(attributes);
3283        self
3284    }
3285
3286    #[must_use]
3287    pub fn named(self, name: impl Into<String>) -> Self {
3288        self.add_attributes(Attributes::named(name))
3289    }
3290
3291    #[must_use]
3292    pub fn join<Mat2>(self, flow: Flow<O1, I2, Mat2>) -> Flow<I1, O2, NotUsed>
3293    where
3294        Mat2: Send + 'static,
3295    {
3296        self.top
3297            .via(flow)
3298            .via(self.bottom)
3299            .map_materialized_value(|_| NotUsed)
3300            .with_attributes(self.attributes)
3301    }
3302
3303    #[must_use]
3304    pub fn atop<OO1: Send + 'static, II2: Send + 'static>(
3305        self,
3306        bidi: BidiFlow<O1, OO1, II2, I2>,
3307    ) -> BidiFlow<I1, OO1, II2, O2> {
3308        BidiFlow {
3309            top: self.top.via(bidi.top).map_materialized_value(|_| NotUsed),
3310            bottom: bidi
3311                .bottom
3312                .via(self.bottom)
3313                .map_materialized_value(|_| NotUsed),
3314            attributes: self.attributes.and(bidi.attributes),
3315        }
3316    }
3317
3318    #[must_use]
3319    pub fn reversed(self) -> BidiFlow<I2, O2, I1, O1> {
3320        BidiFlow {
3321            top: self.bottom,
3322            bottom: self.top.map_materialized_value(|_| NotUsed),
3323            attributes: self.attributes,
3324        }
3325    }
3326}
3327
3328impl<In, Out, Resource, Create, F, Close> Drop
3329    for MapWithResourceStream<In, Out, Resource, Create, F, Close>
3330where
3331    Create: Fn() -> StreamResult<Resource>,
3332    F: Fn(&mut Resource, In) -> StreamResult<Out>,
3333    Close: Fn(Resource) -> StreamResult<Option<Out>>,
3334{
3335    fn drop(&mut self) {
3336        let _ = self.close_resource();
3337    }
3338}
3339
3340fn materialize_inner_flow<In, Out, InnerMat>(
3341    flow: Flow<In, Out, InnerMat>,
3342    input: BoxStream<In>,
3343    materializer: &Materializer,
3344) -> StreamResult<(BoxStream<Out>, InnerMat)>
3345where
3346    In: Send + 'static,
3347    Out: Send + 'static,
3348    InnerMat: Send + 'static,
3349{
3350    let mat = (flow.materialize)()?;
3351    let stream = match flow.transform {
3352        FlowTransform::Pure(transform) => transform(input),
3353        FlowTransform::Runtime(transform) => transform(input, materializer)?,
3354    };
3355    Ok((stream, mat))
3356}
3357
3358struct FutureFlowStream<In, Out, InnerMat, F, Fut> {
3359    future: Arc<F>,
3360    materializer: Materializer,
3361    input: Option<BoxStream<In>>,
3362    current: Option<BoxStream<Out>>,
3363    mat_sender: Option<oneshot::Sender<StreamResult<InnerMat>>>,
3364    initialized: bool,
3365    terminated: bool,
3366    _marker: PhantomData<fn() -> Fut>,
3367}
3368
3369impl<In, Out, InnerMat, F, Fut> FutureFlowStream<In, Out, InnerMat, F, Fut>
3370where
3371    In: Send + 'static,
3372    Out: Send + 'static,
3373    InnerMat: Send + 'static,
3374    F: Fn() -> Fut,
3375    Fut: Future<Output = StreamResult<Flow<In, Out, InnerMat>>> + Send + 'static,
3376{
3377    fn complete_mat(&mut self, result: StreamResult<InnerMat>) {
3378        if let Some(sender) = self.mat_sender.take() {
3379            let _ = sender.send(result);
3380        }
3381    }
3382
3383    fn initialize(&mut self) -> StreamResult<()> {
3384        if self.initialized {
3385            return Ok(());
3386        }
3387        self.initialized = true;
3388        let flow = match catch_unwind_failed("future_flow factory", || (self.future)())
3389            .and_then(run_future_inline_or_spawn)
3390        {
3391            Ok(flow) => flow,
3392            Err(error) => {
3393                self.complete_mat(Err(error.clone()));
3394                return Err(error);
3395            }
3396        };
3397        let input = self.input.take().expect("future_flow input available");
3398        match materialize_inner_flow(flow, input, &self.materializer) {
3399            Ok((stream, mat)) => {
3400                self.current = Some(stream);
3401                self.complete_mat(Ok(mat));
3402                Ok(())
3403            }
3404            Err(error) => {
3405                self.complete_mat(Err(error.clone()));
3406                Err(error)
3407            }
3408        }
3409    }
3410}
3411
3412impl<In, Out, InnerMat, F, Fut> Iterator for FutureFlowStream<In, Out, InnerMat, F, Fut>
3413where
3414    In: Send + 'static,
3415    Out: Send + 'static,
3416    InnerMat: Send + 'static,
3417    F: Fn() -> Fut,
3418    Fut: Future<Output = StreamResult<Flow<In, Out, InnerMat>>> + Send + 'static,
3419{
3420    type Item = StreamResult<Out>;
3421
3422    fn next(&mut self) -> Option<Self::Item> {
3423        if self.terminated {
3424            return None;
3425        }
3426        if let Err(error) = self.initialize() {
3427            self.terminated = true;
3428            return Some(Err(error));
3429        }
3430        match self
3431            .current
3432            .as_mut()
3433            .expect("future_flow current stream initialized")
3434            .next()
3435        {
3436            Some(Ok(item)) => Some(Ok(item)),
3437            Some(Err(error)) => {
3438                self.terminated = true;
3439                Some(Err(error))
3440            }
3441            None => {
3442                self.terminated = true;
3443                None
3444            }
3445        }
3446    }
3447}
3448
3449impl<In, Out, InnerMat, F, Fut> Drop for FutureFlowStream<In, Out, InnerMat, F, Fut> {
3450    fn drop(&mut self) {
3451        if !self.initialized
3452            && let Some(sender) = self.mat_sender.take()
3453        {
3454            let _ = sender.send(Err(StreamError::Failed(
3455                "future flow was never materialized".into(),
3456            )));
3457        }
3458    }
3459}
3460
3461struct LazyFutureFlowStream<In, Out, InnerMat, F, Fut> {
3462    create: Arc<F>,
3463    materializer: Materializer,
3464    input: Option<BoxStream<In>>,
3465    current: Option<BoxStream<Out>>,
3466    mat_sender: Option<oneshot::Sender<StreamResult<InnerMat>>>,
3467    initialized: bool,
3468    terminated: bool,
3469    _marker: PhantomData<fn() -> Fut>,
3470}
3471
3472impl<In, Out, InnerMat, F, Fut> LazyFutureFlowStream<In, Out, InnerMat, F, Fut>
3473where
3474    In: Send + 'static,
3475    Out: Send + 'static,
3476    InnerMat: Send + 'static,
3477    F: Fn() -> Fut,
3478    Fut: Future<Output = StreamResult<Flow<In, Out, InnerMat>>> + Send + 'static,
3479{
3480    fn complete_mat(&mut self, result: StreamResult<InnerMat>) {
3481        if let Some(sender) = self.mat_sender.take() {
3482            let _ = sender.send(result);
3483        }
3484    }
3485
3486    fn initialize(&mut self) -> Result<bool, StreamError> {
3487        if self.initialized {
3488            return Ok(true);
3489        }
3490        self.initialized = true;
3491        let first = match self
3492            .input
3493            .as_mut()
3494            .expect("lazy_future_flow input available")
3495            .next()
3496        {
3497            Some(Ok(item)) => item,
3498            Some(Err(error)) => {
3499                self.complete_mat(Err(error.clone()));
3500                return Err(error);
3501            }
3502            None => {
3503                self.complete_mat(Err(StreamError::Failed(
3504                    "lazy flow was never materialized".into(),
3505                )));
3506                self.terminated = true;
3507                return Ok(false);
3508            }
3509        };
3510
3511        let flow = match catch_unwind_failed("lazy_future_flow factory", || (self.create)())
3512            .and_then(run_future_inline_or_spawn)
3513        {
3514            Ok(flow) => flow,
3515            Err(error) => {
3516                self.complete_mat(Err(error.clone()));
3517                return Err(error);
3518            }
3519        };
3520        let input = prepend_first_stream(
3521            first,
3522            self.input
3523                .take()
3524                .expect("lazy_future_flow input available after first element"),
3525        );
3526        match materialize_inner_flow(flow, input, &self.materializer) {
3527            Ok((stream, mat)) => {
3528                self.current = Some(stream);
3529                self.complete_mat(Ok(mat));
3530                Ok(true)
3531            }
3532            Err(error) => {
3533                self.complete_mat(Err(error.clone()));
3534                Err(error)
3535            }
3536        }
3537    }
3538}
3539
3540impl<In, Out, InnerMat, F, Fut> Iterator for LazyFutureFlowStream<In, Out, InnerMat, F, Fut>
3541where
3542    In: Send + 'static,
3543    Out: Send + 'static,
3544    InnerMat: Send + 'static,
3545    F: Fn() -> Fut,
3546    Fut: Future<Output = StreamResult<Flow<In, Out, InnerMat>>> + Send + 'static,
3547{
3548    type Item = StreamResult<Out>;
3549
3550    fn next(&mut self) -> Option<Self::Item> {
3551        if self.terminated {
3552            return None;
3553        }
3554        match self.initialize() {
3555            Ok(true) => {}
3556            Ok(false) => return None,
3557            Err(error) => {
3558                self.terminated = true;
3559                return Some(Err(error));
3560            }
3561        }
3562        match self
3563            .current
3564            .as_mut()
3565            .expect("lazy_future_flow current stream initialized")
3566            .next()
3567        {
3568            Some(Ok(item)) => Some(Ok(item)),
3569            Some(Err(error)) => {
3570                self.terminated = true;
3571                Some(Err(error))
3572            }
3573            None => {
3574                self.terminated = true;
3575                None
3576            }
3577        }
3578    }
3579}
3580
3581impl<In, Out, InnerMat, F, Fut> Drop for LazyFutureFlowStream<In, Out, InnerMat, F, Fut> {
3582    fn drop(&mut self) {
3583        if !self.initialized
3584            && let Some(sender) = self.mat_sender.take()
3585        {
3586            let _ = sender.send(Err(StreamError::Failed(
3587                "lazy flow was never materialized".into(),
3588            )));
3589        }
3590    }
3591}
3592
3593fn prepend_first_stream<In>(first: In, mut rest: BoxStream<In>) -> BoxStream<In>
3594where
3595    In: Send + 'static,
3596{
3597    let mut first = Some(first);
3598    Box::new(std::iter::from_fn(move || {
3599        if let Some(item) = first.take() {
3600            Some(Ok(item))
3601        } else {
3602            rest.next()
3603        }
3604    }))
3605}
3606
3607pub(super) fn poll_once_or_pending<Fut, T>(future: Fut) -> Result<StreamResult<T>, Pin<Box<Fut>>>
3608where
3609    Fut: Future<Output = StreamResult<T>>,
3610{
3611    let mut future = Box::pin(future);
3612    let _guard = stream_tokio_runtime().enter();
3613    let waker = noop_waker();
3614    let mut cx = Context::from_waker(&waker);
3615    match catch_unwind(AssertUnwindSafe(|| future.as_mut().poll(&mut cx))) {
3616        Ok(Poll::Ready(output)) => Ok(output),
3617        Ok(Poll::Pending) => Err(future),
3618        Err(_) => Ok(Err(StreamError::Failed("future task panicked".into()))),
3619    }
3620}
3621
3622pub(super) fn spawn_completion_task<Fut, T, Msg, Map>(
3623    task_id: usize,
3624    future: Pin<Box<Fut>>,
3625    sender: std::sync::mpsc::Sender<(usize, Msg)>,
3626    map: Map,
3627) -> AbortOnDropHandle<()>
3628where
3629    Fut: Future<Output = StreamResult<T>> + Send + 'static,
3630    T: Send + 'static,
3631    Msg: Send + 'static,
3632    Map: FnOnce(StreamResult<T>) -> Msg + Send + 'static,
3633{
3634    spawn_tokio_task(async move {
3635        let result = AssertUnwindSafe(future).catch_unwind().await;
3636        let message = match result {
3637            Ok(output) => map(output),
3638            Err(_) => map(Err(StreamError::Failed("future task panicked".into()))),
3639        };
3640        let _ = sender.send((task_id, message));
3641    })
3642}
3643
3644pub(super) fn recv_completion<Msg>(
3645    receiver: &std::sync::mpsc::Receiver<(usize, Msg)>,
3646) -> Option<(usize, Msg)> {
3647    let mut idle_spins = 0;
3648    loop {
3649        match receiver.try_recv() {
3650            Ok(message) => return Some(message),
3651            Err(std::sync::mpsc::TryRecvError::Disconnected) => return None,
3652            Err(std::sync::mpsc::TryRecvError::Empty) if idle_spins < STREAM_READY_SPINS => {
3653                idle_spins += STREAM_SPIN_BACKOFF;
3654                for _ in 0..STREAM_SPIN_BACKOFF {
3655                    std::hint::spin_loop();
3656                }
3657            }
3658            Err(std::sync::mpsc::TryRecvError::Empty) => {
3659                idle_spins = 0;
3660                match receiver.recv_timeout(STREAM_MAX_PARK) {
3661                    Ok(message) => return Some(message),
3662                    Err(std::sync::mpsc::RecvTimeoutError::Timeout) => {}
3663                    Err(std::sync::mpsc::RecvTimeoutError::Disconnected) => return None,
3664                }
3665            }
3666        }
3667    }
3668}
3669
3670pub(super) fn run_future_inline_or_spawn<Fut, T>(future: Fut) -> StreamResult<T>
3671where
3672    Fut: Future<Output = StreamResult<T>> + Send + 'static,
3673    T: Send + 'static,
3674{
3675    match poll_once_or_pending(future) {
3676        Ok(result) => result,
3677        Err(future) => {
3678            let (sender, receiver) = std::sync::mpsc::channel::<(usize, StreamResult<T>)>();
3679            let _task = spawn_completion_task(0, future, sender, |result| result);
3680            recv_completion(&receiver)
3681                .map(|(_, result)| result)
3682                .unwrap_or_else(|| Err(StreamError::Failed("future task dropped".into())))
3683        }
3684    }
3685}
3686
3687fn map_async_ordered<Out, Next, F, Fut>(
3688    mut input: BoxStream<Out>,
3689    parallelism: usize,
3690    stage: Arc<F>,
3691) -> BoxStream<Next>
3692where
3693    Out: Send + 'static,
3694    Next: Send + 'static,
3695    F: Fn(Out) -> Fut + Send + Sync + 'static,
3696    Fut: Future<Output = StreamResult<Next>> + Send + 'static,
3697{
3698    let (sender, receiver) = std::sync::mpsc::channel::<(usize, StreamResult<Next>)>();
3699    let mut tasks = HashMap::<usize, AbortOnDropHandle<()>>::with_capacity(parallelism);
3700    let mut next_index = 0_usize;
3701    let mut next_to_emit = 0_usize;
3702    let mut completed = BTreeMap::new();
3703    let mut input_done = false;
3704
3705    Box::new(std::iter::from_fn(move || {
3706        loop {
3707            if let Some(result) = completed.remove(&next_to_emit) {
3708                next_to_emit += 1;
3709                return Some(result);
3710            }
3711
3712            while tasks.len() + completed.len() < parallelism && !input_done {
3713                match input.next() {
3714                    Some(Ok(item)) => {
3715                        let index = next_index;
3716                        next_index += 1;
3717                        match poll_once_or_pending(stage(item)) {
3718                            Ok(result) => {
3719                                if index == next_to_emit {
3720                                    next_to_emit += 1;
3721                                    return Some(result);
3722                                }
3723                                completed.insert(index, result);
3724                            }
3725                            Err(future) => {
3726                                tasks.insert(
3727                                    index,
3728                                    spawn_completion_task(
3729                                        index,
3730                                        future,
3731                                        sender.clone(),
3732                                        |result| result,
3733                                    ),
3734                                );
3735                            }
3736                        }
3737                    }
3738                    Some(Err(error)) => {
3739                        completed.insert(next_index, Err(error));
3740                        next_index += 1;
3741                        input_done = true;
3742                    }
3743                    None => input_done = true,
3744                }
3745            }
3746
3747            if let Some(result) = completed.remove(&next_to_emit) {
3748                next_to_emit += 1;
3749                return Some(result);
3750            }
3751
3752            if tasks.is_empty() {
3753                return None;
3754            }
3755
3756            if let Some((index, result)) = recv_completion(&receiver) {
3757                tasks.remove(&index);
3758                if index == next_to_emit {
3759                    next_to_emit += 1;
3760                    return Some(result);
3761                }
3762                completed.insert(index, result);
3763            }
3764        }
3765    }))
3766}
3767
3768fn supervise_async_result<Next>(
3769    result: StreamResult<Next>,
3770    decider: &SupervisionDecider,
3771) -> Option<StreamResult<Next>> {
3772    match result {
3773        Ok(item) => Some(Ok(item)),
3774        Err(error) => match decide_supervision(decider, &error) {
3775            SupervisionDirective::Stop => Some(Err(error)),
3776            SupervisionDirective::Resume | SupervisionDirective::Restart => None,
3777        },
3778    }
3779}
3780
3781fn map_async_ordered_supervised<Out, Next, F, Fut>(
3782    mut input: BoxStream<Out>,
3783    parallelism: usize,
3784    stage: Arc<F>,
3785    decider: SupervisionDecider,
3786) -> BoxStream<Next>
3787where
3788    Out: Send + 'static,
3789    Next: Send + 'static,
3790    F: Fn(Out) -> Fut + Send + Sync + 'static,
3791    Fut: Future<Output = StreamResult<Next>> + Send + 'static,
3792{
3793    let (sender, receiver) = std::sync::mpsc::channel::<(usize, StreamResult<Next>)>();
3794    let mut tasks = HashMap::<usize, AbortOnDropHandle<()>>::with_capacity(parallelism);
3795    let mut next_index = 0_usize;
3796    let mut next_to_emit = 0_usize;
3797    let mut completed = BTreeMap::<usize, Option<StreamResult<Next>>>::new();
3798    let mut input_done = false;
3799
3800    Box::new(std::iter::from_fn(move || {
3801        loop {
3802            while let Some(result) = completed.remove(&next_to_emit) {
3803                next_to_emit += 1;
3804                if let Some(result) = result {
3805                    return Some(result);
3806                }
3807            }
3808
3809            while tasks.len() + completed.len() < parallelism && !input_done {
3810                match input.next() {
3811                    Some(Ok(item)) => {
3812                        let index = next_index;
3813                        next_index += 1;
3814                        match catch_unwind(AssertUnwindSafe(|| poll_once_or_pending(stage(item)))) {
3815                            Ok(Ok(result)) => {
3816                                let result = supervise_async_result(result, &decider);
3817                                if index == next_to_emit {
3818                                    next_to_emit += 1;
3819                                    if let Some(result) = result {
3820                                        return Some(result);
3821                                    }
3822                                } else {
3823                                    completed.insert(index, result);
3824                                }
3825                            }
3826                            Ok(Err(future)) => {
3827                                tasks.insert(
3828                                    index,
3829                                    spawn_completion_task(
3830                                        index,
3831                                        future,
3832                                        sender.clone(),
3833                                        |result| result,
3834                                    ),
3835                                );
3836                            }
3837                            Err(_) => {
3838                                let error = panic_stream_error("map_async callback");
3839                                let result = supervise_async_result(Err(error), &decider);
3840                                if index == next_to_emit {
3841                                    next_to_emit += 1;
3842                                    if let Some(result) = result {
3843                                        return Some(result);
3844                                    }
3845                                } else {
3846                                    completed.insert(index, result);
3847                                }
3848                            }
3849                        }
3850                    }
3851                    Some(Err(error)) => {
3852                        completed.insert(next_index, Some(Err(error)));
3853                        next_index += 1;
3854                        input_done = true;
3855                    }
3856                    None => input_done = true,
3857                }
3858            }
3859
3860            while let Some(result) = completed.remove(&next_to_emit) {
3861                next_to_emit += 1;
3862                if let Some(result) = result {
3863                    return Some(result);
3864                }
3865            }
3866
3867            if tasks.is_empty() {
3868                return None;
3869            }
3870
3871            if let Some((index, result)) = recv_completion(&receiver) {
3872                tasks.remove(&index);
3873                let result = supervise_async_result(result, &decider);
3874                if index == next_to_emit {
3875                    next_to_emit += 1;
3876                    if let Some(result) = result {
3877                        return Some(result);
3878                    }
3879                } else {
3880                    completed.insert(index, result);
3881                }
3882            }
3883        }
3884    }))
3885}
3886
3887fn concat_streams<Out>(streams: Vec<BoxStream<Out>>) -> BoxStream<Out>
3888where
3889    Out: Send + 'static,
3890{
3891    let mut streams: VecDeque<_> = streams.into();
3892    let mut current = streams.pop_front();
3893    Box::new(std::iter::from_fn(move || {
3894        loop {
3895            match current.as_mut() {
3896                Some(stream) => match stream.next() {
3897                    Some(item) => return Some(item),
3898                    None => current = streams.pop_front(),
3899                },
3900                None => return None,
3901            }
3902        }
3903    }))
3904}
3905
3906fn concat_streams_lazy<Out, Mat>(
3907    initial: BoxStream<Out>,
3908    factories: Vec<Arc<dyn SourceFactory<Out, Mat>>>,
3909    materializer: &Materializer,
3910) -> BoxStream<Out>
3911where
3912    Out: Send + 'static,
3913    Mat: Send + 'static,
3914{
3915    let mut current = Some(initial);
3916    let mut remaining: VecDeque<_> = factories.into();
3917    let materializer = materializer.with_name_prefix(materializer.name_prefix().to_owned());
3918    Box::new(std::iter::from_fn(move || {
3919        loop {
3920            match current.as_mut() {
3921                Some(stream) => match stream.next() {
3922                    Some(item) => return Some(item),
3923                    None => {
3924                        current = remaining.pop_front().map(|factory| {
3925                            match factory.create(&materializer) {
3926                                Ok((stream, _)) => stream,
3927                                Err(error) => {
3928                                    Box::new(std::iter::once(Err(error))) as BoxStream<Out>
3929                                }
3930                            }
3931                        });
3932                    }
3933                },
3934                None => return None,
3935            }
3936        }
3937    }))
3938}
3939
3940fn or_else_stream<Out>(mut primary: BoxStream<Out>, mut secondary: BoxStream<Out>) -> BoxStream<Out>
3941where
3942    Out: Send + 'static,
3943{
3944    let mut primary_emitted = false;
3945    let mut using_secondary = false;
3946    Box::new(std::iter::from_fn(move || {
3947        loop {
3948            if using_secondary {
3949                return secondary.next();
3950            }
3951
3952            match primary.next() {
3953                Some(Ok(item)) => {
3954                    primary_emitted = true;
3955                    return Some(Ok(item));
3956                }
3957                Some(Err(error)) => return Some(Err(error)),
3958                None if primary_emitted => return None,
3959                None => using_secondary = true,
3960            }
3961        }
3962    }))
3963}
3964
3965fn interleave_streams<Out>(
3966    streams: Vec<BoxStream<Out>>,
3967    segment_size: usize,
3968    eager_close: bool,
3969) -> BoxStream<Out>
3970where
3971    Out: Send + 'static,
3972{
3973    if segment_size == 0 {
3974        return Box::new(std::iter::once(Err(StreamError::GraphValidation(
3975            "interleave segment size must be greater than zero".into(),
3976        ))));
3977    }
3978
3979    let mut streams: Vec<Option<BoxStream<Out>>> = streams.into_iter().map(Some).collect();
3980    let mut pending: Vec<Option<StreamResult<Out>>> = (0..streams.len()).map(|_| None).collect();
3981    let mut current = 0usize;
3982    let mut emitted = 0usize;
3983    Box::new(std::iter::from_fn(move || {
3984        loop {
3985            if streams.iter().all(Option::is_none) {
3986                return None;
3987            }
3988            if streams[current].is_none() {
3989                match next_active_stream(&streams, current) {
3990                    Some(next) => {
3991                        current = next;
3992                        emitted = 0;
3993                    }
3994                    None => return None,
3995                }
3996            }
3997
3998            let Some(stream) = streams[current].as_mut() else {
3999                continue;
4000            };
4001            let next_item = pending[current].take().or_else(|| stream.next());
4002            match next_item {
4003                Some(Ok(item)) => {
4004                    emitted += 1;
4005                    if emitted == segment_size {
4006                        emitted = 0;
4007                        if let Some(next) = next_active_stream(&streams, current) {
4008                            current = next;
4009                        }
4010                    }
4011                    return Some(Ok(item));
4012                }
4013                Some(Err(error)) => return Some(Err(error)),
4014                None => {
4015                    streams[current] = None;
4016                    emitted = 0;
4017                    if eager_close {
4018                        return None;
4019                    }
4020                    match next_active_stream(&streams, current) {
4021                        Some(next) => current = next,
4022                        None => return None,
4023                    }
4024                }
4025            }
4026        }
4027    }))
4028}
4029
4030fn next_active_stream<Out>(streams: &[Option<BoxStream<Out>>], current: usize) -> Option<usize>
4031where
4032    Out: Send + 'static,
4033{
4034    if streams.is_empty() {
4035        return None;
4036    }
4037    for offset in 1..=streams.len() {
4038        let index = (current + offset) % streams.len();
4039        if streams[index].is_some() {
4040            return Some(index);
4041        }
4042    }
4043    None
4044}
4045
4046fn materialize_side_sink<Out, Mat>(
4047    sink: &Sink<Out, Mat>,
4048    materializer: &Materializer,
4049    buffer: usize,
4050) -> StreamResult<(std::sync::mpsc::SyncSender<StreamResult<Out>>, Mat)>
4051where
4052    Out: Send + 'static,
4053    Mat: Send + 'static,
4054{
4055    let (sender, receiver) = std::sync::mpsc::sync_channel(buffer);
4056    let mat = sink.run(side_receiver_stream(receiver), materializer)?;
4057    Ok((sender, mat))
4058}
4059
4060fn side_receiver_stream<Out>(
4061    receiver: std::sync::mpsc::Receiver<StreamResult<Out>>,
4062) -> BoxStream<Out>
4063where
4064    Out: Send + 'static,
4065{
4066    Box::new(std::iter::from_fn(move || receiver.recv().ok()))
4067}
4068
4069#[derive(Clone)]
4070enum LiveSubstreamTerminal {
4071    Complete,
4072    Error(StreamError),
4073}
4074
4075const LIVE_SUBSTREAM_CAPACITY: usize = 256;
4076const LIVE_SUBSTREAM_BATCH: usize = 64;
4077const FLAT_MAP_MERGE_SUBSTREAM_WINDOW: usize = 64;
4078
4079struct LiveSubstreamShared<T> {
4080    state: Mutex<LiveSubstreamState<T>>,
4081    available: Condvar,
4082    cancelled: Arc<AtomicBool>,
4083    capacity: usize,
4084    batch_size: usize,
4085}
4086
4087struct LiveSubstreamState<T> {
4088    buffered: usize,
4089    batches: VecDeque<VecDeque<T>>,
4090    terminal: Option<LiveSubstreamTerminal>,
4091}
4092
4093impl<T> LiveSubstreamShared<T> {
4094    fn new() -> Arc<Self> {
4095        Self::with_capacity(LIVE_SUBSTREAM_CAPACITY)
4096    }
4097
4098    fn with_capacity(capacity: usize) -> Arc<Self> {
4099        Self::with_batching(capacity, LIVE_SUBSTREAM_BATCH)
4100    }
4101
4102    fn with_batching(capacity: usize, batch_size: usize) -> Arc<Self> {
4103        Arc::new(Self {
4104            state: Mutex::new(LiveSubstreamState {
4105                buffered: 0,
4106                batches: VecDeque::new(),
4107                terminal: None,
4108            }),
4109            available: Condvar::new(),
4110            cancelled: Arc::new(AtomicBool::new(false)),
4111            capacity,
4112            batch_size: batch_size.max(1),
4113        })
4114    }
4115}
4116
4117struct LiveSubstreamStream<T> {
4118    shared: Arc<LiveSubstreamShared<T>>,
4119    completion: Option<StreamCompletion<NotUsed>>,
4120    local_batch: VecDeque<T>,
4121}
4122
4123impl<T> Iterator for LiveSubstreamStream<T> {
4124    type Item = StreamResult<T>;
4125
4126    fn next(&mut self) -> Option<Self::Item> {
4127        if let Some(item) = self.local_batch.pop_front() {
4128            return Some(Ok(item));
4129        }
4130
4131        let mut state = self
4132            .shared
4133            .state
4134            .lock()
4135            .unwrap_or_else(|poison| poison.into_inner());
4136        loop {
4137            if let Some(mut batch) = state.batches.pop_front() {
4138                state.buffered -= batch.len();
4139                drop(state);
4140                self.shared.available.notify_all();
4141                let item = batch.pop_front().expect("live substream batch has an item");
4142                self.local_batch = batch;
4143                return Some(Ok(item));
4144            }
4145            if let Some(terminal) = state.terminal.clone() {
4146                return match terminal {
4147                    LiveSubstreamTerminal::Complete => None,
4148                    LiveSubstreamTerminal::Error(error) => Some(Err(error)),
4149                };
4150            }
4151            state = self
4152                .shared
4153                .available
4154                .wait(state)
4155                .unwrap_or_else(|poison| poison.into_inner());
4156        }
4157    }
4158}
4159
4160impl<T> Drop for LiveSubstreamStream<T> {
4161    fn drop(&mut self) {
4162        self.shared.cancelled.store(true, Ordering::SeqCst);
4163        self.shared.available.notify_all();
4164        let _ = self.completion.take();
4165    }
4166}
4167
4168/// Push a pre-assembled batch of items in a single lock acquisition.
4169/// Items are moved out of `batch`; on success `batch` is empty.
4170/// Returns `Err(())` if the channel is cancelled — in that case `batch` is also cleared.
4171fn push_live_substream_batch<T>(
4172    shared: &Arc<LiveSubstreamShared<T>>,
4173    batch: &mut VecDeque<T>,
4174) -> Result<(), ()> {
4175    while !batch.is_empty() {
4176        let mut state = shared
4177            .state
4178            .lock()
4179            .unwrap_or_else(|poison| poison.into_inner());
4180        while state.buffered >= shared.capacity && state.terminal.is_none() {
4181            if shared.cancelled.load(Ordering::SeqCst) {
4182                batch.clear();
4183                return Err(());
4184            }
4185            state = shared
4186                .available
4187                .wait(state)
4188                .unwrap_or_else(|poison| poison.into_inner());
4189        }
4190        if shared.cancelled.load(Ordering::SeqCst) || state.terminal.is_some() {
4191            batch.clear();
4192            return Err(());
4193        }
4194        let was_empty = state.buffered == 0;
4195        // Push as many items as the channel has capacity for in one lock hold.
4196        while state.buffered < shared.capacity && !batch.is_empty() {
4197            let item = batch.pop_front().expect("batch non-empty");
4198            if let Some(back) = state.batches.back_mut()
4199                && back.len() < shared.batch_size
4200            {
4201                back.push_back(item);
4202            } else {
4203                let mut new_batch =
4204                    VecDeque::with_capacity(shared.batch_size.min(shared.capacity.max(1)));
4205                new_batch.push_back(item);
4206                state.batches.push_back(new_batch);
4207            }
4208            state.buffered += 1;
4209        }
4210        drop(state);
4211        if was_empty {
4212            shared.available.notify_all();
4213        }
4214    }
4215    Ok(())
4216}
4217
4218fn push_live_substream<T>(shared: &Arc<LiveSubstreamShared<T>>, item: T) -> Result<(), T> {
4219    let mut state = shared
4220        .state
4221        .lock()
4222        .unwrap_or_else(|poison| poison.into_inner());
4223    while state.buffered >= shared.capacity && state.terminal.is_none() {
4224        if shared.cancelled.load(Ordering::SeqCst) {
4225            return Err(item);
4226        }
4227        state = shared
4228            .available
4229            .wait(state)
4230            .unwrap_or_else(|poison| poison.into_inner());
4231    }
4232    if shared.cancelled.load(Ordering::SeqCst) || state.terminal.is_some() {
4233        return Err(item);
4234    }
4235    let was_empty = state.buffered == 0;
4236    if let Some(batch) = state.batches.back_mut()
4237        && batch.len() < shared.batch_size
4238    {
4239        batch.push_back(item);
4240    } else {
4241        let mut batch = VecDeque::with_capacity(shared.batch_size.min(shared.capacity.max(1)));
4242        batch.push_back(item);
4243        state.batches.push_back(batch);
4244    }
4245    state.buffered += 1;
4246    drop(state);
4247    if was_empty {
4248        shared.available.notify_all();
4249    }
4250    Ok(())
4251}
4252
4253fn complete_live_substream<T>(shared: &Arc<LiveSubstreamShared<T>>) {
4254    let mut state = shared
4255        .state
4256        .lock()
4257        .unwrap_or_else(|poison| poison.into_inner());
4258    if state.terminal.is_none() {
4259        state.terminal = Some(LiveSubstreamTerminal::Complete);
4260    }
4261    drop(state);
4262    shared.available.notify_all();
4263}
4264
4265fn fail_live_substream<T>(shared: &Arc<LiveSubstreamShared<T>>, error: StreamError) {
4266    let mut state = shared
4267        .state
4268        .lock()
4269        .unwrap_or_else(|poison| poison.into_inner());
4270    if state.terminal.is_none() {
4271        state.terminal = Some(LiveSubstreamTerminal::Error(error));
4272    }
4273    drop(state);
4274    shared.available.notify_all();
4275}
4276
4277fn cancel_live_substream<T>(shared: &Arc<LiveSubstreamShared<T>>) {
4278    fail_live_substream(shared, StreamError::Cancelled);
4279}
4280
4281fn source_from_live_substream<T>(shared: Arc<LiveSubstreamShared<T>>) -> Source<T>
4282where
4283    T: Send + 'static,
4284{
4285    let claimed = Arc::new(AtomicBool::new(false));
4286    Source::from_materialized_factory(move |_materializer| {
4287        if claimed.swap(true, Ordering::SeqCst) {
4288            return Err(StreamError::Failed(
4289                "substream source cannot be materialized more than once".into(),
4290            ));
4291        }
4292        Ok((
4293            Box::new(LiveSubstreamStream {
4294                shared: Arc::clone(&shared),
4295                completion: None,
4296                local_batch: VecDeque::new(),
4297            }) as BoxStream<T>,
4298            NotUsed,
4299        ))
4300    })
4301}
4302
4303fn source_from_once_stream<T>(stream: BoxStream<T>) -> Source<T>
4304where
4305    T: Send + 'static,
4306{
4307    let stream = Arc::new(Mutex::new(Some(stream)));
4308    Source::from_materialized_factory(move |_materializer| {
4309        let mut slot = stream.lock().unwrap_or_else(|poison| poison.into_inner());
4310        let stream = slot.take().ok_or_else(|| {
4311            StreamError::Failed("substream source cannot be materialized more than once".into())
4312        })?;
4313        Ok((stream, NotUsed))
4314    })
4315}
4316
4317fn prefix_and_tail_stream<Out>(
4318    input: BoxStream<Out>,
4319    n: usize,
4320) -> BoxStream<(Vec<Out>, Source<Out>)>
4321where
4322    Out: Send + 'static,
4323{
4324    let mut input = Some(input);
4325    let mut emitted = false;
4326    Box::new(std::iter::from_fn(move || {
4327        if emitted {
4328            return None;
4329        }
4330        emitted = true;
4331
4332        let mut prefix = Vec::with_capacity(n);
4333        while prefix.len() < n {
4334            match input
4335                .as_mut()
4336                .expect("prefix_and_tail input available")
4337                .next()
4338            {
4339                Some(Ok(item)) => prefix.push(item),
4340                Some(Err(error)) => return Some(Err(error)),
4341                None => return Some(Ok((prefix, Source::empty()))),
4342            }
4343        }
4344
4345        Some(Ok((
4346            prefix,
4347            source_from_once_stream(input.take().expect("tail input available")),
4348        )))
4349    }))
4350}
4351
4352struct GroupByWorkerGuard<Key, Out> {
4353    outer: Arc<LiveSubstreamShared<Source<Out>>>,
4354    active: HashMap<Key, Arc<LiveSubstreamShared<Out>>>,
4355    closed: HashSet<Key>,
4356    armed: bool,
4357}
4358
4359impl<Key, Out> GroupByWorkerGuard<Key, Out>
4360where
4361    Key: Eq + Hash,
4362{
4363    fn new(outer: Arc<LiveSubstreamShared<Source<Out>>>) -> Self {
4364        Self {
4365            outer,
4366            active: HashMap::new(),
4367            closed: HashSet::new(),
4368            armed: true,
4369        }
4370    }
4371
4372    fn disarm(&mut self) {
4373        self.armed = false;
4374    }
4375
4376    fn fail_all(&self, error: StreamError)
4377    where
4378        Out: Send + 'static,
4379    {
4380        fail_live_substream(&self.outer, error.clone());
4381        for substream in self.active.values() {
4382            fail_live_substream(substream, error.clone());
4383        }
4384    }
4385
4386    fn complete_all(&self)
4387    where
4388        Out: Send + 'static,
4389    {
4390        complete_live_substream(&self.outer);
4391        for substream in self.active.values() {
4392            complete_live_substream(substream);
4393        }
4394    }
4395
4396    fn cancel_all(&self)
4397    where
4398        Out: Send + 'static,
4399    {
4400        for substream in self.active.values() {
4401            cancel_live_substream(substream);
4402        }
4403    }
4404}
4405
4406impl<Key, Out> Drop for GroupByWorkerGuard<Key, Out> {
4407    fn drop(&mut self) {
4408        if self.armed {
4409            fail_live_substream(&self.outer, StreamError::AbruptTermination);
4410            for substream in self.active.values() {
4411                fail_live_substream(substream, StreamError::AbruptTermination);
4412            }
4413        }
4414    }
4415}
4416
4417fn group_by_flush_write_batch<Key, Out>(
4418    guard: &mut GroupByWorkerGuard<Key, Out>,
4419    wb_key: &mut Option<Key>,
4420    wb_sub: &mut Option<Arc<LiveSubstreamShared<Out>>>,
4421    wb_items: &mut VecDeque<Out>,
4422    allow_closed_substream_recreation: bool,
4423) where
4424    Key: Clone + Eq + Hash,
4425    Out: Send + 'static,
4426{
4427    if wb_items.is_empty() {
4428        return;
4429    }
4430    let key = wb_key.take().expect("wb_key set when wb_items non-empty");
4431    if let Some(ref sub) = *wb_sub {
4432        if push_live_substream_batch(sub, wb_items).is_err() {
4433            guard.active.remove(&key);
4434            if !allow_closed_substream_recreation {
4435                guard.closed.insert(key);
4436            }
4437        }
4438    } else {
4439        wb_items.clear();
4440    }
4441    *wb_sub = None;
4442}
4443
4444fn group_by_stream<Out, Key, F>(
4445    mut input: BoxStream<Out>,
4446    max_substreams: usize,
4447    allow_closed_substream_recreation: bool,
4448    key_fn: Arc<F>,
4449    batch_mode: GroupByBatchMode,
4450    materializer: &Materializer,
4451) -> BoxStream<Source<Out>>
4452where
4453    Out: Clone + Send + 'static,
4454    Key: Clone + Eq + Hash + Send + 'static,
4455    F: Fn(&Out) -> Key + Send + Sync + 'static,
4456{
4457    let outer = LiveSubstreamShared::new();
4458    let worker_outer = Arc::clone(&outer);
4459    let batch_repeated_keys = batch_mode == GroupByBatchMode::FiniteEagerNoRecreate;
4460    let completion = materializer.spawn_stream(move |cancelled| {
4461        let mut guard = GroupByWorkerGuard::new(worker_outer);
4462
4463        // Writer-side batch: accumulate items for the current key before a
4464        // single lock acquisition.  This reduces per-element mutex traffic by
4465        // LIVE_SUBSTREAM_BATCH (64x) on finite eager single-key sources.
4466        // wb_key/wb_sub track which substream owns the pending items.
4467        let mut wb_key: Option<Key> = None;
4468        let mut wb_sub: Option<Arc<LiveSubstreamShared<Out>>> = None;
4469        let mut wb_items: VecDeque<Out> = VecDeque::with_capacity(LIVE_SUBSTREAM_BATCH);
4470
4471        while !cancelled.load(Ordering::SeqCst) {
4472            if guard.outer.cancelled.load(Ordering::SeqCst) {
4473                guard.disarm();
4474                return Ok(NotUsed);
4475            }
4476
4477            match input.next() {
4478                Some(Ok(item)) => {
4479                    let key = match catch_unwind(AssertUnwindSafe(|| key_fn(&item))) {
4480                        Ok(key) => key,
4481                        Err(_panic) => {
4482                            wb_items.clear();
4483                            guard.fail_all(StreamError::AbruptTermination);
4484                            guard.disarm();
4485                            return Ok(NotUsed);
4486                        }
4487                    };
4488
4489                    // Proactive cancel check — remove the substream from active
4490                    // before we attempt a push, so the batching fast-path below
4491                    // only ever sees live substreams.
4492                    if let Some(current) = guard.active.get(&key)
4493                        && current.cancelled.load(Ordering::SeqCst)
4494                    {
4495                        // If this key had a pending write batch, discard it —
4496                        // the consumer already gave up.
4497                        if wb_key.as_ref() == Some(&key) {
4498                            wb_items.clear();
4499                            wb_key = None;
4500                            wb_sub = None;
4501                        }
4502                        guard.active.remove(&key);
4503                        if !allow_closed_substream_recreation {
4504                            guard.closed.insert(key.clone());
4505                        }
4506                    }
4507
4508                    let mut item = item;
4509                    if let Some(current) = guard.active.get(&key).cloned() {
4510                        if !batch_repeated_keys {
4511                            item = match push_live_substream(&current, item) {
4512                                Ok(()) => {
4513                                    continue;
4514                                }
4515                                Err(item) => item,
4516                            };
4517                            guard.active.remove(&key);
4518                            if !allow_closed_substream_recreation {
4519                                guard.closed.insert(key.clone());
4520                                continue;
4521                            }
4522                        } else {
4523                            // Hot path for finite eager sources: active
4524                            // substream for this key exists, and the upstream
4525                            // is known not to block between local batch flushes.
4526                            if wb_key.as_ref() != Some(&key) {
4527                                group_by_flush_write_batch(
4528                                    &mut guard,
4529                                    &mut wb_key,
4530                                    &mut wb_sub,
4531                                    &mut wb_items,
4532                                    allow_closed_substream_recreation,
4533                                );
4534                                wb_key = Some(key.clone());
4535                                wb_sub = Some(current);
4536                            }
4537                            wb_items.push_back(item);
4538                            if wb_items.len() >= LIVE_SUBSTREAM_BATCH {
4539                                group_by_flush_write_batch(
4540                                    &mut guard,
4541                                    &mut wb_key,
4542                                    &mut wb_sub,
4543                                    &mut wb_items,
4544                                    allow_closed_substream_recreation,
4545                                );
4546                            }
4547                            continue;
4548                        }
4549                    }
4550
4551                    // Key not currently active: flush any pending batch for a
4552                    // different key, then proceed with substream creation below.
4553                    if !wb_items.is_empty() {
4554                        group_by_flush_write_batch(
4555                            &mut guard,
4556                            &mut wb_key,
4557                            &mut wb_sub,
4558                            &mut wb_items,
4559                            allow_closed_substream_recreation,
4560                        );
4561                    }
4562
4563                    if guard.closed.contains(&key) {
4564                        continue;
4565                    }
4566
4567                    if guard.active.len() + guard.closed.len() == max_substreams {
4568                        let error = StreamError::Failed(format!(
4569                            "group_by reached max_substreams ({max_substreams})"
4570                        ));
4571                        guard.fail_all(error.clone());
4572                        guard.disarm();
4573                        return Err(error);
4574                    }
4575
4576                    let substream = LiveSubstreamShared::with_capacity(LIVE_SUBSTREAM_CAPACITY);
4577                    push_live_substream(&substream, item)
4578                        .unwrap_or_else(|_| unreachable!("fresh group_by substream"));
4579                    guard.active.insert(key.clone(), Arc::clone(&substream));
4580                    if push_live_substream(
4581                        &guard.outer,
4582                        source_from_live_substream(Arc::clone(&substream)),
4583                    )
4584                    .is_err()
4585                    {
4586                        guard.cancel_all();
4587                        cancel_live_substream(&substream);
4588                        guard.disarm();
4589                        return Ok(NotUsed);
4590                    }
4591                }
4592                Some(Err(error)) => {
4593                    wb_items.clear();
4594                    guard.fail_all(error.clone());
4595                    guard.disarm();
4596                    return Err(error);
4597                }
4598                None => {
4599                    group_by_flush_write_batch(
4600                        &mut guard,
4601                        &mut wb_key,
4602                        &mut wb_sub,
4603                        &mut wb_items,
4604                        allow_closed_substream_recreation,
4605                    );
4606                    guard.complete_all();
4607                    guard.disarm();
4608                    return Ok(NotUsed);
4609                }
4610            }
4611        }
4612
4613        group_by_flush_write_batch(
4614            &mut guard,
4615            &mut wb_key,
4616            &mut wb_sub,
4617            &mut wb_items,
4618            allow_closed_substream_recreation,
4619        );
4620        guard.complete_all();
4621        guard.disarm();
4622        Ok(NotUsed)
4623    });
4624
4625    Box::new(LiveSubstreamStream {
4626        shared: outer,
4627        completion: Some(completion),
4628        local_batch: VecDeque::new(),
4629    })
4630}
4631
4632#[derive(Clone, Copy, Debug)]
4633enum SplitMode {
4634    When,
4635    After,
4636}
4637
4638#[cfg(test)]
4639struct SplitWorkerGuard<Out> {
4640    outer: Arc<LiveSubstreamShared<Source<Out>>>,
4641    current: Option<Arc<LiveSubstreamShared<Out>>>,
4642    armed: bool,
4643    // Writer-side batch: accumulate items before a single lock acquisition.
4644    // Items are flushed when the batch is full (LIVE_SUBSTREAM_BATCH) or at
4645    // every segment boundary, so downstream is never starved.
4646    pending: VecDeque<Out>,
4647}
4648
4649#[cfg(test)]
4650impl<Out> SplitWorkerGuard<Out> {
4651    fn new(outer: Arc<LiveSubstreamShared<Source<Out>>>) -> Self {
4652        Self {
4653            outer,
4654            current: None,
4655            armed: true,
4656            pending: VecDeque::with_capacity(LIVE_SUBSTREAM_BATCH),
4657        }
4658    }
4659
4660    fn disarm(&mut self) {
4661        self.armed = false;
4662    }
4663
4664    /// Open a new inner substream and publish it to the outer channel.
4665    /// Returns `Err(())` if the outer channel was cancelled/closed.
4666    fn open_segment(&mut self) -> Result<(), ()>
4667    where
4668        Out: Send + 'static,
4669    {
4670        let substream = LiveSubstreamShared::new();
4671        self.current = Some(Arc::clone(&substream));
4672        push_live_substream(&self.outer, source_from_live_substream(substream)).map_err(|_| ())
4673    }
4674
4675    /// Flush the pending write batch to the current inner substream.
4676    /// Returns `Err(())` if the substream was cancelled (batch cleared).
4677    fn flush_pending(&mut self) -> Result<(), ()>
4678    where
4679        Out: Send + 'static,
4680    {
4681        if self.pending.is_empty() {
4682            return Ok(());
4683        }
4684        match self.current {
4685            Some(ref current) => push_live_substream_batch(current, &mut self.pending),
4686            None => {
4687                self.pending.clear();
4688                Ok(())
4689            }
4690        }
4691    }
4692
4693    /// Push an item into the currently open inner substream.
4694    /// Items are accumulated locally and flushed in batches for fewer mutex
4695    /// acquisitions.  Returns `Err(())` if the inner substream was cancelled.
4696    fn push_item(&mut self, item: Out) -> Result<(), ()>
4697    where
4698        Out: Send + 'static,
4699    {
4700        if self.current.is_none() {
4701            // No open segment — outer was cancelled.
4702            return Ok(());
4703        }
4704        self.pending.push_back(item);
4705        if self.pending.len() >= LIVE_SUBSTREAM_BATCH {
4706            self.flush_pending()
4707        } else {
4708            Ok(())
4709        }
4710    }
4711
4712    /// Complete the current inner substream (boundary reached).
4713    /// Flushes any pending items before signalling completion.
4714    fn close_segment(&mut self)
4715    where
4716        Out: Send + 'static,
4717    {
4718        let _ = self.flush_pending();
4719        if let Some(current) = self.current.take() {
4720            complete_live_substream(&current);
4721        }
4722    }
4723
4724    fn fail_current(&mut self, error: StreamError)
4725    where
4726        Out: Send + 'static,
4727    {
4728        self.pending.clear();
4729        if let Some(current) = self.current.take() {
4730            fail_live_substream(&current, error);
4731        }
4732    }
4733
4734    fn fail_all(&mut self, error: StreamError)
4735    where
4736        Out: Send + 'static,
4737    {
4738        self.fail_current(error.clone());
4739        fail_live_substream(&self.outer, error);
4740    }
4741
4742    fn complete_all(&mut self)
4743    where
4744        Out: Send + 'static,
4745    {
4746        self.close_segment();
4747        complete_live_substream(&self.outer);
4748    }
4749}
4750
4751#[cfg(test)]
4752impl<Out> Drop for SplitWorkerGuard<Out> {
4753    fn drop(&mut self) {
4754        if self.armed {
4755            self.pending.clear();
4756            if let Some(current) = self.current.take() {
4757                fail_live_substream(&current, StreamError::AbruptTermination);
4758            }
4759            fail_live_substream(&self.outer, StreamError::AbruptTermination);
4760        }
4761    }
4762}
4763
4764fn split_streams<Out, F>(
4765    input: BoxStream<Out>,
4766    mode: SplitMode,
4767    predicate: Arc<F>,
4768    materializer: &Materializer,
4769) -> BoxStream<Source<Out>>
4770where
4771    Out: Clone + Send + 'static,
4772    F: Fn(&Out) -> bool + Send + Sync + 'static,
4773{
4774    #[cfg(test)]
4775    if current_substream_mode() == SubstreamExecutorMode::LegacyOnly {
4776        return split_streams_legacy(input, mode, predicate, materializer);
4777    }
4778    let parent_cancelled = Arc::new(AtomicBool::new(false));
4779    split_streams_fast(input, mode, predicate, parent_cancelled, materializer)
4780}
4781
4782#[cfg(test)]
4783fn split_streams_legacy<Out, F>(
4784    mut input: BoxStream<Out>,
4785    mode: SplitMode,
4786    predicate: Arc<F>,
4787    materializer: &Materializer,
4788) -> BoxStream<Source<Out>>
4789where
4790    Out: Clone + Send + 'static,
4791    F: Fn(&Out) -> bool + Send + Sync + 'static,
4792{
4793    let outer = LiveSubstreamShared::new();
4794    let worker_outer = Arc::clone(&outer);
4795    let completion = materializer.spawn_stream(move |cancelled| {
4796        let mut guard = SplitWorkerGuard::new(Arc::clone(&worker_outer));
4797
4798        while !cancelled.load(Ordering::SeqCst) {
4799            if worker_outer.cancelled.load(Ordering::SeqCst) {
4800                guard.disarm();
4801                return Ok(NotUsed);
4802            }
4803
4804            match input.next() {
4805                Some(Ok(item)) => {
4806                    let split = match catch_unwind(AssertUnwindSafe(|| predicate(&item))) {
4807                        Ok(split) => split,
4808                        Err(_panic) => {
4809                            guard.fail_all(StreamError::AbruptTermination);
4810                            guard.disarm();
4811                            return Ok(NotUsed);
4812                        }
4813                    };
4814
4815                    match mode {
4816                        SplitMode::When => {
4817                            if split && guard.current.is_some() {
4818                                // Close the previous segment before opening this one.
4819                                guard.close_segment();
4820                            }
4821                            // Lazily open a segment on the first item (or after a boundary).
4822                            if guard.current.is_none() && guard.open_segment().is_err() {
4823                                guard.disarm();
4824                                return Ok(NotUsed);
4825                            }
4826                            if guard.push_item(item).is_err() {
4827                                guard.disarm();
4828                                return Ok(NotUsed);
4829                            }
4830                        }
4831                        SplitMode::After => {
4832                            if guard.current.is_none() && guard.open_segment().is_err() {
4833                                guard.disarm();
4834                                return Ok(NotUsed);
4835                            }
4836                            if guard.push_item(item).is_err() {
4837                                guard.disarm();
4838                                return Ok(NotUsed);
4839                            }
4840                            if split {
4841                                guard.close_segment();
4842                            }
4843                        }
4844                    }
4845                }
4846                Some(Err(error)) => {
4847                    guard.fail_all(error.clone());
4848                    guard.disarm();
4849                    return Err(error);
4850                }
4851                None => {
4852                    guard.complete_all();
4853                    guard.disarm();
4854                    return Ok(NotUsed);
4855                }
4856            }
4857        }
4858
4859        guard.complete_all();
4860        guard.disarm();
4861        Ok(NotUsed)
4862    });
4863
4864    Box::new(LiveSubstreamStream {
4865        shared: outer,
4866        completion: Some(completion),
4867        local_batch: VecDeque::new(),
4868    })
4869}
4870
4871// ── Split-sink fast path ──────────────────────────────────────────────────────
4872
4873/// Drives a segment consumer inline in the split worker thread.
4874/// The worker calls `push_item` for each item in the segment, then `complete`/`fail`.
4875trait SplitConsumer<T: Send + 'static>: Send + 'static {
4876    fn push_item(&mut self, item: T) -> StreamResult<()>;
4877    fn complete(self: Box<Self>);
4878    fn fail(self: Box<Self>, error: StreamError);
4879}
4880
4881struct FoldConsumer<T, Acc> {
4882    acc: Option<Acc>,
4883    f: Arc<dyn Fn(Acc, T) -> Acc + Send + Sync + 'static>,
4884    tx: futures::channel::oneshot::Sender<StreamResult<Acc>>,
4885}
4886
4887impl<T: Send + 'static, Acc: Send + 'static> SplitConsumer<T> for FoldConsumer<T, Acc> {
4888    fn push_item(&mut self, item: T) -> StreamResult<()> {
4889        let acc = self.acc.take().expect("FoldConsumer: push after done");
4890        self.acc = Some((self.f)(acc, item));
4891        Ok(())
4892    }
4893    fn complete(mut self: Box<Self>) {
4894        let acc = self
4895            .acc
4896            .take()
4897            .expect("FoldConsumer: complete called twice");
4898        let _ = self.tx.send(Ok(acc));
4899    }
4900    fn fail(mut self: Box<Self>, error: StreamError) {
4901        self.acc = None;
4902        let _ = self.tx.send(Err(error));
4903    }
4904}
4905
4906struct FoldResultConsumer<T, Acc> {
4907    acc: Option<Acc>,
4908    f: Arc<dyn Fn(Acc, T) -> StreamResult<Acc> + Send + Sync + 'static>,
4909    tx: futures::channel::oneshot::Sender<StreamResult<Acc>>,
4910}
4911
4912impl<T: Send + 'static, Acc: Send + 'static> SplitConsumer<T> for FoldResultConsumer<T, Acc> {
4913    fn push_item(&mut self, item: T) -> StreamResult<()> {
4914        let acc = self
4915            .acc
4916            .take()
4917            .expect("FoldResultConsumer: push after done");
4918        match (self.f)(acc, item) {
4919            Ok(new_acc) => {
4920                self.acc = Some(new_acc);
4921                Ok(())
4922            }
4923            Err(e) => Err(e),
4924        }
4925    }
4926    fn complete(mut self: Box<Self>) {
4927        let acc = self
4928            .acc
4929            .take()
4930            .expect("FoldResultConsumer: complete called twice");
4931        let _ = self.tx.send(Ok(acc));
4932    }
4933    fn fail(mut self: Box<Self>, error: StreamError) {
4934        self.acc = None;
4935        let _ = self.tx.send(Err(error));
4936    }
4937}
4938
4939struct CollectConsumer<T> {
4940    items: Vec<T>,
4941    tx: futures::channel::oneshot::Sender<StreamResult<Vec<T>>>,
4942}
4943
4944impl<T: Send + 'static> SplitConsumer<T> for CollectConsumer<T> {
4945    fn push_item(&mut self, item: T) -> StreamResult<()> {
4946        self.items.push(item);
4947        Ok(())
4948    }
4949    fn complete(self: Box<Self>) {
4950        let _ = self.tx.send(Ok(self.items));
4951    }
4952    fn fail(self: Box<Self>, error: StreamError) {
4953        let _ = self.tx.send(Err(error));
4954    }
4955}
4956
4957struct IgnoreConsumer<T> {
4958    tx: futures::channel::oneshot::Sender<StreamResult<NotUsed>>,
4959    _phantom: std::marker::PhantomData<fn(T)>,
4960}
4961
4962impl<T: Send + 'static> SplitConsumer<T> for IgnoreConsumer<T> {
4963    fn push_item(&mut self, _item: T) -> StreamResult<()> {
4964        Ok(())
4965    }
4966    fn complete(self: Box<Self>) {
4967        let _ = self.tx.send(Ok(NotUsed));
4968    }
4969    fn fail(self: Box<Self>, error: StreamError) {
4970        let _ = self.tx.send(Err(error));
4971    }
4972}
4973
4974struct TerminalDrainCancelGuard<T: 'static> {
4975    hook: Option<Arc<dyn TerminalSourceHookDyn<T>>>,
4976}
4977
4978impl<T: 'static> TerminalDrainCancelGuard<T> {
4979    fn new(hook: Arc<dyn TerminalSourceHookDyn<T>>) -> Self {
4980        Self { hook: Some(hook) }
4981    }
4982
4983    fn disarm(&mut self) {
4984        self.hook = None;
4985    }
4986}
4987
4988impl<T: 'static> Drop for TerminalDrainCancelGuard<T> {
4989    fn drop(&mut self) {
4990        if let Some(hook) = self.hook.take() {
4991            hook.cancel_terminal();
4992        }
4993    }
4994}
4995
4996fn terminal_drain_status<T: 'static>(
4997    hook: &Arc<dyn TerminalSourceHookDyn<T>>,
4998    materializer: &Materializer,
4999    cancelled: &Arc<AtomicBool>,
5000) -> StreamResult<()> {
5001    if materializer.is_shutdown() {
5002        hook.cancel_terminal();
5003        Err(StreamError::AbruptTermination)
5004    } else if cancelled.load(Ordering::SeqCst) {
5005        hook.cancel_terminal();
5006        Err(StreamError::Cancelled)
5007    } else {
5008        Ok(())
5009    }
5010}
5011
5012fn wait_direct_terminal_result<T: 'static, Acc: Send + 'static>(
5013    hook: &Arc<dyn TerminalSourceHookDyn<T>>,
5014    materializer: &Materializer,
5015    cancelled: &Arc<AtomicBool>,
5016    receiver: std::sync::mpsc::Receiver<StreamResult<Acc>>,
5017) -> StreamResult<Acc> {
5018    loop {
5019        if materializer.is_shutdown() {
5020            hook.cancel_terminal();
5021            return Err(StreamError::AbruptTermination);
5022        }
5023        if cancelled.load(Ordering::SeqCst) {
5024            hook.cancel_terminal();
5025            return Err(StreamError::Cancelled);
5026        }
5027        match receiver.recv_timeout(Duration::from_millis(1)) {
5028            Ok(result) => return result,
5029            Err(std::sync::mpsc::RecvTimeoutError::Timeout) => {}
5030            Err(std::sync::mpsc::RecvTimeoutError::Disconnected) => {
5031                return Err(StreamError::AbruptTermination);
5032            }
5033        }
5034    }
5035}
5036
5037fn register_direct_terminal<T: Send + 'static, Acc: Send + 'static>(
5038    hook: Arc<dyn TerminalSourceHookDyn<T>>,
5039    materializer: &Materializer,
5040    consumer: Box<dyn TerminalSinkConsumerDyn<T>>,
5041    receiver: std::sync::mpsc::Receiver<StreamResult<Acc>>,
5042) -> Option<StreamResult<Box<dyn Any + Send>>> {
5043    if !hook.supports_direct_terminal() {
5044        return None;
5045    }
5046    let cancellation = StreamCancellation::for_external_completion();
5047    let cancelled = cancellation.cancelled();
5048    if let Err(error) = hook
5049        .try_register_direct_terminal(consumer, Arc::clone(&cancelled))
5050        .expect("direct terminal support advertised")
5051    {
5052        return Some(Err(error));
5053    }
5054    let worker_materializer = materializer.with_name_prefix(materializer.name_prefix().to_owned());
5055    let (tx, rx) = futures::channel::oneshot::channel();
5056    thread::spawn(move || {
5057        let result = wait_direct_terminal_result(&hook, &worker_materializer, &cancelled, receiver);
5058        let _ = tx.send(result);
5059    });
5060    let completion = StreamCompletion::from_receiver(rx, Some(cancellation));
5061    Some(Ok(Box::new(completion)))
5062}
5063
5064struct DirectFoldConsumer<T, Acc> {
5065    acc: Option<Acc>,
5066    f: Arc<dyn Fn(Acc, T) -> Acc + Send + Sync + 'static>,
5067    tx: std::sync::mpsc::Sender<StreamResult<Acc>>,
5068}
5069
5070impl<T: Send + 'static, Acc: Send + 'static> TerminalSinkConsumerDyn<T>
5071    for DirectFoldConsumer<T, Acc>
5072{
5073    fn on_item(&mut self, item: T) -> StreamResult<()> {
5074        let previous = self.acc.take().expect("fold accumulator present");
5075        match catch_unwind(AssertUnwindSafe(|| (self.f)(previous, item))) {
5076            Ok(next) => {
5077                self.acc = Some(next);
5078                Ok(())
5079            }
5080            Err(_) => Err(StreamError::AbruptTermination),
5081        }
5082    }
5083
5084    fn finish(mut self: Box<Self>, result: StreamResult<()>) {
5085        let result = result.map(|()| self.acc.take().expect("fold accumulator present"));
5086        let _ = self.tx.send(result);
5087    }
5088}
5089
5090struct DirectFoldResultConsumer<T, Acc> {
5091    acc: Option<Acc>,
5092    f: Arc<dyn Fn(Acc, T) -> StreamResult<Acc> + Send + Sync + 'static>,
5093    tx: std::sync::mpsc::Sender<StreamResult<Acc>>,
5094}
5095
5096impl<T: Send + 'static, Acc: Send + 'static> TerminalSinkConsumerDyn<T>
5097    for DirectFoldResultConsumer<T, Acc>
5098{
5099    fn on_item(&mut self, item: T) -> StreamResult<()> {
5100        let previous = self.acc.take().expect("fold accumulator present");
5101        let result = catch_unwind(AssertUnwindSafe(|| (self.f)(previous, item)))
5102            .unwrap_or(Err(StreamError::AbruptTermination));
5103        match result {
5104            Ok(next) => {
5105                self.acc = Some(next);
5106                Ok(())
5107            }
5108            Err(error) => Err(error),
5109        }
5110    }
5111
5112    fn finish(mut self: Box<Self>, result: StreamResult<()>) {
5113        let result = result.map(|()| self.acc.take().expect("fold accumulator present"));
5114        let _ = self.tx.send(result);
5115    }
5116}
5117
5118struct DirectCollectConsumer<T> {
5119    items: Vec<T>,
5120    tx: std::sync::mpsc::Sender<StreamResult<Vec<T>>>,
5121}
5122
5123impl<T: Send + 'static> TerminalSinkConsumerDyn<T> for DirectCollectConsumer<T> {
5124    fn on_item(&mut self, item: T) -> StreamResult<()> {
5125        self.items.push(item);
5126        Ok(())
5127    }
5128
5129    fn finish(self: Box<Self>, result: StreamResult<()>) {
5130        let Self { items, tx } = *self;
5131        let result = result.map(|()| items);
5132        let _ = tx.send(result);
5133    }
5134}
5135
5136struct DirectIgnoreConsumer {
5137    tx: std::sync::mpsc::Sender<StreamResult<NotUsed>>,
5138}
5139
5140impl<T: Send + 'static> TerminalSinkConsumerDyn<T> for DirectIgnoreConsumer {
5141    fn on_item(&mut self, _item: T) -> StreamResult<()> {
5142        Ok(())
5143    }
5144
5145    fn finish(self: Box<Self>, result: StreamResult<()>) {
5146        let result = result.map(|()| NotUsed);
5147        let _ = self.tx.send(result);
5148    }
5149}
5150
5151pub(super) struct ForeachDescriptor<T> {
5152    pub(super) f: Arc<dyn Fn(T) + Send + Sync + 'static>,
5153}
5154
5155struct DirectForeachConsumer<T> {
5156    f: Arc<dyn Fn(T) + Send + Sync + 'static>,
5157    tx: std::sync::mpsc::Sender<StreamResult<NotUsed>>,
5158}
5159
5160impl<T: Send + 'static> TerminalSinkConsumerDyn<T> for DirectForeachConsumer<T> {
5161    fn on_item(&mut self, item: T) -> StreamResult<()> {
5162        catch_unwind(AssertUnwindSafe(|| (self.f)(item)))
5163            .map_err(|_| StreamError::AbruptTermination)
5164    }
5165
5166    fn finish(self: Box<Self>, result: StreamResult<()>) {
5167        let result = result.map(|()| NotUsed);
5168        let _ = self.tx.send(result);
5169    }
5170}
5171
5172// ── FoldFastPathDyn implementations (descriptors) ────────────────────────────
5173
5174pub(super) struct FoldDescriptor<T, Acc> {
5175    pub(super) zero: Acc,
5176    pub(super) f: Arc<dyn Fn(Acc, T) -> Acc + Send + Sync + 'static>,
5177}
5178
5179impl<T: Send + 'static, Acc: Clone + Send + Sync + 'static> FoldFastPathDyn<T>
5180    for FoldDescriptor<T, Acc>
5181{
5182    fn try_register(
5183        &self,
5184        hook: Arc<dyn SplitSegmentHookDyn>,
5185    ) -> Option<StreamResult<Box<dyn Any + Send>>> {
5186        let hook_any = hook.as_any_arc();
5187        let slot = hook_any.downcast::<SegmentConsumerSlot<T>>().ok()?;
5188        if slot.claimed.swap(true, Ordering::SeqCst) {
5189            return Some(Err(StreamError::Failed(
5190                "substream source cannot be materialized more than once".into(),
5191            )));
5192        }
5193        let (tx, rx) = futures::channel::oneshot::channel::<StreamResult<Acc>>();
5194        {
5195            let mut state = slot.state.lock().unwrap_or_else(|p| p.into_inner());
5196
5197            // Check if segment is already complete (non-blocking close_segment set terminal).
5198            if let Some(terminal) = state.terminal.take() {
5199                let buffer = std::mem::take(&mut state.buffer);
5200                state.consumer = SegmentConsumer::DirectTaken;
5201                drop(state);
5202                // Process inline in user thread (no cross-thread sync needed).
5203                let mut acc = self.zero.clone();
5204                for item in buffer {
5205                    acc = (self.f)(acc, item);
5206                }
5207                let result = match terminal {
5208                    LiveSubstreamTerminal::Complete => Ok(acc),
5209                    LiveSubstreamTerminal::Error(e) => Err(e),
5210                };
5211                let _ = tx.send(result);
5212                let completion = StreamCompletion::from_receiver(rx, None);
5213                return Some(Ok(Box::new(completion)));
5214            }
5215
5216            // Terminal not set yet; register consumer for later.
5217            let consumer: Box<dyn SplitConsumer<T>> = Box::new(FoldConsumer {
5218                acc: Some(self.zero.clone()),
5219                f: Arc::clone(&self.f),
5220                tx,
5221            });
5222            state.consumer = SegmentConsumer::Direct(consumer);
5223        }
5224        slot.available.notify_all();
5225        let completion = StreamCompletion::from_receiver(rx, None);
5226        Some(Ok(Box::new(completion)))
5227    }
5228
5229    fn supports_terminal_drain(&self) -> bool {
5230        true
5231    }
5232
5233    fn try_register_direct_terminal(
5234        &self,
5235        hook: Arc<dyn TerminalSourceHookDyn<T>>,
5236        materializer: &Materializer,
5237    ) -> Option<StreamResult<Box<dyn Any + Send>>> {
5238        let (tx, rx) = std::sync::mpsc::channel();
5239        let consumer = Box::new(DirectFoldConsumer {
5240            acc: Some(self.zero.clone()),
5241            f: Arc::clone(&self.f),
5242            tx,
5243        });
5244        register_direct_terminal(hook, materializer, consumer, rx)
5245    }
5246
5247    fn try_register_terminal_drain(
5248        &self,
5249        hook: Arc<dyn TerminalSourceHookDyn<T>>,
5250        materializer: &Materializer,
5251    ) -> Option<StreamResult<Box<dyn Any + Send>>> {
5252        let zero = self.zero.clone();
5253        let f = Arc::clone(&self.f);
5254        let worker_materializer =
5255            materializer.with_name_prefix(materializer.name_prefix().to_owned());
5256        let completion = materializer.spawn_stream(move |cancelled| {
5257            let mut guard = TerminalDrainCancelGuard::new(Arc::clone(&hook));
5258            let mut acc = Some(zero);
5259            let mut batch = Vec::new();
5260            loop {
5261                let status =
5262                    hook.drain_terminal_batch(&worker_materializer, &cancelled, &mut batch)?;
5263                for (index, item) in batch.drain(..).enumerate() {
5264                    let previous = acc.take().expect("fold accumulator present");
5265                    acc = Some(f(previous, item));
5266                    if (index + 1) % 64 == 0 {
5267                        terminal_drain_status(&hook, &worker_materializer, &cancelled)?;
5268                    }
5269                }
5270                if matches!(status, TerminalSourceStatus::Completed) {
5271                    guard.disarm();
5272                    return Ok(acc.expect("fold accumulator present"));
5273                }
5274            }
5275        });
5276        Some(Ok(Box::new(completion)))
5277    }
5278}
5279
5280pub(super) struct FoldResultDescriptor<T, Acc> {
5281    pub(super) zero: Acc,
5282    pub(super) f: Arc<dyn Fn(Acc, T) -> StreamResult<Acc> + Send + Sync + 'static>,
5283}
5284
5285impl<T: Send + 'static, Acc: Clone + Send + Sync + 'static> FoldFastPathDyn<T>
5286    for FoldResultDescriptor<T, Acc>
5287{
5288    fn try_register(
5289        &self,
5290        hook: Arc<dyn SplitSegmentHookDyn>,
5291    ) -> Option<StreamResult<Box<dyn Any + Send>>> {
5292        let hook_any = hook.as_any_arc();
5293        let slot = hook_any.downcast::<SegmentConsumerSlot<T>>().ok()?;
5294        if slot.claimed.swap(true, Ordering::SeqCst) {
5295            return Some(Err(StreamError::Failed(
5296                "substream source cannot be materialized more than once".into(),
5297            )));
5298        }
5299        let (tx, rx) = futures::channel::oneshot::channel::<StreamResult<Acc>>();
5300        {
5301            let mut state = slot.state.lock().unwrap_or_else(|p| p.into_inner());
5302
5303            // Check if segment is already complete.
5304            if let Some(terminal) = state.terminal.take() {
5305                let buffer = std::mem::take(&mut state.buffer);
5306                state.consumer = SegmentConsumer::DirectTaken;
5307                drop(state);
5308                let acc = self.zero.clone();
5309                let result = match terminal {
5310                    LiveSubstreamTerminal::Complete => buffer
5311                        .into_iter()
5312                        .try_fold(acc, |a, item| (self.f)(a, item)),
5313                    LiveSubstreamTerminal::Error(e) => Err(e),
5314                };
5315                let _ = tx.send(result);
5316                let completion = StreamCompletion::from_receiver(rx, None);
5317                return Some(Ok(Box::new(completion)));
5318            }
5319
5320            // Terminal not set yet; register consumer for later.
5321            let consumer: Box<dyn SplitConsumer<T>> = Box::new(FoldResultConsumer {
5322                acc: Some(self.zero.clone()),
5323                f: Arc::clone(&self.f),
5324                tx,
5325            });
5326            state.consumer = SegmentConsumer::Direct(consumer);
5327        }
5328        slot.available.notify_all();
5329        let completion = StreamCompletion::from_receiver(rx, None);
5330        Some(Ok(Box::new(completion)))
5331    }
5332
5333    fn supports_terminal_drain(&self) -> bool {
5334        true
5335    }
5336
5337    fn try_register_direct_terminal(
5338        &self,
5339        hook: Arc<dyn TerminalSourceHookDyn<T>>,
5340        materializer: &Materializer,
5341    ) -> Option<StreamResult<Box<dyn Any + Send>>> {
5342        let (tx, rx) = std::sync::mpsc::channel();
5343        let consumer = Box::new(DirectFoldResultConsumer {
5344            acc: Some(self.zero.clone()),
5345            f: Arc::clone(&self.f),
5346            tx,
5347        });
5348        register_direct_terminal(hook, materializer, consumer, rx)
5349    }
5350
5351    fn try_register_terminal_drain(
5352        &self,
5353        hook: Arc<dyn TerminalSourceHookDyn<T>>,
5354        materializer: &Materializer,
5355    ) -> Option<StreamResult<Box<dyn Any + Send>>> {
5356        let zero = self.zero.clone();
5357        let f = Arc::clone(&self.f);
5358        let worker_materializer =
5359            materializer.with_name_prefix(materializer.name_prefix().to_owned());
5360        let completion = materializer.spawn_stream(move |cancelled| {
5361            let mut guard = TerminalDrainCancelGuard::new(Arc::clone(&hook));
5362            let mut acc = Some(zero);
5363            let mut batch = Vec::new();
5364            loop {
5365                let status =
5366                    hook.drain_terminal_batch(&worker_materializer, &cancelled, &mut batch)?;
5367                for (index, item) in batch.drain(..).enumerate() {
5368                    let previous = acc.take().expect("fold accumulator present");
5369                    match f(previous, item) {
5370                        Ok(next) => {
5371                            acc = Some(next);
5372                        }
5373                        Err(error) => return Err(error),
5374                    }
5375                    if (index + 1) % 64 == 0 {
5376                        terminal_drain_status(&hook, &worker_materializer, &cancelled)?;
5377                    }
5378                }
5379                if matches!(status, TerminalSourceStatus::Completed) {
5380                    guard.disarm();
5381                    return Ok(acc.expect("fold accumulator present"));
5382                }
5383            }
5384        });
5385        Some(Ok(Box::new(completion)))
5386    }
5387}
5388
5389pub(super) struct CollectDescriptor<T> {
5390    pub(super) _phantom: std::marker::PhantomData<fn(T)>,
5391}
5392
5393impl<T: Send + 'static> FoldFastPathDyn<T> for CollectDescriptor<T> {
5394    fn try_register(
5395        &self,
5396        hook: Arc<dyn SplitSegmentHookDyn>,
5397    ) -> Option<StreamResult<Box<dyn Any + Send>>> {
5398        let hook_any = hook.as_any_arc();
5399        let slot = hook_any.downcast::<SegmentConsumerSlot<T>>().ok()?;
5400        if slot.claimed.swap(true, Ordering::SeqCst) {
5401            return Some(Err(StreamError::Failed(
5402                "substream source cannot be materialized more than once".into(),
5403            )));
5404        }
5405        let (tx, rx) = futures::channel::oneshot::channel::<StreamResult<Vec<T>>>();
5406        {
5407            let mut state = slot.state.lock().unwrap_or_else(|p| p.into_inner());
5408
5409            // Check if segment is already complete.
5410            if let Some(terminal) = state.terminal.take() {
5411                let buffer = std::mem::take(&mut state.buffer);
5412                state.consumer = SegmentConsumer::DirectTaken;
5413                drop(state);
5414                let items: Vec<T> = buffer.into_iter().collect();
5415                let result = match terminal {
5416                    LiveSubstreamTerminal::Complete => Ok(items),
5417                    LiveSubstreamTerminal::Error(e) => Err(e),
5418                };
5419                let _ = tx.send(result);
5420                let completion = StreamCompletion::from_receiver(rx, None);
5421                return Some(Ok(Box::new(completion)));
5422            }
5423
5424            // Terminal not set yet; register consumer for later.
5425            let consumer: Box<dyn SplitConsumer<T>> = Box::new(CollectConsumer {
5426                items: Vec::new(),
5427                tx,
5428            });
5429            state.consumer = SegmentConsumer::Direct(consumer);
5430        }
5431        slot.available.notify_all();
5432        let completion = StreamCompletion::from_receiver(rx, None);
5433        Some(Ok(Box::new(completion)))
5434    }
5435
5436    fn supports_terminal_drain(&self) -> bool {
5437        true
5438    }
5439
5440    fn try_register_direct_terminal(
5441        &self,
5442        hook: Arc<dyn TerminalSourceHookDyn<T>>,
5443        materializer: &Materializer,
5444    ) -> Option<StreamResult<Box<dyn Any + Send>>> {
5445        let (tx, rx) = std::sync::mpsc::channel();
5446        let consumer = Box::new(DirectCollectConsumer {
5447            items: Vec::new(),
5448            tx,
5449        });
5450        register_direct_terminal(hook, materializer, consumer, rx)
5451    }
5452
5453    fn try_register_terminal_drain(
5454        &self,
5455        hook: Arc<dyn TerminalSourceHookDyn<T>>,
5456        materializer: &Materializer,
5457    ) -> Option<StreamResult<Box<dyn Any + Send>>> {
5458        let worker_materializer =
5459            materializer.with_name_prefix(materializer.name_prefix().to_owned());
5460        let completion = materializer.spawn_stream(move |cancelled| {
5461            let mut guard = TerminalDrainCancelGuard::new(Arc::clone(&hook));
5462            let mut items = Vec::new();
5463            let mut batch = Vec::new();
5464            loop {
5465                let status =
5466                    hook.drain_terminal_batch(&worker_materializer, &cancelled, &mut batch)?;
5467                for (index, item) in batch.drain(..).enumerate() {
5468                    items.push(item);
5469                    if (index + 1) % 64 == 0 {
5470                        terminal_drain_status(&hook, &worker_materializer, &cancelled)?;
5471                    }
5472                }
5473                if matches!(status, TerminalSourceStatus::Completed) {
5474                    guard.disarm();
5475                    return Ok(items);
5476                }
5477            }
5478        });
5479        Some(Ok(Box::new(completion)))
5480    }
5481}
5482
5483pub(super) struct IgnoreDescriptor<T> {
5484    pub(super) _phantom: std::marker::PhantomData<fn(T)>,
5485}
5486
5487impl<T: Send + 'static> FoldFastPathDyn<T> for IgnoreDescriptor<T> {
5488    fn try_register(
5489        &self,
5490        hook: Arc<dyn SplitSegmentHookDyn>,
5491    ) -> Option<StreamResult<Box<dyn Any + Send>>> {
5492        let hook_any = hook.as_any_arc();
5493        let slot = hook_any.downcast::<SegmentConsumerSlot<T>>().ok()?;
5494        if slot.claimed.swap(true, Ordering::SeqCst) {
5495            return Some(Err(StreamError::Failed(
5496                "substream source cannot be materialized more than once".into(),
5497            )));
5498        }
5499        let (tx, rx) = futures::channel::oneshot::channel::<StreamResult<NotUsed>>();
5500        {
5501            let mut state = slot.state.lock().unwrap_or_else(|p| p.into_inner());
5502
5503            // Check if segment is already complete.
5504            if let Some(terminal) = state.terminal.take() {
5505                state.consumer = SegmentConsumer::DirectTaken;
5506                drop(state);
5507                let result = match terminal {
5508                    LiveSubstreamTerminal::Complete => Ok(NotUsed),
5509                    LiveSubstreamTerminal::Error(e) => Err(e),
5510                };
5511                let _ = tx.send(result);
5512                let completion = StreamCompletion::from_receiver(rx, None);
5513                return Some(Ok(Box::new(completion)));
5514            }
5515
5516            // Terminal not set yet; register consumer for later.
5517            let consumer: Box<dyn SplitConsumer<T>> = Box::new(IgnoreConsumer {
5518                tx,
5519                _phantom: std::marker::PhantomData,
5520            });
5521            state.consumer = SegmentConsumer::Direct(consumer);
5522        }
5523        slot.available.notify_all();
5524        let completion = StreamCompletion::from_receiver(rx, None);
5525        Some(Ok(Box::new(completion)))
5526    }
5527
5528    fn supports_terminal_drain(&self) -> bool {
5529        true
5530    }
5531
5532    fn try_register_direct_terminal(
5533        &self,
5534        hook: Arc<dyn TerminalSourceHookDyn<T>>,
5535        materializer: &Materializer,
5536    ) -> Option<StreamResult<Box<dyn Any + Send>>> {
5537        let (tx, rx) = std::sync::mpsc::channel();
5538        let consumer = Box::new(DirectIgnoreConsumer { tx });
5539        register_direct_terminal(hook, materializer, consumer, rx)
5540    }
5541
5542    fn try_register_terminal_drain(
5543        &self,
5544        hook: Arc<dyn TerminalSourceHookDyn<T>>,
5545        materializer: &Materializer,
5546    ) -> Option<StreamResult<Box<dyn Any + Send>>> {
5547        let worker_materializer =
5548            materializer.with_name_prefix(materializer.name_prefix().to_owned());
5549        let completion = materializer.spawn_stream(move |cancelled| {
5550            let mut guard = TerminalDrainCancelGuard::new(Arc::clone(&hook));
5551            let mut batch = Vec::new();
5552            loop {
5553                let status =
5554                    hook.drain_terminal_batch(&worker_materializer, &cancelled, &mut batch)?;
5555                batch.clear();
5556                if matches!(status, TerminalSourceStatus::Completed) {
5557                    guard.disarm();
5558                    return Ok(NotUsed);
5559                }
5560            }
5561        });
5562        Some(Ok(Box::new(completion)))
5563    }
5564}
5565
5566impl<T: Send + 'static> FoldFastPathDyn<T> for ForeachDescriptor<T> {
5567    fn try_register(
5568        &self,
5569        _hook: Arc<dyn SplitSegmentHookDyn>,
5570    ) -> Option<StreamResult<Box<dyn Any + Send>>> {
5571        None
5572    }
5573
5574    fn supports_terminal_drain(&self) -> bool {
5575        true
5576    }
5577
5578    fn try_register_direct_terminal(
5579        &self,
5580        hook: Arc<dyn TerminalSourceHookDyn<T>>,
5581        materializer: &Materializer,
5582    ) -> Option<StreamResult<Box<dyn Any + Send>>> {
5583        let (tx, rx) = std::sync::mpsc::channel();
5584        let consumer = Box::new(DirectForeachConsumer {
5585            f: Arc::clone(&self.f),
5586            tx,
5587        });
5588        register_direct_terminal(hook, materializer, consumer, rx)
5589    }
5590
5591    fn try_register_terminal_drain(
5592        &self,
5593        hook: Arc<dyn TerminalSourceHookDyn<T>>,
5594        materializer: &Materializer,
5595    ) -> Option<StreamResult<Box<dyn Any + Send>>> {
5596        let f = Arc::clone(&self.f);
5597        let worker_materializer =
5598            materializer.with_name_prefix(materializer.name_prefix().to_owned());
5599        let completion = materializer.spawn_stream(move |cancelled| {
5600            let mut guard = TerminalDrainCancelGuard::new(Arc::clone(&hook));
5601            let mut batch = Vec::new();
5602            loop {
5603                let status =
5604                    hook.drain_terminal_batch(&worker_materializer, &cancelled, &mut batch)?;
5605                for (index, item) in batch.drain(..).enumerate() {
5606                    f(item);
5607                    if (index + 1) % 64 == 0 {
5608                        terminal_drain_status(&hook, &worker_materializer, &cancelled)?;
5609                    }
5610                }
5611                if matches!(status, TerminalSourceStatus::Completed) {
5612                    guard.disarm();
5613                    return Ok(NotUsed);
5614                }
5615            }
5616        });
5617        Some(Ok(Box::new(completion)))
5618    }
5619}
5620
5621// ── SegmentConsumerSlot ───────────────────────────────────────────────────────
5622
5623enum SegmentConsumer<T: Send + 'static> {
5624    Pending,
5625    Direct(Box<dyn SplitConsumer<T>>),
5626    DirectTaken,
5627    Fallback,
5628}
5629
5630struct SegmentSlotState<T: Send + 'static> {
5631    buffer: VecDeque<T>,
5632    consumer: SegmentConsumer<T>,
5633    terminal: Option<LiveSubstreamTerminal>,
5634}
5635
5636pub(super) struct SegmentConsumerSlot<T: Send + 'static> {
5637    claimed: Arc<AtomicBool>,
5638    state: Mutex<SegmentSlotState<T>>,
5639    available: Condvar,
5640    parent_cancelled: Arc<AtomicBool>,
5641}
5642
5643impl<T: Clone + Send + 'static> SegmentConsumerSlot<T> {
5644    fn new(parent_cancelled: Arc<AtomicBool>) -> Arc<Self> {
5645        Arc::new(Self {
5646            claimed: Arc::new(AtomicBool::new(false)),
5647            state: Mutex::new(SegmentSlotState {
5648                buffer: VecDeque::new(),
5649                consumer: SegmentConsumer::Pending,
5650                terminal: None,
5651            }),
5652            available: Condvar::new(),
5653            parent_cancelled,
5654        })
5655    }
5656}
5657
5658impl<T: Clone + Send + 'static> SplitSegmentHookDyn for SegmentConsumerSlot<T> {
5659    fn as_any_arc(self: Arc<Self>) -> Arc<dyn Any + Send + Sync> {
5660        self
5661    }
5662}
5663
5664impl<T: Clone + Send + 'static> SourceFactory<T, NotUsed> for SegmentConsumerSlot<T> {
5665    fn create(
5666        self: Arc<Self>,
5667        _materializer: &Materializer,
5668    ) -> StreamResult<(BoxStream<T>, NotUsed)> {
5669        if self.claimed.swap(true, Ordering::SeqCst) {
5670            return Err(StreamError::Failed(
5671                "substream source cannot be materialized more than once".into(),
5672            ));
5673        }
5674        {
5675            let mut state = self.state.lock().unwrap_or_else(|p| p.into_inner());
5676            state.consumer = SegmentConsumer::Fallback;
5677        }
5678        self.available.notify_all();
5679        let stream = FallbackSegmentStream {
5680            slot: Arc::clone(&self),
5681        };
5682        Ok((Box::new(stream), NotUsed))
5683    }
5684}
5685
5686struct FallbackSegmentStream<T: Send + 'static> {
5687    slot: Arc<SegmentConsumerSlot<T>>,
5688}
5689
5690impl<T: Clone + Send + 'static> Iterator for FallbackSegmentStream<T> {
5691    type Item = StreamResult<T>;
5692
5693    fn next(&mut self) -> Option<Self::Item> {
5694        let mut state = self.slot.state.lock().unwrap_or_else(|p| p.into_inner());
5695        loop {
5696            if let Some(item) = state.buffer.pop_front() {
5697                drop(state);
5698                self.slot.available.notify_all();
5699                return Some(Ok(item));
5700            }
5701            if let Some(terminal) = &state.terminal {
5702                return match terminal {
5703                    LiveSubstreamTerminal::Complete => None,
5704                    LiveSubstreamTerminal::Error(e) => Some(Err(e.clone())),
5705                };
5706            }
5707            if self.slot.parent_cancelled.load(Ordering::SeqCst) {
5708                return Some(Err(StreamError::AbruptTermination));
5709            }
5710            state = self
5711                .slot
5712                .available
5713                .wait(state)
5714                .unwrap_or_else(|p| p.into_inner());
5715        }
5716    }
5717}
5718
5719impl<T: Send + 'static> Drop for FallbackSegmentStream<T> {
5720    fn drop(&mut self) {
5721        let mut state = self.slot.state.lock().unwrap_or_else(|p| p.into_inner());
5722        if state.terminal.is_none() {
5723            state.terminal = Some(LiveSubstreamTerminal::Error(StreamError::Cancelled));
5724        }
5725        drop(state);
5726        self.slot.available.notify_all();
5727    }
5728}
5729
5730// ── SplitFastWorkerGuard ──────────────────────────────────────────────────────
5731
5732struct SplitFastWorkerGuard<Out: Send + 'static> {
5733    outer: Arc<LiveSubstreamShared<Source<Out>>>,
5734    current_slot: Option<Arc<SegmentConsumerSlot<Out>>>,
5735    current_consumer: Option<Box<dyn SplitConsumer<Out>>>,
5736    armed: bool,
5737    parent_cancelled: Arc<AtomicBool>,
5738    local_pending: VecDeque<Out>,
5739}
5740
5741impl<Out: Clone + Send + 'static> SplitFastWorkerGuard<Out> {
5742    fn new(
5743        outer: Arc<LiveSubstreamShared<Source<Out>>>,
5744        parent_cancelled: Arc<AtomicBool>,
5745    ) -> Self {
5746        Self {
5747            outer,
5748            current_slot: None,
5749            current_consumer: None,
5750            armed: true,
5751            parent_cancelled,
5752            local_pending: VecDeque::with_capacity(LIVE_SUBSTREAM_BATCH),
5753        }
5754    }
5755
5756    fn disarm(&mut self) {
5757        self.armed = false;
5758    }
5759
5760    fn open_segment(&mut self) -> Result<(), ()> {
5761        let slot = SegmentConsumerSlot::new(Arc::clone(&self.parent_cancelled));
5762        let factory: Arc<dyn SourceFactory<Out, NotUsed>> =
5763            Arc::clone(&slot) as Arc<dyn SourceFactory<Out, NotUsed>>;
5764        let hook: Arc<dyn SplitSegmentHookDyn> = Arc::clone(&slot) as Arc<dyn SplitSegmentHookDyn>;
5765        let source = Source {
5766            factory,
5767            terminal_factory: None,
5768            hints: SourceHints::default(),
5769            attributes: Attributes::default(),
5770            split_hook: Some(hook),
5771        };
5772        self.current_slot = Some(slot);
5773        push_live_substream(&self.outer, source).map_err(|_| ())
5774    }
5775
5776    fn push_item(&mut self, item: Out) -> Result<(), ()> {
5777        // If we already own the consumer, push directly (lock-free).
5778        if let Some(ref mut consumer) = self.current_consumer {
5779            if let Err(e) = consumer.push_item(item) {
5780                let c = self.current_consumer.take().unwrap();
5781                c.fail(e);
5782                return Err(());
5783            }
5784            return Ok(());
5785        }
5786
5787        // Buffer locally; flush in batches via flush_pending.
5788        self.local_pending.push_back(item);
5789        if self.local_pending.len() >= LIVE_SUBSTREAM_BATCH {
5790            self.flush_pending()
5791        } else {
5792            Ok(())
5793        }
5794    }
5795
5796    /// Flush local_pending to the current slot's buffer, or to the direct consumer
5797    /// if one has registered. Blocks only if buffer is full with no consumer.
5798    fn flush_pending(&mut self) -> Result<(), ()> {
5799        if self.local_pending.is_empty() {
5800            return Ok(());
5801        }
5802
5803        // Fast path: consumer already taken
5804        if let Some(ref mut consumer) = self.current_consumer {
5805            for item in self.local_pending.drain(..) {
5806                if let Err(e) = consumer.push_item(item) {
5807                    let c = self.current_consumer.take().unwrap();
5808                    c.fail(e);
5809                    return Err(());
5810                }
5811            }
5812            return Ok(());
5813        }
5814
5815        let slot = match &self.current_slot {
5816            Some(s) => Arc::clone(s),
5817            None => {
5818                self.local_pending.clear();
5819                return Ok(());
5820            }
5821        };
5822
5823        loop {
5824            if self.parent_cancelled.load(Ordering::SeqCst) {
5825                self.local_pending.clear();
5826                return Err(());
5827            }
5828
5829            let mut state = slot.state.lock().unwrap_or_else(|p| p.into_inner());
5830
5831            if matches!(state.consumer, SegmentConsumer::Direct(_)) {
5832                let consumer =
5833                    match std::mem::replace(&mut state.consumer, SegmentConsumer::DirectTaken) {
5834                        SegmentConsumer::Direct(c) => c,
5835                        _ => unreachable!(),
5836                    };
5837                let mut drain_buf = std::mem::take(&mut state.buffer);
5838                drop(state);
5839                let mut consumer_box = consumer;
5840                for item in drain_buf.drain(..) {
5841                    if let Err(e) = consumer_box.push_item(item) {
5842                        consumer_box.fail(e);
5843                        self.local_pending.clear();
5844                        return Err(());
5845                    }
5846                }
5847                for item in self.local_pending.drain(..) {
5848                    if let Err(e) = consumer_box.push_item(item) {
5849                        consumer_box.fail(e);
5850                        return Err(());
5851                    }
5852                }
5853                self.current_consumer = Some(consumer_box);
5854                return Ok(());
5855            }
5856
5857            if matches!(
5858                state.consumer,
5859                SegmentConsumer::Pending | SegmentConsumer::Fallback
5860            ) {
5861                let is_fallback = matches!(state.consumer, SegmentConsumer::Fallback);
5862                let cap = LIVE_SUBSTREAM_CAPACITY.saturating_sub(state.buffer.len());
5863                if cap > 0 {
5864                    let to_flush = cap.min(self.local_pending.len());
5865                    let items: Vec<Out> = self.local_pending.drain(..to_flush).collect();
5866                    state.buffer.extend(items);
5867                    let has_more = !self.local_pending.is_empty();
5868                    drop(state);
5869                    if is_fallback {
5870                        slot.available.notify_all();
5871                    }
5872                    if !has_more {
5873                        return Ok(());
5874                    }
5875                    // Still have local items. Recheck capacity under the lock
5876                    // instead of waiting unconditionally: the fallback consumer
5877                    // may drain and notify in the gap after our notify.
5878                    continue;
5879                } else {
5880                    // Buffer full; wait.
5881                    state = slot
5882                        .available
5883                        .wait(state)
5884                        .unwrap_or_else(|p| p.into_inner());
5885                    continue;
5886                }
5887            }
5888
5889            // DirectTaken: nothing to do.
5890            return Ok(());
5891        }
5892    }
5893
5894    fn close_segment(&mut self) {
5895        // Flush any locally buffered items before closing.
5896        let _ = self.flush_pending();
5897
5898        // If we own the consumer, complete it directly.
5899        if let Some(consumer) = self.current_consumer.take() {
5900            consumer.complete();
5901            self.current_slot = None;
5902            return;
5903        }
5904
5905        let slot = match self.current_slot.take() {
5906            Some(s) => s,
5907            None => return,
5908        };
5909
5910        if self.parent_cancelled.load(Ordering::SeqCst) {
5911            let mut state = slot.state.lock().unwrap_or_else(|p| p.into_inner());
5912            if state.terminal.is_none() {
5913                state.terminal = Some(LiveSubstreamTerminal::Error(StreamError::AbruptTermination));
5914            }
5915            drop(state);
5916            slot.available.notify_all();
5917            return;
5918        }
5919
5920        let mut state = slot.state.lock().unwrap_or_else(|p| p.into_inner());
5921        match std::mem::replace(&mut state.consumer, SegmentConsumer::DirectTaken) {
5922            SegmentConsumer::Direct(mut consumer_box) => {
5923                let mut drain_buf = std::mem::take(&mut state.buffer);
5924                drop(state);
5925                for item in drain_buf.drain(..) {
5926                    if let Err(e) = consumer_box.push_item(item) {
5927                        consumer_box.fail(e);
5928                        return;
5929                    }
5930                }
5931                consumer_box.complete();
5932            }
5933            SegmentConsumer::DirectTaken => {
5934                // Consumer was already taken (shouldn't happen here).
5935            }
5936            SegmentConsumer::Fallback => {
5937                state.consumer = SegmentConsumer::DirectTaken;
5938                if state.terminal.is_none() {
5939                    state.terminal = Some(LiveSubstreamTerminal::Complete);
5940                }
5941                drop(state);
5942                slot.available.notify_all();
5943            }
5944            SegmentConsumer::Pending => {
5945                // Non-blocking: set terminal so consumer processes when it registers.
5946                state.consumer = SegmentConsumer::Pending;
5947                if state.terminal.is_none() {
5948                    state.terminal = Some(LiveSubstreamTerminal::Complete);
5949                }
5950                drop(state);
5951                // Notify so fallback stream / waiting consumer wakes up.
5952                slot.available.notify_all();
5953            }
5954        }
5955    }
5956
5957    fn fail_segment(&mut self, error: StreamError) {
5958        // Drop any locally buffered items on failure.
5959        self.local_pending.clear();
5960
5961        if let Some(consumer) = self.current_consumer.take() {
5962            consumer.fail(error);
5963            self.current_slot = None;
5964            return;
5965        }
5966        let slot = match self.current_slot.take() {
5967            Some(s) => s,
5968            None => return,
5969        };
5970        let mut state = slot.state.lock().unwrap_or_else(|p| p.into_inner());
5971        match std::mem::replace(&mut state.consumer, SegmentConsumer::DirectTaken) {
5972            SegmentConsumer::Direct(c) => {
5973                drop(state);
5974                c.fail(error);
5975            }
5976            _ => {
5977                if state.terminal.is_none() {
5978                    state.terminal = Some(LiveSubstreamTerminal::Error(error));
5979                }
5980                drop(state);
5981                slot.available.notify_all();
5982            }
5983        }
5984    }
5985
5986    fn fail_all(&mut self, error: StreamError) {
5987        self.fail_segment(error.clone());
5988        fail_live_substream(&self.outer, error);
5989    }
5990
5991    fn complete_all(&mut self) {
5992        self.close_segment();
5993        complete_live_substream(&self.outer);
5994    }
5995}
5996
5997impl<Out: Send + 'static> Drop for SplitFastWorkerGuard<Out> {
5998    fn drop(&mut self) {
5999        if self.armed {
6000            self.local_pending.clear(); // Drop pending items on unclean shutdown
6001            if let Some(consumer) = self.current_consumer.take() {
6002                consumer.fail(StreamError::AbruptTermination);
6003            } else if let Some(slot) = self.current_slot.take() {
6004                let mut state = slot.state.lock().unwrap_or_else(|p| p.into_inner());
6005                match std::mem::replace(&mut state.consumer, SegmentConsumer::DirectTaken) {
6006                    SegmentConsumer::Direct(c) => {
6007                        drop(state);
6008                        c.fail(StreamError::AbruptTermination);
6009                    }
6010                    _ => {
6011                        if state.terminal.is_none() {
6012                            state.terminal =
6013                                Some(LiveSubstreamTerminal::Error(StreamError::AbruptTermination));
6014                        }
6015                        drop(state);
6016                        slot.available.notify_all();
6017                    }
6018                }
6019            }
6020            fail_live_substream(&self.outer, StreamError::AbruptTermination);
6021        }
6022    }
6023}
6024
6025fn split_streams_fast<Out, F>(
6026    mut input: BoxStream<Out>,
6027    mode: SplitMode,
6028    predicate: Arc<F>,
6029    parent_cancelled: Arc<AtomicBool>,
6030    materializer: &Materializer,
6031) -> BoxStream<Source<Out>>
6032where
6033    Out: Clone + Send + 'static,
6034    F: Fn(&Out) -> bool + Send + Sync + 'static,
6035{
6036    let outer = LiveSubstreamShared::new();
6037    let worker_outer = Arc::clone(&outer);
6038    let completion = materializer.spawn_stream(move |cancelled| {
6039        let mut guard =
6040            SplitFastWorkerGuard::new(Arc::clone(&worker_outer), Arc::clone(&parent_cancelled));
6041
6042        while !cancelled.load(Ordering::SeqCst) {
6043            if worker_outer.cancelled.load(Ordering::SeqCst) {
6044                guard.disarm();
6045                return Ok(NotUsed);
6046            }
6047
6048            match input.next() {
6049                Some(Ok(item)) => {
6050                    let split = match catch_unwind(AssertUnwindSafe(|| predicate(&item))) {
6051                        Ok(s) => s,
6052                        Err(_) => {
6053                            guard.fail_all(StreamError::AbruptTermination);
6054                            guard.disarm();
6055                            return Ok(NotUsed);
6056                        }
6057                    };
6058
6059                    match mode {
6060                        SplitMode::When => {
6061                            if split
6062                                && (guard.current_slot.is_some()
6063                                    || guard.current_consumer.is_some())
6064                            {
6065                                guard.close_segment();
6066                            }
6067                            if guard.current_slot.is_none()
6068                                && guard.current_consumer.is_none()
6069                                && guard.open_segment().is_err()
6070                            {
6071                                guard.disarm();
6072                                return Ok(NotUsed);
6073                            }
6074                            if guard.push_item(item).is_err() {
6075                                guard.disarm();
6076                                return Ok(NotUsed);
6077                            }
6078                        }
6079                        SplitMode::After => {
6080                            if guard.current_slot.is_none()
6081                                && guard.current_consumer.is_none()
6082                                && guard.open_segment().is_err()
6083                            {
6084                                guard.disarm();
6085                                return Ok(NotUsed);
6086                            }
6087                            if guard.push_item(item).is_err() {
6088                                guard.disarm();
6089                                return Ok(NotUsed);
6090                            }
6091                            if split {
6092                                guard.close_segment();
6093                            }
6094                        }
6095                    }
6096                }
6097                Some(Err(error)) => {
6098                    guard.fail_all(error.clone());
6099                    guard.disarm();
6100                    return Err(error);
6101                }
6102                None => {
6103                    guard.complete_all();
6104                    guard.disarm();
6105                    return Ok(NotUsed);
6106                }
6107            }
6108        }
6109        guard.complete_all();
6110        guard.disarm();
6111        Ok(NotUsed)
6112    });
6113
6114    Box::new(LiveSubstreamStream {
6115        shared: outer,
6116        completion: Some(completion),
6117        local_batch: VecDeque::new(),
6118    })
6119}
6120
6121// ── Legacy flat_map_merge (oracle; only compiled in test builds) ──────────────
6122#[cfg(test)]
6123struct FlatMapMergeShared<T> {
6124    state: Mutex<FlatMapMergeState<T>>,
6125    available: Condvar,
6126    cancelled: Arc<AtomicBool>,
6127}
6128
6129#[cfg(test)]
6130struct FlatMapMergeState<T> {
6131    queued: VecDeque<(usize, T)>,
6132    window: HashMap<usize, usize>,
6133    active_streams: usize,
6134    input_done: bool,
6135    terminal: Option<StreamError>,
6136}
6137
6138#[cfg(test)]
6139impl<T> FlatMapMergeShared<T> {
6140    fn new() -> Arc<Self> {
6141        Arc::new(Self {
6142            state: Mutex::new(FlatMapMergeState {
6143                queued: VecDeque::new(),
6144                window: HashMap::new(),
6145                active_streams: 0,
6146                input_done: false,
6147                terminal: None,
6148            }),
6149            available: Condvar::new(),
6150            cancelled: Arc::new(AtomicBool::new(false)),
6151        })
6152    }
6153}
6154
6155#[cfg(test)]
6156struct FlatMapMergeStream<T> {
6157    shared: Arc<FlatMapMergeShared<T>>,
6158    completion: Option<StreamCompletion<NotUsed>>,
6159}
6160
6161#[cfg(test)]
6162impl<T> Iterator for FlatMapMergeStream<T> {
6163    type Item = StreamResult<T>;
6164
6165    fn next(&mut self) -> Option<Self::Item> {
6166        let mut state = self
6167            .shared
6168            .state
6169            .lock()
6170            .unwrap_or_else(|poison| poison.into_inner());
6171        loop {
6172            if let Some((stream_id, item)) = state.queued.pop_front() {
6173                // Decrement per-stream in-flight count inside the same lock —
6174                // no second mutex needed.
6175                if let Some(count) = state.window.get_mut(&stream_id) {
6176                    *count -= 1;
6177                    if *count == 0 {
6178                        state.window.remove(&stream_id);
6179                    }
6180                }
6181                drop(state);
6182                self.shared.available.notify_all();
6183                return Some(Ok(item));
6184            }
6185            if let Some(error) = state.terminal.clone() {
6186                return Some(Err(error));
6187            }
6188            if state.input_done && state.active_streams == 0 {
6189                return None;
6190            }
6191            state = self
6192                .shared
6193                .available
6194                .wait(state)
6195                .unwrap_or_else(|poison| poison.into_inner());
6196        }
6197    }
6198}
6199
6200#[cfg(test)]
6201impl<T> Drop for FlatMapMergeStream<T> {
6202    fn drop(&mut self) {
6203        self.shared.cancelled.store(true, Ordering::SeqCst);
6204        self.shared.available.notify_all();
6205        let _ = self.completion.take();
6206    }
6207}
6208
6209#[cfg(test)]
6210fn flat_map_merge_register_stream<T>(shared: &Arc<FlatMapMergeShared<T>>) {
6211    let mut state = shared
6212        .state
6213        .lock()
6214        .unwrap_or_else(|poison| poison.into_inner());
6215    state.active_streams += 1;
6216    drop(state);
6217    shared.available.notify_all();
6218}
6219
6220#[cfg(test)]
6221fn flat_map_merge_finish_stream<T>(
6222    shared: &Arc<FlatMapMergeShared<T>>,
6223    stream_id: usize,
6224    terminal: Result<(), StreamError>,
6225) {
6226    let mut state = shared
6227        .state
6228        .lock()
6229        .unwrap_or_else(|poison| poison.into_inner());
6230    state.active_streams = state.active_streams.saturating_sub(1);
6231    if let Err(error) = terminal
6232        && state.terminal.is_none()
6233    {
6234        state.terminal = Some(error);
6235    }
6236    // Ensure the window entry exists so the consumer can decrement it.
6237    state.window.entry(stream_id).or_default();
6238    drop(state);
6239    shared.available.notify_all();
6240}
6241
6242#[cfg(test)]
6243fn flat_map_merge_push<T>(
6244    shared: &Arc<FlatMapMergeShared<T>>,
6245    stream_id: usize,
6246    item: T,
6247) -> Result<(), T> {
6248    // Single lock: the window map lives inside FlatMapMergeState, so both the
6249    // per-lane backpressure check and the enqueue happen under one mutex.
6250    let mut state = shared
6251        .state
6252        .lock()
6253        .unwrap_or_else(|poison| poison.into_inner());
6254    while state.window.get(&stream_id).copied().unwrap_or(0) >= FLAT_MAP_MERGE_SUBSTREAM_WINDOW
6255        && state.terminal.is_none()
6256    {
6257        if shared.cancelled.load(Ordering::SeqCst) {
6258            return Err(item);
6259        }
6260        state = shared
6261            .available
6262            .wait(state)
6263            .unwrap_or_else(|poison| poison.into_inner());
6264    }
6265    if shared.cancelled.load(Ordering::SeqCst) || state.terminal.is_some() {
6266        return Err(item);
6267    }
6268    *state.window.entry(stream_id).or_insert(0) += 1;
6269    let was_empty = state.queued.is_empty();
6270    state.queued.push_back((stream_id, item));
6271    drop(state);
6272    if was_empty {
6273        shared.available.notify_all();
6274    }
6275    Ok(())
6276}
6277
6278#[cfg(test)]
6279struct FlatMapMergeCoordinatorGuard<T> {
6280    shared: Arc<FlatMapMergeShared<T>>,
6281    armed: bool,
6282}
6283
6284#[cfg(test)]
6285impl<T> FlatMapMergeCoordinatorGuard<T> {
6286    fn new(shared: Arc<FlatMapMergeShared<T>>) -> Self {
6287        Self {
6288            shared,
6289            armed: true,
6290        }
6291    }
6292
6293    fn finish(&mut self) {
6294        let mut state = self
6295            .shared
6296            .state
6297            .lock()
6298            .unwrap_or_else(|poison| poison.into_inner());
6299        state.input_done = true;
6300        drop(state);
6301        self.shared.available.notify_all();
6302        self.armed = false;
6303    }
6304}
6305
6306#[cfg(test)]
6307impl<T> Drop for FlatMapMergeCoordinatorGuard<T> {
6308    fn drop(&mut self) {
6309        if self.armed {
6310            let mut state = self
6311                .shared
6312                .state
6313                .lock()
6314                .unwrap_or_else(|poison| poison.into_inner());
6315            if state.terminal.is_none() {
6316                state.terminal = Some(StreamError::AbruptTermination);
6317            }
6318            state.input_done = true;
6319            drop(state);
6320            self.shared.cancelled.store(true, Ordering::SeqCst);
6321            self.shared.available.notify_all();
6322        }
6323    }
6324}
6325
6326#[cfg(test)]
6327struct FlatMapMergeWorkerGuard<T> {
6328    shared: Arc<FlatMapMergeShared<T>>,
6329    stream_id: usize,
6330    armed: bool,
6331}
6332
6333#[cfg(test)]
6334impl<T> FlatMapMergeWorkerGuard<T> {
6335    fn new(shared: Arc<FlatMapMergeShared<T>>, stream_id: usize) -> Self {
6336        Self {
6337            shared,
6338            stream_id,
6339            armed: true,
6340        }
6341    }
6342
6343    fn finish(&mut self, terminal: Result<(), StreamError>) {
6344        flat_map_merge_finish_stream(&self.shared, self.stream_id, terminal);
6345        self.armed = false;
6346    }
6347}
6348
6349#[cfg(test)]
6350impl<T> Drop for FlatMapMergeWorkerGuard<T> {
6351    fn drop(&mut self) {
6352        if self.armed {
6353            flat_map_merge_finish_stream(
6354                &self.shared,
6355                self.stream_id,
6356                Err(StreamError::AbruptTermination),
6357            );
6358            self.shared.cancelled.store(true, Ordering::SeqCst);
6359        }
6360    }
6361}
6362
6363// ── Mode hook (test-only) ────────────────────────────────────────────────────
6364
6365#[cfg(test)]
6366#[derive(Clone, Copy, PartialEq, Eq, Debug)]
6367pub(crate) enum SubstreamExecutorMode {
6368    Auto,
6369    LegacyOnly,
6370    ReadyRingOnly,
6371    SplitSinkOnly,
6372}
6373
6374#[cfg(test)]
6375thread_local! {
6376    static SUBSTREAM_EXECUTOR_MODE: std::cell::Cell<SubstreamExecutorMode> =
6377        const { std::cell::Cell::new(SubstreamExecutorMode::Auto) };
6378}
6379
6380#[cfg(test)]
6381pub(crate) fn with_substream_mode<R>(mode: SubstreamExecutorMode, f: impl FnOnce() -> R) -> R {
6382    SUBSTREAM_EXECUTOR_MODE.with(|m| {
6383        let old = m.get();
6384        m.set(mode);
6385        let result = f();
6386        m.set(old);
6387        result
6388    })
6389}
6390
6391#[cfg(test)]
6392fn current_substream_mode() -> SubstreamExecutorMode {
6393    SUBSTREAM_EXECUTOR_MODE.with(|m| m.get())
6394}
6395
6396// ── Dispatcher ───────────────────────────────────────────────────────────────
6397
6398fn flat_map_merge_stream<Out, Next, NextMat, F>(
6399    input: BoxStream<Out>,
6400    breadth: usize,
6401    stage: Arc<F>,
6402    materializer: &Materializer,
6403) -> BoxStream<Next>
6404where
6405    Out: Send + 'static,
6406    Next: Send + 'static,
6407    NextMat: Send + 'static,
6408    F: Fn(Out) -> Source<Next, NextMat> + Send + Sync + 'static,
6409{
6410    #[cfg(test)]
6411    if current_substream_mode() == SubstreamExecutorMode::LegacyOnly {
6412        return flat_map_merge_stream_legacy(input, breadth, stage, materializer);
6413    }
6414    flat_map_merge_stream_ready(input, breadth, stage, materializer)
6415}
6416
6417// ── Legacy implementation (oracle / fallback) ─────────────────────────────────
6418
6419#[cfg(test)]
6420fn flat_map_merge_stream_legacy<Out, Next, NextMat, F>(
6421    mut input: BoxStream<Out>,
6422    breadth: usize,
6423    stage: Arc<F>,
6424    materializer: &Materializer,
6425) -> BoxStream<Next>
6426where
6427    Out: Send + 'static,
6428    Next: Send + 'static,
6429    NextMat: Send + 'static,
6430    F: Fn(Out) -> Source<Next, NextMat> + Send + Sync + 'static,
6431{
6432    let worker_materializer = materializer.with_name_prefix(materializer.name_prefix().to_owned());
6433    let shared = FlatMapMergeShared::new();
6434    let worker_shared = Arc::clone(&shared);
6435    let completion = materializer.spawn_stream(move |cancelled| {
6436        let mut next_id = 0usize;
6437        let mut guard = FlatMapMergeCoordinatorGuard::new(worker_shared);
6438        let mut workers = HashMap::<usize, StreamCompletion<NotUsed>>::with_capacity(breadth);
6439
6440        while !cancelled.load(Ordering::SeqCst) && !guard.shared.cancelled.load(Ordering::SeqCst) {
6441            workers.retain(|_, completion| completion.try_wait().is_none());
6442            {
6443                let mut state = guard
6444                    .shared
6445                    .state
6446                    .lock()
6447                    .unwrap_or_else(|poison| poison.into_inner());
6448                while state.active_streams >= breadth
6449                    && state.terminal.is_none()
6450                    && !cancelled.load(Ordering::SeqCst)
6451                    && !guard.shared.cancelled.load(Ordering::SeqCst)
6452                {
6453                    state = guard
6454                        .shared
6455                        .available
6456                        .wait(state)
6457                        .unwrap_or_else(|poison| poison.into_inner());
6458                }
6459                if state.terminal.is_some() {
6460                    let error = state.terminal.clone().expect("terminal checked above");
6461                    drop(state);
6462                    guard.finish();
6463                    return Err(error);
6464                }
6465            }
6466
6467            match input.next() {
6468                Some(Ok(item)) => {
6469                    let source = stage(item);
6470                    let (mut stream, _) =
6471                        match Arc::clone(&source.factory).create(&worker_materializer) {
6472                            Ok(parts) => parts,
6473                            Err(error) => {
6474                                let mut state = guard
6475                                    .shared
6476                                    .state
6477                                    .lock()
6478                                    .unwrap_or_else(|poison| poison.into_inner());
6479                                if state.terminal.is_none() {
6480                                    state.terminal = Some(error.clone());
6481                                }
6482                                drop(state);
6483                                guard.shared.available.notify_all();
6484                                guard.finish();
6485                                return Err(error);
6486                            }
6487                        };
6488                    let stream_id = next_id;
6489                    next_id += 1;
6490                    flat_map_merge_register_stream(&guard.shared);
6491                    let worker_shared = Arc::clone(&guard.shared);
6492                    workers.insert(
6493                        stream_id,
6494                        worker_materializer.spawn_stream(move |inner_cancelled| {
6495                            let mut worker_guard =
6496                                FlatMapMergeWorkerGuard::new(Arc::clone(&worker_shared), stream_id);
6497                            while !inner_cancelled.load(Ordering::SeqCst)
6498                                && !worker_shared.cancelled.load(Ordering::SeqCst)
6499                            {
6500                                match stream.next() {
6501                                    Some(Ok(item)) => {
6502                                        if flat_map_merge_push(&worker_shared, stream_id, item)
6503                                            .is_err()
6504                                        {
6505                                            worker_guard.finish(Ok(()));
6506                                            return Ok(NotUsed);
6507                                        }
6508                                    }
6509                                    Some(Err(error)) => {
6510                                        worker_guard.finish(Err(error.clone()));
6511                                        return Err(error);
6512                                    }
6513                                    None => {
6514                                        worker_guard.finish(Ok(()));
6515                                        return Ok(NotUsed);
6516                                    }
6517                                }
6518                            }
6519                            worker_guard.finish(Ok(()));
6520                            Ok(NotUsed)
6521                        }),
6522                    );
6523                }
6524                Some(Err(error)) => {
6525                    let mut state = guard
6526                        .shared
6527                        .state
6528                        .lock()
6529                        .unwrap_or_else(|poison| poison.into_inner());
6530                    if state.terminal.is_none() {
6531                        state.terminal = Some(error.clone());
6532                    }
6533                    drop(state);
6534                    guard.shared.available.notify_all();
6535                    guard.finish();
6536                    return Err(error);
6537                }
6538                None => {
6539                    guard.finish();
6540                    break;
6541                }
6542            }
6543        }
6544
6545        if guard.armed {
6546            guard.finish();
6547        }
6548        // Drain workers.  The critical invariant: a worker calls
6549        // `flat_map_merge_finish_stream` (decrementing `active_streams` and
6550        // notifying `available`) *before* its spawned task fully exits.  That
6551        // means `try_wait()` may return `None` for a worker that has already
6552        // decremented `active_streams`.  We must therefore check
6553        // `active_streams == 0` *before* sleeping — not only when `workers` is
6554        // empty — otherwise we can miss the final notification and hang forever.
6555        while !cancelled.load(Ordering::SeqCst) && !guard.shared.cancelled.load(Ordering::SeqCst) {
6556            workers.retain(|_, completion| completion.try_wait().is_none());
6557            // Lock state once, check done-condition, and only sleep if we are not
6558            // done yet.  Checking inside the lock prevents the lost-wakeup: any
6559            // worker that fires after we take the lock but before `wait` is called
6560            // will be seen either through `active_streams` here or will produce a
6561            // spurious wakeup that re-evaluates the loop condition.
6562            let state = guard
6563                .shared
6564                .state
6565                .lock()
6566                .unwrap_or_else(|poison| poison.into_inner());
6567            if state.active_streams == 0 {
6568                return Ok(NotUsed);
6569            }
6570            drop(
6571                guard
6572                    .shared
6573                    .available
6574                    .wait(state)
6575                    .unwrap_or_else(|poison| poison.into_inner()),
6576            );
6577        }
6578        Ok(NotUsed)
6579    });
6580
6581    Box::new(FlatMapMergeStream {
6582        shared,
6583        completion: Some(completion),
6584    })
6585}
6586
6587// ── Ready-ring implementation ─────────────────────────────────────────────────
6588
6589const FLAT_MAP_MERGE_READY_BATCH: usize = 16;
6590
6591/// Inline-drain budget: sources whose `max_success_items` does not exceed this
6592/// bound are drained in the coordinator thread without spawning a worker.
6593/// Must be <= FLAT_MAP_MERGE_SUBSTREAM_WINDOW.
6594const FLAT_MAP_MERGE_INLINE_MICRO_MAX: usize = FLAT_MAP_MERGE_READY_BATCH;
6595
6596struct FlatMapMergeReadyShared<T> {
6597    coordinator: Mutex<FlatMapMergeReadyState<T>>,
6598    available: Condvar,
6599    cancelled: Arc<AtomicBool>,
6600}
6601
6602struct FlatMapMergeReadyState<T> {
6603    lanes: HashMap<usize, Arc<FlatMapMergeLane<T>>>,
6604    ready: std::collections::VecDeque<usize>,
6605    queued_items: usize,
6606    active_streams: usize,
6607    input_done: bool,
6608    terminal: Option<StreamError>,
6609    generation: u64,
6610}
6611
6612struct FlatMapMergeLane<T> {
6613    state: Mutex<FlatMapMergeLaneState<T>>,
6614    space_available: Condvar,
6615}
6616
6617struct FlatMapMergeLaneState<T> {
6618    buffer: std::collections::VecDeque<T>,
6619    in_ready_ring: bool,
6620    publishing: bool,
6621    closed: bool,
6622}
6623
6624struct FlatMapMergeReadyStream<T> {
6625    shared: Arc<FlatMapMergeReadyShared<T>>,
6626    completion: Option<StreamCompletion<NotUsed>>,
6627    local_batch: std::collections::VecDeque<T>,
6628}
6629
6630impl<T: Send + 'static> Iterator for FlatMapMergeReadyStream<T> {
6631    type Item = StreamResult<T>;
6632
6633    fn next(&mut self) -> Option<Self::Item> {
6634        // Step 1: return from local batch first.
6635        if let Some(item) = self.local_batch.pop_front() {
6636            return Some(Ok(item));
6637        }
6638
6639        // Step 2: lock coordinator and loop.
6640        let mut coord = self
6641            .shared
6642            .coordinator
6643            .lock()
6644            .unwrap_or_else(|p| p.into_inner());
6645        loop {
6646            // Step 3: terminal wins — fail-fast, do not drain remaining lane items.
6647            if let Some(error) = coord.terminal.clone() {
6648                return Some(Err(error));
6649            }
6650
6651            // Step 4: pop a ready lane id and drop coordinator lock.
6652            if let Some(lane_id) = coord.ready.pop_front() {
6653                let lane = coord.lanes.get(&lane_id).cloned();
6654                drop(coord);
6655
6656                if let Some(lane) = lane {
6657                    // Step 5: lock lane, drain up to BATCH items.
6658                    let mut lane_state = lane.state.lock().unwrap_or_else(|p| p.into_inner());
6659                    while lane_state.publishing {
6660                        lane_state = lane
6661                            .space_available
6662                            .wait(lane_state)
6663                            .unwrap_or_else(|p| p.into_inner());
6664                    }
6665                    let drain_n = lane_state.buffer.len().min(FLAT_MAP_MERGE_READY_BATCH);
6666                    let mut batch = std::collections::VecDeque::with_capacity(drain_n);
6667                    for _ in 0..drain_n {
6668                        if let Some(item) = lane_state.buffer.pop_front() {
6669                            batch.push_back(item);
6670                        }
6671                    }
6672                    let still_has_items = !lane_state.buffer.is_empty();
6673                    let is_closed = lane_state.closed;
6674                    if !still_has_items {
6675                        lane_state.in_ready_ring = false;
6676                    }
6677                    // Step 5 end: drop lane lock.
6678                    drop(lane_state);
6679                    // Notify workers blocked on a full lane.
6680                    if !batch.is_empty() {
6681                        lane.space_available.notify_all();
6682                    }
6683
6684                    // Step 7: re-lock coordinator, update bookkeeping.
6685                    let freed = batch.len();
6686                    let mut coord = self
6687                        .shared
6688                        .coordinator
6689                        .lock()
6690                        .unwrap_or_else(|p| p.into_inner());
6691                    coord.queued_items = coord.queued_items.saturating_sub(freed);
6692                    if still_has_items {
6693                        coord.ready.push_back(lane_id);
6694                    } else if is_closed {
6695                        coord.lanes.remove(&lane_id);
6696                    }
6697                    coord.generation += 1;
6698                    drop(coord);
6699                    self.shared.available.notify_all();
6700
6701                    // Step 8: return first item, stash rest in local_batch.
6702                    let mut iter = batch.into_iter();
6703                    if let Some(first) = iter.next() {
6704                        self.local_batch.extend(iter);
6705                        return Some(Ok(first));
6706                    }
6707                }
6708                // Lane vanished or was empty — re-acquire and retry.
6709                coord = self
6710                    .shared
6711                    .coordinator
6712                    .lock()
6713                    .unwrap_or_else(|p| p.into_inner());
6714                continue;
6715            }
6716
6717            // Step 9: EOS — only once all work has been accounted for.
6718            if coord.terminal.is_none()
6719                && coord.input_done
6720                && coord.active_streams == 0
6721                && coord.queued_items == 0
6722            {
6723                return None;
6724            }
6725
6726            // Step 10: wait with generation-counter protocol.
6727            let seen = coord.generation;
6728            coord = self
6729                .shared
6730                .available
6731                .wait_while(coord, |s| {
6732                    s.generation == seen
6733                        && s.terminal.is_none()
6734                        && s.ready.is_empty()
6735                        && !(s.input_done && s.active_streams == 0 && s.queued_items == 0)
6736                })
6737                .unwrap_or_else(|p| p.into_inner());
6738        }
6739    }
6740}
6741
6742impl<T> Drop for FlatMapMergeReadyStream<T> {
6743    fn drop(&mut self) {
6744        let lanes: Vec<Arc<FlatMapMergeLane<T>>> = {
6745            let mut coord = self
6746                .shared
6747                .coordinator
6748                .lock()
6749                .unwrap_or_else(|p| p.into_inner());
6750            coord.generation += 1;
6751            coord.lanes.values().cloned().collect()
6752        };
6753        self.shared.cancelled.store(true, Ordering::SeqCst);
6754        self.shared.available.notify_all();
6755        for lane in lanes {
6756            lane.space_available.notify_all();
6757        }
6758        let _ = self.completion.take();
6759    }
6760}
6761
6762// Decrement active_streams, mark lane closed, and notify waiters.
6763// Must be called AFTER publishing all batches to the coordinator.
6764fn finish_lane_ready<T>(
6765    shared: &Arc<FlatMapMergeReadyShared<T>>,
6766    stream_id: usize,
6767    terminal: Result<(), StreamError>,
6768) {
6769    let is_error = terminal.is_err();
6770
6771    // Step 1: mark lane closed (lane lock only, no coordinator held).
6772    let lane_is_empty = {
6773        let coord = shared.coordinator.lock().unwrap_or_else(|p| p.into_inner());
6774        let lane_opt = coord.lanes.get(&stream_id).cloned();
6775        drop(coord);
6776        if let Some(lane) = lane_opt {
6777            let mut ls = lane.state.lock().unwrap_or_else(|p| p.into_inner());
6778            ls.closed = true;
6779            ls.buffer.is_empty()
6780        } else {
6781            true
6782        }
6783    };
6784
6785    // Step 2: update coordinator (no lane lock held).
6786    let lanes_to_notify: Vec<Arc<FlatMapMergeLane<T>>> = {
6787        let mut coord = shared.coordinator.lock().unwrap_or_else(|p| p.into_inner());
6788        coord.active_streams = coord.active_streams.saturating_sub(1);
6789        if let Err(ref error) = terminal
6790            && coord.terminal.is_none()
6791        {
6792            coord.terminal = Some(error.clone());
6793        }
6794        if lane_is_empty {
6795            coord.lanes.remove(&stream_id);
6796        }
6797        coord.generation += 1;
6798        if is_error {
6799            coord.lanes.values().cloned().collect()
6800        } else {
6801            vec![]
6802        }
6803    };
6804
6805    if is_error {
6806        shared.cancelled.store(true, Ordering::SeqCst);
6807        for lane in &lanes_to_notify {
6808            lane.space_available.notify_all();
6809        }
6810    }
6811    shared.available.notify_all();
6812}
6813
6814struct FlatMapMergeReadyCoordinatorGuard<T> {
6815    shared: Arc<FlatMapMergeReadyShared<T>>,
6816    armed: bool,
6817}
6818
6819impl<T> FlatMapMergeReadyCoordinatorGuard<T> {
6820    fn new(shared: Arc<FlatMapMergeReadyShared<T>>) -> Self {
6821        Self {
6822            shared,
6823            armed: true,
6824        }
6825    }
6826
6827    fn finish(&mut self) {
6828        let mut coord = self
6829            .shared
6830            .coordinator
6831            .lock()
6832            .unwrap_or_else(|p| p.into_inner());
6833        coord.input_done = true;
6834        coord.generation += 1;
6835        drop(coord);
6836        self.shared.available.notify_all();
6837        self.armed = false;
6838    }
6839}
6840
6841impl<T> Drop for FlatMapMergeReadyCoordinatorGuard<T> {
6842    fn drop(&mut self) {
6843        if self.armed {
6844            let lanes: Vec<Arc<FlatMapMergeLane<T>>> = {
6845                let mut coord = self
6846                    .shared
6847                    .coordinator
6848                    .lock()
6849                    .unwrap_or_else(|p| p.into_inner());
6850                if coord.terminal.is_none() {
6851                    coord.terminal = Some(StreamError::AbruptTermination);
6852                }
6853                coord.input_done = true;
6854                coord.generation += 1;
6855                coord.lanes.values().cloned().collect()
6856            };
6857            self.shared.cancelled.store(true, Ordering::SeqCst);
6858            self.shared.available.notify_all();
6859            for lane in lanes {
6860                lane.space_available.notify_all();
6861            }
6862        }
6863    }
6864}
6865
6866struct FlatMapMergeReadyWorkerGuard<T> {
6867    shared: Arc<FlatMapMergeReadyShared<T>>,
6868    stream_id: usize,
6869    armed: bool,
6870}
6871
6872impl<T> FlatMapMergeReadyWorkerGuard<T> {
6873    fn new(shared: Arc<FlatMapMergeReadyShared<T>>, stream_id: usize) -> Self {
6874        Self {
6875            shared,
6876            stream_id,
6877            armed: true,
6878        }
6879    }
6880
6881    fn finish(&mut self, terminal: Result<(), StreamError>) {
6882        finish_lane_ready(&self.shared, self.stream_id, terminal);
6883        self.armed = false;
6884    }
6885}
6886
6887impl<T> Drop for FlatMapMergeReadyWorkerGuard<T> {
6888    fn drop(&mut self) {
6889        if self.armed {
6890            finish_lane_ready(
6891                &self.shared,
6892                self.stream_id,
6893                Err(StreamError::AbruptTermination),
6894            );
6895            self.shared.cancelled.store(true, Ordering::SeqCst);
6896        }
6897    }
6898}
6899
6900/// Shared publish protocol used by both workers and the coordinator inline path.
6901///
6902/// Steps (per the WP-19 lock contract):
6903/// 1. Lock the lane, mark publication in progress, append `batch` to
6904///    `lane.buffer`, set `in_ready_ring = true` if the transition happens
6905///    (was false → now non-empty).
6906/// 2. Drop the lane lock. Consumers that pop this lane wait for the publication
6907///    gate before draining, so they cannot consume items before `queued_items`
6908///    accounts for them.
6909/// 3. Lock coordinator, increment `queued_items`, enqueue the lane id in `ready`
6910///    iff the ring-transition happened, increment `generation`, drop.
6911/// 4. Clear the publication gate and notify `available`.
6912///
6913/// No-op if `batch` is empty (no transition, no coordinator wakeup needed).
6914fn publish_ready_batch<T>(
6915    shared: &Arc<FlatMapMergeReadyShared<T>>,
6916    stream_id: usize,
6917    lane: &Arc<FlatMapMergeLane<T>>,
6918    batch: std::collections::VecDeque<T>,
6919) {
6920    let batch_len = batch.len();
6921    if batch_len == 0 {
6922        return;
6923    }
6924
6925    // Step 1+2: lane lock, append, decide ready-ring transition.
6926    let was_not_in_ready = {
6927        let mut ls = lane.state.lock().unwrap_or_else(|p| p.into_inner());
6928        let was_not = !ls.in_ready_ring;
6929        ls.publishing = true;
6930        ls.buffer.extend(batch);
6931        if !ls.buffer.is_empty() {
6932            ls.in_ready_ring = true;
6933        }
6934        was_not
6935    };
6936
6937    // Step 3+4: coordinator lock, bookkeeping, notify.
6938    {
6939        let mut coord = shared.coordinator.lock().unwrap_or_else(|p| p.into_inner());
6940        coord.queued_items += batch_len;
6941        if was_not_in_ready {
6942            coord.ready.push_back(stream_id);
6943        }
6944        coord.generation += 1;
6945    }
6946    {
6947        let mut ls = lane.state.lock().unwrap_or_else(|p| p.into_inner());
6948        ls.publishing = false;
6949    }
6950    lane.space_available.notify_all();
6951    shared.available.notify_all();
6952}
6953
6954/// RAII guard for an admitted inline lane.
6955///
6956/// After `active_streams += 1`, exactly one of these outcomes must happen:
6957/// - `hand_off()` is called before spawning a worker (worker will call
6958///   `finish_lane_ready` exactly once).
6959/// - `finish(terminal)` is called after the inline drain completes.
6960///
6961/// If neither is called (e.g. on panic), Drop calls `finish_lane_ready` with
6962/// `AbruptTermination` to prevent an active-lane leak.
6963struct FlatMapMergeReadyInlineGuard<T> {
6964    shared: Arc<FlatMapMergeReadyShared<T>>,
6965    stream_id: usize,
6966    armed: bool,
6967}
6968
6969impl<T> FlatMapMergeReadyInlineGuard<T> {
6970    fn new(shared: Arc<FlatMapMergeReadyShared<T>>, stream_id: usize) -> Self {
6971        Self {
6972            shared,
6973            stream_id,
6974            armed: true,
6975        }
6976    }
6977
6978    fn finish(&mut self, terminal: Result<(), StreamError>) {
6979        finish_lane_ready(&self.shared, self.stream_id, terminal);
6980        self.armed = false;
6981    }
6982
6983    fn hand_off(&mut self) {
6984        self.armed = false;
6985    }
6986}
6987
6988impl<T> Drop for FlatMapMergeReadyInlineGuard<T> {
6989    fn drop(&mut self) {
6990        if self.armed {
6991            finish_lane_ready(
6992                &self.shared,
6993                self.stream_id,
6994                Err(StreamError::AbruptTermination),
6995            );
6996        }
6997    }
6998}
6999
7000fn flat_map_merge_stream_ready<Out, Next, NextMat, F>(
7001    mut input: BoxStream<Out>,
7002    breadth: usize,
7003    stage: Arc<F>,
7004    materializer: &Materializer,
7005) -> BoxStream<Next>
7006where
7007    Out: Send + 'static,
7008    Next: Send + 'static,
7009    NextMat: Send + 'static,
7010    F: Fn(Out) -> Source<Next, NextMat> + Send + Sync + 'static,
7011{
7012    let worker_materializer = materializer.with_name_prefix(materializer.name_prefix().to_owned());
7013    let shared: Arc<FlatMapMergeReadyShared<Next>> = Arc::new(FlatMapMergeReadyShared {
7014        coordinator: Mutex::new(FlatMapMergeReadyState {
7015            lanes: HashMap::new(),
7016            ready: std::collections::VecDeque::new(),
7017            queued_items: 0,
7018            active_streams: 0,
7019            input_done: false,
7020            terminal: None,
7021            generation: 0,
7022        }),
7023        available: Condvar::new(),
7024        cancelled: Arc::new(AtomicBool::new(false)),
7025    });
7026
7027    let worker_shared = Arc::clone(&shared);
7028    let completion = materializer.spawn_stream(move |cancelled| {
7029        let mut next_id = 0usize;
7030        let mut guard = FlatMapMergeReadyCoordinatorGuard::new(worker_shared);
7031        let mut workers = HashMap::<usize, StreamCompletion<NotUsed>>::with_capacity(breadth);
7032
7033        while !cancelled.load(Ordering::SeqCst) && !guard.shared.cancelled.load(Ordering::SeqCst) {
7034            workers.retain(|_, c| c.try_wait().is_none());
7035
7036            // Wait for a free breadth slot using the generation protocol.
7037            {
7038                let mut coord = guard
7039                    .shared
7040                    .coordinator
7041                    .lock()
7042                    .unwrap_or_else(|p| p.into_inner());
7043                loop {
7044                    if coord.terminal.is_some()
7045                        || coord.active_streams < breadth
7046                        || cancelled.load(Ordering::SeqCst)
7047                        || guard.shared.cancelled.load(Ordering::SeqCst)
7048                    {
7049                        break;
7050                    }
7051                    let seen = coord.generation;
7052                    coord = guard
7053                        .shared
7054                        .available
7055                        .wait_while(coord, |s| {
7056                            s.generation == seen
7057                                && s.terminal.is_none()
7058                                && s.active_streams >= breadth
7059                                && !cancelled.load(Ordering::SeqCst)
7060                                && !guard.shared.cancelled.load(Ordering::SeqCst)
7061                        })
7062                        .unwrap_or_else(|p| p.into_inner());
7063                }
7064                if coord.terminal.is_some() {
7065                    let error = coord.terminal.clone().expect("terminal checked");
7066                    drop(coord);
7067                    guard.finish();
7068                    return Err(error);
7069                }
7070            }
7071
7072            match input.next() {
7073                Some(Ok(item)) => {
7074                    // f and factory called outside all locks.
7075                    let source = stage(item);
7076                    // Read hint before factory.create() consumes source.
7077                    let inline_hint = source.hints.inline_micro;
7078                    let (stream, _) = match Arc::clone(&source.factory).create(&worker_materializer)
7079                    {
7080                        Ok(parts) => parts,
7081                        Err(error) => {
7082                            let lanes: Vec<Arc<FlatMapMergeLane<Next>>> = {
7083                                let mut coord = guard
7084                                    .shared
7085                                    .coordinator
7086                                    .lock()
7087                                    .unwrap_or_else(|p| p.into_inner());
7088                                if coord.terminal.is_none() {
7089                                    coord.terminal = Some(error.clone());
7090                                }
7091                                coord.generation += 1;
7092                                coord.lanes.values().cloned().collect()
7093                            };
7094                            guard.shared.cancelled.store(true, Ordering::SeqCst);
7095                            guard.shared.available.notify_all();
7096                            for lane in lanes {
7097                                lane.space_available.notify_all();
7098                            }
7099                            guard.finish();
7100                            return Err(error);
7101                        }
7102                    };
7103
7104                    let stream_id = next_id;
7105                    next_id += 1;
7106                    let mut stream = stream;
7107
7108                    // Create lane and register it.
7109                    let lane: Arc<FlatMapMergeLane<Next>> = Arc::new(FlatMapMergeLane {
7110                        state: Mutex::new(FlatMapMergeLaneState {
7111                            buffer: std::collections::VecDeque::new(),
7112                            in_ready_ring: false,
7113                            publishing: false,
7114                            closed: false,
7115                        }),
7116                        space_available: Condvar::new(),
7117                    });
7118                    {
7119                        let mut coord = guard
7120                            .shared
7121                            .coordinator
7122                            .lock()
7123                            .unwrap_or_else(|p| p.into_inner());
7124                        coord.lanes.insert(stream_id, Arc::clone(&lane));
7125                        coord.active_streams += 1;
7126                        // No generation increment: admission alone does not wake consumer.
7127                    }
7128
7129                    // RAII guard: ensures finish_lane_ready is called exactly once
7130                    // regardless of which path (inline or worker) handles this lane.
7131                    let mut inline_guard =
7132                        FlatMapMergeReadyInlineGuard::new(Arc::clone(&guard.shared), stream_id);
7133
7134                    let is_inline = inline_hint
7135                        .is_some_and(|h| h.max_success_items <= FLAT_MAP_MERGE_INLINE_MICRO_MAX);
7136
7137                    if is_inline {
7138                        // Inline drain path: drain the entire micro-source in the
7139                        // coordinator thread without spawning a worker task.
7140                        // NO locks are held during any stream.next() call.
7141                        let max_items = inline_hint.expect("checked").max_success_items;
7142
7143                        // Check cancellation before starting the drain.
7144                        if guard.shared.cancelled.load(Ordering::SeqCst)
7145                            || cancelled.load(Ordering::SeqCst)
7146                        {
7147                            inline_guard.finish(Ok(()));
7148                            continue;
7149                        }
7150
7151                        // Pull at most max_items successful items, then one terminal
7152                        // probe to distinguish normal EOS from error.
7153                        let max_pulls = max_items.saturating_add(1);
7154                        let mut local_batch = std::collections::VecDeque::with_capacity(max_items);
7155                        let mut terminal_result: Option<Result<(), StreamError>> = None;
7156
7157                        for _ in 0..max_pulls {
7158                            if guard.shared.cancelled.load(Ordering::SeqCst)
7159                                || cancelled.load(Ordering::SeqCst)
7160                            {
7161                                break;
7162                            }
7163                            match stream.next() {
7164                                Some(Ok(item)) => local_batch.push_back(item),
7165                                Some(Err(e)) => {
7166                                    terminal_result = Some(Err(e));
7167                                    break;
7168                                }
7169                                None => {
7170                                    terminal_result = Some(Ok(()));
7171                                    break;
7172                                }
7173                            }
7174                        }
7175
7176                        // Re-check cancellation after pull, before publish.
7177                        if guard.shared.cancelled.load(Ordering::SeqCst)
7178                            || cancelled.load(Ordering::SeqCst)
7179                        {
7180                            inline_guard.finish(Ok(()));
7181                            continue;
7182                        }
7183
7184                        // Publish-before-finish: items enter the coordinator BEFORE
7185                        // active_streams is decremented by finish_lane_ready.
7186                        publish_ready_batch(&guard.shared, stream_id, &lane, local_batch);
7187
7188                        match terminal_result {
7189                            Some(Err(e)) => {
7190                                // Inner error: shared.cancelled will be set by
7191                                // finish_lane_ready, breaking the outer loop next iter.
7192                                inline_guard.finish(Err(e));
7193                            }
7194                            Some(Ok(())) | None => {
7195                                inline_guard.finish(Ok(()));
7196                            }
7197                        }
7198                    } else {
7199                        // Worker path: spawn the ready-ring worker (unchanged).
7200                        inline_guard.hand_off();
7201                        let worker_shared = Arc::clone(&guard.shared);
7202                        let worker_lane = Arc::clone(&lane);
7203                        workers.insert(
7204                            stream_id,
7205                            worker_materializer.spawn_stream(move |inner_cancelled| {
7206                                let mut worker_guard = FlatMapMergeReadyWorkerGuard::new(
7207                                    Arc::clone(&worker_shared),
7208                                    stream_id,
7209                                );
7210
7211                                loop {
7212                                    if inner_cancelled.load(Ordering::SeqCst)
7213                                        || worker_shared.cancelled.load(Ordering::SeqCst)
7214                                    {
7215                                        worker_guard.finish(Ok(()));
7216                                        return Ok(NotUsed);
7217                                    }
7218
7219                                    // Step 1: wait for ring capacity under lane lock.
7220                                    let capacity;
7221                                    {
7222                                        let mut ls = worker_lane
7223                                            .state
7224                                            .lock()
7225                                            .unwrap_or_else(|p| p.into_inner());
7226                                        while ls.buffer.len() >= FLAT_MAP_MERGE_SUBSTREAM_WINDOW
7227                                            && !worker_shared.cancelled.load(Ordering::SeqCst)
7228                                            && !ls.closed
7229                                            && !inner_cancelled.load(Ordering::SeqCst)
7230                                        {
7231                                            ls = worker_lane
7232                                                .space_available
7233                                                .wait(ls)
7234                                                .unwrap_or_else(|p| p.into_inner());
7235                                        }
7236                                        if worker_shared.cancelled.load(Ordering::SeqCst)
7237                                            || ls.closed
7238                                            || inner_cancelled.load(Ordering::SeqCst)
7239                                        {
7240                                            drop(ls);
7241                                            worker_guard.finish(Ok(()));
7242                                            return Ok(NotUsed);
7243                                        }
7244                                        capacity =
7245                                            FLAT_MAP_MERGE_SUBSTREAM_WINDOW - ls.buffer.len();
7246                                    }
7247                                    // Step 2: lane lock dropped.
7248
7249                                    // Step 3: pull items outside all locks.
7250                                    let batch_size = capacity.min(FLAT_MAP_MERGE_READY_BATCH);
7251                                    let mut local_batch =
7252                                        std::collections::VecDeque::with_capacity(batch_size);
7253                                    let mut terminal_result: Option<Result<(), StreamError>> = None;
7254                                    for _ in 0..batch_size {
7255                                        if inner_cancelled.load(Ordering::SeqCst)
7256                                            || worker_shared.cancelled.load(Ordering::SeqCst)
7257                                        {
7258                                            break;
7259                                        }
7260                                        match stream.next() {
7261                                            Some(Ok(item)) => local_batch.push_back(item),
7262                                            Some(Err(e)) => {
7263                                                terminal_result = Some(Err(e));
7264                                                break;
7265                                            }
7266                                            None => {
7267                                                terminal_result = Some(Ok(()));
7268                                                break;
7269                                            }
7270                                        }
7271                                    }
7272                                    let batch_len = local_batch.len();
7273
7274                                    // Step 4+5: re-lock lane, append batch; use
7275                                    // publish_ready_batch for the shared protocol.
7276                                    publish_ready_batch(
7277                                        &worker_shared,
7278                                        stream_id,
7279                                        &worker_lane,
7280                                        local_batch,
7281                                    );
7282
7283                                    match terminal_result {
7284                                        Some(Ok(())) => {
7285                                            worker_guard.finish(Ok(()));
7286                                            return Ok(NotUsed);
7287                                        }
7288                                        Some(Err(e)) => {
7289                                            worker_guard.finish(Err(e.clone()));
7290                                            return Err(e);
7291                                        }
7292                                        None => {
7293                                            if batch_len == 0 {
7294                                                worker_guard.finish(Ok(()));
7295                                                return Ok(NotUsed);
7296                                            }
7297                                        }
7298                                    }
7299                                }
7300                            }),
7301                        );
7302                    }
7303                }
7304                Some(Err(error)) => {
7305                    let lanes: Vec<Arc<FlatMapMergeLane<Next>>> = {
7306                        let mut coord = guard
7307                            .shared
7308                            .coordinator
7309                            .lock()
7310                            .unwrap_or_else(|p| p.into_inner());
7311                        if coord.terminal.is_none() {
7312                            coord.terminal = Some(error.clone());
7313                        }
7314                        coord.generation += 1;
7315                        coord.lanes.values().cloned().collect()
7316                    };
7317                    guard.shared.cancelled.store(true, Ordering::SeqCst);
7318                    guard.shared.available.notify_all();
7319                    for lane in lanes {
7320                        lane.space_available.notify_all();
7321                    }
7322                    guard.finish();
7323                    return Err(error);
7324                }
7325                None => {
7326                    guard.finish();
7327                    break;
7328                }
7329            }
7330        }
7331
7332        if guard.armed {
7333            guard.finish();
7334        }
7335
7336        // Drain workers.  Same invariant as legacy: a worker may decrement
7337        // active_streams before its StreamCompletion reflects completion.
7338        // Always check active_streams == 0 under the lock before sleeping.
7339        while !cancelled.load(Ordering::SeqCst) && !guard.shared.cancelled.load(Ordering::SeqCst) {
7340            workers.retain(|_, c| c.try_wait().is_none());
7341            let coord = guard
7342                .shared
7343                .coordinator
7344                .lock()
7345                .unwrap_or_else(|p| p.into_inner());
7346            if coord.active_streams == 0 {
7347                return Ok(NotUsed);
7348            }
7349            let seen = coord.generation;
7350            drop(
7351                guard
7352                    .shared
7353                    .available
7354                    .wait_while(coord, |s| s.generation == seen && s.active_streams > 0)
7355                    .unwrap_or_else(|p| p.into_inner()),
7356            );
7357        }
7358        Ok(NotUsed)
7359    });
7360
7361    Box::new(FlatMapMergeReadyStream {
7362        shared,
7363        completion: Some(completion),
7364        local_batch: std::collections::VecDeque::new(),
7365    })
7366}
7367
7368fn flat_map_concat_stream<Out, Next, NextMat, F>(
7369    mut input: BoxStream<Out>,
7370    stage: Arc<F>,
7371    materializer: &Materializer,
7372) -> BoxStream<Next>
7373where
7374    Out: Send + 'static,
7375    Next: Send + 'static,
7376    NextMat: Send + 'static,
7377    F: Fn(Out) -> Source<Next, NextMat> + Send + Sync + 'static,
7378{
7379    let materializer = materializer.with_name_prefix(materializer.name_prefix().to_owned());
7380    let mut current: Option<BoxStream<Next>> = None;
7381    Box::new(std::iter::from_fn(move || {
7382        loop {
7383            if let Some(stream) = current.as_mut() {
7384                match stream.next() {
7385                    Some(item) => return Some(item),
7386                    None => current = None,
7387                }
7388            }
7389
7390            match input.next() {
7391                Some(Ok(item)) => {
7392                    let source = stage(item);
7393                    current = Some(match Arc::clone(&source.factory).create(&materializer) {
7394                        Ok((stream, _)) => stream,
7395                        Err(error) => Box::new(std::iter::once(Err(error))) as BoxStream<Next>,
7396                    });
7397                }
7398                Some(Err(error)) => return Some(Err(error)),
7399                None => return None,
7400            }
7401        }
7402    }))
7403}
7404
7405fn map_async_unordered<Out, Next, F, Fut>(
7406    mut input: BoxStream<Out>,
7407    parallelism: usize,
7408    stage: Arc<F>,
7409) -> BoxStream<Next>
7410where
7411    Out: Send + 'static,
7412    Next: Send + 'static,
7413    F: Fn(Out) -> Fut + Send + Sync + 'static,
7414    Fut: Future<Output = StreamResult<Next>> + Send + 'static,
7415{
7416    let (sender, receiver) = std::sync::mpsc::channel::<(usize, StreamResult<Next>)>();
7417    let mut tasks = HashMap::<usize, AbortOnDropHandle<()>>::with_capacity(parallelism);
7418    let mut next_task_id = 0_usize;
7419    let mut input_done = false;
7420
7421    Box::new(std::iter::from_fn(move || {
7422        loop {
7423            while tasks.len() < parallelism && !input_done {
7424                match input.next() {
7425                    Some(Ok(item)) => match poll_once_or_pending(stage(item)) {
7426                        Ok(result) => return Some(result),
7427                        Err(future) => {
7428                            let task_id = next_task_id;
7429                            next_task_id += 1;
7430                            tasks.insert(
7431                                task_id,
7432                                spawn_completion_task(task_id, future, sender.clone(), |result| {
7433                                    result
7434                                }),
7435                            );
7436                        }
7437                    },
7438                    Some(Err(error)) => {
7439                        input_done = true;
7440                        return Some(Err(error));
7441                    }
7442                    None => input_done = true,
7443                }
7444            }
7445
7446            if tasks.is_empty() {
7447                return None;
7448            }
7449
7450            if let Some((task_id, result)) = recv_completion(&receiver) {
7451                tasks.remove(&task_id);
7452                return Some(result);
7453            }
7454        }
7455    }))
7456}
7457
7458fn map_async_unordered_supervised<Out, Next, F, Fut>(
7459    mut input: BoxStream<Out>,
7460    parallelism: usize,
7461    stage: Arc<F>,
7462    decider: SupervisionDecider,
7463) -> BoxStream<Next>
7464where
7465    Out: Send + 'static,
7466    Next: Send + 'static,
7467    F: Fn(Out) -> Fut + Send + Sync + 'static,
7468    Fut: Future<Output = StreamResult<Next>> + Send + 'static,
7469{
7470    let (sender, receiver) = std::sync::mpsc::channel::<(usize, StreamResult<Next>)>();
7471    let mut tasks = HashMap::<usize, AbortOnDropHandle<()>>::with_capacity(parallelism);
7472    let mut next_task_id = 0_usize;
7473    let mut input_done = false;
7474
7475    Box::new(std::iter::from_fn(move || {
7476        loop {
7477            while tasks.len() < parallelism && !input_done {
7478                match input.next() {
7479                    Some(Ok(item)) => {
7480                        match catch_unwind(AssertUnwindSafe(|| poll_once_or_pending(stage(item)))) {
7481                            Ok(Ok(result)) => {
7482                                if let Some(result) = supervise_async_result(result, &decider) {
7483                                    return Some(result);
7484                                }
7485                            }
7486                            Ok(Err(future)) => {
7487                                let task_id = next_task_id;
7488                                next_task_id += 1;
7489                                tasks.insert(
7490                                    task_id,
7491                                    spawn_completion_task(
7492                                        task_id,
7493                                        future,
7494                                        sender.clone(),
7495                                        |result| result,
7496                                    ),
7497                                );
7498                            }
7499                            Err(_) => {
7500                                let error = panic_stream_error("map_async_unordered callback");
7501                                if let Some(result) = supervise_async_result(Err(error), &decider) {
7502                                    return Some(result);
7503                                }
7504                            }
7505                        }
7506                    }
7507                    Some(Err(error)) => {
7508                        input_done = true;
7509                        return Some(Err(error));
7510                    }
7511                    None => input_done = true,
7512                }
7513            }
7514
7515            if tasks.is_empty() {
7516                return None;
7517            }
7518
7519            if let Some((task_id, result)) = recv_completion(&receiver) {
7520                tasks.remove(&task_id);
7521                if let Some(result) = supervise_async_result(result, &decider) {
7522                    return Some(result);
7523                }
7524            }
7525        }
7526    }))
7527}
7528
7529#[inline(always)]
7530fn map_async_partitioned_serial<Out, Key, Next, Partition, F, Fut>(
7531    mut input: BoxStream<Out>,
7532    partition: Arc<Partition>,
7533    stage: Arc<F>,
7534) -> BoxStream<Next>
7535where
7536    Out: Send + 'static,
7537    Key: Send + 'static,
7538    Next: Send + 'static,
7539    Partition: Fn(&Out) -> Key + Send + Sync + 'static,
7540    F: Fn(Out) -> Fut + Send + Sync + 'static,
7541    Fut: Future<Output = StreamResult<Next>> + Send + 'static,
7542{
7543    let (sender, receiver) = std::sync::mpsc::channel::<(usize, StreamResult<Next>)>();
7544    let mut tasks = HashMap::<usize, AbortOnDropHandle<()>>::with_capacity(1);
7545    let mut next_index = 0_usize;
7546    Box::new(std::iter::from_fn(move || {
7547        // With parallelism = 1 any pending future is awaited inline at spawn, so
7548        // the receiver can only have a message when a task is outstanding; an
7549        // unguarded recv here busy-loops forever on the all-inline-ready path.
7550        if !tasks.is_empty()
7551            && let Some((task_id, result)) = recv_completion(&receiver)
7552        {
7553            tasks.remove(&task_id);
7554            return Some(result);
7555        }
7556
7557        let item = input.next()?;
7558        match item {
7559            Ok(item) => {
7560                let _ = partition(&item);
7561                let index = next_index;
7562                next_index += 1;
7563                Some(match poll_once_or_pending(stage(item)) {
7564                    Ok(result) => result,
7565                    Err(future) => {
7566                        tasks.insert(
7567                            index,
7568                            spawn_completion_task(index, future, sender.clone(), |result| result),
7569                        );
7570                        let (task_id, result) =
7571                            recv_completion(&receiver).expect("pending map_async task completion");
7572                        tasks.remove(&task_id);
7573                        result
7574                    }
7575                })
7576            }
7577            Err(error) => Some(Err(error)),
7578        }
7579    }))
7580}
7581
7582#[inline(always)]
7583fn map_async_partitioned_scanning<Out, Key, Next, Partition, F, Fut>(
7584    mut input: BoxStream<Out>,
7585    parallelism: usize,
7586    per_partition: usize,
7587    partition: Arc<Partition>,
7588    stage: Arc<F>,
7589) -> BoxStream<Next>
7590where
7591    Out: Send + 'static,
7592    Key: Clone + Eq + Hash + Send + 'static,
7593    Next: Send + 'static,
7594    Partition: Fn(&Out) -> Key + Send + Sync + 'static,
7595    F: Fn(Out) -> Fut + Send + Sync + 'static,
7596    Fut: Future<Output = StreamResult<Next>> + Send + 'static,
7597{
7598    let (sender, receiver) = std::sync::mpsc::channel::<(usize, (Key, StreamResult<Next>))>();
7599    let mut tasks = HashMap::<usize, AbortOnDropHandle<()>>::with_capacity(parallelism);
7600    let mut active_by_key = HashMap::<Key, usize>::with_capacity(parallelism);
7601    let mut pending = VecDeque::<(usize, Key, Out)>::with_capacity(parallelism);
7602    let mut completed = BTreeMap::<usize, StreamResult<Next>>::new();
7603    let mut next_index = 0_usize;
7604    let mut next_to_emit = 0_usize;
7605    let mut input_done = false;
7606
7607    Box::new(std::iter::from_fn(move || {
7608        loop {
7609            if let Some(result) = completed.remove(&next_to_emit) {
7610                next_to_emit += 1;
7611                return Some(result);
7612            }
7613
7614            while tasks.len() + completed.len() < parallelism {
7615                let next = pending
7616                    .iter()
7617                    .position(|(_, key, _)| {
7618                        active_by_key.get(key).copied().unwrap_or(0) < per_partition
7619                    })
7620                    .and_then(|index| pending.remove(index))
7621                    .or_else(|| {
7622                        if input_done {
7623                            return None;
7624                        }
7625                        match input.next() {
7626                            Some(Ok(item)) => {
7627                                let key = partition(&item);
7628                                let index = next_index;
7629                                next_index += 1;
7630                                Some((index, key, item))
7631                            }
7632                            Some(Err(error)) => {
7633                                completed.insert(next_index, Err(error));
7634                                next_index += 1;
7635                                input_done = true;
7636                                None
7637                            }
7638                            None => {
7639                                input_done = true;
7640                                None
7641                            }
7642                        }
7643                    });
7644
7645                let Some((index, key, item)) = next else {
7646                    break;
7647                };
7648                if active_by_key.get(&key).copied().unwrap_or(0) >= per_partition {
7649                    pending.push_back((index, key, item));
7650                    if input_done || pending.len() >= parallelism {
7651                        break;
7652                    }
7653                    continue;
7654                }
7655                *active_by_key.entry(key.clone()).or_default() += 1;
7656                match poll_once_or_pending(stage(item)) {
7657                    Ok(result) => {
7658                        if let Some(count) = active_by_key.get_mut(&key) {
7659                            *count -= 1;
7660                            if *count == 0 {
7661                                active_by_key.remove(&key);
7662                            }
7663                        }
7664                        completed.insert(index, result);
7665                    }
7666                    Err(future) => {
7667                        tasks.insert(
7668                            index,
7669                            spawn_completion_task(index, future, sender.clone(), move |result| {
7670                                (key, result)
7671                            }),
7672                        );
7673                    }
7674                }
7675            }
7676
7677            if let Some(result) = completed.remove(&next_to_emit) {
7678                next_to_emit += 1;
7679                return Some(result);
7680            }
7681
7682            if tasks.is_empty() {
7683                return None;
7684            }
7685
7686            if let Some((index, (key, result))) = recv_completion(&receiver) {
7687                tasks.remove(&index);
7688                if let Some(count) = active_by_key.get_mut(&key) {
7689                    *count -= 1;
7690                    if *count == 0 {
7691                        active_by_key.remove(&key);
7692                    }
7693                }
7694                if index == next_to_emit {
7695                    next_to_emit += 1;
7696                    return Some(result);
7697                }
7698                completed.insert(index, result);
7699            }
7700        }
7701    }))
7702}
7703
7704fn map_async_partitioned<Out, Key, Next, Partition, F, Fut>(
7705    mut input: BoxStream<Out>,
7706    parallelism: usize,
7707    per_partition: usize,
7708    partition: Arc<Partition>,
7709    stage: Arc<F>,
7710) -> BoxStream<Next>
7711where
7712    Out: Send + 'static,
7713    Key: Clone + Eq + Hash + Send + 'static,
7714    Next: Send + 'static,
7715    Partition: Fn(&Out) -> Key + Send + Sync + 'static,
7716    F: Fn(Out) -> Fut + Send + Sync + 'static,
7717    Fut: Future<Output = StreamResult<Next>> + Send + 'static,
7718{
7719    if parallelism == 1 {
7720        return map_async_partitioned_serial(input, partition, stage);
7721    }
7722    // At small parallelism the bounded scan is cheaper than maintaining slot queues.
7723    if parallelism <= 4 {
7724        return map_async_partitioned_scanning(input, parallelism, per_partition, partition, stage);
7725    }
7726
7727    let (sender, receiver) = std::sync::mpsc::channel::<(usize, (usize, StreamResult<Next>))>();
7728    let mut tasks = HashMap::<usize, AbortOnDropHandle<()>>::with_capacity(parallelism);
7729    let mut slots_by_key = HashMap::<Key, usize>::with_capacity(parallelism);
7730    let mut slots = Vec::<PartitionSlot<Key, Out>>::with_capacity(parallelism);
7731    let mut free_slots = Vec::<usize>::new();
7732    let mut ready_slots = VecDeque::<usize>::with_capacity(parallelism);
7733    let mut completed = BTreeMap::<usize, StreamResult<Next>>::new();
7734    let mut next_index = 0_usize;
7735    let mut next_to_emit = 0_usize;
7736    let mut input_done = false;
7737    Box::new(std::iter::from_fn(move || {
7738        loop {
7739            if let Some(result) = completed.remove(&next_to_emit) {
7740                next_to_emit += 1;
7741                return Some(result);
7742            }
7743
7744            while tasks.len() + completed.len() < parallelism {
7745                if let Some((index, slot, item)) =
7746                    pop_ready_partition_slot(&mut slots, &mut ready_slots, per_partition)
7747                {
7748                    match poll_once_or_pending(stage(item)) {
7749                        Ok(result) => {
7750                            let mut remove_empty = false;
7751                            if let Some(state) = slots.get_mut(slot) {
7752                                state.active -= 1;
7753                                remove_empty = state.active == 0
7754                                    && state.queued.is_empty()
7755                                    && !state.in_ready_queue;
7756                            }
7757                            if remove_empty {
7758                                retire_partition_slot(
7759                                    slot,
7760                                    &mut slots_by_key,
7761                                    &mut slots,
7762                                    &mut free_slots,
7763                                );
7764                            } else {
7765                                ready_partition_slot(
7766                                    &mut slots,
7767                                    &mut ready_slots,
7768                                    slot,
7769                                    per_partition,
7770                                );
7771                            }
7772                            if index == next_to_emit {
7773                                next_to_emit += 1;
7774                                return Some(result);
7775                            }
7776                            completed.insert(index, result);
7777                        }
7778                        Err(future) => {
7779                            tasks.insert(
7780                                index,
7781                                spawn_completion_task(
7782                                    index,
7783                                    future,
7784                                    sender.clone(),
7785                                    move |result| (slot, result),
7786                                ),
7787                            );
7788                        }
7789                    }
7790                    continue;
7791                }
7792
7793                if input_done {
7794                    break;
7795                }
7796
7797                match input.next() {
7798                    Some(Ok(item)) => {
7799                        let key = partition(&item);
7800                        let index = next_index;
7801                        next_index += 1;
7802                        let slot =
7803                            partition_slot_for(key, &mut slots_by_key, &mut slots, &mut free_slots);
7804                        let state = &mut slots[slot];
7805                        if state.active < per_partition {
7806                            match poll_once_or_pending(stage(item)) {
7807                                Ok(result) => {
7808                                    if index == next_to_emit {
7809                                        next_to_emit += 1;
7810                                        if state.queued.is_empty() && !state.in_ready_queue {
7811                                            retire_partition_slot(
7812                                                slot,
7813                                                &mut slots_by_key,
7814                                                &mut slots,
7815                                                &mut free_slots,
7816                                            );
7817                                        }
7818                                        return Some(result);
7819                                    }
7820                                    completed.insert(index, result);
7821                                    if state.queued.is_empty() && !state.in_ready_queue {
7822                                        retire_partition_slot(
7823                                            slot,
7824                                            &mut slots_by_key,
7825                                            &mut slots,
7826                                            &mut free_slots,
7827                                        );
7828                                    }
7829                                }
7830                                Err(future) => {
7831                                    state.active += 1;
7832                                    tasks.insert(
7833                                        index,
7834                                        spawn_completion_task(
7835                                            index,
7836                                            future,
7837                                            sender.clone(),
7838                                            move |result| (slot, result),
7839                                        ),
7840                                    );
7841                                }
7842                            }
7843                        } else {
7844                            state.queued.push_back((index, item));
7845                        }
7846                    }
7847                    Some(Err(error)) => {
7848                        completed.insert(next_index, Err(error));
7849                        next_index += 1;
7850                        input_done = true;
7851                        break;
7852                    }
7853                    None => {
7854                        input_done = true;
7855                        break;
7856                    }
7857                }
7858            }
7859
7860            if let Some(result) = completed.remove(&next_to_emit) {
7861                next_to_emit += 1;
7862                return Some(result);
7863            }
7864
7865            if tasks.is_empty() {
7866                return None;
7867            }
7868
7869            if let Some((index, (slot, result))) = recv_completion(&receiver) {
7870                tasks.remove(&index);
7871                let mut remove_empty = false;
7872                if let Some(state) = slots.get_mut(slot) {
7873                    state.active -= 1;
7874                    remove_empty =
7875                        state.active == 0 && state.queued.is_empty() && !state.in_ready_queue;
7876                }
7877                if remove_empty {
7878                    retire_partition_slot(slot, &mut slots_by_key, &mut slots, &mut free_slots);
7879                } else {
7880                    ready_partition_slot(&mut slots, &mut ready_slots, slot, per_partition);
7881                }
7882                if index == next_to_emit {
7883                    next_to_emit += 1;
7884                    return Some(result);
7885                }
7886                completed.insert(index, result);
7887            }
7888        }
7889    }))
7890}
7891
7892#[cfg(test)]
7893mod flat_map_merge_ready_ring_tests {
7894    use super::*;
7895    use std::sync::mpsc;
7896    use std::time::Duration;
7897
7898    fn run_sorted<T: Ord + Send + 'static>(source: crate::Source<T>) -> Vec<T> {
7899        let mut v = source.run_collect().unwrap();
7900        v.sort_unstable();
7901        v
7902    }
7903
7904    #[test]
7905    fn ready_ring_empty_upstream() {
7906        let legacy = with_substream_mode(SubstreamExecutorMode::LegacyOnly, || {
7907            run_sorted(crate::Source::<i32>::empty().flat_map_merge(4, crate::Source::single))
7908        });
7909        let ring = with_substream_mode(SubstreamExecutorMode::ReadyRingOnly, || {
7910            run_sorted(crate::Source::<i32>::empty().flat_map_merge(4, crate::Source::single))
7911        });
7912        assert_eq!(legacy, ring);
7913        assert!(ring.is_empty());
7914    }
7915
7916    #[test]
7917    fn ready_ring_single_lane() {
7918        let legacy = with_substream_mode(SubstreamExecutorMode::LegacyOnly, || {
7919            run_sorted(crate::Source::single(42_i32).flat_map_merge(4, crate::Source::single))
7920        });
7921        let ring = with_substream_mode(SubstreamExecutorMode::ReadyRingOnly, || {
7922            run_sorted(crate::Source::single(42_i32).flat_map_merge(4, crate::Source::single))
7923        });
7924        assert_eq!(legacy, ring);
7925        assert_eq!(ring, vec![42]);
7926    }
7927
7928    #[test]
7929    fn ready_ring_breadth_one_exact_order() {
7930        // breadth=1 + single-item inner: output must be same set as legacy.
7931        let make = || {
7932            crate::Source::from_iter(0_i32..5)
7933                .flat_map_merge(1, |x| crate::Source::single(x * 10))
7934                .run_collect()
7935                .unwrap()
7936        };
7937        let mut legacy = with_substream_mode(SubstreamExecutorMode::LegacyOnly, make);
7938        let mut ring = with_substream_mode(SubstreamExecutorMode::ReadyRingOnly, make);
7939        legacy.sort_unstable();
7940        ring.sort_unstable();
7941        assert_eq!(legacy, ring);
7942    }
7943
7944    #[test]
7945    fn ready_ring_breadth_gt_input() {
7946        let make = || {
7947            run_sorted(
7948                crate::Source::from_iter(0_i32..3)
7949                    .flat_map_merge(100, |x| crate::Source::from_iter([x, x + 1, x + 2])),
7950            )
7951        };
7952        let legacy = with_substream_mode(SubstreamExecutorMode::LegacyOnly, make);
7953        let ring = with_substream_mode(SubstreamExecutorMode::ReadyRingOnly, make);
7954        assert_eq!(legacy, ring);
7955        assert_eq!(ring.len(), 9);
7956    }
7957
7958    #[test]
7959    fn ready_ring_mixed_short_long() {
7960        let make = || {
7961            run_sorted(crate::Source::from_iter(0_i32..8).flat_map_merge(4, |x| {
7962                if x % 3 == 0 {
7963                    crate::Source::from_iter(0..20_i32).map(move |i| x * 100 + i)
7964                } else {
7965                    crate::Source::single(x)
7966                }
7967            }))
7968        };
7969        let legacy = with_substream_mode(SubstreamExecutorMode::LegacyOnly, make);
7970        let ring = with_substream_mode(SubstreamExecutorMode::ReadyRingOnly, make);
7971        assert_eq!(legacy, ring);
7972    }
7973
7974    #[test]
7975    fn ready_ring_respects_breadth_bound() {
7976        use std::sync::atomic::{AtomicUsize, Ordering as Ord};
7977        let active = Arc::new(AtomicUsize::new(0));
7978        let max_active = Arc::new(AtomicUsize::new(0));
7979        let a2 = Arc::clone(&active);
7980        let m2 = Arc::clone(&max_active);
7981        let mut values = with_substream_mode(SubstreamExecutorMode::ReadyRingOnly, || {
7982            crate::Source::from_iter(0_i32..6)
7983                .flat_map_merge(2, move |item| {
7984                    let a = Arc::clone(&a2);
7985                    let m = Arc::clone(&m2);
7986                    crate::Source::future(move || {
7987                        let a = Arc::clone(&a);
7988                        let m = Arc::clone(&m);
7989                        async move {
7990                            let now = a.fetch_add(1, Ord::SeqCst) + 1;
7991                            let mut seen = m.load(Ord::SeqCst);
7992                            while now > seen {
7993                                match m.compare_exchange(seen, now, Ord::SeqCst, Ord::SeqCst) {
7994                                    Ok(_) => break,
7995                                    Err(v) => seen = v,
7996                                }
7997                            }
7998                            thread::sleep(Duration::from_millis(20));
7999                            a.fetch_sub(1, Ord::SeqCst);
8000                            Ok(item)
8001                        }
8002                    })
8003                })
8004                .run_collect()
8005                .unwrap()
8006        });
8007        values.sort_unstable();
8008        assert_eq!(values, vec![0, 1, 2, 3, 4, 5]);
8009        assert!(max_active.load(std::sync::atomic::Ordering::SeqCst) <= 2);
8010    }
8011
8012    #[test]
8013    fn ready_ring_fairness_slow_lane_not_starved() {
8014        use std::sync::atomic::{AtomicBool, Ordering as Ord};
8015        let slow_emitted = Arc::new(AtomicBool::new(false));
8016        let slow_flag = Arc::clone(&slow_emitted);
8017        let results = with_substream_mode(SubstreamExecutorMode::ReadyRingOnly, || {
8018            crate::Source::from_iter(0_i32..3)
8019                .flat_map_merge(4, move |lane_id| {
8020                    let flag = Arc::clone(&slow_flag);
8021                    match lane_id {
8022                        0 => crate::Source::from_iter(0_i32..50),
8023                        1 => crate::Source::from_iter(100_i32..150),
8024                        _ => crate::Source::future(move || {
8025                            let flag = Arc::clone(&flag);
8026                            async move {
8027                                thread::sleep(Duration::from_millis(10));
8028                                flag.store(true, Ord::SeqCst);
8029                                Ok(999_i32)
8030                            }
8031                        }),
8032                    }
8033                })
8034                .run_collect()
8035                .unwrap()
8036        });
8037        assert!(slow_emitted.load(std::sync::atomic::Ordering::SeqCst));
8038        assert!(results.contains(&999));
8039        assert_eq!(results.len(), 101);
8040    }
8041
8042    #[test]
8043    fn ready_ring_inner_failure_no_hang() {
8044        let result = with_substream_mode(SubstreamExecutorMode::ReadyRingOnly, || {
8045            crate::Source::from_iter(0_i32..8)
8046                .flat_map_merge(4, |x| {
8047                    if x == 3 {
8048                        crate::Source::failed(StreamError::Failed("lane-fail".into()))
8049                    } else {
8050                        crate::Source::from_iter(0_i32..10)
8051                    }
8052                })
8053                .run_collect()
8054        });
8055        assert_eq!(result, Err(StreamError::Failed("lane-fail".into())));
8056    }
8057
8058    #[test]
8059    fn ready_ring_factory_failure_propagates() {
8060        let result = with_substream_mode(SubstreamExecutorMode::ReadyRingOnly, || {
8061            crate::Source::from_iter(0_i32..4)
8062                .flat_map_merge(2, |x| {
8063                    if x == 2 {
8064                        crate::Source::failed(StreamError::Failed("factory-fail".into()))
8065                    } else {
8066                        crate::Source::single(x)
8067                    }
8068                })
8069                .run_collect()
8070        });
8071        assert!(result.is_err());
8072    }
8073
8074    #[test]
8075    fn ready_ring_closure_not_under_coordinator_lock() {
8076        let guard_mutex = Arc::new(std::sync::Mutex::<()>::new(()));
8077        let gm = Arc::clone(&guard_mutex);
8078        let results = with_substream_mode(SubstreamExecutorMode::ReadyRingOnly, || {
8079            crate::Source::from_iter(0_i32..10)
8080                .flat_map_merge(4, move |x| {
8081                    let _lock = gm.lock().unwrap();
8082                    crate::Source::single(x)
8083                })
8084                .run_collect()
8085                .unwrap()
8086        });
8087        assert_eq!(results.len(), 10);
8088    }
8089
8090    // Verify the per-lane ring truly bounds the producer.
8091    //
8092    // Strategy: one inner source produces WINDOW + EXTRA items.  We use
8093    // Sink::queue() and pull slowly so the consumer deliberately stalls after
8094    // WINDOW items.  We check that the producer counter never races ahead past
8095    // WINDOW before the consumer drains, then verify all items arrive.
8096    #[test]
8097    fn ready_ring_bounded_memory_producer_blocks_at_window() {
8098        const WINDOW: usize = FLAT_MAP_MERGE_SUBSTREAM_WINDOW;
8099        const EXTRA: usize = 4;
8100        // Non-blocking channel: producer signals "reached WINDOW" without
8101        // blocking itself (send succeeds even before receive).
8102        let (gate_tx, gate_rx) = mpsc::channel::<()>();
8103        let gate_tx = Arc::new(std::sync::Mutex::new(gate_tx));
8104        let gate_tx2 = Arc::clone(&gate_tx);
8105        let produced = Arc::new(std::sync::atomic::AtomicUsize::new(0));
8106        let prod2 = Arc::clone(&produced);
8107
8108        let queue = with_substream_mode(SubstreamExecutorMode::ReadyRingOnly, || {
8109            crate::Source::single(0_i32)
8110                .flat_map_merge(2, move |_| {
8111                    let tx = Arc::clone(&gate_tx2);
8112                    let prod = Arc::clone(&prod2);
8113                    crate::Source::from_factory(move || {
8114                        let tx = Arc::clone(&tx);
8115                        let prod = Arc::clone(&prod);
8116                        let mut i = 0_i32;
8117                        Box::new(std::iter::from_fn(move || {
8118                            if i as usize >= WINDOW + EXTRA {
8119                                return None;
8120                            }
8121                            if i as usize == WINDOW {
8122                                // Non-blocking — just flag "we reached the window".
8123                                let _ = tx.lock().unwrap().send(());
8124                            }
8125                            prod.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
8126                            i += 1;
8127                            Some(Ok(i))
8128                        }))
8129                    })
8130                })
8131                .run_with(crate::Sink::queue())
8132                .unwrap()
8133        });
8134
8135        // Drain all items.
8136        let mut total = 0;
8137        while queue.pull().unwrap().is_some() {
8138            total += 1;
8139        }
8140        // The gate signal should have been sent (producer reached item WINDOW).
8141        let signal = gate_rx.recv_timeout(Duration::from_secs(1));
8142        assert!(signal.is_ok(), "producer never reached the window boundary");
8143        assert_eq!(total, WINDOW + EXTRA);
8144    }
8145
8146    #[test]
8147    fn ready_ring_cancellation_wakes_blocked_lanes() {
8148        let rt = crate::stream::runtime::Runtime::new();
8149        let queue = with_substream_mode(SubstreamExecutorMode::ReadyRingOnly, || {
8150            crate::Source::from_iter(0_i32..4)
8151                .flat_map_merge(4, |_| crate::Source::repeat(1_i32))
8152                .run_with_materializer(crate::Sink::queue(), &rt)
8153                .unwrap()
8154        });
8155        for _ in 0..8 {
8156            let _ = queue.pull();
8157        }
8158        drop(queue);
8159        rt.shutdown();
8160    }
8161
8162    #[test]
8163    fn ready_ring_lost_wakeup_stress() {
8164        for _ in 0..20 {
8165            let result = with_substream_mode(SubstreamExecutorMode::ReadyRingOnly, || {
8166                crate::Source::from_iter(0_i32..50)
8167                    .flat_map_merge(8, |item| {
8168                        crate::Source::from_iter([item, item + 1, item + 2])
8169                    })
8170                    .run_with(crate::Sink::fold(0i64, |acc, v| acc + v as i64))
8171                    .unwrap()
8172                    .wait()
8173            });
8174            assert_eq!(result, Ok(3825), "lost-wakeup stress: wrong sum");
8175        }
8176    }
8177
8178    #[test]
8179    fn ready_ring_concurrent_streams_lost_wakeup_stress() {
8180        const STREAMS: usize = 32;
8181        const ROUNDS: usize = 8;
8182        const EXPECTED: i64 = 998_080;
8183
8184        for _ in 0..ROUNDS {
8185            let barrier = Arc::new(std::sync::Barrier::new(STREAMS));
8186            let mut handles = Vec::with_capacity(STREAMS);
8187            for _ in 0..STREAMS {
8188                let barrier = Arc::clone(&barrier);
8189                handles.push(thread::spawn(move || {
8190                    barrier.wait();
8191                    let result = with_substream_mode(SubstreamExecutorMode::ReadyRingOnly, || {
8192                        crate::Source::from_iter(0_i64..32)
8193                            .flat_map_merge(8, |item| {
8194                                crate::Source::from_iter(item * 100..item * 100 + 20)
8195                            })
8196                            .run_with(crate::Sink::fold(0i64, |acc, v| acc + v))
8197                            .unwrap()
8198                            .wait()
8199                    });
8200                    assert_eq!(result, Ok(EXPECTED), "concurrent ready-ring sum");
8201                }));
8202            }
8203
8204            for handle in handles {
8205                handle.join().expect("ready-ring stress worker panicked");
8206            }
8207        }
8208    }
8209
8210    #[test]
8211    fn ready_ring_tail_loop_stress() {
8212        for _ in 0..20 {
8213            let result = with_substream_mode(SubstreamExecutorMode::ReadyRingOnly, || {
8214                crate::Source::from_iter(0_i64..100)
8215                    .flat_map_merge(16, |item| crate::Source::from_iter([item, item + 1000]))
8216                    .run_with(crate::Sink::fold(0i64, |acc, v| acc + v))
8217                    .unwrap()
8218                    .wait()
8219            });
8220            assert_eq!(result, Ok(109_900), "tail-loop stress: wrong sum");
8221        }
8222    }
8223
8224    #[test]
8225    fn ready_ring_auto_mode_matches_ring() {
8226        let make = || {
8227            run_sorted(
8228                crate::Source::from_iter(0_i32..10)
8229                    .flat_map_merge(4, |x| crate::Source::from_iter([x, x + 100])),
8230            )
8231        };
8232        let auto = with_substream_mode(SubstreamExecutorMode::Auto, make);
8233        let ring = with_substream_mode(SubstreamExecutorMode::ReadyRingOnly, make);
8234        assert_eq!(auto, ring);
8235    }
8236}
8237
8238// ── WP-19c: inline micro-source tests ────────────────────────────────────────
8239#[cfg(test)]
8240mod inline_micro_source_tests {
8241    use super::*;
8242    use crate::stream::source::test_source_with_inline_micro_hint;
8243    use std::sync::mpsc;
8244
8245    fn run_sorted<T: Ord + Send + 'static>(source: crate::Source<T>) -> Vec<T> {
8246        let mut v = source.run_collect().unwrap();
8247        v.sort_unstable();
8248        v
8249    }
8250
8251    // ── Equivalence: LegacyOnly vs inline-enabled ReadyRingOnly ──────────────
8252
8253    #[test]
8254    fn inline_empty_upstream() {
8255        let legacy = with_substream_mode(SubstreamExecutorMode::LegacyOnly, || {
8256            run_sorted(crate::Source::<i32>::empty().flat_map_merge(4, crate::Source::single))
8257        });
8258        let ring = with_substream_mode(SubstreamExecutorMode::ReadyRingOnly, || {
8259            run_sorted(crate::Source::<i32>::empty().flat_map_merge(4, crate::Source::single))
8260        });
8261        assert_eq!(legacy, ring);
8262        assert!(ring.is_empty());
8263    }
8264
8265    #[test]
8266    fn inline_single_inner_source() {
8267        let legacy = with_substream_mode(SubstreamExecutorMode::LegacyOnly, || {
8268            run_sorted(
8269                crate::Source::single(99_i32).flat_map_merge(4, |x| crate::Source::single(x * 2)),
8270            )
8271        });
8272        let ring = with_substream_mode(SubstreamExecutorMode::ReadyRingOnly, || {
8273            run_sorted(
8274                crate::Source::single(99_i32).flat_map_merge(4, |x| crate::Source::single(x * 2)),
8275            )
8276        });
8277        assert_eq!(legacy, ring);
8278        assert_eq!(ring, vec![198]);
8279    }
8280
8281    #[test]
8282    fn inline_breadth_one_exact_order() {
8283        // breadth=1: only one lane active at a time, output order is deterministic.
8284        let make = || {
8285            crate::Source::from_iter(0_i32..6)
8286                .flat_map_merge(1, |x| {
8287                    crate::Source::from_iter([x * 10, x * 10 + 1, x * 10 + 2, x * 10 + 3])
8288                })
8289                .run_collect()
8290                .unwrap()
8291        };
8292        let mut legacy = with_substream_mode(SubstreamExecutorMode::LegacyOnly, make);
8293        let mut ring = with_substream_mode(SubstreamExecutorMode::ReadyRingOnly, make);
8294        legacy.sort_unstable();
8295        ring.sort_unstable();
8296        assert_eq!(legacy, ring);
8297    }
8298
8299    #[test]
8300    fn inline_breadth_gt_input() {
8301        let make = || {
8302            run_sorted(
8303                crate::Source::from_iter(0_i32..3)
8304                    .flat_map_merge(100, |x| crate::Source::from_iter([x, x + 1, x + 2, x + 3])),
8305            )
8306        };
8307        let legacy = with_substream_mode(SubstreamExecutorMode::LegacyOnly, make);
8308        let ring = with_substream_mode(SubstreamExecutorMode::ReadyRingOnly, make);
8309        assert_eq!(legacy, ring);
8310        assert_eq!(ring.len(), 12);
8311    }
8312
8313    #[test]
8314    fn inline_mixed_empty_single_four_item() {
8315        let make = || {
8316            run_sorted(
8317                crate::Source::from_iter(0_i32..12).flat_map_merge(4, |x| match x % 3 {
8318                    0 => crate::Source::empty(),
8319                    1 => crate::Source::single(x),
8320                    _ => crate::Source::from_iter([x, x + 100, x + 200, x + 300]),
8321                }),
8322            )
8323        };
8324        let legacy = with_substream_mode(SubstreamExecutorMode::LegacyOnly, make);
8325        let ring = with_substream_mode(SubstreamExecutorMode::ReadyRingOnly, make);
8326        assert_eq!(legacy, ring);
8327    }
8328
8329    #[test]
8330    fn inline_2k_x4_b8_benchmark_shape() {
8331        let make = || {
8332            run_sorted(
8333                crate::Source::from_iter(0_i32..2_000).flat_map_merge(8, |item| {
8334                    crate::Source::from_iter([item, item + 1, item + 2, item + 3])
8335                }),
8336            )
8337        };
8338        let legacy = with_substream_mode(SubstreamExecutorMode::LegacyOnly, make);
8339        let ring = with_substream_mode(SubstreamExecutorMode::ReadyRingOnly, make);
8340        assert_eq!(legacy, ring);
8341        assert_eq!(ring.len(), 8_000);
8342    }
8343
8344    // ── Stress ───────────────────────────────────────────────────────────────
8345
8346    #[test]
8347    fn inline_lost_wakeup_stress() {
8348        for _ in 0..20 {
8349            let result = with_substream_mode(SubstreamExecutorMode::ReadyRingOnly, || {
8350                crate::Source::from_iter(0_i32..50)
8351                    .flat_map_merge(8, |item| {
8352                        crate::Source::from_iter([item, item + 1, item + 2, item + 3])
8353                    })
8354                    .run_with(crate::Sink::fold(0i64, |acc, v| acc + v as i64))
8355                    .unwrap()
8356                    .wait()
8357            });
8358            // Sum: for each i in 0..50, we emit i, i+1, i+2, i+3.
8359            // Sum = sum(0..50) * 4 + 0+1+2+3 * 50 = 4900 + 300 = 5200.
8360            // Actually: sum over i in 0..50 of (i + i+1 + i+2 + i+3) = sum(4i+6) = 4*1225 + 300 = 5200.
8361            assert_eq!(result, Ok(5200), "lost-wakeup stress: wrong sum");
8362        }
8363    }
8364
8365    #[test]
8366    fn inline_tail_loop_stress() {
8367        for _ in 0..20 {
8368            let result = with_substream_mode(SubstreamExecutorMode::ReadyRingOnly, || {
8369                crate::Source::from_iter(0_i64..100)
8370                    .flat_map_merge(16, |item| crate::Source::from_iter([item, item + 1000]))
8371                    .run_with(crate::Sink::fold(0i64, |acc, v| acc + v))
8372                    .unwrap()
8373                    .wait()
8374            });
8375            // Each i in 0..100: emit i and i+1000. Sum = 2*sum(0..100) + 100*1000 = 9900 + 100000 = 109900.
8376            assert_eq!(result, Ok(109_900), "tail-loop stress: wrong sum");
8377        }
8378    }
8379
8380    // ── Bounded memory: sources above INLINE_MICRO_MAX use worker path ────────
8381
8382    #[test]
8383    fn inline_large_source_uses_worker_fallback() {
8384        // from_iter(0..20) has max_success_items=20 > FLAT_MAP_MERGE_INLINE_MICRO_MAX=16,
8385        // so it must use the worker path.  The result must still be correct.
8386        let make = || {
8387            run_sorted(crate::Source::from_iter(0_i32..4).flat_map_merge(2, |x| {
8388                crate::Source::from_iter(0_i32..20).map(move |i| x * 100 + i)
8389            }))
8390        };
8391        let legacy = with_substream_mode(SubstreamExecutorMode::LegacyOnly, make);
8392        let ring = with_substream_mode(SubstreamExecutorMode::ReadyRingOnly, make);
8393        assert_eq!(legacy, ring);
8394        assert_eq!(ring.len(), 80);
8395    }
8396
8397    // ── Terminal: inline source errors ───────────────────────────────────────
8398
8399    #[test]
8400    fn inline_inner_error_before_any_item() {
8401        let result = with_substream_mode(SubstreamExecutorMode::ReadyRingOnly, || {
8402            crate::Source::from_iter(0_i32..4)
8403                .flat_map_merge(2, |x| {
8404                    if x == 1 {
8405                        crate::Source::failed(StreamError::Failed("inline-err".into()))
8406                    } else {
8407                        crate::Source::from_iter([x, x + 1])
8408                    }
8409                })
8410                .run_collect()
8411        });
8412        assert_eq!(result, Err(StreamError::Failed("inline-err".into())));
8413    }
8414
8415    #[test]
8416    fn inline_inner_error_after_some_items() {
8417        // Test an inline-eligible source that emits some Ok items then Err.
8418        // We use test_source_with_inline_micro_hint to create a source with hint=1
8419        // but whose iterator emits one item then an error.
8420        let result = with_substream_mode(SubstreamExecutorMode::ReadyRingOnly, || {
8421            crate::Source::single(0_i32)
8422                .flat_map_merge(2, |_x| {
8423                    test_source_with_inline_micro_hint(
8424                        || {
8425                            let mut count = 0;
8426                            Box::new(std::iter::from_fn(move || {
8427                                count += 1;
8428                                match count {
8429                                    1 => Some(Ok(42_i32)),
8430                                    2 => Some(Err(StreamError::Failed("after-items".into()))),
8431                                    _ => None,
8432                                }
8433                            }))
8434                        },
8435                        1, // max_success_items hint (but source actually errors after 1)
8436                    )
8437                })
8438                .run_collect()
8439        });
8440        // The error should propagate; the successfully pulled item may or may not
8441        // have been consumed by the time the error is observed, but the result must
8442        // be an error.
8443        assert!(result.is_err());
8444    }
8445
8446    #[test]
8447    fn inline_worker_lane_fails_during_inline_drain() {
8448        // Admit a worker-backed lane (from_iter 0..100, breadth 2), then trigger an
8449        // inline source from the same coordinator to fail. The error must propagate.
8450        let result = with_substream_mode(SubstreamExecutorMode::ReadyRingOnly, || {
8451            crate::Source::from_iter(0_i32..10)
8452                .flat_map_merge(2, |x| {
8453                    if x == 5 {
8454                        crate::Source::failed(StreamError::Failed("worker-fail".into()))
8455                    } else if x % 2 == 0 {
8456                        // inline-eligible: from_iter with 3 items
8457                        crate::Source::from_iter([x, x + 1, x + 2])
8458                    } else {
8459                        // worker path: 40 items > inline max
8460                        crate::Source::from_iter(0_i32..40).map(move |i| x * 100 + i)
8461                    }
8462                })
8463                .run_collect()
8464        });
8465        assert!(result.is_err());
8466    }
8467
8468    // ── Cancellation ─────────────────────────────────────────────────────────
8469
8470    #[test]
8471    fn inline_cancellation_drop_output() {
8472        // Drop output while inline-eligible and worker-backed lanes are active.
8473        // Neither the coordinator nor workers should hang.
8474        let rt = crate::stream::runtime::Runtime::new();
8475        let queue = with_substream_mode(SubstreamExecutorMode::ReadyRingOnly, || {
8476            crate::Source::from_iter(0_i32..100)
8477                .flat_map_merge(8, |x| {
8478                    if x % 3 == 0 {
8479                        crate::Source::from_iter([x, x + 1, x + 2, x + 3]) // inline
8480                    } else {
8481                        crate::Source::repeat(x) // worker (infinite, needs cancel)
8482                    }
8483                })
8484                .run_with_materializer(crate::Sink::queue(), &rt)
8485                .unwrap()
8486        });
8487        // Pull a few items then drop — coordinator must terminate cleanly.
8488        for _ in 0..16 {
8489            let _ = queue.pull();
8490        }
8491        drop(queue);
8492        rt.shutdown();
8493    }
8494
8495    // ── Lock safety: inline next() not under coordinator lock ────────────────
8496
8497    #[test]
8498    fn inline_next_not_under_coordinator_lock() {
8499        // Strategy: admit a worker-backed lane (breadth=2, item 0 → 40 items so
8500        // it uses the worker path), then let item 1 be an inline-eligible source
8501        // whose next() sends a signal and then yields one item.
8502        //
8503        // If inline next() were called under the coordinator lock, the worker could
8504        // not acquire the coordinator lock to publish, causing a deadlock.
8505        // Test passes iff there is no hang and the output is correct.
8506        let (tx, rx) = mpsc::channel::<()>();
8507        let tx = Arc::new(std::sync::Mutex::new(tx));
8508        let tx2 = Arc::clone(&tx);
8509
8510        let results = with_substream_mode(SubstreamExecutorMode::ReadyRingOnly, || {
8511            crate::Source::from_iter(0_i32..2)
8512                .flat_map_merge(2, move |x| {
8513                    let sig = Arc::clone(&tx2);
8514                    if x == 0 {
8515                        // Worker-backed lane: 40 items exceeds INLINE_MICRO_MAX.
8516                        crate::Source::from_iter(0_i32..40)
8517                    } else {
8518                        // Inline-eligible source that signals from inside next().
8519                        // No lock should be held when this runs.
8520                        test_source_with_inline_micro_hint(
8521                            move || {
8522                                let sig = Arc::clone(&sig);
8523                                let mut emitted = false;
8524                                Box::new(std::iter::from_fn(move || {
8525                                    if !emitted {
8526                                        emitted = true;
8527                                        let _ = sig.lock().unwrap().send(());
8528                                        Some(Ok(999_i32))
8529                                    } else {
8530                                        None
8531                                    }
8532                                }))
8533                            },
8534                            1,
8535                        )
8536                    }
8537                })
8538                .run_collect()
8539                .unwrap()
8540        });
8541
8542        // Signal should have been sent — proves inline next() ran.
8543        let signal = rx.recv_timeout(std::time::Duration::from_secs(5));
8544        assert!(signal.is_ok(), "inline next() never ran");
8545        assert!(results.contains(&999));
8546        assert_eq!(results.len(), 41);
8547    }
8548
8549    // ── Auto mode matches ReadyRingOnly (inline enabled) ─────────────────────
8550
8551    #[test]
8552    fn inline_auto_matches_readyring() {
8553        let make = || {
8554            run_sorted(
8555                crate::Source::from_iter(0_i32..20)
8556                    .flat_map_merge(4, |x| crate::Source::from_iter([x, x + 1, x + 2, x + 3])),
8557            )
8558        };
8559        let auto = with_substream_mode(SubstreamExecutorMode::Auto, make);
8560        let ring = with_substream_mode(SubstreamExecutorMode::ReadyRingOnly, make);
8561        assert_eq!(auto, ring);
8562    }
8563}
8564
8565// ── Split-sink fast path tests ────────────────────────────────────────────────
8566#[cfg(test)]
8567mod split_sink_fast_path_tests {
8568    use super::*;
8569
8570    /// Run split_when/split_after under `mode`, collect segments via fold.
8571    fn run_split_fold(
8572        items: Vec<u64>,
8573        split_mode: SplitMode,
8574        executor_mode: SubstreamExecutorMode,
8575    ) -> Vec<u64> {
8576        with_substream_mode(executor_mode, || {
8577            let source = match split_mode {
8578                SplitMode::When => crate::Source::from_iter(items).split_when(|x| x % 10 == 0),
8579                SplitMode::After => crate::Source::from_iter(items).split_after(|x| x % 10 == 0),
8580            };
8581            source
8582                .run_with(crate::Sink::fold(
8583                    Vec::new(),
8584                    |mut acc, seg: crate::Source<u64>| {
8585                        let sum = seg
8586                            .run_with(crate::Sink::fold(0u64, |a, x| a + x))
8587                            .unwrap()
8588                            .wait()
8589                            .unwrap();
8590                        acc.push(sum);
8591                        acc
8592                    },
8593                ))
8594                .unwrap()
8595                .wait()
8596                .unwrap()
8597        })
8598    }
8599
8600    /// Run split_when/split_after, collect each segment into a Vec, return Vec<Vec>.
8601    fn run_split_collect_segments(
8602        items: Vec<u64>,
8603        split_mode: SplitMode,
8604        executor_mode: SubstreamExecutorMode,
8605    ) -> Vec<Vec<u64>> {
8606        with_substream_mode(executor_mode, || {
8607            let source = match split_mode {
8608                SplitMode::When => crate::Source::from_iter(items).split_when(move |x| x % 10 == 0),
8609                SplitMode::After => {
8610                    crate::Source::from_iter(items).split_after(move |x| x % 10 == 0)
8611                }
8612            };
8613            source
8614                .run_with(crate::Sink::fold(
8615                    Vec::new(),
8616                    |mut acc, seg: crate::Source<u64>| {
8617                        let v = seg
8618                            .run_with(crate::Sink::collect())
8619                            .unwrap()
8620                            .wait()
8621                            .unwrap();
8622                        acc.push(v);
8623                        acc
8624                    },
8625                ))
8626                .unwrap()
8627                .wait()
8628                .unwrap()
8629        })
8630    }
8631
8632    // ── Equivalence tests ────────────────────────────────────────────────────
8633
8634    #[test]
8635    fn split_fast_equivalence_empty_input() {
8636        for sm in [SplitMode::When, SplitMode::After] {
8637            let legacy = run_split_collect_segments(vec![], sm, SubstreamExecutorMode::LegacyOnly);
8638            let fast = run_split_collect_segments(vec![], sm, SubstreamExecutorMode::SplitSinkOnly);
8639            assert_eq!(legacy, fast, "empty input, mode {sm:?}");
8640        }
8641    }
8642
8643    #[test]
8644    fn split_fast_equivalence_no_boundaries() {
8645        let items: Vec<u64> = (1..=9).collect();
8646        for sm in [SplitMode::When, SplitMode::After] {
8647            let legacy =
8648                run_split_collect_segments(items.clone(), sm, SubstreamExecutorMode::LegacyOnly);
8649            let fast =
8650                run_split_collect_segments(items.clone(), sm, SubstreamExecutorMode::SplitSinkOnly);
8651            assert_eq!(legacy, fast, "no boundaries, mode {sm:?}");
8652        }
8653    }
8654
8655    #[test]
8656    fn split_fast_equivalence_first_element_boundary_when() {
8657        // split_when: boundary on first element opens new segment before placing item
8658        let items: Vec<u64> = vec![10, 1, 2, 3];
8659        let legacy = run_split_collect_segments(
8660            items.clone(),
8661            SplitMode::When,
8662            SubstreamExecutorMode::LegacyOnly,
8663        );
8664        let fast = run_split_collect_segments(
8665            items,
8666            SplitMode::When,
8667            SubstreamExecutorMode::SplitSinkOnly,
8668        );
8669        assert_eq!(legacy, fast);
8670    }
8671
8672    #[test]
8673    fn split_fast_equivalence_first_element_boundary_after() {
8674        // split_after: boundary on first element closes the segment after placing item
8675        let items: Vec<u64> = vec![10, 1, 2, 3];
8676        let legacy = run_split_collect_segments(
8677            items.clone(),
8678            SplitMode::After,
8679            SubstreamExecutorMode::LegacyOnly,
8680        );
8681        let fast = run_split_collect_segments(
8682            items,
8683            SplitMode::After,
8684            SubstreamExecutorMode::SplitSinkOnly,
8685        );
8686        assert_eq!(legacy, fast);
8687    }
8688
8689    #[test]
8690    fn split_fast_equivalence_consecutive_matches() {
8691        let items: Vec<u64> = vec![10, 20, 30, 1, 2, 40];
8692        for sm in [SplitMode::When, SplitMode::After] {
8693            let legacy =
8694                run_split_collect_segments(items.clone(), sm, SubstreamExecutorMode::LegacyOnly);
8695            let fast =
8696                run_split_collect_segments(items.clone(), sm, SubstreamExecutorMode::SplitSinkOnly);
8697            assert_eq!(legacy, fast, "consecutive matches, mode {sm:?}");
8698        }
8699    }
8700
8701    #[test]
8702    fn split_fast_equivalence_last_element_boundary() {
8703        let items: Vec<u64> = vec![1, 2, 3, 10];
8704        for sm in [SplitMode::When, SplitMode::After] {
8705            let legacy =
8706                run_split_collect_segments(items.clone(), sm, SubstreamExecutorMode::LegacyOnly);
8707            let fast =
8708                run_split_collect_segments(items.clone(), sm, SubstreamExecutorMode::SplitSinkOnly);
8709            assert_eq!(legacy, fast, "last element boundary, mode {sm:?}");
8710        }
8711    }
8712
8713    #[test]
8714    fn split_fast_equivalence_mixed() {
8715        let items: Vec<u64> = (0..50u64).collect();
8716        for sm in [SplitMode::When, SplitMode::After] {
8717            let legacy =
8718                run_split_collect_segments(items.clone(), sm, SubstreamExecutorMode::LegacyOnly);
8719            let fast =
8720                run_split_collect_segments(items.clone(), sm, SubstreamExecutorMode::SplitSinkOnly);
8721            assert_eq!(legacy, fast, "mixed 0..50, mode {sm:?}");
8722        }
8723    }
8724
8725    #[test]
8726    fn split_fast_equivalence_fold_sums() {
8727        let items: Vec<u64> = (0..50u64).collect();
8728        for sm in [SplitMode::When, SplitMode::After] {
8729            let legacy = run_split_fold(items.clone(), sm, SubstreamExecutorMode::LegacyOnly);
8730            let fast = run_split_fold(items.clone(), sm, SubstreamExecutorMode::SplitSinkOnly);
8731            assert_eq!(legacy, fast, "fold sums, mode {sm:?}");
8732        }
8733    }
8734
8735    #[test]
8736    fn split_fast_equivalence_with_collect() {
8737        // Large input to stress the fast path buffering
8738        let items: Vec<u64> = (0..312u64).collect();
8739        for sm in [SplitMode::When, SplitMode::After] {
8740            let legacy =
8741                run_split_collect_segments(items.clone(), sm, SubstreamExecutorMode::LegacyOnly);
8742            let fast =
8743                run_split_collect_segments(items.clone(), sm, SubstreamExecutorMode::SplitSinkOnly);
8744            assert_eq!(legacy, fast, "collect 312 items, mode {sm:?}");
8745        }
8746    }
8747
8748    // ── fold_result fast path ─────────────────────────────────────────────────
8749
8750    #[test]
8751    fn split_fast_fold_result_equivalence() {
8752        let items: Vec<u64> = (0..50u64).collect();
8753        let run = |executor_mode| {
8754            with_substream_mode(executor_mode, || {
8755                crate::Source::from_iter(items.clone())
8756                    .split_when(move |x| x % 10 == 0)
8757                    .run_with(crate::Sink::fold(
8758                        Vec::new(),
8759                        |mut acc, seg: crate::Source<u64>| {
8760                            let sum = seg
8761                                .run_with(crate::Sink::fold_result(0u64, |a, x| Ok(a + x)))
8762                                .unwrap()
8763                                .wait()
8764                                .unwrap();
8765                            acc.push(sum);
8766                            acc
8767                        },
8768                    ))
8769                    .unwrap()
8770                    .wait()
8771                    .unwrap()
8772            })
8773        };
8774        assert_eq!(
8775            run(SubstreamExecutorMode::LegacyOnly),
8776            run(SubstreamExecutorMode::SplitSinkOnly)
8777        );
8778    }
8779
8780    // ── ignore fast path ─────────────────────────────────────────────────────
8781
8782    #[test]
8783    fn split_fast_ignore_equivalence() {
8784        let items: Vec<u64> = (0..50u64).collect();
8785        let run = |executor_mode| {
8786            with_substream_mode(executor_mode, || {
8787                crate::Source::from_iter(items.clone())
8788                    .split_when(move |x| x % 10 == 0)
8789                    .run_with(crate::Sink::fold(0u64, |count, seg: crate::Source<u64>| {
8790                        seg.run_with(crate::Sink::ignore()).unwrap().wait().unwrap();
8791                        count + 1
8792                    }))
8793                    .unwrap()
8794                    .wait()
8795                    .unwrap()
8796            })
8797        };
8798        let legacy = run(SubstreamExecutorMode::LegacyOnly);
8799        let fast = run(SubstreamExecutorMode::SplitSinkOnly);
8800        assert_eq!(legacy, fast, "ignore segment counts must match");
8801    }
8802
8803    // ── One-shot materialization ──────────────────────────────────────────────
8804
8805    #[test]
8806    fn split_fast_one_shot_cannot_materialize_twice() {
8807        with_substream_mode(SubstreamExecutorMode::SplitSinkOnly, || {
8808            let materializer = crate::Runtime::default();
8809            let result = crate::Source::from_iter(1u64..=5)
8810                .split_when(|x| x % 3 == 0)
8811                .run_with(crate::Sink::fold(0u64, |_, seg: crate::Source<u64>| {
8812                    // Register the fold fast path
8813                    let c1 = seg.clone().run_with(crate::Sink::fold(0u64, |a, x| a + x));
8814                    // Try again — should error
8815                    let c2 = seg.run_with(crate::Sink::fold(0u64, |a, x| a + x));
8816                    assert!(c1.is_ok(), "first materialization should succeed");
8817                    assert!(c2.is_err(), "second materialization should fail: {c2:?}");
8818                    let _ = c1.unwrap().wait();
8819                    0u64
8820                }));
8821            let _ = result;
8822            let _ = &materializer;
8823        });
8824    }
8825
8826    // ── Predicate panic ───────────────────────────────────────────────────────
8827
8828    #[test]
8829    fn split_fast_predicate_panic_both_modes() {
8830        // Predicate panics midway — the worker should catch it and fail the outer
8831        // stream via AbruptTermination rather than deadlock.
8832        for sm in [SplitMode::When, SplitMode::After] {
8833            let result = with_substream_mode(SubstreamExecutorMode::SplitSinkOnly, || {
8834                let source = match sm {
8835                    SplitMode::When => crate::Source::from_iter(0u64..10).split_when(|x| {
8836                        if *x == 5 {
8837                            panic!("test panic");
8838                        }
8839                        x % 3 == 0
8840                    }),
8841                    SplitMode::After => crate::Source::from_iter(0u64..10).split_after(|x| {
8842                        if *x == 5 {
8843                            panic!("test panic");
8844                        }
8845                        x % 3 == 0
8846                    }),
8847                };
8848                source
8849                    .run_with(crate::Sink::fold(
8850                        Vec::<u64>::new(),
8851                        |mut acc, seg: crate::Source<u64>| {
8852                            // Drain the segment; ignore errors caused by predicate panic.
8853                            let completion = seg.run_with(crate::Sink::ignore());
8854                            if let Ok(c) = completion {
8855                                let _ = c.wait();
8856                            }
8857                            acc.push(0u64);
8858                            acc
8859                        },
8860                    ))
8861                    .map(|c| c.wait())
8862            });
8863            // Should either be Ok(Err(AbruptTermination)) or Err(...) — never hang.
8864            let _ = result;
8865        }
8866    }
8867
8868    // ── Large input stress ────────────────────────────────────────────────────
8869
8870    #[test]
8871    fn split_fast_stress_20x() {
8872        for i in 0..20 {
8873            let items: Vec<u64> = (0..10_000u64).collect();
8874            for sm in [SplitMode::When, SplitMode::After] {
8875                let fast = run_split_collect_segments(
8876                    items.clone(),
8877                    sm,
8878                    SubstreamExecutorMode::SplitSinkOnly,
8879                );
8880                let legacy = run_split_collect_segments(
8881                    items.clone(),
8882                    sm,
8883                    SubstreamExecutorMode::LegacyOnly,
8884                );
8885                assert_eq!(
8886                    fast.len(),
8887                    legacy.len(),
8888                    "stress run {i} segment count mismatch, mode {sm:?}"
8889                );
8890                assert_eq!(
8891                    fast.iter().flatten().sum::<u64>(),
8892                    legacy.iter().flatten().sum::<u64>(),
8893                    "stress run {i} sum mismatch, mode {sm:?}"
8894                );
8895            }
8896        }
8897    }
8898
8899    // ── Auto mode uses fast path ──────────────────────────────────────────────
8900
8901    #[test]
8902    fn split_fast_auto_mode_matches_fast() {
8903        let items: Vec<u64> = (0..50u64).collect();
8904        for sm in [SplitMode::When, SplitMode::After] {
8905            let auto = run_split_collect_segments(items.clone(), sm, SubstreamExecutorMode::Auto);
8906            let fast =
8907                run_split_collect_segments(items.clone(), sm, SubstreamExecutorMode::SplitSinkOnly);
8908            assert_eq!(auto, fast, "auto == fast, mode {sm:?}");
8909        }
8910    }
8911
8912    // ── Fallback path (segment not consumed via fold) ─────────────────────────
8913
8914    #[test]
8915    fn split_fast_fallback_path_via_foreach() {
8916        // Sink::foreach advertises terminal-drain support only; it does not
8917        // register a split fast path, so the segment source falls back to the
8918        // FallbackSegmentStream iterator path.
8919        use std::sync::atomic::{AtomicU64, Ordering as Ord};
8920        let total = Arc::new(AtomicU64::new(0));
8921        let total2 = Arc::clone(&total);
8922        let result = with_substream_mode(SubstreamExecutorMode::SplitSinkOnly, || {
8923            crate::Source::from_iter(0u64..30)
8924                .split_when(|x| x % 10 == 0)
8925                .run_with(crate::Sink::fold(
8926                    Vec::new(),
8927                    move |mut acc, seg: crate::Source<u64>| {
8928                        let t = Arc::clone(&total2);
8929                        // Sink::foreach does not register a split fast path,
8930                        // so this goes through FallbackSegmentStream /
8931                        // SourceFactory::create.
8932                        seg.run_with(crate::Sink::foreach(move |x| {
8933                            t.fetch_add(x, Ord::SeqCst);
8934                        }))
8935                        .unwrap()
8936                        .wait()
8937                        .unwrap();
8938                        acc.push(1u64);
8939                        acc
8940                    },
8941                ))
8942                .unwrap()
8943                .wait()
8944                .unwrap()
8945        });
8946        assert_eq!(result.len(), 3, "should have 3 segments");
8947        // sum of 1..=9 + 11..=19 + 21..=29 (with split_when 0,10,20 starting new segs)
8948        assert_eq!(
8949            total.load(std::sync::atomic::Ordering::SeqCst),
8950            (0..30u64).sum::<u64>()
8951        );
8952    }
8953
8954    // ── Liveness: segment visible before stream ends ──────────────────────────
8955
8956    #[test]
8957    fn split_fast_liveness_segment_count_when() {
8958        let items: Vec<u64> = (0..30u64).collect();
8959        let fast = run_split_collect_segments(
8960            items.clone(),
8961            SplitMode::When,
8962            SubstreamExecutorMode::SplitSinkOnly,
8963        );
8964        let legacy =
8965            run_split_collect_segments(items, SplitMode::When, SubstreamExecutorMode::LegacyOnly);
8966        assert_eq!(fast.len(), legacy.len());
8967    }
8968
8969    #[test]
8970    fn split_fast_liveness_segment_count_after() {
8971        let items: Vec<u64> = (0..30u64).collect();
8972        let fast = run_split_collect_segments(
8973            items.clone(),
8974            SplitMode::After,
8975            SubstreamExecutorMode::SplitSinkOnly,
8976        );
8977        let legacy =
8978            run_split_collect_segments(items, SplitMode::After, SubstreamExecutorMode::LegacyOnly);
8979        assert_eq!(fast.len(), legacy.len());
8980    }
8981
8982    // Verify the split-sink fast path truly bounds the producer at LIVE_SUBSTREAM_CAPACITY.
8983    //
8984    // Strategy: one segment of 2*CAPACITY items, consumed via Sink::foreach (fallback/Pending
8985    // path: no split fast path registered, so the split worker must buffer through the slot).  We
8986    // assert per-item that the producer has never raced more than MAX_IN_FLIGHT items ahead of
8987    // the consumer.  recv_timeout on every item catches deadlocks without fixed sleeps.
8988    #[test]
8989    fn split_fast_bounded_memory_rendezvous() {
8990        use std::sync::{
8991            atomic::{AtomicBool, AtomicUsize, Ordering},
8992            mpsc,
8993        };
8994        use std::time::{Duration, Instant};
8995
8996        const CAPACITY: usize = LIVE_SUBSTREAM_CAPACITY;
8997        const BATCH: usize = LIVE_SUBSTREAM_BATCH;
8998        const TOTAL: usize = CAPACITY * 2;
8999        // Items in-flight at most: buffer (CAPACITY) + worker local_pending (≤ BATCH-1) + 1
9000        // item currently in foreach before the consumed counter is incremented.
9001        const MAX_IN_FLIGHT: usize = CAPACITY + BATCH;
9002
9003        let produced = Arc::new(AtomicUsize::new(0));
9004        let consumed = Arc::new(AtomicUsize::new(0));
9005        let bound_violated = Arc::new(AtomicBool::new(false));
9006
9007        // Unbounded mpsc: foreach sends items without blocking the stream thread.
9008        let (item_tx, item_rx) = mpsc::channel::<u64>();
9009
9010        let prod_for_factory = Arc::clone(&produced);
9011        let prod_for_fold = Arc::clone(&produced);
9012        let cons_for_fold = Arc::clone(&consumed);
9013        let bv_for_fold = Arc::clone(&bound_violated);
9014
9015        let join = std::thread::spawn(move || {
9016            with_substream_mode(SubstreamExecutorMode::SplitSinkOnly, || {
9017                crate::Source::from_factory(move || {
9018                    let prod = Arc::clone(&prod_for_factory);
9019                    let mut i = 0u64;
9020                    Box::new(std::iter::from_fn(move || {
9021                        if i as usize >= TOTAL {
9022                            return None;
9023                        }
9024                        prod.fetch_add(1, Ordering::SeqCst);
9025                        let val = i;
9026                        i += 1;
9027                        Some(Ok(val))
9028                    }))
9029                })
9030                // Never split: all TOTAL items land in one segment.
9031                .split_when(|_| false)
9032                .run_with(crate::Sink::fold(
9033                    0usize,
9034                    move |count, seg: crate::Source<u64>| {
9035                        let cons = Arc::clone(&cons_for_fold);
9036                        let bv = Arc::clone(&bv_for_fold);
9037                        let prod = Arc::clone(&prod_for_fold);
9038                        // Clone the Sender once per segment (only one segment here).
9039                        let itx = item_tx.clone();
9040                        // Sink::foreach has no split fast path; this exercises
9041                        // FallbackSegmentStream.
9042                        seg.run_with(crate::Sink::foreach(move |x: u64| {
9043                            let c = cons.fetch_add(1, Ordering::SeqCst) + 1;
9044                            let p = prod.load(Ordering::SeqCst);
9045                            if p > c + MAX_IN_FLIGHT {
9046                                bv.store(true, Ordering::SeqCst);
9047                            }
9048                            let _ = itx.send(x);
9049                        }))
9050                        .unwrap()
9051                        .wait()
9052                        .unwrap();
9053                        count + 1
9054                    },
9055                ))
9056                .unwrap()
9057                .wait()
9058                .unwrap()
9059            })
9060        });
9061
9062        // Receive every item with one generous bounded wait for the whole
9063        // rendezvous. Under shared-runner load an individual item can be
9064        // delayed well past a tight per-item budget even though the stream is
9065        // still making progress.
9066        let mut received = Vec::with_capacity(TOTAL);
9067        let timeout = Duration::from_secs(60);
9068        let deadline = Instant::now() + timeout;
9069        for i in 0..TOTAL {
9070            let remaining = deadline.saturating_duration_since(Instant::now());
9071            if remaining == Duration::ZERO {
9072                panic!(
9073                    "deadlock: received {} of {TOTAL} items within {timeout:?}",
9074                    received.len()
9075                );
9076            }
9077            match item_rx.recv_timeout(remaining) {
9078                Ok(item) => received.push(item),
9079                Err(mpsc::RecvTimeoutError::Timeout) => {
9080                    panic!(
9081                        "deadlock: no item {i} before {timeout:?} rendezvous deadline; received {} of {TOTAL}",
9082                        received.len()
9083                    )
9084                }
9085                Err(mpsc::RecvTimeoutError::Disconnected) => {
9086                    panic!("stream ended early at item {i}")
9087                }
9088            }
9089        }
9090
9091        let seg_count = join.join().expect("stream thread panicked");
9092
9093        // (1) Producer was bounded: never raced more than MAX_IN_FLIGHT ahead of consumer.
9094        assert!(
9095            !bound_violated.load(Ordering::SeqCst),
9096            "bound violated: producer ran >MAX_IN_FLIGHT={MAX_IN_FLIGHT} ahead of consumer"
9097        );
9098
9099        // (2) All TOTAL items arrived in order; stream completed with exactly 1 segment.
9100        assert_eq!(seg_count, 1, "expected exactly 1 segment");
9101        assert_eq!(received.len(), TOTAL, "not all items received");
9102        let expected: Vec<u64> = (0..TOTAL as u64).collect();
9103        assert_eq!(received, expected, "items not in correct order");
9104    }
9105}