trinket-streamdeck/pico/src/main.rs

#![no_std]
#![no_main]

use panic_halt as _;

use common::{handle_input, Handshake, UnsafeSync};
use core::cell::RefCell;
use cortex_m::{asm::delay as cycle_delay, interrupt::Mutex, peripheral::NVIC};
use once_cell::unsync::OnceCell;
use rp2040_hal::{
    pac::{interrupt, Interrupt, Peripherals},
    usb::UsbBus,
    Sio, Timer, Watchdog,
};
use usb_device::{class_prelude::*, prelude::*};
use usbd_serial::{SerialPort, USB_CLASS_CDC};

mod bsp;
mod buttons;
mod device_signature;

use bsp::{entry, Pins};
use buttons::Buttons;

static USB_ALLOCATOR: UnsafeSync<OnceCell<UsbBusAllocator<UsbBus>>> =
    UnsafeSync::new(OnceCell::new());
// static RTC_COUNT: Mutex<RefCell<Option<Rtc<Count32Mode>>>> = Mutex::new(RefCell::new(None));
static TIMER: Mutex<RefCell<Option<Timer>>> = Mutex::new(RefCell::new(None));

static USB_BUS: Mutex<RefCell<Option<UsbDevice<UsbBus>>>> = Mutex::new(RefCell::new(None));
static USB_SERIAL: Mutex<RefCell<Option<SerialPort<UsbBus>>>> = Mutex::new(RefCell::new(None));

static HANDSHAKE: Mutex<RefCell<Handshake>> = Mutex::new(RefCell::new(Handshake::new()));

fn get_current_time() -> u32 {
    let micros =
        cortex_m::interrupt::free(|cs| TIMER.borrow(cs).borrow().as_ref().unwrap().get_counter());

    let millis = micros / 1_000;

    millis as u32
}

fn is_handshake_complete() -> bool {
    cortex_m::interrupt::free(|cs| HANDSHAKE.borrow(cs).borrow().is_complete())
}

#[entry]
fn main() -> ! {
    let mut peripherals = Peripherals::take().unwrap();

    let mut watchdog = Watchdog::new(peripherals.WATCHDOG);

    let clocks = rp2040_hal::clocks::init_clocks_and_plls(
        bsp::XOSC_CRYSTAL_FREQ,
        peripherals.XOSC,
        peripherals.CLOCKS,
        peripherals.PLL_SYS,
        peripherals.PLL_USB,
        &mut peripherals.RESETS,
        &mut watchdog,
    )
    .ok()
    .unwrap();

    let timer = Timer::new(peripherals.TIMER, &mut peripherals.RESETS);

    cortex_m::interrupt::free(|cs| {
        *TIMER.borrow(cs).borrow_mut() = Some(timer);
    });

    let sio = Sio::new(peripherals.SIO);

    let pins = Pins::new(
        peripherals.IO_BANK0,
        peripherals.PADS_BANK0,
        sio.gpio_bank0,
        &mut peripherals.RESETS,
    );

    let usb_allocator = UsbBusAllocator::new(UsbBus::new(
        peripherals.USBCTRL_REGS,
        peripherals.USBCTRL_DPRAM,
        clocks.usb_clock,
        true,
        &mut peripherals.RESETS,
    ));

    // SAFETY: No interrupt will access USB_ALLOCATOR before this point, since it is only
    // referenced from inside the USB_SERIAL mutex and the USB_BUS mutex, which have not been set
    // yet.
    //
    // SAFETY: This is the only line of the program that modifies USB_ALLOCATOR
    let allocator = unsafe {
        USB_ALLOCATOR.get().set(usb_allocator).ok().unwrap();
        USB_ALLOCATOR.get().get().unwrap()
    };

    cortex_m::interrupt::free(|cs| {
        let serial = USB_SERIAL.borrow(cs);
        let mut serial = serial.borrow_mut();
        *serial = Some(SerialPort::new(allocator));
    });

    cortex_m::interrupt::free(|cs| {
        let bus = USB_BUS.borrow(cs);
        let mut bus = bus.borrow_mut();
        *bus = Some(
            UsbDeviceBuilder::new(allocator, UsbVidPid(0x16c0, 0x27dd))
                .manufacturer("Aode Lion")
                .product("Pico Streamdeck")
                .serial_number("AAAB")
                .device_class(USB_CLASS_CDC)
                .build(),
        );
    });

    unsafe {
        NVIC::unmask(Interrupt::USBCTRL_IRQ);
    }

    let mut buttons = Buttons {
        d0: pins.gpio0,
        d1: pins.gpio2,
        d2: pins.gpio4,
        d3: pins.gpio6,
        d4: pins.gpio8,
        d5: pins.gpio14,
        d6: pins.gpio16,
        d7: pins.gpio21,
        d8: pins.gpio22,
        d9: pins.gpio27,
    }
    .init();

    loop {
        let current_time = get_current_time();

        if !is_handshake_complete() {
            cycle_delay(1024 * 15);
            continue;
        }

        let mut output = [0; 16];

        if let Some(len) = buttons.tick(current_time, &mut output) {
            let _ = cortex_m::interrupt::free(|cs| {
                let mut serial = USB_SERIAL.borrow(cs).borrow_mut();
                let serial = serial.as_mut().unwrap();
                let _ = serial.write(&output[0..len])?;
                serial.flush()
            });
        }
    }
}

