kaptan/libs/kaptan-runtime/src/mailbox.rs

//! Lock-free single-producer, single-consumer latest-value mailboxes.

use core::{
    marker::PhantomData,
    sync::atomic::{AtomicU32, Ordering},
};

use kaptan_sensor::{imu::ImuSnapshot, range::RangeSnapshot};
use kaptan_time::Timestamp;

/// Fixed-width snapshot representation for [`LatestMailbox`].
///
/// The representation consists only of atomic words. Implementations must
/// encode and decode every field without loss so an unread older observation
/// can be safely replaced by a newer one.
pub trait MailboxValue<const WORDS: usize>: Copy {
    /// Converts a snapshot into its atomic-word representation.
    fn into_words(self) -> [u32; WORDS];

    /// Reconstructs a snapshot from its atomic-word representation.
    fn from_words(words: [u32; WORDS]) -> Self;
}

/// Lock-free latest-value SPSC mailbox.
///
/// Publishing a newer value intentionally replaces an unread older value: the
/// output path is best-effort and must never delay acquisition. The sequence
/// counter is a seqlock; odd values indicate a writer and even values identify
/// a stable snapshot. Exactly one producer and one consumer may use an
/// instance concurrently. Sequence operations are sequentially consistent so
/// a reader cannot validate an old even sequence around concurrently written
/// atomic payload words; the payload words themselves need only be relaxed.
pub struct LatestMailbox<Value, const WORDS: usize> {
    sequence: AtomicU32,
    words: [AtomicU32; WORDS],
    value: PhantomData<Value>,
}

impl<Value: MailboxValue<WORDS>, const WORDS: usize> LatestMailbox<Value, WORDS> {
    /// Creates an empty mailbox.
    pub const fn new() -> Self {
        Self {
            sequence: AtomicU32::new(0),
            words: [const { AtomicU32::new(0) }; WORDS],
            value: PhantomData,
        }
    }

    /// Publishes a snapshot, replacing any unread older snapshot.
    pub fn publish(&self, value: Value) {
        let sequence = self.sequence.fetch_add(1, Ordering::SeqCst);
        debug_assert_eq!(sequence & 1, 0, "SPSC producer must not reenter");

        for (destination, word) in self.words.iter().zip(value.into_words()) {
            destination.store(word, Ordering::Relaxed);
        }

        self.sequence
            .store(sequence.wrapping_add(2), Ordering::SeqCst);
    }

    /// Returns the newest stable snapshot different from `observed_sequence`.
    pub fn take_newer(&self, observed_sequence: &mut u32) -> Option<Value> {
        let before = self.sequence.load(Ordering::SeqCst);
        if before == *observed_sequence || before & 1 != 0 {
            return None;
        }

        let mut words = [0; WORDS];
        for (destination, source) in words.iter_mut().zip(&self.words) {
            *destination = source.load(Ordering::Relaxed);
        }

        if self.sequence.load(Ordering::SeqCst) != before {
            return None;
        }

        *observed_sequence = before;
        Some(Value::from_words(words))
    }
}

impl<Value: MailboxValue<WORDS>, const WORDS: usize> Default for LatestMailbox<Value, WORDS> {
    fn default() -> Self {
        Self::new()
    }
}

impl MailboxValue<9> for ImuSnapshot {
    fn into_words(self) -> [u32; 9] {
        let timestamp = self.timestamp.as_microseconds();
        [
            timestamp as u32,
            (timestamp >> 32) as u32,
            self.acceleration_milligravity[0] as u16 as u32,
            self.acceleration_milligravity[1] as u16 as u32,
            self.acceleration_milligravity[2] as u16 as u32,
            self.angular_rate_milliradians_per_second[0] as u16 as u32,
            self.angular_rate_milliradians_per_second[1] as u16 as u32,
            self.angular_rate_milliradians_per_second[2] as u16 as u32,
            self.temperature_centidegrees_celsius as u16 as u32,
        ]
    }

    fn from_words(words: [u32; 9]) -> Self {
        ImuSnapshot::new(
            timestamp_from_words(words[0], words[1]),
            [
                words[2] as u16 as i16,
                words[3] as u16 as i16,
                words[4] as u16 as i16,
            ],
            [
                words[5] as u16 as i16,
                words[6] as u16 as i16,
                words[7] as u16 as i16,
            ],
            words[8] as u16 as i16,
        )
    }
}

impl MailboxValue<3> for RangeSnapshot {
    fn into_words(self) -> [u32; 3] {
        let timestamp = self.timestamp.as_microseconds();
        [
            timestamp as u32,
            (timestamp >> 32) as u32,
            self.distance_millimetres,
        ]
    }

    fn from_words(words: [u32; 3]) -> Self {
        RangeSnapshot::new(timestamp_from_words(words[0], words[1]), words[2])
    }
}

/// Latest-value mailbox for IMU telemetry snapshots.
pub type LatestImuMailbox = LatestMailbox<ImuSnapshot, 9>;
/// Latest-value mailbox for range telemetry snapshots.
pub type LatestRangeMailbox = LatestMailbox<RangeSnapshot, 3>;

const fn timestamp_from_words(low: u32, high: u32) -> Timestamp {
    Timestamp::from_microseconds((low as u64) | ((high as u64) << 32))
}

#[cfg(test)]
mod tests {
    use super::{LatestImuMailbox, LatestRangeMailbox};
    use kaptan_sensor::{imu::ImuSnapshot, range::RangeSnapshot};
    use kaptan_time::Timestamp;

    #[test]
    fn publishes_only_the_newest_unread_imu_sample() {
        let mailbox = LatestImuMailbox::new();
        let mut observed = 0;
        mailbox.publish(ImuSnapshot::new(
            Timestamp::from_microseconds(1),
            [1; 3],
            [2; 3],
            3,
        ));
        mailbox.publish(ImuSnapshot::new(
            Timestamp::from_microseconds(2),
            [4; 3],
            [5; 3],
            6,
        ));

        let sample = mailbox.take_newer(&mut observed).unwrap();
        assert_eq!(sample.timestamp, Timestamp::from_microseconds(2));
        assert_eq!(sample.acceleration_milligravity, [4; 3]);
        assert!(mailbox.take_newer(&mut observed).is_none());
    }

    #[test]
    fn publishes_only_the_newest_unread_range_sample() {
        let mailbox = LatestRangeMailbox::new();
        let mut observed = 0;
        mailbox.publish(RangeSnapshot::new(Timestamp::from_microseconds(1), 100));
        mailbox.publish(RangeSnapshot::new(Timestamp::from_microseconds(2), 200));

        assert_eq!(
            mailbox.take_newer(&mut observed),
            Some(RangeSnapshot::new(Timestamp::from_microseconds(2), 200))
        );
        assert!(mailbox.take_newer(&mut observed).is_none());
    }
}