cheapsdo2.0/firmware/src/main.rs

#![no_std]
#![no_main]
#![feature(type_alias_impl_trait)]
#![feature(slice_split_once)]
use defmt_rtt as _; // global logger

use panic_probe as _;
use stm32f1xx_hal as _;

// same panicking *behavior* as `panic-probe` but doesn't print a panic message
// this prevents the panic message being printed *twice* when `defmt::panic` is invoked
#[defmt::panic_handler]
fn panic() -> ! {
    cortex_m::asm::udf()
}

use rtic::app;

#[app(device = stm32f1xx_hal::pac, peripherals = true, dispatchers = [SPI3])]
mod app {

    use cortex_m::asm::delay;
    use rtic_monotonics::systick::*;

    use stm32f1xx_hal::{
        gpio::{self, gpioa, gpioc, Alternate, Output, PushPull},
        pac,
        pac::{RCC, TIM2, TIM3, TIM4},
        prelude::*,
        rcc::Enable,
        rcc::Reset,
        timer::{self, Channel, PwmHz, Tim4NoRemap},
    };

    use stm32f1xx_hal::usb::{Peripheral, UsbBus, UsbBusType};
    use usb_device::prelude::*;

    use heapless::Vec;
    use postcard::{from_bytes_cobs, to_vec_cobs};

    use cheapsdo_protocol::{*, DeviceMessage};

    const USB_BUFFER_SIZE : usize = 64;

    #[local]
    struct Local {
        board_led: gpioc::PC13<Output<PushPull>>,
        tim2: TIM2,
        tim3: TIM3,
        pwm: PwmHz<TIM4, Tim4NoRemap, timer::Ch<0>, gpio::Pin<'B', 6, Alternate>>,
    }

    #[shared]
    struct Shared {
        usb_dev: UsbDevice<'static, UsbBusType>,
        serial: usbd_serial::SerialPort<'static, UsbBusType>,
        current_freq: f64,
        short_avg: f64,
        buffer: Vec<u8, USB_BUFFER_SIZE>,
    }

    const TARGET_FREQ: f64 = 10.0f64;

