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);