RKTK API Docs RKTK Home Repo

rktk_drivers_nrf/rgb/
ws2812_pwm.rs

1//! WS2812 PWM driver for nRF52
2//! Based on https://github.com/embassy-rs/embassy/blob/main/examples/nrf52840/src/bin/pwm_sequence_ws2812b.rs
3
4use core::convert::Infallible;
5
6use bitvec::prelude::*;
7use embassy_nrf::{
8    Peripheral,
9    gpio::{OutputDrive, Pin},
10    pwm::{
11        Config, Instance, Prescaler, SequenceLoad, SequencePwm, SingleSequenceMode, SingleSequencer,
12    },
13};
14use embassy_time::Timer;
15use rktk::drivers::interface::rgb::{LedRgb, RgbDriver};
16
17/// WS2812 NeoPixel driver using PWM on nRF.
18///
19/// ## Note about `BUFFER_SIZE`
20/// Buffer size should be large enough to hold the number of pixels you want to control.
21/// It should be larger than `led_count * 24 + 100` to ensure enough space for the sequence.
22///
23/// Even buffer size is not enough, pixels will be written, but the driver will log a warning.
24pub struct Ws2812Pwm<
25    const BUFFER_SIZE: usize,
26    PWM: Instance,
27    PWMP: Peripheral<P = PWM>,
28    DATA: Pin,
29    DATAP: Peripheral<P = DATA>,
30> {
31    pwm: PWMP,
32    pin: DATAP,
33}
34
35const T1H: u16 = 0x8000 | 13;
36const T0H: u16 = 0x8000 | 7;
37const RES: u16 = 0x8000;
38
39impl<
40    const BUFFER_SIZE: usize,
41    PWM: Instance + 'static,
42    PWMP: Peripheral<P = PWM> + 'static,
43    DATA: Pin,
44    DATAP: Peripheral<P = DATA> + 'static,
45> Ws2812Pwm<BUFFER_SIZE, PWM, PWMP, DATA, DATAP>
46{
47    pub fn new(pwm: PWMP, data: DATAP) -> Self {
48        Self { pwm, pin: data }
49    }
50}
51
52impl<
53    const BUFFER_SIZE: usize,
54    PWM: Instance + 'static,
55    PWMP: Peripheral<P = PWM> + 'static,
56    DATA: Pin,
57    DATAP: Peripheral<P = DATA> + 'static,
58> RgbDriver for Ws2812Pwm<BUFFER_SIZE, PWM, PWMP, DATA, DATAP>
59{
60    type Error = Infallible;
61
62    async fn write<I: IntoIterator<Item = LedRgb<u8>>>(
63        &mut self,
64        pixels: I,
65    ) -> Result<(), Self::Error> {
66        let mut pwm_config = Config::default();
67        pwm_config.sequence_load = SequenceLoad::Common;
68        pwm_config.prescaler = Prescaler::Div1; // 16MHz
69        pwm_config.max_duty = 20; // 1.25us (1s / 16Mhz * 20)
70        pwm_config.ch0_drive = OutputDrive::HighDrive;
71        let Ok(mut seq_pwm) = SequencePwm::new_1ch(&mut self.pwm, &mut self.pin, pwm_config) else {
72            // TODO: Handle error
73            return Ok(());
74        };
75
76        let mut words = heapless::Vec::<u16, BUFFER_SIZE>::from_slice(&[RES; 100]).unwrap();
77
78        'outer: for color in pixels {
79            for bit in color[1]
80                .view_bits::<Msb0>()
81                .iter()
82                .chain(color[0].view_bits())
83                .chain(color[2].view_bits())
84            {
85                if words.push(if *bit { T1H } else { T0H }).is_err() {
86                    rktk_log::warn!("WS2812Pwm buffer size is not enough. Increase BUFFER_SIZE.");
87                    break 'outer;
88                }
89            }
90        }
91
92        let seq_config = embassy_nrf::pwm::SequenceConfig::default();
93        let sequencer = SingleSequencer::new(&mut seq_pwm, words.as_slice(), seq_config);
94        let _ = sequencer.start(SingleSequenceMode::Times(1));
95
96        // Wait for a long time is important. Otherwise, the PWM will be stopped before the
97        // sequence is finished.
98        Timer::after_millis(5).await;
99
100        Ok(())
101    }
102}