    #[init]
    fn init(cx: init::Context) -> (Shared, Local) {
        let rcc = cx.device.RCC.constrain();
        let mut flash = cx.device.FLASH.constrain();

        let clocks = rcc
            .cfgr
            .use_hse(8.MHz())
            .sysclk(48.MHz())
            .pclk1(24.MHz())
            .freeze(&mut flash.acr);

        assert!(clocks.usbclk_valid());
        defmt::info!("Clock Setup done");

        // Initialize the systick interrupt & obtain the token to prove that we did
        let systick_mono_token = rtic_monotonics::create_systick_token!();
        Systick::start(cx.core.SYST, clocks.sysclk().to_Hz(), systick_mono_token);

        let mut gpioc = cx.device.GPIOC.split();

        // Configure gpio C pin 13 as a push-pull output. The `crh` register is passed to the function
        // in order to configure the port. For pins 0-7, crl should be passed instead.
        let board_led = gpioc.pc13.into_push_pull_output(&mut gpioc.crh);

        let mut afio = cx.device.AFIO.constrain();
        let mut gpiob = cx.device.GPIOB.split();
        let pwm_pin = gpiob.pb6.into_alternate_push_pull(&mut gpiob.crl);
        let mut pwm =
            cx.device
                .TIM4
                .pwm_hz::<Tim4NoRemap, _, _>(pwm_pin, &mut afio.mapr, 32.kHz(), &clocks);
        pwm.enable(Channel::C1);

        defmt::info!("PWM Setup done");

        let tim2 = cx.device.TIM2;
        unsafe {
            let rcc = &*RCC::ptr();
            TIM2::enable(rcc);
            TIM2::reset(rcc);
        }

        // Enable external clocking
        tim2.smcr.write(|w| {
            w.etf().no_filter(); // No filter for to 10Mhz clock
            w.etps().div1(); // No divider
            w.etp().not_inverted(); // on rising edege at ETR pin
            w.ece().enabled() // mode 2 (use ETR pin)
        });

        tim2.ccmr1_input().write(|w| {
            w.cc1s().ti2(); // Input capture using TI2 input
            w.ic1f().no_filter() // No filter on input capture input
                                 //w.ic1psc().bits(0) // Disable prescaler, not safely implement by HAL yet
        });

        tim2.ccer.write(|w| {
            w.cc1p().set_bit(); // Use rising edge on TI
            w.cc1e().set_bit() // Enable input capture
        });

        tim2.cr2.write(|w| {
            w.mms().update() // Trigger output on update/overflow
        });

        tim2.ccer.write(|w| {
            w.cc1p().set_bit(); // Use rising edge on TI
            w.cc1e().set_bit() // Enable input capture
        });

        tim2.cr2.write(|w| {
            w.mms().update() // Trigger output on update/overflow
        });

        // Counting up to 10^7 should need 24 bits
        // Clock tim2 by tim1s overflow to make a 32bit timer
        let tim3 = cx.device.TIM3;
        unsafe {
            let rcc = &*RCC::ptr();
            TIM3::enable(rcc);
            TIM3::reset(rcc);
        }

        tim3.smcr.write(|w| {
            w.ts().itr1(); // Trigger from internal trigger 1
            w.sms().ext_clock_mode() // Use trigger as clock
        });

        tim3.ccmr1_input().write(|w| {
            w.cc1s().ti1(); // Input capture using TI1 input
            w.ic1f().no_filter() // No filter on input capture input
                                 //w.ic1psc().bits(0) // Disable prescaler, not safely implement by HAL yet
        });

        tim3.ccer.write(|w| {
            w.cc1p().set_bit(); // Use rising edge on TI
            w.cc1e().set_bit() // Enable input capture
        });

        tim2.cr1.write(|w| w.cen().enabled());
        tim3.cr1.write(|w| w.cen().enabled());

        defmt::info!("Timer Setup done");

        static mut USB_BUS: Option<usb_device::bus::UsbBusAllocator<UsbBusType>> = None;

        let mut gpioa = cx.device.GPIOA.split();

        let mut usb_dp = gpioa.pa12.into_push_pull_output(&mut gpioa.crh);
        usb_dp.set_low();
        delay(clocks.sysclk().raw() / 100);

        let usb_dm = gpioa.pa11;
        let usb_dp = usb_dp.into_floating_input(&mut gpioa.crh);

        let usb = Peripheral {
            usb: cx.device.USB,
            pin_dm: usb_dm,
            pin_dp: usb_dp,
        };

        unsafe {
            USB_BUS.replace(UsbBus::new(usb));
        }

        let serial = usbd_serial::SerialPort::new(unsafe { USB_BUS.as_ref().unwrap() });

        let usb_dev = UsbDeviceBuilder::new(
            unsafe { USB_BUS.as_ref().unwrap() },
            UsbVidPid(0x16c0, 0x27dd),
        )
        .manufacturer("Arbitrary Precision Instruments")
        .product("cheapsdo")
        .serial_number("1337")
        .device_class(usbd_serial::USB_CLASS_CDC)
        .build();

        update_pwm::spawn().unwrap();

        (
            Shared {
                serial,
                usb_dev,
                current_freq: 0.0f64,
                short_avg: 0.0f64,
                buffer: Vec::new(),
            },
            Local {
                board_led,
                tim2,
                tim3,
                pwm,
            },
        )
    }


    const WINDOW_LEN: usize = 100;

    #[task(local=[tim2, tim3, pwm, board_led], shared=[current_freq, short_avg])]
    async fn update_pwm(mut cx: update_pwm::Context) {
        defmt::info!("Update Task started");

        let tim2 = cx.local.tim2;
        let tim3 = cx.local.tim3;
        let pwm = cx.local.pwm;
        let board_led = cx.local.board_led;

        let max_pwm = pwm.get_max_duty() as i32;
        let mut cur_pwm = max_pwm / 2;

        // Inialize last_ic
        while !tim2.sr.read().cc1if().bit_is_set() || !tim3.sr.read().cc1if().bit_is_set() {
            Systick::delay(10.millis()).await;
        }
        let ic1 = tim2.ccr1().read().bits();
        let ic2 = tim3.ccr1().read().bits();

        let mut last_ic = ic2 << 16 | ic1;

        loop {
            let mut short_avg = 0.0f64;
            let mut last_freq = 10.0f64;

            let mut count = 0;
            while count < WINDOW_LEN {
                while !tim3.sr.read().cc1if().bit_is_set() || !tim3.sr.read().cc1if().bit_is_set() {
                    Systick::delay(10.millis()).await;
                }
                let ic1 = tim2.ccr1().read().bits();
                let ic2 = tim3.ccr1().read().bits();

                let sum_ic = ic2 << 16 | ic1;

                let diff_ic = if sum_ic > last_ic {
                    sum_ic - last_ic
                } else {
                    u32::MAX - last_ic + sum_ic
                };

                defmt::info!("ic1:\t{}", ic1);
                defmt::info!("ic2:\t{}", ic2);
                defmt::info!("sum_ic:\t{}", sum_ic);
                defmt::info!("diff_ic:\t{}", diff_ic);

                last_ic = sum_ic;

                let freq = (diff_ic as f64) / 1_000_000f64;
                defmt::info!("freq:\t{} MHz", freq);

                let diff = freq - last_freq;
                last_freq = freq;
                if diff > 0.000_100 || diff < -0.000_100 {
                    defmt::info!("Out of range, dropping sample.");
                    continue;
                }

                cx.shared.current_freq.lock(|current_freq| {
                    *current_freq = freq;
                });

                short_avg += freq / WINDOW_LEN as f64;

                count += 1;
                board_led.toggle();
            }

            defmt::info!("short_avg:\t{} MHz", short_avg);
            cx.shared.short_avg.lock(|avg| {
                *avg = short_avg;
            });

            let diff = (TARGET_FREQ - short_avg) * 1_000_000.0;
            if diff > 1.0 || diff < -1.0 {
                cur_pwm += diff as i32 * 15;
            } else if diff < -0.001 {
                cur_pwm -= 1;
            } else if diff > 0.001 {
                cur_pwm += 1;
            }

            cur_pwm = if cur_pwm < 0 { 0 } else { cur_pwm };
            cur_pwm = if cur_pwm > max_pwm { max_pwm } else { cur_pwm };

            pwm.set_duty(Channel::C1, cur_pwm as u16);
            defmt::info!("pwm:\t{}", cur_pwm);

            Systick::delay(500.millis()).await;
        }
    }

