../writing-keyboard-firmware-in-rust

Oxidating all the microcontrollers

Huh

Most user-programmable mechanical keyboards use an embedded OS named QMK, which has support for both AVR and ARM microcontrollers. Some keyboards alternatively use ZMK which adds support for Bluetooth.

However as it turns out, the Rust-on-ARM story is pretty fleshed out with the following projects:

For this I built an ergo keyboard (the corne v3) with RGB LEDs and OLED displays. For the controller(s) I went with nice!nanos, which use the nRF52840 SoC. The nRF52840 supports Bluetooth, but as I don’t move the keyboard often, I decided to assume the keyboard will be always wired.

This was my first time writing for a microcontroller using something other than Arduino and it was an incredibly fluid experience. I spent a grand total of zero seconds debugging memory related problems, all bugs I had ended up being logic issues that were easily fixed by logging with defmt.

Programming

Embassy

Initially I attempted using RTIC to provide support for concurrent tasks, however RTIC requires you to manage setting up interrupts (for UART, I2C, etc) yourself, and quickly this got annoying.

After finding embassy I quickly ported my code over to use it. Embassy uses rust’s async machinery to allow for true tasks that run forever (in rtic each task is a function triggered by an interrupt). Embassy provides async compatible interfaces for USB, UART, and I2C too.

Creating a UART interface in Embassy is as simple as:

let uart_config = uarte::Config::default();
let irq = interrupt::take!(UARTE0_UART0);
let uart = uarte::Uarte::new(p.UARTE0, irq, p.P0_08, p.P1_04, uart_config);

uart.write(b"hello world").await?;
let mut buf = [0u8; 1];
uart.read(&mut buf).await?;

To communicate between tasks, I use the channel provided by embassy to have the producer task wait for the consumer task to process messages if the queue is full.

Keyberon

The most important part of the keyboard is having it, well, work as a keyboard. Luckily I don’t have to implement all the state management of a keyboard myself, as the TeXitoi/keyberon project conveniently exists.

Keyberon was incredibly easy to work with as it splits itself up into the few logically separated components needed for a keyboard:

To get Keyberon working I simply had to:

Specify matrix pins and construct the matrix struct

I used a macro for this as it requires that the gpio pin values are partially moved from the peripherals struct (this allows the compiler to ensure there is only one user of each pin at compile time)

macro_rules! build_matrix {
    ($p:ident) => {{
        use embassy_nrf::gpio::{Input, Level, OutputDrive, Pin, Pull};
        use keyberon::matrix::Matrix;
        Matrix::new(
            [
                Input::new($p.P0_31.degrade(), Pull::Up),
                Input::new($p.P0_29.degrade(), Pull::Up),
                Input::new($p.P0_02.degrade(), Pull::Up),
                Input::new($p.P1_15.degrade(), Pull::Up),
                Input::new($p.P1_13.degrade(), Pull::Up),
                Input::new($p.P1_11.degrade(), Pull::Up),
            ],
            [
                Output::new($p.P0_22.degrade(), Level::High, OutputDrive::Standard),
                Output::new($p.P0_24.degrade(), Level::High, OutputDrive::Standard),
                Output::new($p.P1_00.degrade(), Level::High, OutputDrive::Standard),
                Output::new($p.P0_11.degrade(), Level::High, OutputDrive::Standard),
            ],
        )
        .unwrap()
    }};
}

Specify my keyboard layout

Keyberon provides a nice macro for doing this. In the same file I also declare the hold-taps and chords I want to use.