#[allow(non_snake_case)]
#[interrupt]
fn USBCTRL_IRQ() {
    poll_usb();
}

fn poll_usb() -> Option<()> {
    let ready = cortex_m::interrupt::free(|cs1| {
        cortex_m::interrupt::free(|cs2| {
            let mut usb_dev = USB_BUS.borrow(cs1).borrow_mut();
            let usb_dev = usb_dev.as_mut()?;
            let mut serial = USB_SERIAL.borrow(cs2).borrow_mut();
            let serial = serial.as_mut()?;

            Some(usb_dev.poll(&mut [&mut *serial]))
        })
    })?;

    if !ready {
        return Some(());
    }

    let mut input = [0u8; 32];

    let count = cortex_m::interrupt::free(|cs| {
        USB_SERIAL
            .borrow(cs)
            .borrow_mut()
            .as_mut()?
            .read(&mut input)
            .ok()
    })?;

    let current_time = get_current_time();

    cortex_m::interrupt::free(|cs| {
        let mut handshake = HANDSHAKE.borrow(cs).borrow_mut();

        if count == 32 {
            if let common::HandshakeResponse::Write { bytes } =
                handshake.perform(&input, current_time)?
            {
                cortex_m::interrupt::free(|cs| {
                    let mut serial = USB_SERIAL.borrow(cs).borrow_mut();
                    let serial = serial.as_mut()?;
                    serial.write(bytes).ok()?;
                    serial.flush().ok()
                })?;
            }

            return None;
        }

        Some(())
    })?;

    let mut output_buffer = [0; 256];

    let n = handle_input(&mut Device, &input, &mut output_buffer, count)?;

    if n == 0 {
        return None;
    }

    cortex_m::interrupt::free(|cs| {
        let mut serial = USB_SERIAL.borrow(cs).borrow_mut();
        let serial = serial.as_mut()?;
        serial.write(&output_buffer[0..n]).ok()?;
        serial.flush().ok()
    })?;

    Some(())
}

struct Device;

impl common::Device<8, 1> for Device {
    fn serial_number(&mut self) -> [u8; 8] {
        device_signature::read_uid()
    }

    fn extras() -> [(&'static str, &'static str); 1] {
        [("version", env!("FIRMWARE_VERSION"))]
    }

    fn reset(&mut self) {
        cortex_m::interrupt::disable();
        rp2040_hal::rom_data::reset_to_usb_boot(0, 0);
    }
}