RKTK API Docs RKTK Home Repo

rktk_drivers_rp/rgb/
ws2812_pio.rs

1use core::convert::Infallible;
2
3use embassy_rp::dma::{AnyChannel, Channel};
4use embassy_rp::pio::{
5    Config, FifoJoin, Instance, Pio, PioPin, ShiftConfig, ShiftDirection, StateMachine,
6};
7use embassy_rp::{Peri, clocks};
8use embassy_time::Timer;
9use fixed::types::U24F8;
10use fixed_macro::fixed;
11use rktk::drivers::interface::rgb::{LedRgb, RgbDriver};
12
13pub struct Ws2812Pio<'a, const MAX_LED_COUNT: usize, I: Instance> {
14    dma: Peri<'a, AnyChannel>,
15    sm: StateMachine<'a, I, 0>,
16}
17
18impl<'a, const MAX_LED_COUNT: usize, I: Instance> Ws2812Pio<'a, MAX_LED_COUNT, I> {
19    pub fn new<'b: 'a>(
20        mut pio: Pio<'a, I>,
21        data_pin: Peri<'b, impl PioPin>,
22        dma: Peri<'a, impl Channel>,
23    ) -> Self {
24        // Setup sm0
25
26        // prepare the PIO program
27        let side_set = pio::SideSet::new(false, 1, false);
28        let mut a: pio::Assembler<32> = pio::Assembler::new_with_side_set(side_set);
29
30        const T1: u8 = 2; // start bit
31        const T2: u8 = 5; // data bit
32        const T3: u8 = 3; // stop bit
33        const CYCLES_PER_BIT: u32 = (T1 + T2 + T3) as u32;
34
35        let mut wrap_target = a.label();
36        let mut wrap_source = a.label();
37        let mut do_zero = a.label();
38        a.set_with_side_set(pio::SetDestination::PINDIRS, 1, 0);
39        a.bind(&mut wrap_target);
40        // Do stop bit
41        a.out_with_delay_and_side_set(pio::OutDestination::X, 1, T3 - 1, 0);
42        // Do start bit
43        a.jmp_with_delay_and_side_set(pio::JmpCondition::XIsZero, &mut do_zero, T1 - 1, 1);
44        // Do data bit = 1
45        a.jmp_with_delay_and_side_set(pio::JmpCondition::Always, &mut wrap_target, T2 - 1, 1);
46        a.bind(&mut do_zero);
47        // Do data bit = 0
48        a.nop_with_delay_and_side_set(T2 - 1, 0);
49        a.bind(&mut wrap_source);
50
51        let prg = a.assemble_with_wrap(wrap_source, wrap_target);
52        let mut cfg = Config::default();
53
54        // Pin config
55        let out_pin = pio.common.make_pio_pin(data_pin);
56        cfg.set_out_pins(&[&out_pin]);
57        cfg.set_set_pins(&[&out_pin]);
58
59        cfg.use_program(&pio.common.load_program(&prg), &[&out_pin]);
60
61        // Clock config, measured in kHz to avoid overflows
62        // TODO CLOCK_FREQ should come from embassy_rp
63        let clock_freq = U24F8::from_num(clocks::clk_sys_freq() / 1000);
64        let ws2812_freq = fixed!(800: U24F8);
65        let bit_freq = ws2812_freq * CYCLES_PER_BIT;
66        cfg.clock_divider = clock_freq / bit_freq;
67
68        // FIFO config
69        cfg.fifo_join = FifoJoin::TxOnly;
70        cfg.shift_out = ShiftConfig {
71            auto_fill: true,
72            threshold: 24,
73            direction: ShiftDirection::Left,
74        };
75
76        pio.sm0.set_config(&cfg);
77        pio.sm0.set_enable(true);
78
79        Self {
80            dma: dma.into(),
81            sm: pio.sm0,
82        }
83    }
84}
85
86impl<const MAX_LED_COUNT: usize, I: Instance + 'static> RgbDriver
87    for Ws2812Pio<'static, MAX_LED_COUNT, I>
88{
89    type Error = Infallible;
90
91    async fn write<IT: IntoIterator<Item = LedRgb<u8>>>(
92        &mut self,
93        pixels: IT,
94    ) -> Result<(), Self::Error> {
95        let mut words = [0u32; MAX_LED_COUNT];
96        for (p, w) in pixels.into_iter().zip(words.iter_mut()) {
97            *w = (u32::from(p[1]) << 24) | (u32::from(p[0]) << 16) | (u32::from(p[2]) << 8);
98        }
99
100        // DMA transfer
101        self.sm
102            .tx()
103            .dma_push(self.dma.reborrow(), &words, false)
104            .await;
105
106        Timer::after_micros(55).await;
107
108        Ok(())
109    }
110}