logic.rs

use crate::{
    encoders::Encoders,
    tb6612fng::{Movement, Tb6612fng},
};
use core::sync::atomic::{AtomicU16, AtomicU32, AtomicU8, Ordering};
use embassy_executor::Spawner;
use embassy_stm32::{
    pac::{self, wwdg::vals::Wdgtb},
    peripherals::WWDG,
    time::{hz, Hertz},
};
use embassy_sync::{
    blocking_mutex::raw::CriticalSectionRawMutex, channel::Channel, signal::Signal,
};
use embassy_time::{Duration, Instant, Ticker};
use futures::FutureExt;
use heapless::Deque;
use i2c2_target::{I2C2, MESSAGE_SIZE};

const STANDBY: u8 = 0x80;
const STANDBY_BOTH: u16 = ((STANDBY as u16) << 8) | STANDBY as u16;

static MAX_MOTOR_PERCENTAGE: AtomicU8 = AtomicU8::new(100);
static MOTOR_SPEED: AtomicU16 = AtomicU16::new(STANDBY_BOTH);
static PWM_FREQUENCY: AtomicU32 = AtomicU32::new(0);
static WATCHDOG_TICKS: AtomicU8 = AtomicU8::new(5);

struct State {
    i2c: I2C2,
    encoders: Encoders<'static>,
}

impl State {
    fn new(i2c2: I2C2, encoders: Encoders<'static>) -> Self {
        Self {
            i2c: i2c2,
            encoders,
        }
    }
}

static mut STATE: Option<State> = None;

fn wwdg_configure() {
    // Enable WWDG clock
    pac::RCC.apb1enr().modify(|w| w.set_wwdgen(true));
    // With PCLK1 at 36MHz, 63 as a timeout value with wdgtb at /8
    // means ~58ms.  Comparing with a value of 45 means that at
    // least ~42ms must have passed before the watchdog can be
    // petted, so we have 8ms on each side.
    pac::WWDG.cfr().write(|w| {
        w.set_wdgtb(Wdgtb::DIV8);
        w.set_w((1 << 6) | 31);
    });
    wwdg_pet();
}

fn wwdg_pet() {
    // Reset watchdog and start it if not started yet
    pac::WWDG.cr().write(|w| {
        w.set_t((1 << 6) | 63);
        w.set_wdga(true);
    });
}

#[allow(clippy::needless_pass_by_value, clippy::used_underscore_binding)] // Clippy bug
pub fn start_i2c_target(
    spawner: Spawner,
    i2c2: I2C2,
    motors: Tb6612fng<'static>,
    encoders: Encoders<'static>,
    _wwdg: WWDG,
) {
    unsafe {
        STATE = Some(State::new(i2c2, encoders));
        STATE.as_mut().unwrap().i2c.start(&i2c_callback);
    }
    spawner.spawn(watchdog()).unwrap();
    spawner.spawn(motor_control(motors)).unwrap();
}

static WATCHDOG_PING: Signal<CriticalSectionRawMutex, Instant> = Signal::new();

#[embassy_executor::task]
async fn watchdog() {
    wwdg_configure();
    let mut ticker = Ticker::every(Duration::from_millis(50));
    let mut watchdog_expiration: Instant = Instant::now();
    loop {
        ticker.next().await;
        wwdg_pet();
        if let Some(ping) = WATCHDOG_PING.try_take() {
            watchdog_expiration = ping
                + Duration::from_millis(100 * u64::from(WATCHDOG_TICKS.load(Ordering::Relaxed)));
        }
        if is_moving() && Instant::now() >= watchdog_expiration {
            defmt::debug!("stopping motors after watchdog has expired");
            standby().await;
        }
    }
}

async fn standby() {
    MOTOR_CONTROL
        .send(MotorCommand::SetSpeed(STANDBY, STANDBY))
        .await;
}

fn is_moving() -> bool {
    let motor_speed = MOTOR_SPEED.load(Ordering::Relaxed);
    motor_speed != 0 && motor_speed != STANDBY_BOTH
}

enum MotorCommand {
    SetSpeed(u8, u8),
    SetFrequency(Hertz),
}

// Set channel size to 4 for:
//  - two regular motor commands
//  - one standby command sent by I²C
//  - one standby command sent by the software watchdog
static MOTOR_CONTROL: Channel<CriticalSectionRawMutex, MotorCommand, 4> = Channel::new();

#[embassy_executor::task]
#[allow(clippy::cast_possible_wrap)]
async fn motor_control(mut motors: Tb6612fng<'static>) {
    let (mut left, mut right, mut standby) = (0u8, 0u8, true);
    motors.standby_enter();
    PWM_FREQUENCY.store(motors.get_frequency().0, Ordering::Relaxed);
    loop {
        match MOTOR_CONTROL.receive().await {
            MotorCommand::SetSpeed(new_left, new_right) => {
                let (max_pwm, mmp) = (
                    motors.max_pwm() as i32,
                    i32::from(MAX_MOTOR_PERCENTAGE.load(Ordering::Relaxed)),
                );
                let scaled_speed = |speed| match speed as i8 {
                    -128 => None,
                    speed @ (-100..=100) => Some(i32::from(speed) * max_pwm * mmp / (100 * 100)),
                    speed => {
                        defmt::warn!("speed {} out of range [-100, 100]", speed);
                        None
                    }
                };
                let (speed_left, speed_right) = (scaled_speed(new_left), scaled_speed(new_right));
                if (speed_left, speed_right) == (None, None) {
                    motors.standby_enter();
                    standby = true;
                    (left, right) = (new_left, new_right);
                } else {
                    if standby {
                        motors.standby_leave();
                        standby = false;
                    }
                    if let Some(speed_left) = speed_left {
                        motors.move_left(Movement::Advance(speed_left));
                        left = new_left;
                    }
                    if let Some(speed_right) = speed_right {
                        motors.move_right(Movement::Advance(speed_right));
                        right = new_right;
                    }
                }
                MOTOR_SPEED.store((u16::from(left) << 8) | u16::from(right), Ordering::Relaxed);
            }
            MotorCommand::SetFrequency(freq) => {
                motors.set_frequency(freq);
                PWM_FREQUENCY.store(freq.0, Ordering::Relaxed);
            }
        }
    }
}