pub static LAYERS: Layers  = keyberon::layout::layout! {
    {
        ['`' Q W E R T Y U I O P '\''],
        [LShift A S D F G H J K L ; RShift],
        [LCtrl Z X C V B N M , . / RCtrl],
        [n n n LGui {ALT_TAB} {L1_SP} {L2_SP} Enter BSpace n n n],
        [Escape {m(&[KeyCode::LAlt, KeyCode::X])} {m(&[KeyCode::Space, KeyCode::Grave])} Delete < {m(&[KeyCode::LShift, KeyCode::SColon])} > '\\' / '"' '\'' '_'],
    }
    {
        ['`' ! @ '{' '}' | '`' ~ '\\' n '"'  n],
        [ t  # $ '(' ')' n  +  -  /   * '\'' t],
        [ t  % ^ '[' ']' n  &  =  ,   . '_'  t],
        [n n n LGui LAlt =  = Tab BSpace n n n],
        [n n n n    n    n  n n   n      n n n],
    }
    {
        [n Kb1 Kb2 Kb3 Kb4 Kb5 Kb6 Kb7 Kb8 Kb9 Kb0 n],
        [t F1  F2  F3  F4  F5  Left Down Up Right VolUp t],
        [t F6  F7  F8  F9  F10 PgDown {m(&[KeyCode::LCtrl, KeyCode::Down])} {m(&[KeyCode::LCtrl, KeyCode::Up])} PgUp VolDown t],
        [n n n F11 F12 t t RAlt End n n n],
        [n n n n   n   n n n    n   n n n],
    }
};

Poll the matrix

Keyberon provides the matrix struct, but won’t handle polling the matrix at an interval for us. For that I use an embassy task to poll the matrix every POLL_PERIOD and feed the results through the debouncer and chording engine. The processed key events are then pushed to a channel for processing by the layout task.

#[embassy::task]
async fn keyboard_poll_task(
    mut matrix: Matrix<Input<'static, AnyPin>, Output<'static, AnyPin>, COLS_PER_SIDE, ROWS>,
    mut debouncer: Debouncer<[[bool; COLS_PER_SIDE]; ROWS]>,
    mut chording: Chording<{ keyboard_thing::layout::NUM_CHORDS }>,
) {
    loop {
        let events = debouncer
            .events(matrix.get().unwrap())
            .collect::<heapless::Vec<_, 8>>();

        for event in &events {
            for chan in KEY_EVENT_CHANS {
                let _ = chan.try_send(*event);
            }
        }

        let events = chording.tick(events);

        let count = events.iter().filter(|e| e.is_press()).count() as u32;
        TOTAL_LHS_KEYPRESSES.fetch_add(count, core::sync::atomic::Ordering::Relaxed);

        for event in events {
            PROCESSED_KEY_CHAN.send(event).await;
        }

        Timer::after(POLL_PERIOD).await;
    }
}

Push events through the layout

As before with the matrix, we decide when the layout state should be updated. For that another task is used to update the matrix state when a key is pressed or released. This task also receives the key events from the other half.

#[embassy::task]
async fn keyboard_event_task(layout: &'static Mutex<ThreadModeRawMutex, Layout>) {
    loop {
        let event = PROCESSED_KEY_CHAN.recv().await;
        let mut count = if event.is_press() { 1 } else { 0 };
        if event.is_press() {
            KEYPRESS_EVENT.set();
        }
        interacted();
        {
            let mut layout = layout.lock().await;
            layout.event(event);
            while let Ok(event) = PROCESSED_KEY_CHAN.try_recv() {
                layout.event(event);
                count += if event.is_press() { 1 } else { 0 };
            }
        }
        TOTAL_KEYPRESSES.fetch_add(count, core::sync::atomic::Ordering::Relaxed);
    }
}

Extract keycode events and submit to the computer

To extract keycode events we use another task that extracts which keys are currently pressed every 1ms and submits the event to the task handling the USB HID messaging.

#[embassy::task]
async fn layout_task(layout: &'static Mutex<ThreadModeRawMutex, Layout>) {
    let mut last_report = None;
    loop {
        {
            let mut layout = layout.lock().await;
            layout.tick();

            let collect = layout
                .keycodes()
                .filter_map(|k| Keyboard::try_from_primitive(k as u8).ok())
                .collect::<heapless::Vec<_, 24>>();

            if last_report.as_ref() != Some(&collect) {
                last_report = Some(collect.clone());
                HID_CHAN.send(NKROBootKeyboardReport::new(&collect)).await;
            }
        }

        Timer::after(Duration::from_millis(1)).await;
    }
}

Tie everything together

First the required things are initialized, then we can start the tasks that perform all the keyboard processing.

let matrix = keyboard_thing::build_matrix!(p);
let debouncer = Debouncer::new(
    [[false; COLS_PER_SIDE]; ROWS],
    [[false; COLS_PER_SIDE]; ROWS],
    DEBOUNCER_TICKS,
);
let chording = Chording::new(&keyboard_thing::layout::CHORDS);

let layout = forever!(Mutex::new(Layout::new(&keyboard_thing::layout::LAYERS)));

spawner
    .spawn(keyboard_poll_task(matrix, debouncer, chording))
    .unwrap();
spawner.spawn(keyboard_event_task(layout)).unwrap();
spawner.spawn(layout_task(layout)).unwrap();

NeoPixels

My Corne kit came with per-key and under-glow neopixels, so why not use them!

Fortunately the jamesmunns/nrf-smartled library exists for controlling neopixels by ~~ab~~​using the PWM peripheral on the nRF52840. Now to take advantage of all 64Mhz we just need to define the positions of each LED (they’re connected serially):

// underglow LEDs are left to right
#[rustfmt::skip]
pub const UNDERGLOW_LED_POSITIONS: [(u8, u8); UNDERGLOW_LEDS] = [
    // top row: 1, 2, 3
    (0, 1), (2, 1), (4, 1),
    // bottom row: 4, 5, 6
    (4, 2), (2, 3), (0, 3),
];

// switch leds are bottom to top
#[rustfmt::skip]
pub const SWITCH_LED_POSITIONS: [(u8, u8); SWITCH_LEDS] = [
    // first column: 7, 8, 9, 10
    (3, 5), (2, 5), (1, 5), (0, 5),
    // second column: 11, 12, 13, 14
    (0, 4), (1, 4), (2, 4), (3, 4),
    // third column: 15, 16, 17, 18
    (3, 3), (2, 3), (1, 3), (0, 3),
    // fourth column: 19, 20, 21
    (0, 2), (1, 2), (2, 2),
    // fifth column: 22, 23, 24
    (2, 1), (1, 1), (0, 1),
    // sixth column: 25, 26, 27
    (0, 0), (1, 0), (2, 0)
];

And now we can create a fancy rainbow effect:

pub fn rainbow_single(x: u8, y: u8, offset: u8) -> Hsv {
    Hsv {
        hue: x
            .wrapping_mul(6)
            .wrapping_add(y.wrapping_mul(2))
            .wrapping_add(offset),
        sat: 255,
        val: 127,
    }
}

pub fn rainbow(offset: u8) -> impl Iterator<Item = RGB8> {
    colour_gen(move |x, y| hsv2rgb(rainbow_single(x, y, offset)))
}

And use an embassy task to update and render it at 30fps.

#[embassy::task]
async fn led_task(mut leds: Leds) {
    let fps = 30;
    let mut tapwaves = TapWaves::new();
    let mut ticker = Ticker::every(Duration::from_millis(1000 / fps));
    let mut counter = WrappingID::<u16>::new(0);

    loop {
        while let Ok(event) = LED_KEY_LISTEN_CHAN.try_recv() {
            tapwaves.update(event);
        }

        tapwaves.tick();

        leds.send(tapwaves.render(|x, y| rainbow_single(x, y, counter.get() as u8)));

        counter.inc();

        if (counter.get() % 128) == 0 {
            let _ = COMMAND_CHAN.try_send((
                DomToSub::ResyncLeds(counter.get()),
                Duration::from_millis(5),
            ));
        }

        ticker.next().await;
    }
}

I also added in animated waves that emanate from each key when pressed, I’m really making full use of the nRF’s FP unit here.

OLEDs

The Corne has support for a 128x32 display on each side, enough space that I’m struggling to decide what to put on each.

For the right side I have some metrics displayed: the total number of keypresses, the current keypresses per second, and the number of seconds the keyboard has been on. I also have a sliding display of keypresses at the bottom:

And for the left side I currently have a badly drawn bongo cat, I’m still thinking of what to use the remaining 96x32 pixels for.

To control the OLED displays, I use jamwaffles/ssd1306 which handles writing out a buffer the display, and embedded-graphics to draw text, images, and other geometry.

I use the following to initialize the display and handle turning the display off during periods of inactivity:

type OledDisplay<'a, T> =
    Ssd1306<I2CInterface<Twim<'a, T>>, DisplaySize128x32, BufferedGraphicsMode<DisplaySize128x32>>;

pub struct Oled<'a, T: Instance> {
    status: bool,
    display: OledDisplay<'a, T>,
}

