3 files changed, 34 insertions, 36 deletions

diff --git a/embassy-nrf/src/gpiote.rs b/embassy-nrf/src/gpiote.rs
index 2b1ca40b1..9626d3299 100644
--- a/embassy-nrf/src/gpiote.rs
+++ b/embassy-nrf/src/gpiote.rs
@@ -1,6 +1,8 @@
 use anyfmt::{panic, *};
 use core::cell::Cell;
+use core::future::Future;
 use core::ptr;
+use core::task::{Context, Poll};
 use embassy::util::Signal;
 
 use crate::hal::gpio::{Input, Level, Output, Pin, Port};
@@ -201,10 +203,18 @@ impl Gpiote {
         })
     }
 
-    pub fn new_port_input<'a, T>(&'a self, pin: Pin<Input<T>>) -> PortInput<'a, T> {
+    pub fn wait_port_input<'a, T>(
+        &'a self,
+        pin: &'a Pin<Input<T>>,
+        polarity: PortInputPolarity,
+    ) -> PortInputFuture<'a, T> {
         interrupt::free(|_| {
             unsafe { INSTANCE = self };
-            PortInput { gpiote: self, pin }
+            PortInputFuture {
+                gpiote: self,
+                pin,
+                polarity,
+            }
         })
     }
 
@@ -285,29 +295,28 @@ impl Gpiote {
     }
 }
 
-pub struct PortInput<'a, T> {
+pub struct PortInputFuture<'a, T> {
     gpiote: &'a Gpiote,
-    pin: Pin<Input<T>>,
+    pin: &'a Pin<Input<T>>,
+    polarity: PortInputPolarity,
 }
 
-impl<'a, T> Drop for PortInput<'a, T> {
+impl<'a, T> Drop for PortInputFuture<'a, T> {
     fn drop(&mut self) {
         pin_conf(&self.pin).modify(|_, w| w.sense().disabled());
         self.gpiote.port_signals[pin_num(&self.pin)].reset();
     }
 }
 
-impl<'a, T> PortInput<'a, T> {
-    pub async fn wait(&self, polarity: PortInputPolarity) {
-        pin_conf(&self.pin).modify(|_, w| match polarity {
+impl<'a, T> Future for PortInputFuture<'a, T> {
+    type Output = ();
+
+    fn poll(self: core::pin::Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        pin_conf(&self.pin).modify(|_, w| match self.polarity {
             PortInputPolarity::Low => w.sense().low(),
             PortInputPolarity::High => w.sense().high(),
         });
-        self.gpiote.port_signals[pin_num(&self.pin)].wait().await;
-    }
-
-    pub fn pin(&self) -> &Pin<Input<T>> {
-        &self.pin
+        self.gpiote.port_signals[pin_num(&self.pin)].poll_wait(cx)
     }
 }
 
diff --git a/embassy/src/util/signal.rs b/embassy/src/util/signal.rs
index f28166d3c..e2cd3a80f 100644
--- a/embassy/src/util/signal.rs
+++ b/embassy/src/util/signal.rs
@@ -2,7 +2,6 @@ use anyfmt::panic;
 use core::cell::UnsafeCell;
 use core::future::Future;
 use core::mem;
-use core::pin::Pin;
 use core::task::{Context, Poll, Waker};
 
 pub struct Signal<T> {
@@ -42,21 +41,9 @@ impl<T: Send> Signal<T> {
         })
     }
 
-    pub fn wait<'a>(&'a self) -> impl Future<Output = T> + 'a {
-        WaitFuture { signal: self }
-    }
-}
-
-struct WaitFuture<'a, T> {
-    signal: &'a Signal<T>,
-}
-
-impl<'a, T: Send> Future for WaitFuture<'a, T> {
-    type Output = T;
-
-    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<T> {
+    pub fn poll_wait(&self, cx: &mut Context<'_>) -> Poll<T> {
         cortex_m::interrupt::free(|_| unsafe {
-            let state = &mut *self.signal.state.get();
+            let state = &mut *self.state.get();
             match state {
                 State::None => {
                     *state = State::Waiting(cx.waker().clone());
@@ -71,4 +58,8 @@ impl<'a, T: Send> Future for WaitFuture<'a, T> {
             }
         })
     }
+
+    pub fn wait(&self) -> impl Future<Output = T> + '_ {
+        futures::future::poll_fn(move |cx| self.poll_wait(cx))
+    }
 }
diff --git a/examples/src/bin/gpiote_port.rs b/examples/src/bin/gpiote_port.rs
index bc0cb4367..585317642 100644
--- a/examples/src/bin/gpiote_port.rs
+++ b/examples/src/bin/gpiote_port.rs
@@ -12,15 +12,13 @@ use nrf52840_hal::gpio;
 
 use embassy::executor::{task, Executor};
 use embassy::util::Forever;
-use embassy_nrf::gpiote;
-
-async fn button(g: &gpiote::Gpiote, n: usize, pin: gpio::Pin<gpio::Input<gpio::PullUp>>) {
-    let ch = g.new_port_input(pin);
+use embassy_nrf::gpiote::{Gpiote, PortInputPolarity};
 
+async fn button(g: &Gpiote, n: usize, pin: gpio::Pin<gpio::Input<gpio::PullUp>>) {
     loop {
-        ch.wait(gpiote::PortInputPolarity::Low).await;
+        g.wait_port_input(&pin, PortInputPolarity::Low).await;
         info!("Button {:?} pressed!", n);
-        ch.wait(gpiote::PortInputPolarity::High).await;
+        g.wait_port_input(&pin, PortInputPolarity::High).await;
         info!("Button {:?} released!", n);
     }
 }
@@ -30,12 +28,12 @@ async fn run() {
     let p = unwrap!(embassy_nrf::pac::Peripherals::take());
     let port0 = gpio::p0::Parts::new(p.P0);
 
-    let g = gpiote::Gpiote::new(p.GPIOTE);
+    let g = Gpiote::new(p.GPIOTE);
     info!(
         "sizeof Signal<()> = {:usize}",
         mem::size_of::<embassy::util::Signal<()>>()
     );
-    info!("sizeof gpiote = {:usize}", mem::size_of::<gpiote::Gpiote>());
+    info!("sizeof gpiote = {:usize}", mem::size_of::<Gpiote>());
 
     info!("Starting!");