3 files changed, 23 insertions, 15 deletions

diff --git a/embassy-extras/src/peripheral.rs b/embassy-extras/src/peripheral.rs
index aa23bc978..e2435d63f 100644
--- a/embassy-extras/src/peripheral.rs
+++ b/embassy-extras/src/peripheral.rs
@@ -32,32 +32,33 @@ impl<S: PeripheralState> PeripheralMutex<S> {
         }
     }
 
-    /// safety: self must be pinned.
-    unsafe fn setup(&mut self) {
-        self.irq.disable();
+    pub fn register_interrupt(self: Pin<&mut Self>) {
+        let this = unsafe { self.get_unchecked_mut() };
+        if this.irq_setup_done {
+            return;
+        }
+
+        this.irq.disable();
         compiler_fence(Ordering::SeqCst);
 
-        self.irq.set_handler(|p| {
+        this.irq.set_handler(|p| {
             // Safety: it's OK to get a &mut to the state, since
             // - We're in the IRQ, no one else can't preempt us
             // - We can't have preempted a with() call because the irq is disabled during it.
-            let state = &mut *(p as *mut S);
+            let state = unsafe { &mut *(p as *mut S) };
             state.on_interrupt();
         });
-        self.irq
-            .set_handler_context((&mut self.state) as *mut _ as *mut ());
+        this.irq
+            .set_handler_context((&mut this.state) as *mut _ as *mut ());
 
         compiler_fence(Ordering::SeqCst);
-        self.irq.enable();
+        this.irq.enable();
 
-        self.irq_setup_done = true;
+        this.irq_setup_done = true;
     }
 
     pub fn with<R>(self: Pin<&mut Self>, f: impl FnOnce(&mut S, &mut S::Interrupt) -> R) -> R {
         let this = unsafe { self.get_unchecked_mut() };
-        if !this.irq_setup_done {
-            unsafe { this.setup() }
-        }
 
         this.irq.disable();
         compiler_fence(Ordering::SeqCst);
diff --git a/embassy-nrf/src/buffered_uarte.rs b/embassy-nrf/src/buffered_uarte.rs
index 9b4ec6061..6cc5f1322 100644
--- a/embassy-nrf/src/buffered_uarte.rs
+++ b/embassy-nrf/src/buffered_uarte.rs
@@ -207,7 +207,9 @@ impl<'a, U: Instance, T: TimerInstance, P1: ConfigurablePpi, P2: ConfigurablePpi
     for BufferedUarte<'a, U, T, P1, P2>
 {
     fn poll_fill_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<&[u8]>> {
-        self.inner().with(|state, _irq| {
+        let mut inner = self.inner();
+        inner.as_mut().register_interrupt();
+        inner.with(|state, _irq| {
             // Conservative compiler fence to prevent optimizations that do not
             // take in to account actions by DMA. The fence has been placed here,
             // before any DMA action has started
@@ -230,7 +232,9 @@ impl<'a, U: Instance, T: TimerInstance, P1: ConfigurablePpi, P2: ConfigurablePpi
     }
 
     fn consume(self: Pin<&mut Self>, amt: usize) {
-        self.inner().with(|state, irq| {
+        let mut inner = self.inner();
+        inner.as_mut().register_interrupt();
+        inner.with(|state, irq| {
             trace!("consume {:?}", amt);
             state.rx.pop(amt);
             irq.pend();
@@ -242,7 +246,9 @@ impl<'a, U: Instance, T: TimerInstance, P1: ConfigurablePpi, P2: ConfigurablePpi
     for BufferedUarte<'a, U, T, P1, P2>
 {
     fn poll_write(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8]) -> Poll<Result<usize>> {
-        self.inner().with(|state, irq| {
+        let mut inner = self.inner();
+        inner.as_mut().register_interrupt();
+        inner.with(|state, irq| {
             trace!("poll_write: {:?}", buf.len());
 
             let tx_buf = state.tx.push_buf();
diff --git a/embassy-nrf/src/spim.rs b/embassy-nrf/src/spim.rs
index 4a3adeb3d..ef89afbc6 100644
--- a/embassy-nrf/src/spim.rs
+++ b/embassy-nrf/src/spim.rs
@@ -148,6 +148,7 @@ impl<T: Instance> FullDuplex<u8> for Spim<T> {
             slice_in_ram_or(rx, Error::DMABufferNotInDataMemory)?;
             slice_in_ram_or(tx, Error::DMABufferNotInDataMemory)?;
 
+            self.as_mut().inner().register_interrupt();
             self.as_mut().inner().with(|s, _irq| {
                 // Conservative compiler fence to prevent optimizations that do not
                 // take in to account actions by DMA. The fence has been placed here,