impl<'a, T: Instance> Oled<'a, T> {
    pub fn new(twim: Twim<'a, T>) -> Self {
        let i2c = I2CDisplayInterface::new(twim);
        let display = Ssd1306::new(i2c, DisplaySize128x32, DisplayRotation::Rotate0)
            .into_buffered_graphics_mode();
        Self {
            status: true,
            display,
        }
    }

    pub async fn init(&mut self) -> Result<(), DisplayError> {
        self.display.set_rotation(DisplayRotation::Rotate90).await?;
        self.display.set_brightness(Brightness::BRIGHTEST).await?;
        self.display.init().await?;
        Ok(())
    }

    // ...
}

pub const OLED_TIMEOUT: Duration = Duration::from_secs(30);
static INTERACTED_EVENT: Event = Event::new();

pub fn interacted() {
    INTERACTED_EVENT.set();
}

async fn turn_off(oled: &Mutex<ThreadModeRawMutex, Oled<'_, impl Instance>>) {
    Timer::after(OLED_TIMEOUT).await;

    let _ = oled.lock().await.set_off().await;

    turn_on(oled).await;
}

async fn turn_on(oled: &Mutex<ThreadModeRawMutex, Oled<'_, impl Instance>>) {
    INTERACTED_EVENT.wait().await;

    let _ = oled.lock().await.set_on().await;
}

