4 files changed, 63 insertions, 79 deletions

diff --git a/embassy-rp/src/pio.rs b/embassy-rp/src/pio.rs
index 8492ca31d..7faec10b5 100644
--- a/embassy-rp/src/pio.rs
+++ b/embassy-rp/src/pio.rs
@@ -772,19 +772,6 @@ pub trait PioStateMachine: sealed::PioStateMachine + Sized + Unpin {
         }
     }
 
-    fn write_instr<I>(&mut self, start: usize, instrs: I)
-    where
-        I: Iterator<Item = u16>,
-    {
-        write_instr(
-            Self::Pio::PIO,
-            Self::Pio::PIO_NO,
-            start,
-            instrs,
-            MEM_USED_BY_STATEMACHINE | Self::Sm::SM_NO as u32,
-        );
-    }
-
     fn wait_push<'a>(&'a mut self, value: u32) -> FifoOutFuture<'a, Self::Pio, Self> {
         FifoOutFuture::new(self, value)
     }
@@ -880,71 +867,59 @@ pub trait PioStateMachine: sealed::PioStateMachine + Sized + Unpin {
     }
 }
 
-/*
-This is a bit array containing 4 bits for every word in the PIO instruction memory.
-*/
-// Bit 3-2
-//const MEM_USE_MASK: u32 = 0b1100;
-const MEM_NOT_USED: u32 = 0b0000;
-const MEM_USED_BY_STATEMACHINE: u32 = 0b0100;
-const MEM_USED_BY_COMMON: u32 = 0b1000;
-
-// Bit 1-0 is the number of the state machine
-//const MEM_STATE_MASK: u32 = 0b0011;
-
-// Should use mutex if running on multiple cores
-static mut INSTR_MEM_STATUS: &'static mut [[u32; 4]; 2] = &mut [[0; 4]; 2];
-
-fn instr_mem_get_status(pio_no: u8, addr: u8) -> u32 {
-    ((unsafe { INSTR_MEM_STATUS[pio_no as usize][(addr >> 3) as usize] }) >> ((addr & 0x07) * 4)) & 0xf
-}
-
-fn instr_mem_set_status(pio_no: u8, addr: u8, status: u32) {
-    let w = unsafe { &mut INSTR_MEM_STATUS[pio_no as usize][(addr >> 3) as usize] };
-    let shift = (addr & 0x07) * 4;
-    *w = (*w & !(0xf << shift)) | (status << shift);
-}
-
-fn instr_mem_is_free(pio_no: u8, addr: u8) -> bool {
-    instr_mem_get_status(pio_no, addr) == MEM_NOT_USED
+pub struct PioCommonInstance<PIO: PioInstance> {
+    instructions_used: u32,
+    pio: PhantomData<PIO>,
 }
 
-pub struct PioCommonInstance<PIO: PioInstance> {
+pub struct PioInstanceMemory<PIO: PioInstance> {
+    used_mask: u32,
     pio: PhantomData<PIO>,
 }
 
 impl<PIO: PioInstance> sealed::PioCommon for PioCommonInstance<PIO> {
     type Pio = PIO;
 }
-impl<PIO: PioInstance> PioCommon for PioCommonInstance<PIO> {}
-
-fn write_instr<I>(pio: &pac::pio::Pio, pio_no: u8, start: usize, instrs: I, mem_user: u32)
-where
-    I: Iterator<Item = u16>,
-{
-    for (i, instr) in instrs.enumerate() {
-        let addr = (i + start) as u8;
-        assert!(
-            instr_mem_is_free(pio_no, addr),
-            "Trying to write already used PIO instruction memory at {}",
-            addr
-        );
-        unsafe {
-            pio.instr_mem(addr as usize).write(|w| {
-                w.set_instr_mem(instr);
-            });
-            instr_mem_set_status(pio_no, addr, mem_user);
+impl<PIO: PioInstance> PioCommon for PioCommonInstance<PIO> {
+    fn write_instr<I>(&mut self, start: usize, instrs: I) -> PioInstanceMemory<Self::Pio>
+    where
+        I: Iterator<Item = u16>,
+    {
+        let mut used_mask = 0;
+        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
+            );
+            unsafe {
+                PIO::PIO.instr_mem(addr as usize).write(|w| {
+                    w.set_instr_mem(instr);
+                });
+            }
+            used_mask |= mask;
         }
+        self.instructions_used |= used_mask;
+        PioInstanceMemory {
+            used_mask,
+            pio: PhantomData,
+        }
+    }
+
+    fn free_instr(&mut self, instrs: PioInstanceMemory<Self::Pio>) {
+        self.instructions_used &= !instrs.used_mask;
     }
 }
 
 pub trait PioCommon: sealed::PioCommon + Sized {
-    fn write_instr<I>(&mut self, start: usize, instrs: I)
+    fn write_instr<I>(&mut self, start: usize, instrs: I) -> PioInstanceMemory<Self::Pio>
     where
-        I: Iterator<Item = u16>,
-    {
-        write_instr(Self::Pio::PIO, Self::Pio::PIO_NO, start, instrs, MEM_USED_BY_COMMON);
-    }
+        I: Iterator<Item = u16>;
+
+    // TODO make instruction memory that is currently in use unfreeable
+    fn free_instr(&mut self, instrs: PioInstanceMemory<Self::Pio>);
 
     fn is_irq_set(&self, irq_no: u8) -> bool {
         assert!(irq_no < 8);
@@ -1027,6 +1002,7 @@ pub trait PioPeripheral: sealed::PioPeripheral + Sized {
     ) {
         (
             PioCommonInstance {
+                instructions_used: 0,
                 pio: PhantomData::default(),
             },
             PioStateMachineInstance {
diff --git a/examples/rp/src/bin/pio_async.rs b/examples/rp/src/bin/pio_async.rs
index 3cfeec71f..1b075b8fd 100644
--- a/examples/rp/src/bin/pio_async.rs
+++ b/examples/rp/src/bin/pio_async.rs
@@ -28,7 +28,7 @@ fn setup_pio_task_sm0(pio: &mut PioCommonInstance<Pio0>, sm: &mut PioStateMachin
     let out_pin = pio.make_pio_pin(pin);
     let pio_pins = [&out_pin];
     sm.set_out_pins(&pio_pins);
-    sm.write_instr(relocated.origin() as usize, relocated.code());
+    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);
@@ -51,16 +51,15 @@ async fn pio_task_sm0(mut sm: PioStateMachineInstance<Pio0, Sm0>) {
     }
 }
 
-#[embassy_executor::task]
-async fn pio_task_sm1(mut sm: PioStateMachineInstance<Pio0, Sm1>) {
+fn setup_pio_task_sm1(pio: &mut PioCommonInstance<Pio0>, sm: &mut PioStateMachineInstance<Pio0, Sm1>) {
     // 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);
-    sm.write_instr(relocated.origin() as usize, relocated.code());
-    pio_instr_util::exec_jmp(&mut sm, relocated.origin());
+    pio.write_instr(relocated.origin() as usize, relocated.code());
+    pio_instr_util::exec_jmp(sm, relocated.origin());
     sm.set_clkdiv((125e6 / 2e3 * 256.0) as u32);
     sm.set_set_range(0, 0);
     let pio::Wrap { source, target } = relocated.wrap();
@@ -68,6 +67,10 @@ async fn pio_task_sm1(mut sm: PioStateMachineInstance<Pio0, Sm1>) {
 
     sm.set_autopush(true);
     sm.set_in_shift_dir(ShiftDirection::Right);
+}
+
+#[embassy_executor::task]
+async fn pio_task_sm1(mut sm: PioStateMachineInstance<Pio0, Sm1>) {
     sm.set_enable(true);
     loop {
         let rx = sm.wait_pull().await;
@@ -75,8 +78,7 @@ async fn pio_task_sm1(mut sm: PioStateMachineInstance<Pio0, Sm1>) {
     }
 }
 
-#[embassy_executor::task]
-async fn pio_task_sm2(mut sm: PioStateMachineInstance<Pio0, Sm2>) {
+fn setup_pio_task_sm2(pio: &mut PioCommonInstance<Pio0>, sm: &mut PioStateMachineInstance<Pio0, Sm2>) {
     // Setup sm2
 
     // Repeatedly trigger IRQ 3
@@ -90,13 +92,17 @@ async fn pio_task_sm2(mut sm: PioStateMachineInstance<Pio0, Sm2>) {
         ".wrap",
     );
     let relocated = RelocatedProgram::new(&prg.program);
-    sm.write_instr(relocated.origin() as usize, relocated.code());
+    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(&mut sm, relocated.origin());
+    pio_instr_util::exec_jmp(sm, relocated.origin());
     sm.set_clkdiv((125e6 / 2e3 * 256.0) as u32);
+}
+
+#[embassy_executor::task]
+async fn pio_task_sm2(mut sm: PioStateMachineInstance<Pio0, Sm2>) {
     sm.set_enable(true);
     loop {
         sm.wait_irq(3).await;
@@ -109,9 +115,11 @@ async fn main(spawner: Spawner) {
     let p = embassy_rp::init(Default::default());
     let pio = p.PIO0;
 
-    let (mut pio0, mut sm0, sm1, sm2, ..) = pio.split();
+    let (mut pio0, mut sm0, mut sm1, mut sm2, ..) = pio.split();
 
     setup_pio_task_sm0(&mut pio0, &mut sm0, p.PIN_0.degrade());
+    setup_pio_task_sm1(&mut pio0, &mut sm1);
+    setup_pio_task_sm2(&mut pio0, &mut sm2);
     spawner.spawn(pio_task_sm0(sm0)).unwrap();
     spawner.spawn(pio_task_sm1(sm1)).unwrap();
     spawner.spawn(pio_task_sm2(sm2)).unwrap();
diff --git a/examples/rp/src/bin/pio_dma.rs b/examples/rp/src/bin/pio_dma.rs
index 145e4a656..7d4919f75 100644
--- a/examples/rp/src/bin/pio_dma.rs
+++ b/examples/rp/src/bin/pio_dma.rs
@@ -4,7 +4,7 @@
 use defmt::info;
 use embassy_executor::Spawner;
 use embassy_futures::join::join;
-use embassy_rp::pio::{PioPeripheral, PioStateMachine, ShiftDirection};
+use embassy_rp::pio::{PioCommon, PioPeripheral, PioStateMachine, ShiftDirection};
 use embassy_rp::relocate::RelocatedProgram;
 use embassy_rp::{pio_instr_util, Peripheral};
 use {defmt_rtt as _, panic_probe as _};
@@ -19,7 +19,7 @@ fn swap_nibbles(v: u32) -> u32 {
 async fn main(_spawner: Spawner) {
     let p = embassy_rp::init(Default::default());
     let pio = p.PIO0;
-    let (_, mut sm, ..) = pio.split();
+    let (mut pio0, mut sm, ..) = pio.split();
 
     let prg = pio_proc::pio_asm!(
         ".origin 0",
@@ -34,7 +34,7 @@ async fn main(_spawner: Spawner) {
     );
 
     let relocated = RelocatedProgram::new(&prg.program);
-    sm.write_instr(relocated.origin() as usize, relocated.code());
+    pio0.write_instr(relocated.origin() as usize, relocated.code());
     pio_instr_util::exec_jmp(&mut sm, relocated.origin());
     sm.set_clkdiv((125e6 / 10e3 * 256.0) as u32);
     let pio::Wrap { source, target } = relocated.wrap();
diff --git a/examples/rp/src/bin/ws2812-pio.rs b/examples/rp/src/bin/ws2812-pio.rs
index 211c60c49..041e8ae11 100644
--- a/examples/rp/src/bin/ws2812-pio.rs
+++ b/examples/rp/src/bin/ws2812-pio.rs
@@ -19,7 +19,7 @@ pub struct Ws2812<P: PioInstance, S: SmInstance> {
 }
 
 impl<P: PioInstance, S: SmInstance> Ws2812<P, S> {
-    pub fn new(pio: PioCommonInstance<P>, mut sm: PioStateMachineInstance<P, S>, pin: gpio::AnyPin) -> Self {
+    pub fn new(mut pio: PioCommonInstance<P>, mut sm: PioStateMachineInstance<P, S>, pin: gpio::AnyPin) -> Self {
         // Setup sm0
 
         // prepare the PIO program
@@ -50,7 +50,7 @@ impl<P: PioInstance, S: SmInstance> Ws2812<P, S> {
         let prg = a.assemble_with_wrap(wrap_source, wrap_target);
 
         let relocated = RelocatedProgram::new(&prg);
-        sm.write_instr(relocated.origin() as usize, relocated.code());
+        pio.write_instr(relocated.origin() as usize, relocated.code());
         pio_instr_util::exec_jmp(&mut sm, relocated.origin());
 
         // Pin config