2 files changed, 52 insertions, 108 deletions

diff --git a/embassy-nrf/src/buffered_uarte.rs b/embassy-nrf/src/buffered_uarte.rs
index c67b6f166..dbdb6f514 100644
--- a/embassy-nrf/src/buffered_uarte.rs
+++ b/embassy-nrf/src/buffered_uarte.rs
@@ -5,7 +5,6 @@
 //! - nrf52832: Section 35
 //! - nrf52840: Section 6.34
 use core::cmp::min;
-use core::marker::PhantomData;
 use core::mem;
 use core::ops::Deref;
 use core::pin::Pin;
@@ -20,7 +19,7 @@ use crate::hal::gpio::Port as GpioPort;
 use crate::interrupt::{self, OwnedInterrupt};
 use crate::pac;
 use crate::pac::uarte0;
-use crate::util::peripheral;
+use crate::util::peripheral::{PeripheralMutex, PeripheralState};
 use crate::util::ring_buffer::RingBuffer;
 
 // Re-export SVD variants to allow user to directly set values
@@ -49,8 +48,7 @@ enum TxState {
 ///     - nrf52832: Section 15.2
 ///     - nrf52840: Section 6.1.2
 pub struct BufferedUarte<'a, T: Instance> {
-    reg: peripheral::Registration<State<'a, T>>,
-    wtf: PhantomData<&'a ()>,
+    inner: PeripheralMutex<T::Interrupt, State<'a, T>>,
 }
 
 impl<'a, T: Instance> Unpin for BufferedUarte<'a, T> {}
