rp/pio: add load_program, use_program

programs contain information we could pull from them directly and use to
validate other configuration of the state machine instead of asking the
user to pull them out and hand them to us bit by bit. unfortunately
programs do not specify how many in or out bits they use, so we can only
handle side-set and wrapping jumps like this. it's still something though.

5 files changed, 74 insertions, 115 deletions

diff --git a/embassy-rp/src/pio.rs b/embassy-rp/src/pio.rs
index ea6814fb8..f835fc8d9 100644
--- a/embassy-rp/src/pio.rs
+++ b/embassy-rp/src/pio.rs
@@ -9,12 +9,14 @@ use embassy_cortex_m::interrupt::{Interrupt, InterruptExt};
 use embassy_hal_common::{into_ref, Peripheral, PeripheralRef};
 use embassy_sync::waitqueue::AtomicWaker;
 use pac::io::vals::Gpio0ctrlFuncsel;
+use pio::{SideSet, Wrap};
 
 use crate::dma::{Channel, Transfer, Word};
 use crate::gpio::sealed::Pin as SealedPin;
 use crate::gpio::{self, AnyPin, Drive, Pull, SlewRate};
 use crate::pac::dma::vals::TreqSel;
-use crate::{interrupt, pac, peripherals, RegExt};
+use crate::relocate::RelocatedProgram;
+use crate::{interrupt, pac, peripherals, pio_instr_util, RegExt};
 
 struct Wakers([AtomicWaker; 12]);
 
@@ -460,6 +462,13 @@ impl<'d, PIO: Instance, const SM: usize> Drop for StateMachine<'d, PIO, SM> {
     }
 }
 
+fn assert_consecutive<'d, PIO: Instance>(pins: &[&Pin<'d, PIO>]) {
+    for (p1, p2) in pins.iter().zip(pins.iter().skip(1)) {
+        // purposely does not allow wrap-around because we can't claim pins 30 and 31.
+        assert!(p1.pin() + 1 == p2.pin(), "pins must be consecutive");
+    }
+}
+
 impl<'d, PIO: Instance + 'd, const SM: usize> StateMachine<'d, PIO, SM> {
     #[inline(always)]
     fn this_sm() -> crate::pac::pio::StateMachine {
@@ -504,26 +513,26 @@ impl<'d, PIO: Instance + 'd, const SM: usize> StateMachine<'d, PIO, SM> {
         }
     }
 
-    pub fn set_side_enable(&self, enable: bool) {
-        unsafe {
-            Self::this_sm().execctrl().modify(|w| w.set_side_en(enable));
-        }
-    }
-
-    pub fn is_side_enabled(&self) -> bool {
-        unsafe { Self::this_sm().execctrl().read().side_en() }
-    }
-
-    pub fn set_side_pindir(&mut self, pindir: bool) {
+    /// Configures this state machine to use the given program, including jumping to the origin
+    /// of the program. The state machine is not started.
+    pub fn use_program(&mut self, prog: &LoadedProgram<'d, PIO>, side_set: &[&Pin<'d, PIO>]) {
+        assert!((prog.side_set.bits() - prog.side_set.optional() as u8) as usize == side_set.len());
+        assert_consecutive(side_set);
         unsafe {
-            Self::this_sm().execctrl().modify(|w| w.set_side_pindir(pindir));
+            Self::this_sm().execctrl().modify(|w| {
+                w.set_side_en(prog.side_set.optional());
+                w.set_side_pindir(prog.side_set.pindirs());
+                w.set_wrap_bottom(prog.wrap.target);
+                w.set_wrap_top(prog.wrap.source);
+            });
+            Self::this_sm().pinctrl().modify(|w| {
+                w.set_sideset_count(prog.side_set.bits());
+                w.set_sideset_base(side_set.first().map_or(0, |p| p.pin()));
+            });
+            pio_instr_util::exec_jmp(self, prog.origin);
         }
     }
 
