Rewriting smart-keymap's Key Storage Implementation

I recently finished a big rewrite of how the keys were stored in my smart-keymap keyboard firmware.

Smart Keymap

In case you’ve found this blogpost and aren’t familiar with smart-keymap:

smart-keymap is a library for building keyboard firmware by declaring keymaps in Nickel, powered by Rust.

A ‘smart keyboard’ is a keyboard with additional bells and whistles, such as alternate key functionality (layers, tap-hold keys, chords, etc.), or RGB effects, etc. – e.g. QMK, ZMK are popular smart keyboard firmware frameworks.

This project provides a library which handles the keymap behaviour part of this. With smart-keymap, it’s easy to write keyboard firmware, in Rust using embassy or RTIC, or quickly provide powerful keymap functionality to EVT example code written in C.

Rewrite Motivation

I was asked in a GitHub issue if was possible to implement the sophisticated artseyio in smart-keymap.

After brushing smart-keymap’s chording engine up a bit so that it supported overlapping chords, I found that the artseyio keymap barely fit on the low-budget CH32X035 MCU. (With a slightly older compiler version, I observed it didn’t even fit).

I considered that the implementation (details discussed below) maybe wasn’t the most efficient; I wanted the artseyio keymap firmware to fit on the CH32X035.

The Numbers

The rewrite can be found at pull request #381 of smart-keymap.

Overall, the summary says +6,783 −8,730.

The number added is smaller than the number deleted. I’m happy with that!

Though, this large number of lines changed does include changes to automatically generated test snapshots.
For an example of code that I touched, the git summary indicates ncl/keymap-codegen.ncl had around 2000 lines of code changed.

As for the goal of reducing firmware size?
For the example firmware built as a CI check, its size reduced from 78% of the CH32X’s flash size, down to 54% of the firmware’s flash size.
For the motivating artseyio keymap, its size reduced from 95% of the CH32X’s flash size, down to 60%.

Rewrite Details: from Tree-Based Key Storage, to Struct-of-Arrays

The inspiration for writing smart-keymap was semickolon’s kirei. Kirei emphasised the somewhat unusual “keys as the main abstraction of keymap behaviour”.
In the project’s discord channel, the author discussed all sorts of inspiring ideas, like how a keymap engine is bears resemblance to a parser.. or how, if you squint, various smart keyboard functionality (like tap hold, layers, chords) can be generalised as varying the functionality depending on recent key presses to the keymap.
(That’s why kirei didn’t directly implement layering).

So.
smart-keymap takes this idea of “the key as the main abstraction of keymap behaviour”, and implements functionality like “tap hold keys” or “layered keys”.
It was really natural to describe these as tree structures: a tap-hold key is a key which acts like one key when tapped, as another key when held. It’s obvious to describe this with a struct like { tap, hold }. Similarly, a layered key is a key which behaves differently depending on which layers are active; so, { base, layered: [...] } is an obvious struct.

In smart-keymap’s case, the keys themselves are conceptually something like:

   Key = Chorded { chord: Key, passthrough: Key }
       | Layered { base: Key, layered: Key }
       | TapHold { tap: Key, hold: Key }
       | HidKeyCode
       | LayerModifier
       | ...

Tree-like Structs on the Stack

There were a couple of problems of directly implementing the structs like this, particularly for embedded firmware:

Similar to how with a linked list List<x> = Node<x> | Empty, enum<T> { Node(T), Empty } doesn’t work in Rust, a naive attempt at writing out these Key structs isn’t going to work.
Roughly, the code needs to know what size the values are in order to call the function (because the values are on the stack, not on the heap). And so it’s a problem if the structure doesn’t have a determined size.

