ProjectNautilus/flight/src/commanded_state.rs

use log::trace;
use std::any::type_name;
use std::fmt::Debug;
use std::ops::Deref;
use std::time::Instant;

#[derive(Debug, Clone)]
struct StateEntry<T> {
    state: T,
    valid_until: Instant,
    priority: u8,
}

#[derive(Debug, Clone)]
pub struct CommandedState<T, const N: usize = 2> {
    entries: [Option<StateEntry<T>>; N],
    default: T,
    changed: bool,
}

impl<T, const N: usize> CommandedState<T, N>
where
    T: Debug,
{
    pub fn new(default: T) -> Self {
        trace!(
            "CommandedState::<{}>::new(default: {default:?})",
            type_name::<T>()
        );
        Self {
            entries: [const { None }; N],
            default,
            changed: true,
        }
    }

    pub fn consume_changed(&mut self) -> bool {
        trace!(
            "CommandedState::<{}>::consume_changed(self: {self:?})",
            type_name::<T>()
        );
        let res = self.changed;
        self.changed = false;
        res
    }

    fn evaluate_change_if<F>(&mut self, current_time: Instant, f: F) -> &T
    where
        F: FnOnce(&T, &T) -> bool,
    {
        trace!(
            "CommandedState::<{}>::evaluate_change_if(self: {self:?}, current_time: {current_time:?}, f: {})",
            type_name::<T>(),
            type_name::<F>()
        );
        let mut original_value: Option<StateEntry<T>> = None;
        for i in 0..N {
            if let Some(entry) = &self.entries[0] {
                if entry.valid_until >= current_time {
                    if let Some(original_value) = original_value {
                        self.changed |= f(&original_value.state, &entry.state);
                    }
                    // SAFETY: we just checked that this index is Some
                    // Unfortunately we can't use entry in order to satisfy the borrow checker
                    return self.entries[0].as_ref().map(|x| &x.state).unwrap();
                }
                // An expired entry - we need to remove it

                // Empty out this entry (keeping it in original value in case it is needed)
                if i == 0 {
                    original_value = self.entries[0].take();
                } else {
                    self.entries[0] = None;
                }
                // Bump up all future entries
                self.entries.rotate_left(1);
            } else {
                break;
            }
        }
        if let Some(original_value) = original_value {
            self.changed |= f(&original_value.state, &self.default);
        }
        &self.default
    }

    fn insert_change_if<F>(&mut self, state: T, valid_until: Instant, priority: u8, f: F)
    where
        F: FnOnce(&T, &T) -> bool,
    {
        trace!(
            "CommandedState::<{}>::insert_change_if(self: {self:?}, state: {state:?}, valid_until: {valid_until:?}, priority: {priority}, f: {})",
            type_name::<T>(),
            type_name::<F>()
        );
        for i in 0..N {
            if let Some(entry) = &mut self.entries[i] {
                // The current entry exists - let's find out it replacing it is an option
                if priority >= entry.priority {
                    // We have enough priority to modify this entry
                    if priority == entry.priority || entry.valid_until <= valid_until {
                        // If we are same priority (always replace) or
                        // we will be valid for longer then we can just replace this entry
                        if i == 0 {
                            self.changed |= f(&state, &entry.state);
                        }
                        *entry = StateEntry {
                            state,
                            valid_until,
                            priority,
                        };
                        // Remove entries which will expire before this one
                        for j in (i + 1)..N {
                            if let Some(later_entry) = &self.entries[j] {
                                // If the later entry exists and expires before this one
                                if later_entry.valid_until <= valid_until {
                                    self.entries[j] = None;
                                    self.entries[j..].rotate_left(1);
                                }
                            }
                        }
                        return;
                    }
                    // We want to add a higher priority entry, but we will expire first
                    if i == 0 {
                        self.changed |= f(&state, &entry.state);
                    }
                    self.entries[i..].rotate_right(1);
                    self.entries[i] = Some(StateEntry {
                        state,
                        valid_until,
                        priority,
                    });
                    return;
                }
            } else {
                // No entry exists so we can just put ourselves there
                if i == 0 {
                    self.changed |= f(&state, &self.default);
                }
                self.entries[i] = Some(StateEntry {
                    state,
                    valid_until,
                    priority,
                });
                return;
            }
        }
        // If we got here then we didn't find a valid place to insert it
    }

    #[allow(unused)]
    #[inline]
    pub fn evaluate_changed(&mut self, current_time: Instant) -> &T {
        self.evaluate_change_if(current_time, |_, _| true)
    }

    #[allow(unused)]
    #[inline]
    pub fn insert_changed(&mut self, state: T, valid_until: Instant, priority: u8) {
        self.insert_change_if(state, valid_until, priority, |_, _| true);
    }
}