pub async fn display_timeout_task<'a, T: Instance>(oled: &Mutex<ThreadModeRawMutex, Oled<'a, T>>)
where
    Twim<'a, T>: I2c<u8>,
{
    loop {
        select(turn_on(oled), turn_off(oled)).await;
    }
}

To then generate the content for the displays, I use the following (for the rhs):

async fn render_normal(&mut self) {
    let character_style = MonoTextStyle::new(&PROFONT_9_POINT, BinaryColor::On);
    let textbox_style = TextBoxStyleBuilder::new()
        .height_mode(embedded_text::style::HeightMode::FitToText)
        .alignment(embedded_text::alignment::HorizontalAlignment::Justified)
        .paragraph_spacing(6)
        .build();

    let bounds = Rectangle::new(Point::zero(), Size::new(32, 0));

    self.buf.clear();

    let kp = TOTAL_KEYPRESSES.load(core::sync::atomic::Ordering::Relaxed);
    let cps = AVERAGE_KEYPRESSES.load(core::sync::atomic::Ordering::Relaxed);
    let cps = f32::trunc(cps * 10.0) / 10.0;
    let mut fp_buf = dtoa::Buffer::new();
    let cps = fp_buf.format_finite(cps);

    let _ = uwriteln!(&mut self.buf, "kp:");
    let _ = uwriteln!(&mut self.buf, "{}", kp);
    let _ = uwriteln!(&mut self.buf, "cps:");
    let _ = uwriteln!(&mut self.buf, "{}/s", cps);
    let _ = uwriteln!(&mut self.buf, "tick:");
    let _ = uwriteln!(&mut self.buf, "{}", self.ticks);

    let text_box =
        TextBox::with_textbox_style(&self.buf, bounds, character_style, textbox_style);

    let lines = {
        let samples = self.sample_buffer.lock().await;
        samples
            .oldest_ordered()
            .enumerate()
            .map(|(idx, height)| {
                Line::new(
                    Point::new(idx as i32, 128 - (*height as i32).clamp(0, 16)),
                    Point::new(idx as i32, 128),
                )
                .into_styled(PrimitiveStyle::with_stroke(BinaryColor::On, 1))
            })
            .collect::<heapless::Vec<_, 32>>()
    };

    let _ = self
        .oled
        .lock()
        .await
        .draw(move |d| {
            let _ = text_box.draw(d);
            for line in lines {
                let _ = line.draw(d);
            }
        })
        .await;
}

Inter-board communication

Since I’m not using Bluetooth, the right side of the split needs to communicate with the left side, there are a few ways to do this but I went with a UART as it easily allows both sides to send messages to the other.

However a UART has no provisions for error checking (outside of a parity bit) or framing so I have to do that myself.

To handle encoding and decoding of messages I use jamesmunns/postcard which conveniently also handles framing and failure recovery through the use of COBS

To ensure message delivery a uuid and checksum is attached to each command, when one side receives a message and validates the checksum, an Ack message with the same uuid is sent back to the other side. After a command is sent, the keyboard waits a period of time for an Ack before considering the message to have not been received. This period is variable depending on the message sent, keypress events have a longer timeout as duplicated keypresses aren’t a good thing, messages that can tolerate duplication are sent with a lower timeout to decrease latency.

The messages sent between sides are defined as plain rust enums that derive serde::Serialize and serde::Deserialize:

#[derive(Serialize, Deserialize, Eq, PartialEq, Format, Hash, Clone)]
pub enum DomToSub {
    ResyncLeds(u16),
    Reset,
    SyncKeypresses(u16),
    WritePixels {
        row: u8,
        data_0: [u8; 4],
        data_1: [u8; 4],
    },
}

#[derive(Serialize, Deserialize, Eq, PartialEq, Debug, Format, Hash, Clone)]
pub enum SubToDom {
    KeyPressed(u8),
    KeyReleased(u8),
}

These messages are then wrapped in the structs defined here:

#[derive(Serialize, Deserialize, defmt::Format, Debug)]
pub struct Command<T> {
    pub uuid: u8,
    pub csum: u8,
    pub cmd: T,
}