@@ -126,10 +124,12 @@ impl<'a, T: Instance> BufferedUarte<'a, T> {
         // Configure frequency
         uarte.baudrate.write(|w| w.baudrate().variant(baudrate));
 
+        // Disable the irq, let the Registration enable it when everything is set up.
+        irq.disable();
         irq.pend();
 
         BufferedUarte {
-            reg: peripheral::Registration::new(
+            inner: PeripheralMutex::new(
                 irq,
                 State {
                     inner: uarte,
@@ -143,7 +143,6 @@ impl<'a, T: Instance> BufferedUarte<'a, T> {
                     tx_waker: WakerRegistration::new(),
                 },
             ),
-            wtf: PhantomData,
         }
     }
 }
@@ -158,7 +157,8 @@ impl<'a, T: Instance> Drop for BufferedUarte<'a, T> {
 impl<'a, T: Instance> AsyncBufRead for BufferedUarte<'a, T> {
     fn poll_fill_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<&[u8]>> {
         let this = unsafe { self.get_unchecked_mut() };
-        this.reg.with(|state, _| {
+        let reg = unsafe { Pin::new_unchecked(&mut this.inner) };
+        reg.with(|_irq, state| {
             let z: Poll<Result<&[u8]>> = state.poll_fill_buf(cx);
             let z: Poll<Result<&[u8]>> = unsafe { mem::transmute(z) };
             z
@@ -167,14 +167,16 @@ impl<'a, T: Instance> AsyncBufRead for BufferedUarte<'a, T> {
 
     fn consume(self: Pin<&mut Self>, amt: usize) {
         let this = unsafe { self.get_unchecked_mut() };
-        this.reg.with(|state, irq| state.consume(irq, amt))
+        let reg = unsafe { Pin::new_unchecked(&mut this.inner) };
+        reg.with(|irq, state| state.consume(irq, amt))
     }
 }
 
 impl<'a, T: Instance> AsyncWrite for BufferedUarte<'a, T> {
     fn poll_write(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8]) -> Poll<Result<usize>> {
         let this = unsafe { self.get_unchecked_mut() };
-        this.reg.with(|state, irq| state.poll_write(irq, cx, buf))
+        let reg = unsafe { Pin::new_unchecked(&mut this.inner) };
+        reg.with(|irq, state| state.poll_write(irq, cx, buf))
     }
 }
 
@@ -262,12 +264,7 @@ impl<'a, T: Instance> State<'a, T> {
     }
 }
 
-impl<'a, T: Instance> peripheral::State for State<'a, T> {
-    type Interrupt = T::Interrupt;
-    fn store<'b>() -> &'b peripheral::Store<Self> {
-        unsafe { mem::transmute(T::storage()) }
-    }
-
+impl<'a, T: Instance> PeripheralState for State<'a, T> {
     fn on_interrupt(&mut self) {
         trace!("irq: start");
         let mut more_work = true;
@@ -414,28 +411,15 @@ mod private {
     impl Sealed for crate::pac::UARTE1 {}
 }
 
-pub trait Instance:
-    Deref<Target = uarte0::RegisterBlock> + Sized + private::Sealed + 'static
-{
+pub trait Instance: Deref<Target = uarte0::RegisterBlock> + private::Sealed {
     type Interrupt: OwnedInterrupt;
-    fn storage() -> &'static peripheral::Store<State<'static, Self>>;
 }
 
 impl Instance for pac::UARTE0 {
     type Interrupt = interrupt::UARTE0_UART0Interrupt;
-    fn storage() -> &'static peripheral::Store<State<'static, Self>> {
-        static STORAGE: peripheral::Store<State<'static, crate::pac::UARTE0>> =
-            peripheral::Store::uninit();
-        &STORAGE
-    }
 }
 
 #[cfg(any(feature = "52833", feature = "52840", feature = "9160"))]
 impl Instance for pac::UARTE1 {
     type Interrupt = interrupt::UARTE1Interrupt;
-    fn storage() -> &'static peripheral::Store<State<'static, Self>> {
-        static STORAGE: peripheral::Store<State<'static, crate::pac::UARTE1>> =
-            peripheral::Store::uninit();
-        &STORAGE
-    }
 }
diff --git a/embassy-nrf/src/util/peripheral.rs b/embassy-nrf/src/util/peripheral.rs
index 9b3384e5d..07dc4a7bc 100644
--- a/embassy-nrf/src/util/peripheral.rs
+++ b/embassy-nrf/src/util/peripheral.rs
@@ -1,110 +1,70 @@
-use core::mem;
-use core::mem::MaybeUninit;
-use core::ptr;
+use core::pin::Pin;
 use core::sync::atomic::{compiler_fence, Ordering};
 use core::{cell::UnsafeCell, marker::PhantomData};
 
+use crate::fmt::*;
 use crate::interrupt::OwnedInterrupt;
 
-pub struct Store<T>(MaybeUninit<UnsafeCell<T>>);
-impl<T> Store<T> {
-    pub const fn uninit() -> Self {
-        Self(MaybeUninit::uninit())
-    }
-
-    unsafe fn as_mut_ptr(&self) -> *mut T {
-        (*self.0.as_ptr()).get()
-    }
-
-    unsafe fn as_mut(&self) -> &mut T {
-        &mut *self.as_mut_ptr()
-    }
-
-    unsafe fn write(&self, val: T) {
-        ptr::write(self.as_mut_ptr(), val)
-    }
-
-    unsafe fn drop_in_place(&self) {
-        ptr::drop_in_place(self.as_mut_ptr())
-    }
-
-    unsafe fn read(&self) -> T {
-        ptr::read(self.as_mut_ptr())
-    }
-}
-unsafe impl<T> Send for Store<T> {}
-unsafe impl<T> Sync for Store<T> {}
-
-pub trait State: Sized {
-    type Interrupt: OwnedInterrupt;
+pub trait PeripheralState {
     fn on_interrupt(&mut self);
-    #[doc(hidden)]
-    fn store<'a>() -> &'a Store<Self>;
 }
 
-pub struct Registration<P: State> {
-    irq: P::Interrupt,
-    not_send: PhantomData<*mut P>,
+pub struct PeripheralMutex<I: OwnedInterrupt, S: PeripheralState> {
+    inner: Option<(I, UnsafeCell<S>)>,
+    not_send: PhantomData<*mut ()>,
 }
 
-impl<P: State> Registration<P> {
-    pub fn new(irq: P::Interrupt, state: P) -> Self {
-        // safety:
-        // - No other PeripheralRegistration can already exist because we have the owned interrupt
-        // - therefore, storage is uninitialized
-        // - therefore it's safe to overwrite it without dropping the previous contents
-        unsafe { P::store().write(state) }
-
-        irq.set_handler(
-            |_| {
-                // safety:
-                // - If a PeripheralRegistration instance exists, P::storage() is initialized.
-                // - It's OK to get a &mut to it since the irq is disabled.
-                unsafe { P::store().as_mut() }.on_interrupt();
-            },
-            core::ptr::null_mut(),
-        );
-
-        compiler_fence(Ordering::SeqCst);
-        irq.enable();
-
+impl<I: OwnedInterrupt, S: PeripheralState> PeripheralMutex<I, S> {
+    pub fn new(irq: I, state: S) -> Self {
         Self {
-            irq,
+            inner: Some((irq, UnsafeCell::new(state))),
             not_send: PhantomData,
         }
     }
 
-    pub fn with<R>(&mut self, f: impl FnOnce(&mut P, &mut P::Interrupt) -> R) -> R {
-        self.irq.disable();
+    pub fn with<R>(self: Pin<&mut Self>, f: impl FnOnce(&mut I, &mut S) -> R) -> R {
+        let this = unsafe { self.get_unchecked_mut() };
+        let (irq, state) = unwrap!(this.inner.as_mut());
+
+        irq.disable();
         compiler_fence(Ordering::SeqCst);
 
-        // safety:
-        // - If a PeripheralRegistration instance exists, P::storage() is initialized.
-        // - It's OK to get a &mut to it since the irq is disabled.
-        let r = f(unsafe { P::store().as_mut() }, &mut self.irq);
+        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 = unsafe { &mut *(p as *mut S) };
+                state.on_interrupt();
+            },
+            state.get() as *mut (),
+        );
+
+        // Safety: it's OK to get a &mut to the state, since the irq is disabled.
+        let state = unsafe { &mut *state.get() };
+
+        let r = f(irq, state);
 
         compiler_fence(Ordering::SeqCst);
-        self.irq.enable();
+        irq.enable();
 
         r
     }
 
-    pub fn free(self) -> (P::Interrupt, P) {
-        let irq = unsafe { ptr::read(&self.irq) };
+    pub fn free(self: Pin<&mut Self>) -> (I, S) {
+        let this = unsafe { self.get_unchecked_mut() };
+        let (irq, state) = unwrap!(this.inner.take());
         irq.disable();
         irq.remove_handler();
-        mem::forget(self);
-        let storage = P::store();
-        (irq, unsafe { storage.read() })
+        (irq, state.into_inner())
     }
 }
 
-impl<P: State> Drop for Registration<P> {
+impl<I: OwnedInterrupt, S: PeripheralState> Drop for PeripheralMutex<I, S> {
     fn drop(&mut self) {
-        self.irq.disable();
-        self.irq.remove_handler();
-
-        let storage = P::store();
-        unsafe { storage.drop_in_place() };
+        if let Some((irq, state)) = &mut self.inner {
+            irq.disable();
+            irq.remove_handler();
+        }
     }
 }