    #[task(binds = USB_HP_CAN_TX, shared = [usb_dev, serial, buffer, current_freq, short_avg])]
    fn usb_tx(cx: usb_tx::Context) {
        let mut usb_dev = cx.shared.usb_dev;
        let mut serial = cx.shared.serial;
        let mut buffer = cx.shared.buffer;
        let mut current_freq = cx.shared.current_freq;
        let mut short_avg = cx.shared.short_avg;

        (
            &mut usb_dev,
            &mut serial,
            &mut buffer,
            &mut current_freq,
            &mut short_avg,
        )
            .lock(|usb_dev, serial, buffer, current_freq, short_avg| {
                usb_poll(usb_dev, serial, buffer, current_freq, short_avg);
            });
    }

    #[task(binds = USB_LP_CAN_RX0, shared = [usb_dev, serial, buffer, current_freq, short_avg])]
    fn usb_rx0(cx: usb_rx0::Context) {
        let mut usb_dev = cx.shared.usb_dev;
        let mut serial = cx.shared.serial;
        let mut buffer = cx.shared.buffer;
        let mut current_freq = cx.shared.current_freq;
        let mut short_avg = cx.shared.short_avg;

        (
            &mut usb_dev,
            &mut serial,
            &mut buffer,
            &mut current_freq,
            &mut short_avg,
        )
            .lock(|usb_dev, serial, buffer, current_freq, short_avg| {
                usb_poll(usb_dev, serial, buffer, current_freq, short_avg);
            });
    }

    fn usb_poll<B: usb_device::bus::UsbBus>(
        usb_dev: &mut usb_device::prelude::UsbDevice<'static, B>,
        serial: &mut usbd_serial::SerialPort<'static, B>,
        buffer: &mut Vec<u8, USB_BUFFER_SIZE>,
        current_freq: &mut f64,
        short_avg: &mut f64,
    ) {
        if !usb_dev.poll(&mut [serial]) {
            return;
        }

        let mut tmp = [0u8; 16];
        match serial.read(&mut tmp) {
            Ok(count) if count > 0 => {
                if buffer.extend_from_slice(&tmp[0..count]).is_err() {
                    buffer.clear();
                    defmt::error!("Buffer overflow while waiting for the end of the packet");
                }
            }
            _ => {}
        }

        loop {
            if let Some((msg, rest)) = buffer.split_once(|&x| x == 0) {
                let mut message = [0u8; 128];
                message[0..msg.len()].clone_from_slice(msg);
                let host_msg = from_bytes_cobs::<HostMessage>(&mut message);

                match host_msg {
                    Ok(host_msg) => match host_msg {
                        HostMessage::RequestStatus => {
                            let device_msg = DeviceMessage::Status(StatusMessage{
                                measured_frequency: *current_freq,
                                average_frequency: *short_avg,
                                pwm: 0,
                            });

                            let bytes = to_vec_cobs::<DeviceMessage, USB_BUFFER_SIZE>(&device_msg).unwrap();
                            serial.write(bytes.as_slice()).unwrap();
                        },
                        HostMessage::SetPLLOutputs => {
                             defmt::error!("PLL output is not implemented yet")
                        }
                    }
                    Err(err) => defmt::error!("Unable to parse host message: {}", err),
                };

                *buffer = Vec::<u8, USB_BUFFER_SIZE>::from_slice(rest).unwrap();
            } else {
                break;
            }
        }
    }
}