mecha10_core/context/receiver/

combinators.rs

//! `map`, `filter`, and `filter_map` combinators.

use super::relay::RelayAction;
use super::types::{Receiver, ReceiverInner};
use tokio::sync::mpsc;

enum MappedSender<T> {
    Unbounded(mpsc::UnboundedSender<T>),
    Bounded(mpsc::Sender<T>),
}

impl<T> Receiver<T> {
    /// Map (transform) messages as they are received
    ///
    /// Creates a new receiver that applies a transformation function to each message.
    /// This is useful for data transformations like resizing images, converting units,
    /// or extracting fields from complex messages.
    ///
    /// # Arguments
    ///
    /// * `f` - Transformation function to apply to each message
    ///
    /// # Example
    ///
    /// ```rust
    /// # use mecha10_core::prelude::*;
    /// # use mecha10_core::topics::{actuator, sensor};
    /// # async fn example(ctx: &Context) -> Result<()> {
    /// // Resize camera images on the fly
    /// let mut resized_images = ctx.subscribe(sensor::CAMERA_RGB)
    ///     .await?
    ///     .map(|img| resize_image(img, 320, 240));
    ///
    /// while let Some(small_image) = resized_images.recv().await {
    ///     // Process resized image
    /// }
    ///
    /// // Extract specific fields from complex messages
    /// let mut speeds = ctx.subscribe(actuator::CMD_VEL)
    ///     .await?
    ///     .map(|twist| twist.linear[0]);
    ///
    /// while let Some(speed) = speeds.recv().await {
    ///     println!("Current speed: {}", speed);
    /// }
    /// # Ok(())
    /// # }
    /// # fn resize_image(img: mecha10_core::messages::Image, w: u32, h: u32) -> mecha10_core::messages::Image { img }
    /// ```
    pub fn map<U, F>(self, f: F) -> Receiver<U>
    where
        F: FnMut(T) -> U + Send + 'static,
        T: Send + 'static,
        U: Send + 'static,
    {
        let (tx, rx) = match &self.inner {
            ReceiverInner::Unbounded(_) => {
                let (tx, rx) = mpsc::unbounded_channel();
                (MappedSender::Unbounded(tx), ReceiverInner::Unbounded(rx))
            }
            ReceiverInner::Bounded(bounded_rx) => {
                let capacity = bounded_rx.max_capacity();
                let (tx, rx) = mpsc::channel(capacity);
                (MappedSender::Bounded(tx), ReceiverInner::Bounded(rx))
            }
            ReceiverInner::DropOldest(drop_oldest_rx) => {
                let capacity = drop_oldest_rx.capacity();
                let (tx, rx) = mpsc::channel(capacity);
                (MappedSender::Bounded(tx), ReceiverInner::Bounded(rx))
            }
        };

        let mut receiver = self;
        let mut transform = f;
        tokio::spawn(async move {
            while let Some(msg) = receiver.recv().await {
                let transformed = transform(msg);

                let send_result = match tx {
                    MappedSender::Unbounded(ref tx) => tx.send(transformed).map_err(|_| ()),
                    MappedSender::Bounded(ref tx) => tx.send(transformed).await.map_err(|_| ()),
                };

                if send_result.is_err() {
                    break;
                }
            }
        });

        Receiver { inner: rx }
    }

    /// Filter messages based on a predicate
    ///
    /// Creates a new receiver that only passes through messages that satisfy the predicate.
    /// This is useful for filtering sensor data, ignoring invalid readings, or implementing
    /// conditional processing.
    ///
    /// # Arguments
    ///
    /// * `predicate` - Function that returns `true` for messages to keep, `false` to drop
    ///
    /// # Example
    ///
    /// ```rust
    /// # use mecha10_core::prelude::*;
    /// # use mecha10_core::topics::{actuator, sensor};
    /// # async fn example(ctx: &Context) -> Result<()> {
    /// // Only process high-speed motor commands
    /// let mut high_speeds = ctx.subscribe(actuator::CMD_VEL)
    ///     .await?
    ///     .filter(|twist| twist.linear[0] > 0.5);
    ///
    /// while let Some(twist) = high_speeds.recv().await {
    ///     println!("High speed detected: {}", twist.linear[0]);
    /// }
    ///
    /// // Filter out invalid lidar readings
    /// let mut valid_scans = ctx.subscribe(sensor::LIDAR_SCAN)
    ///     .await?
    ///     .filter(|scan| scan.ranges.iter().all(|&r| r > 0.0 && r < 100.0));
    ///
    /// while let Some(scan) = valid_scans.recv().await {
    ///     // All ranges are valid
    /// }
    /// # Ok(())
    /// # }
    /// ```
    pub fn filter<F>(self, mut predicate: F) -> Receiver<T>
    where
        F: FnMut(&T) -> bool + Send + 'static,
        T: Send + 'static,
    {
        self.spawn_relay(move |msg| {
            if predicate(&msg) {
                RelayAction::Forward(msg)
            } else {
                RelayAction::Drop
            }
        })
    }

    /// Filter and map messages in one operation
    ///
    /// Combines filtering and mapping into a single efficient operation. The function
    /// returns `Some(U)` for messages to keep (and transform), or `None` to drop.
    /// This is more efficient than chaining `.filter()` and `.map()` separately.
    ///
    /// # Arguments
    ///
    /// * `f` - Function that returns `Some(transformed)` to keep, `None` to drop
    ///
    /// # Example
    ///
    /// ```rust
    /// # use mecha10_core::prelude::*;
    /// # use mecha10_core::topics::{actuator, system};
    /// # async fn example(ctx: &Context) -> Result<()> {
    /// // Extract speed only when robot is moving
    /// let mut active_speeds = ctx.subscribe(actuator::CMD_VEL)
    ///     .await?
    ///     .filter_map(|twist| {
    ///         if twist.linear[0] > 0.1 {
    ///             Some(twist.linear[0])
    ///         } else {
    ///             None
    ///         }
    ///     });
    ///
    /// while let Some(speed) = active_speeds.recv().await {
    ///     println!("Robot moving at: {}", speed);
    /// }
    ///
    /// // Parse valid diagnostics readings
    /// let mut parsed_data = ctx.subscribe(system::DIAGNOSTICS)
    ///     .await?
    ///     .filter_map(|raw| parse_sensor_data(&raw).ok());
    ///
    /// while let Some(data) = parsed_data.recv().await {
    ///     // Process successfully parsed data
    /// }
    /// # Ok(())
    /// # }
    /// # fn parse_sensor_data(raw: &String) -> Result<String> { Ok(raw.clone()) }
    /// ```
    pub fn filter_map<U, F>(self, mut f: F) -> Receiver<U>
    where
        F: FnMut(T) -> Option<U> + Send + 'static,
        T: Send + 'static,
        U: Send + 'static,
    {
        self.spawn_relay(move |msg| match f(msg) {
            Some(transformed) => RelayAction::Forward(transformed),
            None => RelayAction::Drop,
        })
    }
}