fn i2c_callback(command: &[u8], response: &mut Deque<u8, MESSAGE_SIZE>) {
    #[allow(static_mut_refs)]
    let state = unsafe { STATE.as_mut().unwrap() };
    if process_command(state, command, response) {
        WATCHDOG_PING.signal(Instant::now());
    } else {
        // If this does not succeed, this will reset due to
        // a lack of petting the window watchdog.
        while standby().now_or_never().is_none() {}
    }
}

const CMD_FIRMWARE_VERSION: u8 = 0x08;
const CMD_WHO_AM_I: u8 = 0x0f;
const CMD_PWM_FREQUENCY: u8 = 0x10;
const CMD_MAX_MOTOR_PERCENTAGE: u8 = 0x11;
const CMD_MOTOR_SHUTDOWN_TIMEOUT: u8 = 0x28;
const CMD_MOTOR_SPEED: u8 = 0x30;
const CMD_ENCODER_TICKS: u8 = 0x32;
const CMD_STATUS: u8 = 0x36;

include!(concat!(env!("OUT_DIR"), "/version.rs"));

#[allow(clippy::cast_possible_truncation)]
fn process_command(
    state: &mut State,
    command: &[u8],
    response: &mut Deque<u8, MESSAGE_SIZE>,
) -> bool {
    defmt::trace!("Processing command {:?}", command);
    match command {
        [CMD_FIRMWARE_VERSION] => {
            for b in PKG_VERSION {
                response.push_back(b).unwrap();
            }
        }
        [CMD_WHO_AM_I, ..] => {
            response.push_back(0x57).unwrap();
        }
        &[CMD_PWM_FREQUENCY, a, b, c] => {
            let f = u32::from_le_bytes([a, b, c, 0]);
            if (1..=100_000).contains(&f) {
                if MOTOR_CONTROL
                    .try_send(MotorCommand::SetFrequency(hz(f)))
                    .is_err()
                {
                    defmt::error!("unable to queue up motor frequency");
                    return false;
                }
            } else {
                defmt::warn!("incorrect PWM frequency {}", f);
                return false;
            }
        }
        [CMD_PWM_FREQUENCY] => {
            let freq = PWM_FREQUENCY.load(Ordering::Relaxed);
            for &b in &freq.to_le_bytes()[..3] {
                response.push_back(b).unwrap();
            }
        }
        &[CMD_MAX_MOTOR_PERCENTAGE, p] => {
            if (1..=100).contains(&p) {
                MAX_MOTOR_PERCENTAGE.store(p, Ordering::Relaxed);
            } else {
                defmt::warn!("incorrect max percentage {}", p);
                return false;
            }
        }
        [CMD_MAX_MOTOR_PERCENTAGE] => {
            response
                .push_back(MAX_MOTOR_PERCENTAGE.load(Ordering::Relaxed))
                .unwrap();
        }
        &[CMD_MOTOR_SPEED, left, right] => {
            if MOTOR_CONTROL
                .try_send(MotorCommand::SetSpeed(left, right))
                .is_err()
            {
                defmt::error!("unable to queue up motor speeds");
                return false;
            }
        }
        [CMD_MOTOR_SPEED] => {
            let motor_speed = MOTOR_SPEED.load(Ordering::Relaxed);
            response.push_back((motor_speed >> 8) as u8).unwrap();
            response.push_back(motor_speed as u8).unwrap();
        }
        &[CMD_MOTOR_SHUTDOWN_TIMEOUT, ticks] => {
            if (1..=100).contains(&ticks) {
                WATCHDOG_TICKS.store(ticks, Ordering::Relaxed);
            } else {
                defmt::warn!("ticks out of range: {}", ticks);
                return false;
            }
        }
        [CMD_MOTOR_SHUTDOWN_TIMEOUT] => {
            response
                .push_back(WATCHDOG_TICKS.load(Ordering::Relaxed))
                .unwrap();
        }
        [CMD_ENCODER_TICKS] => {
            let (left, right) = state.encoders.ticks();
            let (left, right) = (left.to_le_bytes(), right.to_le_bytes());
            response.push_back(left[0]).unwrap();
            response.push_back(left[1]).unwrap();
            response.push_back(right[0]).unwrap();
            response.push_back(right[1]).unwrap();
        }
        [CMD_STATUS] => {
            response.push_back(u8::from(is_moving())).unwrap();
        }
        _ => {
            defmt::warn!("unknown command or args {:#04x}", command);
            return false;
        }
    }
    true
}