pub fn csum<T: Hash>(v: T) -> u8 {
    let mut hasher = StableHasher::new(fnv::FnvHasher::default());
    v.hash(&mut hasher);
    let checksum = hasher.finish();

    let bytes = checksum.to_le_bytes();

    bytes.iter().fold(0, core::ops::BitXor::bitxor)
}

impl<T: Hash> Command<T> {
    pub fn new(cmd: T) -> Self {
        static UUID_GEN: AtomicU8 = AtomicU8::new(0);
        let uuid = UUID_GEN.fetch_add(1, core::sync::atomic::Ordering::SeqCst);
        let csum = csum((&cmd, uuid));
        Self { uuid, csum, cmd }
    }

    /// validate the data of the command
    pub fn validate(&self) -> bool {
        let csum = csum((&self.cmd, self.uuid));
        csum == self.csum
    }

    pub fn ack(&self) -> Ack {
        let csum = csum(self.uuid);
        Ack {
            uuid: self.uuid,
            csum,
        }
    }
}

#[derive(Serialize, Deserialize, defmt::Format, Debug)]
pub struct Ack {
    pub uuid: u8,
    pub csum: u8,
}

#[derive(Serialize, Deserialize, defmt::Format, Debug)]
#[repr(u8)]
pub enum CmdOrAck<T> {
    Cmd(Command<T>),
    Ack(Ack),
}

impl Ack {
    pub fn validate(self) -> Option<Self> {
        let csum = csum(self.uuid);
        if csum == self.csum {
            Some(self)
        } else {
            None
        }
    }
}

And then are serialized with postcard and transmitted to the other side:

async fn task(self) {
    loop {
        let val = self.mix_chan.recv().await;

        let mut buf = [0u8; BUF_SIZE];
        if let Ok(buf) =
            postcard::serialize_with_flavor(&val, Cobs::try_new(Slice::new(&mut buf)).unwrap())
        {
            let r = self.tx.write(buf).await;
            debug!("Transmitted {:?}, r: {:?}", val, r);
        }
    }
}

Other dumb things

Okay so I have a keyboard running firmware on rust, oh and I can also talk to it over USB serial in the same way each half talks to the other. What can I do?

Metrics

With a little bit of code on the keyboard we can have it reply with the keypress counter when queried:

#[embassy::task]
async fn usb_serial_task(mut class: CdcAcmClass<'static, UsbDriver>) {
    loop {
        let in_chan: &mut Channel<ThreadModeRawMutex, u8, 128> = forever!(Channel::new());
        let out_chan: &mut Channel<ThreadModeRawMutex, u8, 128> = forever!(Channel::new());
        let msg_out_chan: &mut Channel<ThreadModeRawMutex, HostToKeyboard, 16> =
            forever!(Channel::new());
        let msg_in_chan: &mut Channel<ThreadModeRawMutex, (KeyboardToHost, Duration), 16> =
            forever!(Channel::new());
        class.wait_connection().await;
        let mut wrapper = UsbSerialWrapper::new(&mut class, &*in_chan, &*out_chan);
        let mut eventer = Eventer::new(&*in_chan, &*out_chan, msg_out_chan.sender());

        let handle = async {
            loop {
                match msg_out_chan.recv().await {
                    HostToKeyboard::RequestStats => {
                        msg_in_chan
                            .send((
                                KeyboardToHost::Stats {
                                    keypresses: TOTAL_KEYPRESSES
                                        .load(core::sync::atomic::Ordering::Relaxed),
                                },
                                Duration::from_millis(5),
                            ))
                            .await;
                    },
                    // ...
                }
            }
        };

        let (e_a, e_b, e_c) = eventer.split_tasks(msg_in_chan);

        select3(wrapper.run(), select3(e_a, e_b, e_c), handle).await;
    }
}

And then with the help of another rust program to periodically request the number of keys pressed from the keyboard (I could do this by keylogging, but that’s not as fun) and export the count to Prometheus, we get a fancy dashboard:

Video playback

Since the nRF52840 is pretty powerful, we can get away with streaming a video to the displays of the keyboard. The left side handles receiving frames from the computer over USB serial, and then sends the frame to its OLED task if the packet is for the LHS, or forwarded to the RHS otherwise.

The OLED tasks on each side have a channel to receive frames from, when a frame is received the task stops rendering the original content for a second and instead displays the received frame.

The result is this:

If you’re interested, you can find the source code here

/programming/ /rust/ /keyboards/