3 files changed, 37 insertions, 24 deletions

diff --git a/embassy-rp/src/pio.rs b/embassy-rp/src/pio.rs
index 52ef89a16..769641ac9 100644
--- a/embassy-rp/src/pio.rs
+++ b/embassy-rp/src/pio.rs
@@ -180,21 +180,12 @@ impl<'a, 'd, PIO: PioInstance, const SM: usize> Drop for FifoInFuture<'a, 'd, PI
 
 /// Future that waits for IRQ
 #[must_use = "futures do nothing unless you `.await` or poll them"]
-pub struct IrqFuture<PIO: PioInstance> {
-    pio: PhantomData<PIO>,
+pub struct IrqFuture<'a, 'd, PIO: PioInstance> {
+    pio: PhantomData<&'a PioIrq<'d, PIO, 0>>,
     irq_no: u8,
 }
 
-impl<'a, PIO: PioInstance> IrqFuture<PIO> {
-    pub fn new(irq_no: u8) -> Self {
-        IrqFuture {
-            pio: PhantomData::default(),
-            irq_no,
-        }
-    }
-}
-
-impl<'d, PIO: PioInstance> Future for IrqFuture<PIO> {
+impl<'a, 'd, PIO: PioInstance> Future for IrqFuture<'a, 'd, PIO> {
     type Output = ();
     fn poll(self: FuturePin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
         //debug!("Poll {},{}", PIO::PIO_NO, SM);
@@ -224,7 +215,7 @@ impl<'d, PIO: PioInstance> Future for IrqFuture<PIO> {
     }
 }
 