-    pub fn is_side_pindir(&self) -> bool {
-        unsafe { Self::this_sm().execctrl().read().side_pindir() }
-    }
-
     pub fn set_jmp_pin(&mut self, pin: u8) {
         unsafe {
             Self::this_sm().execctrl().modify(|w| w.set_jmp_pin(pin));
@@ -534,23 +543,6 @@ impl<'d, PIO: Instance + 'd, const SM: usize> StateMachine<'d, PIO, SM> {
         unsafe { Self::this_sm().execctrl().read().jmp_pin() }
     }
 
-    pub fn set_wrap(&self, source: u8, target: u8) {
-        unsafe {
-            Self::this_sm().execctrl().modify(|w| {
-                w.set_wrap_top(source);
-                w.set_wrap_bottom(target)
-            });
-        }
-    }
-
-    /// Get wrapping addresses. Returns (source, target).
-    pub fn get_wrap(&self) -> (u8, u8) {
-        unsafe {
-            let r = Self::this_sm().execctrl().read();
-            (r.wrap_top(), r.wrap_bottom())
-        }
-    }
-
     pub fn set_fifo_join(&mut self, join: FifoJoin) {
         let (rx, tx) = match join {
             FifoJoin::Duplex => (false, false),
@@ -667,28 +659,6 @@ impl<'d, PIO: Instance + 'd, const SM: usize> StateMachine<'d, PIO, SM> {
     pub fn get_addr(&self) -> u8 {
         unsafe { Self::this_sm().addr().read().addr() }
     }
-    pub fn set_sideset_count(&mut self, count: u8) {
-        unsafe {
-            Self::this_sm().pinctrl().modify(|w| w.set_sideset_count(count));
-        }
-    }
-
-    pub fn get_sideset_count(&self) -> u8 {
-        unsafe { Self::this_sm().pinctrl().read().sideset_count() }
-    }
-
-    pub fn set_sideset_base_pin(&mut self, base_pin: &Pin<PIO>) {
-        unsafe {
-            Self::this_sm().pinctrl().modify(|w| w.set_sideset_base(base_pin.pin()));
-        }
-    }
-
-    pub fn get_sideset_base(&self) -> u8 {
-        unsafe {
-            let r = Self::this_sm().pinctrl().read();
-            r.sideset_base()
-        }
-    }
 
     /// Set the range of out pins affected by a set instruction.
     pub fn set_set_range(&mut self, base: u8, count: u8) {
@@ -799,8 +769,35 @@ pub struct InstanceMemory<'d, PIO: Instance> {
     pio: PhantomData<&'d mut PIO>,
 }
 
+pub struct LoadedProgram<'d, PIO: Instance> {
+    pub used_memory: InstanceMemory<'d, PIO>,
+    origin: u8,
+    wrap: Wrap,
+    side_set: SideSet,
+}
+
 impl<'d, PIO: Instance> Common<'d, PIO> {
-    pub fn write_instr<I>(&mut self, start: usize, instrs: I) -> InstanceMemory<'d, PIO>
+    pub fn load_program<const SIZE: usize>(&mut self, prog: &RelocatedProgram<SIZE>) -> LoadedProgram<'d, PIO> {
+        match self.try_load_program(prog) {
+            Ok(r) => r,
+            Err(at) => panic!("Trying to write already used PIO instruction memory at {}", at),
+        }
+    }
+
+    pub fn try_load_program<const SIZE: usize>(
+        &mut self,
+        prog: &RelocatedProgram<SIZE>,
+    ) -> Result<LoadedProgram<'d, PIO>, usize> {
+        let used_memory = self.try_write_instr(prog.origin() as _, prog.code())?;
+        Ok(LoadedProgram {
+            used_memory,
+            origin: prog.origin(),
+            wrap: prog.wrap(),
+            side_set: prog.side_set(),
+        })
+    }
+
+    pub fn try_write_instr<I>(&mut self, start: usize, instrs: I) -> Result<InstanceMemory<'d, PIO>, usize>
     where
         I: Iterator<Item = u16>,
     {
@@ -808,11 +805,9 @@ impl<'d, PIO: Instance> Common<'d, PIO> {
         for (i, instr) in instrs.enumerate() {
             let addr = (i + start) as u8;
             let mask = 1 << (addr as usize);
-            assert!(
-                self.instructions_used & mask == 0,
-                "Trying to write already used PIO instruction memory at {}",
-                addr
-            );
+            if self.instructions_used & mask != 0 {
+                return Err(addr as usize);
+            }
             unsafe {
                 PIO::PIO.instr_mem(addr as usize).write(|w| {
                     w.set_instr_mem(instr);
@@ -821,15 +816,14 @@ impl<'d, PIO: Instance> Common<'d, PIO> {
             used_mask |= mask;
         }
         self.instructions_used |= used_mask;
-        InstanceMemory {
+        Ok(InstanceMemory {
             used_mask,
             pio: PhantomData,
-        }
+        })
     }
 
-    /// Free instruction memory previously allocated with [`Common::write_instr`].
-    /// This is always possible but unsafe if any state machine is still using this
-    /// bit of memory.
+    /// Free instruction memory. This is always possible but unsafe if any
+    /// state machine is still using this bit of memory.
     pub unsafe fn free_instr(&mut self, instrs: InstanceMemory<PIO>) {
         self.instructions_used &= !instrs.used_mask;
     }
diff --git a/examples/rp/src/bin/pio_async.rs b/examples/rp/src/bin/pio_async.rs
index 461ea3ff9..9f47c2316 100644
--- a/examples/rp/src/bin/pio_async.rs
+++ b/examples/rp/src/bin/pio_async.rs
@@ -5,11 +5,10 @@ use defmt::info;
 use embassy_executor::Spawner;
 use embassy_rp::peripherals::PIO0;
 use embassy_rp::pio::{Common, Irq, Pio, PioPin, ShiftDirection, StateMachine};
-use embassy_rp::pio_instr_util;
 use embassy_rp::relocate::RelocatedProgram;
 use {defmt_rtt as _, panic_probe as _};
 
-fn setup_pio_task_sm0(pio: &mut Common<PIO0>, sm: &mut StateMachine<PIO0, 0>, pin: impl PioPin) {
+fn setup_pio_task_sm0<'a>(pio: &mut Common<'a, PIO0>, sm: &mut StateMachine<'a, PIO0, 0>, pin: impl PioPin) {
     // Setup sm0
 
     // Send data serially to pin
@@ -22,15 +21,12 @@ fn setup_pio_task_sm0(pio: &mut Common<PIO0>, sm: &mut StateMachine<PIO0, 0>, pi
     );
 
     let relocated = RelocatedProgram::new(&prg.program);
+    sm.use_program(&pio.load_program(&relocated), &[]);
     let out_pin = pio.make_pio_pin(pin);
     let pio_pins = [&out_pin];
     sm.set_out_pins(&pio_pins);
-    pio.write_instr(relocated.origin() as usize, relocated.code());
-    pio_instr_util::exec_jmp(sm, relocated.origin());
     sm.set_clkdiv((125e6 / 20.0 / 2e2 * 256.0) as u32);
     sm.set_set_range(0, 1);
-    let pio::Wrap { source, target } = relocated.wrap();
-    sm.set_wrap(source, target);
 
     sm.set_autopull(true);
     sm.set_out_shift_dir(ShiftDirection::Left);
@@ -48,20 +44,16 @@ async fn pio_task_sm0(mut sm: StateMachine<'static, PIO0, 0>) {
     }
 }
 
-fn setup_pio_task_sm1(pio: &mut Common<PIO0>, sm: &mut StateMachine<PIO0, 1>) {
+fn setup_pio_task_sm1<'a>(pio: &mut Common<'a, PIO0>, sm: &mut StateMachine<'a, PIO0, 1>) {
     // Setupm sm1
 
     // Read 0b10101 repeatedly until ISR is full
     let prg = pio_proc::pio_asm!(".origin 8", "set x, 0x15", ".wrap_target", "in x, 5 [31]", ".wrap",);
 
     let relocated = RelocatedProgram::new(&prg.program);
-    pio.write_instr(relocated.origin() as usize, relocated.code());
-    pio_instr_util::exec_jmp(sm, relocated.origin());
+    sm.use_program(&pio.load_program(&relocated), &[]);
     sm.set_clkdiv((125e6 / 2e3 * 256.0) as u32);
     sm.set_set_range(0, 0);
-    let pio::Wrap { source, target } = relocated.wrap();
-    sm.set_wrap(source, target);
-
     sm.set_autopush(true);
     sm.set_in_shift_dir(ShiftDirection::Right);
 }
@@ -75,7 +67,7 @@ async fn pio_task_sm1(mut sm: StateMachine<'static, PIO0, 1>) {
     }
 }
 
-fn setup_pio_task_sm2(pio: &mut Common<PIO0>, sm: &mut StateMachine<PIO0, 2>) {
+fn setup_pio_task_sm2<'a>(pio: &mut Common<'a, PIO0>, sm: &mut StateMachine<'a, PIO0, 2>) {
     // Setup sm2
 
     // Repeatedly trigger IRQ 3
@@ -89,12 +81,7 @@ fn setup_pio_task_sm2(pio: &mut Common<PIO0>, sm: &mut StateMachine<PIO0, 2>) {
         ".wrap",
     );
     let relocated = RelocatedProgram::new(&prg.program);
-    pio.write_instr(relocated.origin() as usize, relocated.code());
-
-    let pio::Wrap { source, target } = relocated.wrap();
-    sm.set_wrap(source, target);
-
-    pio_instr_util::exec_jmp(sm, relocated.origin());
+    sm.use_program(&pio.load_program(&relocated), &[]);
     sm.set_clkdiv((125e6 / 2e3 * 256.0) as u32);
 }
 
diff --git a/examples/rp/src/bin/pio_dma.rs b/examples/rp/src/bin/pio_dma.rs
index c664482e5..1c4e127c7 100644
--- a/examples/rp/src/bin/pio_dma.rs
+++ b/examples/rp/src/bin/pio_dma.rs
@@ -6,7 +6,7 @@ use embassy_executor::Spawner;
 use embassy_futures::join::join;
 use embassy_rp::pio::{Pio, ShiftDirection};
 use embassy_rp::relocate::RelocatedProgram;
-use embassy_rp::{pio_instr_util, Peripheral};
+use embassy_rp::Peripheral;
 use {defmt_rtt as _, panic_probe as _};
 
 fn swap_nibbles(v: u32) -> u32 {
@@ -38,11 +38,8 @@ async fn main(_spawner: Spawner) {
     );
 
     let relocated = RelocatedProgram::new(&prg.program);
-    common.write_instr(relocated.origin() as usize, relocated.code());
-    pio_instr_util::exec_jmp(&mut sm, relocated.origin());
+    sm.use_program(&common.load_program(&relocated), &[]);
     sm.set_clkdiv((125e6 / 10e3 * 256.0) as u32);
-    let pio::Wrap { source, target } = relocated.wrap();
-    sm.set_wrap(source, target);
     sm.set_autopull(true);
     sm.set_autopush(true);
     sm.set_pull_threshold(32);
diff --git a/examples/rp/src/bin/pio_hd44780.rs b/examples/rp/src/bin/pio_hd44780.rs
index 17b2440cf..c3466b554 100644
--- a/examples/rp/src/bin/pio_hd44780.rs
+++ b/examples/rp/src/bin/pio_hd44780.rs
@@ -123,15 +123,9 @@ impl<'l> HD44780<'l> {
         embassy_rp::pio_instr_util::set_pindir(&mut sm0, 0b11111);
 
         let relocated = RelocatedProgram::new(&prg.program);
-        common.write_instr(relocated.origin() as usize, relocated.code());
-        embassy_rp::pio_instr_util::exec_jmp(&mut sm0, relocated.origin());
+        sm0.use_program(&common.load_program(&relocated), &[&e]);
         sm0.set_clkdiv(125 * 256);
-        let pio::Wrap { source, target } = relocated.wrap();
-        sm0.set_wrap(source, target);
-        sm0.set_side_enable(true);
         sm0.set_out_pins(&[&db4, &db5, &db6, &db7]);
-        sm0.set_sideset_base_pin(&e);
-        sm0.set_sideset_count(2);
         sm0.set_out_shift_dir(ShiftDirection::Left);
         sm0.set_fifo_join(FifoJoin::TxOnly);
         sm0.set_autopull(true);
@@ -199,17 +193,11 @@ impl<'l> HD44780<'l> {
         );
 
         let relocated = RelocatedProgram::new(&prg.program);
-        common.write_instr(relocated.origin() as usize, relocated.code());
-        embassy_rp::pio_instr_util::exec_jmp(&mut sm0, relocated.origin());
-        let pio::Wrap { source, target } = relocated.wrap();
+        sm0.use_program(&common.load_program(&relocated), &[&e]);
         sm0.set_clkdiv(8 * 256); // ~64ns/insn
-        sm0.set_side_enable(false);
         sm0.set_jmp_pin(db7pin);
-        sm0.set_wrap(source, target);
         sm0.set_set_pins(&[&rs, &rw]);
         sm0.set_out_pins(&[&db4, &db5, &db6, &db7]);
-        sm0.set_sideset_base_pin(&e);
-        sm0.set_sideset_count(1);
         sm0.set_out_shift_dir(ShiftDirection::Left);
         sm0.set_fifo_join(FifoJoin::TxOnly);
 
diff --git a/examples/rp/src/bin/ws2812-pio.rs b/examples/rp/src/bin/ws2812-pio.rs
index 2e6860d8b..889970541 100644
--- a/examples/rp/src/bin/ws2812-pio.rs
+++ b/examples/rp/src/bin/ws2812-pio.rs
@@ -5,7 +5,6 @@
 use defmt::*;
 use embassy_executor::Spawner;
 use embassy_rp::pio::{Common, FifoJoin, Instance, Pio, PioPin, ShiftDirection, StateMachine};
-use embassy_rp::pio_instr_util;
 use embassy_rp::relocate::RelocatedProgram;
 use embassy_time::{Duration, Timer};
 use smart_leds::RGB8;
@@ -45,15 +44,11 @@ impl<'d, P: Instance, const S: usize> Ws2812<'d, P, S> {
 
         let prg = a.assemble_with_wrap(wrap_source, wrap_target);
 
-        let relocated = RelocatedProgram::new(&prg);
-        pio.write_instr(relocated.origin() as usize, relocated.code());
-        pio_instr_util::exec_jmp(&mut sm, relocated.origin());
-
         // Pin config
         let out_pin = pio.make_pio_pin(pin);
-        sm.set_set_pins(&[&out_pin]);
-        sm.set_sideset_base_pin(&out_pin);
-        sm.set_sideset_count(1);
+
+        let relocated = RelocatedProgram::new(&prg);
+        sm.use_program(&pio.load_program(&relocated), &[&out_pin]);
 
         // Clock config
         // TODO CLOCK_FREQ should come from embassy_rp
@@ -70,8 +65,6 @@ impl<'d, P: Instance, const S: usize> Ws2812<'d, P, S> {
         }
 
         sm.set_clkdiv((int << 8) | frac);
-        let pio::Wrap { source, target } = relocated.wrap();
-        sm.set_wrap(source, target);
 
         // FIFO config
         sm.set_autopull(true);