Skip to main content

rx_rust/operators/transforming/
buffer_with_count.rs

1use crate::utils::types::MaybeSend;
2use crate::{
3    observable::Observable,
4    observable::Subscription,
5    observer::{Flow, Observer, Termination},
6};
7use educe::Educe;
8use std::num::NonZeroUsize;
9
10/// Periodically gathers items from an Observable into bundles and emits these bundles as `Vec<T>`, when the bundle reaches a specified size.
11/// See <https://reactivex.io/documentation/operators/buffer.html>
12///
13/// # Examples
14/// ```rust
15/// use rx_rust::{
16///     observable::ObservableExt,
17///     observer::Termination,
18///     operators::{
19///         creating::from_iter::FromIter,
20///         transforming::buffer_with_count::BufferWithCount,
21///     },
22/// };
23/// use std::num::NonZeroUsize;
24///
25/// let mut values = Vec::new();
26/// let mut terminations = Vec::new();
27///
28/// let observable = BufferWithCount::new(FromIter::new(vec![1, 2, 3, 4]), NonZeroUsize::new(3).unwrap());
29/// observable.subscribe_with_callback(
30///     |value| values.push(value),
31///     |termination| terminations.push(termination),
32/// );
33///
34/// assert_eq!(values, vec![vec![1, 2, 3], vec![4]]);
35/// assert_eq!(terminations, vec![Termination::Completed]);
36/// ```
37#[derive(Educe)]
38#[educe(Debug, Clone)]
39pub struct BufferWithCount<OE> {
40    source: OE,
41    count: NonZeroUsize,
42}
43
44impl<OE> BufferWithCount<OE> {
45    pub fn new(source: OE, count: NonZeroUsize) -> Self {
46        Self { source, count }
47    }
48}
49
50impl<'or, T, E, OE> Observable<'or, Vec<T>, E> for BufferWithCount<OE>
51where
52    T: MaybeSend + 'or,
53    OE: Observable<'or, T, E>,
54{
55    type D = OE::D;
56
57    fn subscribe(
58        self,
59        observer: impl Observer<Vec<T>, E> + MaybeSend + 'or,
60    ) -> Subscription<Self::D> {
61        let observer = BufferWithCountObserver {
62            observer,
63            values: Vec::default(),
64            count: self.count,
65        };
66        self.source.subscribe(observer)
67    }
68}
69
70struct BufferWithCountObserver<T, OR> {
71    observer: OR,
72    values: Vec<T>,
73    count: NonZeroUsize,
74}
75
76impl<T, E, OR> Observer<T, E> for BufferWithCountObserver<T, OR>
77where
78    OR: Observer<Vec<T>, E>,
79{
80    fn on_next(&mut self, value: T) -> Flow {
81        self.values.push(value);
82        if self.values.len() >= self.count.get() {
83            self.observer.on_next(std::mem::take(&mut self.values))
84        } else {
85            Flow::Continue
86        }
87    }
88
89    fn on_termination(mut self, termination: Termination<E>) {
90        // The final value ends the stream, so a downstream that stopped on it is not completed
91        // on top of that: it has already ended itself.
92        if matches!(termination, Termination::Completed)
93            && !self.values.is_empty()
94            && self
95                .observer
96                .on_next(std::mem::take(&mut self.values))
97                .is_stop()
98        {
99            return;
100        }
101        self.observer.on_termination(termination);
102    }
103}