-impl<'d, PIO: PioInstance> Drop for IrqFuture<PIO> {
+impl<'a, 'd, PIO: PioInstance> Drop for IrqFuture<'a, 'd, PIO> {
     fn drop(&mut self) {
         unsafe {
             PIO::PIO.irqs(0).inte().write_clear(|m| {
@@ -688,10 +679,6 @@ impl<'d, PIO: PioInstance + 'd, const SM: usize> PioStateMachine<'d, PIO, SM> {
         FifoInFuture::new(self)
     }
 
-    pub fn wait_irq(&self, irq_no: u8) -> IrqFuture<PIO> {
-        IrqFuture::new(irq_no)
-    }
-
     pub fn has_tx_stalled(&self) -> bool {
         unsafe {
             let fdebug = PIO::PIO.fdebug();
@@ -862,7 +849,8 @@ impl<'d, PIO: PioInstance> PioCommon<'d, PIO> {
 
     /// Register a pin for PIO usage. Pins will be released from the PIO block
     /// (i.e., have their `FUNCSEL` reset to `NULL`) when the [`PioCommon`] *and*
-    /// all [`PioStateMachine`]s for this block have been dropped.
+    /// all [`PioStateMachine`]s for this block have been dropped. **Other members
+    /// of [`Pio`] do not keep pin registrations alive.**
     pub fn make_pio_pin(&mut self, pin: impl Peripheral<P = impl PioPin + 'd> + 'd) -> Pin<'d, PIO> {
         into_ref!(pin);
         unsafe {
@@ -877,8 +865,25 @@ impl<'d, PIO: PioInstance> PioCommon<'d, PIO> {
     }
 }
 
+pub struct PioIrq<'d, PIO: PioInstance, const N: usize> {
+    pio: PhantomData<&'d PIO>,
+}
+
+impl<'d, PIO: PioInstance, const N: usize> PioIrq<'d, PIO, N> {
+    pub fn wait<'a>(&'a mut self) -> IrqFuture<'a, 'd, PIO> {
+        IrqFuture {
+            pio: PhantomData,
+            irq_no: N as u8,
+        }
+    }
+}
+
 pub struct Pio<'d, PIO: PioInstance> {
     pub common: PioCommon<'d, PIO>,
+    pub irq0: PioIrq<'d, PIO, 0>,
+    pub irq1: PioIrq<'d, PIO, 1>,
+    pub irq2: PioIrq<'d, PIO, 2>,
+    pub irq3: PioIrq<'d, PIO, 3>,
     pub sm0: PioStateMachine<'d, PIO, 0>,
     pub sm1: PioStateMachine<'d, PIO, 1>,
     pub sm2: PioStateMachine<'d, PIO, 2>,
@@ -894,6 +899,10 @@ impl<'d, PIO: PioInstance> Pio<'d, PIO> {
                 instructions_used: 0,
                 pio: PhantomData,
             },
+            irq0: PioIrq { pio: PhantomData },
+            irq1: PioIrq { pio: PhantomData },
+            irq2: PioIrq { pio: PhantomData },
+            irq3: PioIrq { pio: PhantomData },
             sm0: PioStateMachine { pio: PhantomData },
             sm1: PioStateMachine { pio: PhantomData },
             sm2: PioStateMachine { pio: PhantomData },
diff --git a/examples/rp/src/bin/pio_async.rs b/examples/rp/src/bin/pio_async.rs
index 11b290869..3d76a7d7b 100644
--- a/examples/rp/src/bin/pio_async.rs
+++ b/examples/rp/src/bin/pio_async.rs
@@ -4,7 +4,7 @@
 use defmt::info;
 use embassy_executor::Spawner;
 use embassy_rp::peripherals::PIO0;
-use embassy_rp::pio::{Pio, PioCommon, PioPin, PioStateMachine, ShiftDirection};
+use embassy_rp::pio::{Pio, PioCommon, PioIrq, PioPin, PioStateMachine, ShiftDirection};
 use embassy_rp::pio_instr_util;
 use embassy_rp::relocate::RelocatedProgram;
 use {defmt_rtt as _, panic_probe as _};
@@ -99,10 +99,10 @@ fn setup_pio_task_sm2(pio: &mut PioCommon<PIO0>, sm: &mut PioStateMachine<PIO0,
 }
 
 #[embassy_executor::task]
-async fn pio_task_sm2(mut sm: PioStateMachine<'static, PIO0, 2>) {
+async fn pio_task_sm2(mut irq: PioIrq<'static, PIO0, 3>, mut sm: PioStateMachine<'static, PIO0, 2>) {
     sm.set_enable(true);
     loop {
-        sm.wait_irq(3).await;
+        irq.wait().await;
         info!("IRQ trigged");
     }
 }
@@ -114,6 +114,7 @@ async fn main(spawner: Spawner) {
 
     let Pio {
         mut common,
+        irq3,
         mut sm0,
         mut sm1,
         mut sm2,
@@ -125,5 +126,5 @@ async fn main(spawner: Spawner) {
     setup_pio_task_sm2(&mut common, &mut sm2);
     spawner.spawn(pio_task_sm0(sm0)).unwrap();
     spawner.spawn(pio_task_sm1(sm1)).unwrap();
-    spawner.spawn(pio_task_sm2(sm2)).unwrap();
+    spawner.spawn(pio_task_sm2(irq3, sm2)).unwrap();
 }
diff --git a/examples/rp/src/bin/pio_hd44780.rs b/examples/rp/src/bin/pio_hd44780.rs
index bc51d43c4..7c1d7acfe 100644
--- a/examples/rp/src/bin/pio_hd44780.rs
+++ b/examples/rp/src/bin/pio_hd44780.rs
@@ -85,7 +85,10 @@ impl<'l> HD44780<'l> {
 
         let db7pin = db7.pin();
         let Pio {
-            mut common, mut sm0, ..
+            mut common,
+            mut irq0,
+            mut sm0,
+            ..
         } = Pio::new(pio);
 
         // takes command words (<wait:24> <command:4> <0:4>)
@@ -145,7 +148,7 @@ impl<'l> HD44780<'l> {
         sm0.push_tx((50 << 8) | 0x20);
         sm0.push_tx(0b1100_0000);
 
-        sm0.wait_irq(0).await;
+        irq0.wait().await;
         sm0.set_enable(false);
 
         // takes command sequences (<rs:1> <count:7>, data...)