From 4e83970cc3f0b5fc48d4c208ea5f1e256169d215 Mon Sep 17 00:00:00 2001
From: Joakim Hulthe <joakim@hulthe.net>
Date: Thu, 16 Mar 2023 22:59:59 +0100
Subject: [PATCH] Improve ws2812 api

---
 src/main.rs   | 51 ++++++++++++++++++++++++++++++---------------------
 src/ws2812.rs | 30 +++++++++++++++++-------------
 2 files changed, 47 insertions(+), 34 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index 5b2c97a..c15e16f 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -11,7 +11,6 @@
 #![feature(type_alias_impl_trait)]
 
 extern crate cortex_m_rt;
-//extern crate panic_halt;
 
 mod board;
 mod keyboard;
@@ -22,10 +21,7 @@ mod ws2812;
 
 use board::Board;
 use embassy_executor::Spawner;
-use embassy_rp::{
-    gpio::{Level, Output, Pin},
-    pio::PioPeripheral,
-};
+use embassy_rp::gpio::{Level, Output, Pin};
 use embassy_time::{Duration, Timer};
 use ws2812::Rgb;
 
@@ -64,9 +60,8 @@ async fn main(spawner: Spawner) {
     //let mut led = Output::new(board.d13, Level::Low);
     let _neopixel_power = Output::new(board.neopixel_power, Level::High);
 
-    let (_, sm, ..) = p.PIO0.split();
-    let mut neopixel = ws2812::Ws2812::new(sm, p.DMA_CH0, board.neopixel.degrade());
-    //let mut neopixel = ws2812::Ws2812::new(sm, p.DMA_CH0, board.d5.degrade());
+    let mut neopixel = ws2812::Ws2812::new(p.PIO0, p.DMA_CH0, board.neopixel.degrade());
+    let mut neopixels_d5 = ws2812::Ws2812::new(p.PIO1, p.DMA_CH1, board.d5.degrade());
 
     neopixel.write(&[Rgb::new(0xb7, 0x31, 0x2c)]).await;
 
@@ -111,18 +106,32 @@ async fn main(spawner: Spawner) {
 
     //keyboard::test::type_string("Hello there!\n").await;
     //keyboard::test::rollover(['h', 'e', 'l', 'o', 't', 'r', 'a', 'b', 'c', 'd', 'i']).await;
-    loop {
-        //neopixel
-        //    .write(&[
-        //        Rgb::new(0xAA, 0xFF, 0x00),
-        //        Rgb::new(0xAA, 0xFF, 0x00),
-        //        Rgb::new(0xAA, 0xFF, 0x00),
-        //        Rgb::new(0xAA, 0xFF, 0x00),
-        //        Rgb::new(0xAA, 0xFF, 0x00),
-        //        Rgb::new(0xAA, 0xFF, 0x00),
-        //    ])
-        //    .await;
-        //Timer::after(Duration::from_millis(10)).await;
-        Timer::after(Duration::from_secs(10)).await;
+    for w in 0usize.. {
+        neopixel.write(&[wheel(w as u8)]).await;
+        neopixels_d5
+            .write(&[
+                wheel((w + 50) as u8),
+                wheel((w + 100) as u8),
+                wheel((w + 150) as u8),
+                wheel((w + 200) as u8),
+            ])
+            .await;
+        Timer::after(Duration::from_millis(10)).await;
+        //Timer::after(Duration::from_secs(10)).await;
     }
 }
+
+/// Input a value 0 to 255 to get a color value
+// The colours are a transition r - g - b - back to r.
+fn wheel(mut wheel_pos: u8) -> Rgb {
+    wheel_pos = 255 - wheel_pos;
+    if wheel_pos < 85 {
+        return Rgb::new(255 - wheel_pos * 3, 0, wheel_pos * 3);
+    }
+    if wheel_pos < 170 {
+        wheel_pos -= 85;
+        return Rgb::new(0, wheel_pos * 3, 255 - wheel_pos * 3);
+    }
+    wheel_pos -= 170;
+    Rgb::new(wheel_pos * 3, 255 - wheel_pos * 3, 0)
+}
diff --git a/src/ws2812.rs b/src/ws2812.rs
index e662be2..91a5b56 100644
--- a/src/ws2812.rs
+++ b/src/ws2812.rs
@@ -3,27 +3,30 @@ use core::mem::transmute;
 
 use embassy_rp::dma::{self, AnyChannel};
 use embassy_rp::pio::{
-    FifoJoin, PioInstance, PioStateMachine, PioStateMachineInstance, ShiftDirection, SmInstance,
+    FifoJoin, PioInstance, PioPeripheral, PioStateMachine, PioStateMachineInstance, ShiftDirection,
+    SmInstanceBase,
 };
 use embassy_rp::pio_instr_util;
 use embassy_rp::relocate::RelocatedProgram;
 use embassy_rp::{gpio, PeripheralRef};
 
-pub struct Ws2812<P: PioInstance, S: SmInstance> {
-    sm: PioStateMachineInstance<P, S>,
+pub struct Ws2812<P: PioInstance> {
+    sm: PioStateMachineInstance<P, SmInstanceBase<0>>,
     dma: PeripheralRef<'static, AnyChannel>,
 }
 
 /// An Rgb value that can be safely transmuted to u32 for use with Ws2812.
 #[repr(transparent)]
+#[derive(Clone, Copy, PartialEq, Eq)]
 pub struct Rgb(u32);
 
-impl<P: PioInstance, S: SmInstance> Ws2812<P, S> {
-    pub fn new(
-        mut sm: PioStateMachineInstance<P, S>,
+impl<P: PioInstance> Ws2812<P> {
+    pub fn new<PP: PioPeripheral<Pio = P>>(
+        pio: PP,
         dma: impl dma::Channel,
         pin: gpio::AnyPin,
     ) -> Self {
+        let (_, mut sm, ..) = pio.split();
         // prepare the PIO program
         let side_set = pio::SideSet::new(false, 1, false);
         let mut a: pio::Assembler<32> = pio::Assembler::new_with_side_set(side_set);
@@ -94,28 +97,29 @@ impl<P: PioInstance, S: SmInstance> Ws2812<P, S> {
     }
 
     pub async fn write(&mut self, colors: &[Rgb]) {
-        let colors = rgbs_to_u32s(colors);
+        let colors = Rgb::slice_as_u32s(colors);
         self.sm.dma_push(self.dma.reborrow(), colors).await;
     }
 }
 
-#[inline(always)]
-fn rgbs_to_u32s(rgbs: &[Rgb]) -> &[u32] {
-    // SAFETY: Rgb contains only a u32, and is #[repr(transparent)]
-    unsafe { transmute(rgbs) }
-}
-
 impl Rgb {
     #[inline(always)]
     pub const fn new(r: u8, g: u8, b: u8) -> Self {
         Self(u32::from_be_bytes([g, r, b, 0]))
     }
 
+    /// Get the red, green, and blue components of this Rgb.
     #[inline(always)]
     pub const fn components(&self) -> (u8, u8, u8) {
         let [g, r, b, _] = self.0.to_be_bytes();
         (r, g, b)
     }
+
+    #[inline(always)]
+    pub fn slice_as_u32s(rgbs: &[Rgb]) -> &[u32] {
+        // SAFETY: Rgb contains only a u32, and is #[repr(transparent)]
+        unsafe { transmute(rgbs) }
+    }
 }
 
 impl Debug for Rgb {