We can adjust the structure of the tree so that key values can’t be arbitrarily deep:

   Key = ChordedKey;
   ChordedKey = Chorded { chord: LayeredKey, passthrough: LayeredKey } | LayeredKey;
   LayeredKey = Layered { base: TapHoldKey, layered: TapHoldKey } | TapHoldKey;
   TapHoldKey = TapHold { tap: BaseKey, hold: BaseKey } | BaseKey;
   BaseKey = HidKeyCode | LayerModifier | ...`

In order to express ChordedKey = Chorded { chord, passthrough } | Layered without using the heap in Rust, I used a lot of ‘wrapping’ code to coerce the various values to be the same type (& also indicate to the compiler that the struct has finite size).
Chorded<K> { chord: K, passthrough: K } requires that chord and passthrough have the same type; which is a bit annoying if chord is an HidKeyCode, and passthrough is a TapHold<HidKeyCode>.

The size of each key value grows exponentially large, for the depth of this hierarchy of keys.
For the struct above: a TapHold key is going to be twice as large as a base key, a LayeredKey might be NumLayers times that.
A ChordedKey will be twice that. (And, if smart-keymap were to have another kind of smart key that nests ChordedKeys…).

Flattening the Tree using Struct-of-Arrays

The idea to reduce storage size: instead of TapHold<K> { tap: K, hold: K }, we have TapHold { tap: Ref, hold: Ref }, and a struct like Keys { keyboard: [...], tap_hold: [...], layered: [...], ... }.

That is, rather than nesting the key data inside each key, the key definitions are stored in a struct of arrays, and each key definition indirectly refers to other keys using a Ref.
e.g. a Chorded key’s chord might be tap_hold::Ref(3); and looking up keys.tap_hold[3] then refers to some TapHold { tap: Ref, hold: Ref } value.

Since each key definition only nests ‘references’, the size of the key data does not grow exponentially as a result of a deeper hierarchy of key implementations.
This also avoids all of the complicated fuss that was required for the previous nested tree structure implementation.

:o) This is a very straightforward to change to make to the code.

I think if you squint, this is comparable to going from an AST-based interpreter, to compiling the AST to instructions.

Lesson Learned

Things you might be able to take away from this?

• For emdedded firmware, even low cost MCUs like the CH32X have enough flash space that you can run fairly sophisticated code without needing to fuss over having a particularly efficient implementation.

• If firmware size is an issue, having using tree-like structs on the stack is something to avoid.

Rewrite Approach: A Pragmatic All-in-One-Go

I was impatient to implement this change.

It probably would have been possible to add indirection so that the old code still worked, add in the new code (& make use of the common functionality), then shift over to the new code.

Instead, I opted to just rewrite in place.

The upside is: I wasn’t constrained by the old implementation.

One downside to this: this meant most of the functionality wasn’t going to work until the rewrite was finished.

Another downside was that, the rewrite can’t then be split up into discrete commits and have a working project. Because traits core to the implementation were written, the tests would fail (or not even compile) for practically everything more granular than the full rewrite.

Positive About Rust: Frictionless Refactoring

The rewrite basically involved smooshing two different traits (key::Key and key::KeyState) into one (key::System).

Because the associated types and methods were similar, (and because I’d reshaped the code a bit with “use nested Refs, not nested key data” in mind), much of the rewrite was fairly straightforward.

This meant I rarely to re-think the details. And having a comprehensive integration test suite helped provide confidence that the code was working as it should; or catch times when I’d trip over Rust subtleties.

Rust’s enums make for a really clean way of expressing different Rust states some value might have. This made the rewrite flow.

Typo-Preventer

Using Rust’s enums and pattern matching helps prevent typos.

In the smart-keymap codebase, the various smart key implementations (keyboard, tap-hold, layered, etc.) get aggregated by a composite key. This composite type has a Ref type like enum Ref{ Keyboard(keyboard::Ref), TapHold(tap_hold::Ref), .. }.

This sum type make it easier to retain type safety across code which dispatches method calls to the respective key system implementations, helping to catch typos from copy pasting.

Rust’s Enums Help Keep Code Small

The meat of this rewrite, where replacing key::Key and key::KeyState with key::System involved the most thought, was in the keymap::Keymap implementation itself.

I feel that by using enums, this helped reduce the amount of code I had to keep in mind at once.

For example, the new_pressed_key method returns a PressedKeyResult enum (indicating whether the key press resulted in a pending key state, a resolved key state, or the press of another key).
Because this method returns an enum value, as opposed to calling back to a method on Keymap or some other technique, I believe it’s easier to reason about assumptions about what goes into the enum value, and assumptions are made from using the enum value. – It’s easier to use correctly, and more difficult to use incorrectly than an arbitrary method call.

---

---