impl<T, const N: usize> CommandedState<T, N>
where
    T: PartialEq + Debug,
{
    #[inline]
    pub fn evaluate(&mut self, current_time: Instant) -> &T {
        self.evaluate_change_if(current_time, T::ne)
    }

    #[inline]
    pub fn insert(&mut self, state: T, valid_until: Instant, priority: u8) {
        self.insert_change_if(state, valid_until, priority, T::ne);
    }
}

impl<T, const N: usize> Deref for CommandedState<T, N> {
    type Target = T;

    fn deref(&self) -> &Self::Target {
        for i in 0..N {
            if let Some(entry) = &self.entries[i] {
                return &entry.state;
            }
        }
        &self.default
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use log::LevelFilter;
    use std::time::Duration;

    #[test]
    fn commanded_state() {
        let _ = env_logger::builder()
            .is_test(true)
            .filter_level(LevelFilter::Trace)
            .try_init();

        let mut state = CommandedState::<_, 2>::new(4);
        let now = Instant::now();
        assert!(!state.consume_changed());
        assert_eq!(4, *state.evaluate(now));
        assert!(!state.consume_changed());

        state.insert(5, now, 0);
        assert!(state.consume_changed());
        assert!(!state.consume_changed());
        assert_eq!(5, *state.evaluate(now));
        assert!(!state.consume_changed());

        let now = now + Duration::from_secs(1);
        assert_eq!(4, *state.evaluate(now));
        assert!(state.consume_changed());

        state.insert(6, now + Duration::from_secs(2), 0);
        assert!(state.consume_changed());
        assert_eq!(6, *state.evaluate(now));
        assert!(!state.consume_changed());

        state.insert(7, now + Duration::from_secs(1), 1);
        assert!(state.consume_changed());
        assert_eq!(7, *state);

        state.insert(8, now + Duration::from_secs(3), 0);
        assert!(!state.consume_changed());
        assert_eq!(7, *state);

        state.insert(9, now + Duration::from_secs(2), 1);
        assert!(state.consume_changed());
        assert_eq!(9, *state);

        assert_eq!(9, *state.evaluate(now));
        assert!(!state.consume_changed());

        let now = now + Duration::from_secs(2);
        assert_eq!(9, *state.evaluate(now));
        assert!(!state.consume_changed());

        let now = now + Duration::from_secs(1);
        assert_eq!(8, *state.evaluate(now));
        assert!(state.consume_changed());

        let now = now + Duration::from_secs(1);
        assert_eq!(4, *state.evaluate(now));
        assert!(state.consume_changed());

        state.insert(10, now + Duration::from_secs(1), 0);
        state.insert(11, now, 1);
        state.insert(12, now + Duration::from_secs(1), 1);
        assert!(state.consume_changed());
        assert_eq!(12, *state);

        let now = now + Duration::from_secs(2);
        assert_eq!(4, *state.evaluate(now));
        assert!(state.consume_changed());

        state.insert(13, now + Duration::from_secs(1), 0);
        state.insert(14, now, 1);
        assert_eq!(14, *state);
        let now = now + Duration::from_secs(2);
        assert_eq!(4, *state.evaluate(now));
        assert!(state.consume_changed());
    }

    #[test]
    fn changed() {
        let _ = env_logger::builder()
            .is_test(true)
            .filter_level(LevelFilter::Trace)
            .try_init();

        let mut state = CommandedState::<_, 2>::new(4);
        let now = Instant::now();
        assert!(!state.consume_changed());
        assert_eq!(4, *state);

        state.insert(4, now + Duration::from_secs(1), 0);
        assert!(!state.consume_changed());

        let now = Instant::now() + Duration::from_secs(2);
        assert_eq!(4, *state.evaluate_changed(now));
        assert!(state.consume_changed());

        state.insert_changed(4, now + Duration::from_secs(1), 0);
        assert!(state.consume_changed());

        let now = Instant::now() + Duration::from_secs(2);
        assert_eq!(4, *state.evaluate(now));
    }
}