Refactor keyboarg lights

.cargo/config.toml

@@ -9,4 +9,4 @@ rustflags = [
 ]
 
 #runner = "elf2uf2-rs -d"
-runner = "probe-run --chip RP2040"
+runner = "probe-rs run --chip RP2040"

Cargo.lock

@@ -726,6 +726,15 @@ dependencies = [
  "wasi",
 ]
 
+[[package]]
+name = "glam"
+version = "0.25.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "151665d9be52f9bb40fc7966565d39666f2d1e69233571b71b87791c7e0528b3"
+dependencies = [
+ "libm",
+]
+
 [[package]]
 name = "half"
 version = "2.3.1"
@@ -875,6 +884,12 @@ version = "0.2.151"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "302d7ab3130588088d277783b1e2d2e10c9e9e4a16dd9050e6ec93fb3e7048f4"
 
+[[package]]
+name = "libm"
+version = "0.2.8"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "4ec2a862134d2a7d32d7983ddcdd1c4923530833c9f2ea1a44fc5fa473989058"
+
 [[package]]
 name = "libredox"
 version = "0.0.1"
@@ -1518,7 +1533,9 @@ dependencies = [
  "embedded-io-async",
  "fixed",
  "futures",
+ "glam",
  "heapless 0.7.17",
+ "libm",
  "log",
  "once_cell",
  "pio",

lib/Cargo.toml

@@ -37,6 +37,8 @@ critical-section = "1.1.1"
 crc-any = "2.4.3"
 serde = { version = "1.0.163", default-features = false, features = ["derive"] }
 bytemuck = { version = "1.13.1", features = ["derive"] }
+libm = "0.2.8"
+glam = { version = "0.25.0", default-features = false, features = ["libm"] }
 
 [target.'cfg(target_arch = "x86_64")'.dependencies]
 embassy-executor = { version = "0.5.0", features = ["arch-std"] }

lib/src/keyboard/lights.rs

@@ -5,9 +5,9 @@ use futures::{select_biased, FutureExt};
 use tgnt::button::Button;
 
 use crate::{
+    lights::shaders::{LsdHyperspace, OrthoRainbow, PowerOffAnim, PowerOnAnim, Shader, Shaders},
     rgb::Rgb,
     usb::{UsbEvent, USB_EVENTS},
-    util::wheel,
 };
 
 use super::{Event, EventKind, KbEvents, State, SWITCH_COUNT};
@@ -15,83 +15,160 @@ use super::{Event, EventKind, KbEvents, State, SWITCH_COUNT};
 /// Duration until the keyboard starts the idle animation
 const UNTIL_IDLE: Duration = Duration::from_secs(30);
 
-///// Duration from idle until the keyboard goes to sleep
+/// DUration between each animation frame.
-//const UNTIL_SLEEP: Duration = Duration::from_secs(10);
+const FRAMETIME: Duration = Duration::from_millis(16);
 
-const IDLE_ANIMATION_SPEED: u64 = 3;
+#[derive(Default)]
-const IDLE_ANIMATION_KEY_OFFSET: u64 = 10;
+enum LightsState {
-const IDLE_BRIGHTNESS_RAMPUP: Duration = Duration::from_secs(120);
+    Active {
-const MAX_IDLE_BRIGHTESS: f32 = 0.2;
+        keys: [KeyLedState; SWITCH_COUNT],
-const MIN_IDLE_BRIGHTESS: f32 = 0.05;
+        next_frame: Instant,
-
+        idle_at: Instant,
-#[derive(Clone, Copy)]
-enum LightState {
-    Solid(Rgb),
-    #[allow(dead_code)]
-    SolidThenFade {
-        color: Rgb,
-        solid_until: Instant,
-        fade_by: f32,
     },
-    FadeBy(f32),
+    PoweringOff(PowerOffAnim),
+    PoweringOn(PowerOnAnim),
+    #[default]
+    PoweredOff,
+    Idle(Shaders),
+}
+
+impl LightsState {
+    pub fn active_default() -> Self {
+        let now = Instant::now();
+        LightsState::Active {
+            keys: Default::default(),
+            next_frame: now,
+            idle_at: now + UNTIL_IDLE,
+        }
+    }
+}
+
+#[derive(Clone, Copy, Default)]
+enum KeyLedState {
+    #[default]
     None,
+    Solid(Rgb),
+    FadeBy(f32),
 }
 
 #[embassy_executor::task]
 pub(super) async fn task(mut events: KbEvents, state: &'static State) {
-    let mut lights: [LightState; SWITCH_COUNT] = [LightState::None; SWITCH_COUNT];
+    let mut lights = LightsState::default();
-    let mut next_frame = Instant::now();
-    let mut idle_at = Instant::now() + UNTIL_IDLE;
-    let mut usb_enabled: bool = false;
     let mut usb_events = USB_EVENTS
-        .subscriber()
+        .dyn_subscriber()
         .expect("USB_EVENTS: out of subscribers");
 
     loop {
-        let wait_for_idle = async {
+        match &mut lights {
-            if usb_enabled {
+            LightsState::Active {
-                Timer::at(idle_at).await
+                keys,
-            } else {
+                next_frame,
-                // if usb is disabled, we never want to start the idle animation
+                idle_at,
-                pending().await
+            } => {
+                select_biased! {
+                    event = events.recv().fuse() => {
+                        *idle_at = Instant::now() + UNTIL_IDLE;
+                        handle_event(event, state, keys).await;
+                    }
+                    ev = usb_events.next_message_pure().fuse() => {
+                        *idle_at = Instant::now() + UNTIL_IDLE;
+                        handle_usb_event(ev, &mut lights).await;
+                    }
+                    _ = Timer::at(*idle_at).fuse() => lights = LightsState::Idle(Shaders::LsdHyperspace(LsdHyperspace)),
+                    _ = Timer::at(*next_frame).fuse() => {
+                        keypress_tick(state, keys).await;
+                        *next_frame = Instant::now() + FRAMETIME;
+                    }
+                }
+            }
+            LightsState::PoweringOff(anim) => {
+                select_biased! {
+                    ev = usb_events.next_message_pure().fuse() => handle_usb_event(ev, &mut lights).await,
+                    _ = play_shader(state, anim).fuse() => lights = LightsState::PoweredOff,
+                }
+            }
+            LightsState::PoweringOn(anim) => {
+                select_biased! {
+                    ev = usb_events.next_message_pure().fuse() => handle_usb_event(ev, &mut lights).await,
+                    _ = play_shader(state, anim).fuse() => lights = LightsState::active_default(),
+                }
+            }
+            LightsState::PoweredOff => {
+                let ev = usb_events.next_message_pure().await;
+                handle_usb_event(ev, &mut lights).await;
+            }
+            LightsState::Idle(anim) => {
+                select_biased! {
+                    ev = usb_events.next_message_pure().fuse() => handle_usb_event(ev, &mut lights).await,
+                    event = events.recv().fuse() => {
+                        let now = Instant::now();
+                        let mut keys = Default::default();
+                        handle_event(event, state, &mut keys).await;
+                        lights = LightsState::Active { keys, next_frame: now, idle_at: now + UNTIL_IDLE};
+                    }
+                    _ = play_shader(state, anim).fuse() => {}
+                }
             }
         };
-
-        select_biased! {
-            event = events.recv().fuse() => {
-                handle_event(event, state, &mut lights).await;
-                idle_at = Instant::now() + UNTIL_IDLE;
-            }
-            _ = Timer::at(next_frame).fuse() => {
-                tick(state, &mut lights).await;
-                next_frame = Instant::now() + Duration::from_millis(16);
-            }
-            event = usb_events.next_message_pure().fuse() => {
-                handle_usb_event(event, &mut usb_enabled, &mut lights).await;
-                idle_at = Instant::now() + UNTIL_IDLE;
-            }
-            _ = wait_for_idle.fuse() => {
-                select_biased! {
-                    event = events.recv().fuse() => {
-                        state.lights.update(|lights| {
-                            lights.iter_mut().for_each(|rgb| *rgb = Rgb::new(0, 0, 0));
-                        }).await;
-                        handle_event(event, state, &mut lights).await;
-                    }
-                    event = usb_events.next_message_pure().fuse() => {
-                        handle_usb_event(event, &mut usb_enabled, &mut lights).await;
-                    }
-                    _ = idle_animation(state).fuse() => {}
-                }
-                idle_at = Instant::now() + UNTIL_IDLE;
-            }
-        }
     }
 }
 
-async fn tick(state: &'static State, lights: &mut [LightState; SWITCH_COUNT]) {
+async fn play_shader(state: &'static State, shader: &impl Shader) {
+    const SWITCH_COORDS: [(u16, u16); SWITCH_COUNT] = [
+        (0, 1),
+        (1, 1),
+        (2, 1),
+        (3, 1),
+        (4, 1),
+        (4, 2),
+        (3, 2),
+        (2, 2),
+        (1, 2),
+        (0, 2),
+        (0, 3),
+        (1, 3),
+        (2, 3),
+        (3, 3),
+        (4, 3),
+        (2, 4),
+        (3, 4),
+        (4, 4),
+    ];
+
+    const BRIGHTNESS: f32 = 0.15;
+
+    let switch_coords = SWITCH_COORDS.map(|(x, y)| (f32::from(x) / 4.0, f32::from(y) / 4.0));
+
+    let animate_shader = async {
+        loop {
             let now = Instant::now();
 
+            state
+                .lights
+                .update(|rgbs| {
+                    (switch_coords.into_iter().zip(rgbs.iter_mut()))
+                        .for_each(|(uv, rgb)| *rgb = shader.sample(now, uv) * BRIGHTNESS)
+                })
+                .await;
+
+            Timer::after(FRAMETIME).await;
+        }
+    };
+
+    let end_animation = async {
+        match shader.end_time() {
+            Some(end_time) => Timer::at(end_time).await,
+            None => pending().await,
+        };
+    };
+
+    select_biased! {
+        _ = animate_shader.fuse() => {}
+        _ = end_animation.fuse() => {}
+    }
+}
+
+async fn keypress_tick(state: &'static State, lights: &mut [KeyLedState; SWITCH_COUNT]) {
     state
         .lights
         .update(|rgbs| {
@@ -104,114 +181,40 @@ async fn tick(state: &'static State, lights: &mut [LightState; SWITCH_COUNT]) {
                 };
 
                 match &*light {
-                    LightState::None => {}
+                    KeyLedState::None => *rgb = Rgb::new(0, 0, 0),
-                    LightState::FadeBy(fade) => {
+                    KeyLedState::FadeBy(fade) => {
                         let [r, g, b] = rgb
                             .components()
                             .map(|c| ((c as f32) * fade.clamp(0.0, 1.0)) as u8);
                         *rgb = Rgb::new(r, g, b);
                         if *rgb == Rgb::new(0, 0, 0) {
-                            *light = LightState::None;
+                            *light = KeyLedState::None;
-                        }
-                    }
-                    &LightState::Solid(color) => *rgb = color,
-                    &LightState::SolidThenFade {
-                        color,
-                        solid_until,
-                        fade_by,
-                    } => {
-                        *rgb = color;
-                        if now >= solid_until {
-                            *light = LightState::FadeBy(fade_by);
                         }
                     }
+                    &KeyLedState::Solid(color) => *rgb = color,
                 }
             }
         })
         .await;
 }
 
-async fn idle_animation(state: &'static State) {
-    for tick in 0.. {
-        const FRAMETIME: Duration = Duration::from_millis(16);
-
-        state
-            .lights
-            .update(|lights| {
-                const N_MAX: u64 = IDLE_BRIGHTNESS_RAMPUP.as_millis() / FRAMETIME.as_millis();
-
-                let brightness = if tick >= N_MAX {
-                    MAX_IDLE_BRIGHTESS
-                } else {
-                    ((tick as f32) / N_MAX as f32).clamp(MIN_IDLE_BRIGHTESS, MAX_IDLE_BRIGHTESS)
-                };
-
-                for (n, &i) in state.led_map.iter().enumerate() {
-                    let Some(light) = lights.get_mut(i) else {
-                        continue;
-                    };
-                    let rgb = wheel(
-                        (n as u64 * IDLE_ANIMATION_KEY_OFFSET + tick * IDLE_ANIMATION_SPEED) as u8,
-                    );
-                    *light = rgb * brightness;
-                }
-            })
-            .await;
-
-        Timer::after(FRAMETIME).await;
-    }
-}
-
 async fn handle_event(
     event: Event,
     state: &'static State,
-    lights: &mut [LightState; SWITCH_COUNT],
+    lights: &mut [KeyLedState; SWITCH_COUNT],
 ) {
     let rgb = match event.kind {
         EventKind::Press { button } => match button {
-            Button::Key(..) => LightState::Solid(Rgb::new(0, 150, 0)),
+            Button::Key(..) => KeyLedState::Solid(Rgb::new(0, 150, 0)),
-            Button::Mod(..) => LightState::Solid(Rgb::new(0, 0, 150)),
+            Button::Mod(..) => KeyLedState::Solid(Rgb::new(0, 0, 150)),
-            Button::ModTap(..) => LightState::Solid(Rgb::new(0, 0, 150)),
+            Button::ModTap(..) => KeyLedState::Solid(Rgb::new(0, 0, 150)),
-            Button::Compose2(..) | Button::Compose3(..) => LightState::Solid(Rgb::new(0, 100, 100)),
+            Button::Compose2(..) | Button::Compose3(..) => {
-            Button::Layer(..) => LightState::Solid(Rgb::new(120, 0, 120)),
+                KeyLedState::Solid(Rgb::new(0, 100, 100))
-            /*
-            Button::NextLayer | Button::PrevLayer => {
-                yield_now().await; // dirty hack to make sure layer_switch_task gets to run first
-                let layer = state.current_layer.load(Ordering::Relaxed);
-                let layer = min(layer, state.layers.len().saturating_sub(1) as u16);
-                let buttons_to_light_up = if state.layers.len() <= 3 {
-                    match layer {
-                        0 => [0, 1, 2, 3, 4].as_ref(),
-                        1 => &[5, 6, 7, 8, 9],
-                        2 => &[10, 11, 12, 13, 14],
-                        _ => &[],
             }
-                } else {
+            Button::Layer(..) => KeyLedState::Solid(Rgb::new(120, 0, 120)),
-                    match layer {
+            _ => KeyLedState::Solid(Rgb::new(150, 0, 0)),
-                        0 => [0, 5, 10].as_ref(),
-                        1 => &[1, 6, 11],
-                        2 => &[2, 7, 12],
-                        3 => &[3, 8, 13],
-                        4 => &[4, 9, 14],
-                        _ => &[],
-                    }
-                };
-
-                let solid_until = Instant::now() + Duration::from_millis(200);
-                for &button in buttons_to_light_up {
-                    let Some(light) = lights.get_mut(button) else { continue; };
-                    *light = LightState::SolidThenFade {
-                        color: Rgb::new(120, 0, 120),
-                        solid_until,
-                        fade_by: 0.85,
-                    }
-                }
-                LightState::Solid(Rgb::new(100, 0, 100))
-            }
-            */
-            _ => LightState::Solid(Rgb::new(150, 0, 0)),
         },
-        EventKind::Release { .. } => LightState::FadeBy(0.85),
+        EventKind::Release { .. } => KeyLedState::FadeBy(0.85),
     };
 
     if event.source != state.half {
@@ -224,30 +227,19 @@ async fn handle_event(
     *light = rgb;
 }
 
-async fn handle_usb_event(
+async fn handle_usb_event(event: UsbEvent, state: &mut LightsState) {
-    event: UsbEvent,
+    let usb_enabled = match event {
-    is_enabled: &mut bool,
+        UsbEvent::Suspended(false) | UsbEvent::Configured(true) => true,
-    lights: &mut [LightState; SWITCH_COUNT],
+        UsbEvent::Configured(false) | UsbEvent::Suspended(true) | UsbEvent::Reset => false,
-) {
+        _ => return,
-    match event {
-        UsbEvent::Suspended(false) | UsbEvent::Configured(true) => {
-            let new_state = LightState::SolidThenFade {
-                color: Rgb::new(0, 255, 0),
-                solid_until: Instant::now() + Duration::from_millis(200),
-                fade_by: 0.85,
     };
-            lights.iter_mut().for_each(|state| *state = new_state);
+
-            *is_enabled = true;
+    let start = Instant::now();
-        }
+    *state = match (&state, usb_enabled) {
-        UsbEvent::Configured(false) | UsbEvent::Suspended(true) | UsbEvent::Reset => {
+        (LightsState::PoweringOn(..), true) => return,
-            let new_state = LightState::SolidThenFade {
+        (LightsState::PoweringOff(..), false) => return,
-                color: Rgb::new(255, 0, 0),
+        (LightsState::PoweredOff, false) => return,
-                solid_until: Instant::now() + Duration::from_millis(200),
+        (_, true) => LightsState::PoweringOn(PowerOnAnim { start }),
-                fade_by: 0.85,
+        (_, false) => LightsState::PoweringOff(PowerOffAnim { start }),
     };
-            lights.iter_mut().for_each(|state| *state = new_state);
-            *is_enabled = false;
-        }
-        _ => {}
-    }
 }

lib/src/lights.rs

@@ -1,3 +1,5 @@
+pub mod shaders;
+
 use crate::ws2812::Ws2812;
 use embassy_rp::pio;
 use embassy_sync::mutex::Mutex;

lib/src/lights/shaders.rs

@@ -0,0 +1,124 @@
+use core::f32::consts::PI;
+
+use embassy_time::{Duration, Instant};
+use glam::{vec2, vec3, Vec3};
+use libm::cosf;
+
+use crate::rgb::Rgb;
+
+/// A fragment shader.
+pub trait Shader {
+    /// Sample a normalized coordinate (0 to 1) using the shader function and return a color.
+    fn sample(&self, time: Instant, uv: (f32, f32)) -> Rgb;
+
+    fn end_time(&self) -> Option<Instant> {
+        None
+    }
+}
+
+pub enum Shaders {
+    OrthoRainbow(OrthoRainbow),
+    LsdHyperspace(LsdHyperspace),
+}
+
+impl Shader for Shaders {
+    fn sample(&self, time: Instant, uv: (f32, f32)) -> Rgb {
+        match self {
+            Shaders::OrthoRainbow(s) => s.sample(time, uv),
+            Shaders::LsdHyperspace(s) => s.sample(time, uv),
+        }
+    }
+
+    fn end_time(&self) -> Option<Instant> {
+        match self {
+            Shaders::OrthoRainbow(s) => s.end_time(),
+            Shaders::LsdHyperspace(s) => s.end_time(),
+        }
+    }
+}
+
+pub struct OrthoRainbow;
+impl Shader for OrthoRainbow {
+    fn sample(&self, time: Instant, (x, y): (f32, f32)) -> Rgb {
+        let time = time.as_millis() as f32 / 1000.0;
+        let r = 0.5 + 0.5 * cosf(time + x + 0.0);
+        let g = 0.5 + 0.5 * cosf(time + y + 2.0);
+        let b = 0.5 + 0.5 * cosf(time + x + 4.0);
+
+        Rgb::from_f32s(r, g, b)
+    }
+}
+
+pub struct LsdHyperspace;
+impl Shader for LsdHyperspace {
+    fn sample(&self, time: Instant, uv: (f32, f32)) -> Rgb {
+        let time = time.as_millis() as f32 / 1000.0 * 3.0;
+        let uv = vec2(uv.0, uv.1);
+
+        let center = vec2(0.5, 0.5);
+
+        let dist = (uv - center).length();
+
+        let fac = dist * PI + vec3(0.0, 1.0, 4.0) - time;
+        let col = cos3(fac) * cos3(fac * 0.5);
+
+        Rgb::from_f32s(col.x, col.y, col.z)
+    }
+}
+
+pub struct PowerOffAnim {
+    /// Animation starting time.
+    pub start: Instant,
+}
+
+impl PowerOffAnim {
+    const FADE_FACTOR: f32 = 0.5;
+
+    /// Animation duration, in seconds.
+    const DURATION_SEC: u16 = 5;
+}
+
+impl Shader for PowerOffAnim {
+    fn sample(&self, time: Instant, (_, y): (f32, f32)) -> Rgb {
+        let time = time.as_millis().saturating_sub(self.start.as_millis());
+        let time = time as f32 / 1000.0;
+
+        let duration: f32 = Self::DURATION_SEC.into();
+        let r = Self::FADE_FACTOR * (duration - time - 1.0 + y);
+        Rgb::from_f32s(r, 0.0, 0.0)
+    }
+
+    fn end_time(&self) -> Option<Instant> {
+        Some(self.start + Duration::from_secs(Self::DURATION_SEC.into()))
+    }
+}
+
+pub struct PowerOnAnim {
+    /// Animation starting time.
+    pub start: Instant,
+}
+
+impl PowerOnAnim {
+    const FADE_FACTOR: f32 = 1.0;
+
+    /// Animation duration, in seconds.
+    const DURATION_SEC: u16 = 1;
+}
+
+fn cos3(v: Vec3) -> Vec3 {
+    vec3(cosf(v.x), cosf(v.y), cosf(v.z))
+}
+
+impl Shader for PowerOnAnim {
+    fn sample(&self, time: Instant, (_, y): (f32, f32)) -> Rgb {
+        let time = time.as_millis().saturating_sub(self.start.as_millis());
+        let time = time as f32 / 1000.0;
+
+        let duration: f32 = Self::DURATION_SEC.into();
+        let g = Self::FADE_FACTOR * (duration - time - y);
+        Rgb::from_f32s(0.0, g, 0.0)
+    }
+    fn end_time(&self) -> Option<Instant> {
+        Some(self.start + Duration::from_secs(Self::DURATION_SEC.into()))
+    }
+}

lib/src/rgb.rs

@@ -15,6 +15,20 @@ impl Rgb {
         Self(u32::from_be_bytes([g, r, b, 0]))
     }
 
+    pub fn from_f64s(r: f64, g: f64, b: f64) -> Self {
+        let r = r.clamp(0.0, 1.0) * 255.0;
+        let g = g.clamp(0.0, 1.0) * 255.0;
+        let b = b.clamp(0.0, 1.0) * 255.0;
+        Self::new(r as u8, g as u8, b as u8)
+    }
+
+    pub fn from_f32s(r: f32, g: f32, b: f32) -> Self {
+        let r = r.clamp(0.0, 1.0) * 255.0;
+        let g = g.clamp(0.0, 1.0) * 255.0;
+        let b = b.clamp(0.0, 1.0) * 255.0;
+        Self::new(r as u8, g as u8, b as u8)
+    }
+
     /// Get the red, green, and blue components of this Rgb.
     #[inline(always)]
     pub const fn components(&self) -> [u8; 3] {