Skip to main content
RKTK API Docs RKTK Home Repo

rktk_drivers_rp/rgb/
ws2812_pio.rs

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