3 files changed, 338 insertions, 48 deletions

diff --git a/embassy-nrf/src/buffered_uarte.rs b/embassy-nrf/src/buffered_uarte.rs
index 517afa38c..a5a37b982 100644
--- a/embassy-nrf/src/buffered_uarte.rs
+++ b/embassy-nrf/src/buffered_uarte.rs
@@ -15,7 +15,9 @@ use crate::gpio::sealed::Pin as _;
 use crate::gpio::{OptionalPin as GpioOptionalPin, Pin as GpioPin};
 use crate::pac;
 use crate::ppi::{AnyConfigurableChannel, ConfigurableChannel, Event, Ppi, Task};
+use crate::timer::Frequency;
 use crate::timer::Instance as TimerInstance;
+use crate::timer::Timer;
 use crate::uarte::{Config, Instance as UarteInstance};
 
 // Re-export SVD variants to allow user to directly set values
@@ -35,7 +37,7 @@ enum TxState {
 
 struct State<'d, U: UarteInstance, T: TimerInstance> {
     phantom: PhantomData<&'d mut U>,
-    timer: T,
+    timer: Timer<'d, T>,
     _ppi_ch1: Ppi<'d, AnyConfigurableChannel>,
     _ppi_ch2: Ppi<'d, AnyConfigurableChannel>,
 
@@ -76,10 +78,11 @@ impl<'d, U: UarteInstance, T: TimerInstance> BufferedUarte<'d, U, T> {
         rx_buffer: &'d mut [u8],
         tx_buffer: &'d mut [u8],
     ) -> Self {
-        unborrow!(timer, ppi_ch1, ppi_ch2, irq, rxd, txd, cts, rts);
+        unborrow!(ppi_ch1, ppi_ch2, irq, rxd, txd, cts, rts);
 
         let r = U::regs();
-        let rt = timer.regs();
+
+        let mut timer = Timer::new_irqless(timer);
 
         rxd.conf().write(|w| w.input().connect().drive().h0h1());
         r.psel.rxd.write(|w| unsafe { w.bits(rxd.psel_bits()) });
@@ -133,25 +136,19 @@ impl<'d, U: UarteInstance, T: TimerInstance> BufferedUarte<'d, U, T> {
         // This gives us the amount of 16M ticks for 20 bits.
         let timeout = 0x8000_0000 / (config.baudrate as u32 / 40);
 
-        rt.tasks_stop.write(|w| unsafe { w.bits(1) });
-        rt.bitmode.write(|w| w.bitmode()._32bit());
-        rt.prescaler.write(|w| unsafe { w.prescaler().bits(0) });
-        rt.cc[0].write(|w| unsafe { w.bits(timeout) });
-        rt.mode.write(|w| w.mode().timer());
-        rt.shorts.write(|w| {
-            w.compare0_clear().set_bit();
-            w.compare0_stop().set_bit();
-            w
-        });
+        timer.set_frequency(Frequency::F16MHz);
+        timer.cc(0).write(timeout);
+        timer.cc(0).short_compare_clear();
+        timer.cc(0).short_compare_stop();
 
         let mut ppi_ch1 = Ppi::new(ppi_ch1.degrade_configurable());
         ppi_ch1.set_event(Event::from_reg(&r.events_rxdrdy));
-        ppi_ch1.set_task(Task::from_reg(&rt.tasks_clear));
-        ppi_ch1.set_fork_task(Task::from_reg(&rt.tasks_start));
+        ppi_ch1.set_task(timer.task_clear());
+        ppi_ch1.set_fork_task(timer.task_start());
         ppi_ch1.enable();
 
         let mut ppi_ch2 = Ppi::new(ppi_ch2.degrade_configurable());
-        ppi_ch2.set_event(Event::from_reg(&rt.events_compare[0]));
+        ppi_ch2.set_event(timer.cc(0).event_compare());
         ppi_ch2.set_task(Task::from_reg(&r.tasks_stoprx));
         ppi_ch2.enable();
 
@@ -181,11 +178,10 @@ impl<'d, U: UarteInstance, T: TimerInstance> BufferedUarte<'d, U, T> {
         inner.as_mut().register_interrupt();
         inner.with(|state, _irq| {
             let r = U::regs();
-            let rt = state.timer.regs();
 
             let timeout = 0x8000_0000 / (baudrate as u32 / 40);
-            rt.cc[0].write(|w| unsafe { w.bits(timeout) });
-            rt.tasks_clear.write(|w| unsafe { w.bits(1) });
+            state.timer.cc(0).write(timeout);
+            state.timer.clear();
 
             r.baudrate.write(|w| w.baudrate().variant(baudrate));
         });
@@ -268,11 +264,10 @@ impl<'d, U: UarteInstance, T: TimerInstance> AsyncWrite for BufferedUarte<'d, U,
 impl<'a, U: UarteInstance, T: TimerInstance> Drop for State<'a, U, T> {
     fn drop(&mut self) {
         let r = U::regs();
-        let rt = self.timer.regs();
 
         // TODO this probably deadlocks. do like Uarte instead.
 
-        rt.tasks_stop.write(|w| unsafe { w.bits(1) });
+        self.timer.stop();
         if let RxState::Receiving = self.rx_state {
             r.tasks_stoprx.write(|w| unsafe { w.bits(1) });
         }
@@ -293,7 +288,6 @@ impl<'a, U: UarteInstance, T: TimerInstance> PeripheralState for State<'a, U, T>
     fn on_interrupt(&mut self) {
         trace!("irq: start");
         let r = U::regs();
-        let rt = self.timer.regs();
 
         loop {
             match self.rx_state {
@@ -330,7 +324,7 @@ impl<'a, U: UarteInstance, T: TimerInstance> PeripheralState for State<'a, U, T>
                 RxState::Receiving => {
                     trace!("  irq_rx: in state receiving");
                     if r.events_endrx.read().bits() != 0 {
-                        rt.tasks_stop.write(|w| unsafe { w.bits(1) });
+                        self.timer.stop();
 
                         let n: usize = r.rxd.amount.read().amount().bits() as usize;
                         trace!("  irq_rx: endrx {:?}", n);
diff --git a/embassy-nrf/src/timer.rs b/embassy-nrf/src/timer.rs
index 2490bfd93..a6e91f228 100644
--- a/embassy-nrf/src/timer.rs
+++ b/embassy-nrf/src/timer.rs
@@ -1,15 +1,30 @@
 #![macro_use]
 
+use core::marker::PhantomData;
+use core::task::Poll;
+
 use embassy::interrupt::Interrupt;
+use embassy::interrupt::InterruptExt;
+use embassy::util::OnDrop;
 use embassy::util::Unborrow;
+use embassy_extras::unborrow;
+use futures::future::poll_fn;
 
 use crate::pac;
+use crate::ppi::Event;
+use crate::ppi::Task;
 
 pub(crate) mod sealed {
+    use embassy::util::AtomicWaker;
+
     use super::*;
 
     pub trait Instance {
-        fn regs(&self) -> &pac::timer0::RegisterBlock;
+        /// The number of CC registers this instance has.
+        const CCS: usize;
+        fn regs() -> &'static pac::timer0::RegisterBlock;
+        /// Storage for the waker for CC register `n`.
+        fn waker(n: usize) -> &'static AtomicWaker;
     }
     pub trait ExtendedInstance {}
 }
@@ -20,19 +35,306 @@ pub trait Instance: Unborrow<Target = Self> + sealed::Instance + 'static {
 pub trait ExtendedInstance: Instance + sealed::ExtendedInstance {}
 
 macro_rules! impl_timer {
-    ($type:ident, $pac_type:ident, $irq:ident) => {
+    ($type:ident, $pac_type:ident, $irq:ident, $ccs:literal) => {
         impl crate::timer::sealed::Instance for peripherals::$type {
-            fn regs(&self) -> &pac::timer0::RegisterBlock {
+            const CCS: usize = $ccs;
+            fn regs() -> &'static pac::timer0::RegisterBlock {
                 unsafe { &*(pac::$pac_type::ptr() as *const pac::timer0::RegisterBlock) }
             }
+            fn waker(n: usize) -> &'static ::embassy::util::AtomicWaker {
+                use ::embassy::util::AtomicWaker;
+                const NEW_AW: AtomicWaker = AtomicWaker::new();
+                static WAKERS: [AtomicWaker; $ccs] = [NEW_AW; $ccs];
+                &WAKERS[n]
+            }
         }
         impl crate::timer::Instance for peripherals::$type {
             type Interrupt = crate::interrupt::$irq;
         }
     };
+    ($type:ident, $pac_type:ident, $irq:ident) => {
+        impl_timer!($type, $pac_type, $irq, 4);
+    };
     ($type:ident, $pac_type:ident, $irq:ident, extended) => {
-        impl_timer!($type, $pac_type, $irq);
+        impl_timer!($type, $pac_type, $irq, 6);
         impl crate::timer::sealed::ExtendedInstance for peripherals::$type {}
         impl crate::timer::ExtendedInstance for peripherals::$type {}
     };
 }
+
+#[repr(u8)]
+pub enum Frequency {
+    // I'd prefer not to prefix these with `F`, but Rust identifiers can't start with digits.
+    F16MHz = 0,
+    F8MHz = 1,
+    F4MHz = 2,
+    F2MHz = 3,
+    F1MHz = 4,
+    F500kHz = 5,
+    F250kHz = 6,
+    F125kHz = 7,
+    F62500Hz = 8,
+    F31250Hz = 9,
+}
+
+/// nRF Timer driver.
+///
+/// The timer has an internal counter, which is incremented for every tick of the timer.
+/// The counter is 32-bit, so it wraps back to 0 at 4294967296.
+///
+/// It has either 4 or 6 Capture/Compare registers, which can be used to capture the current state of the counter
+/// or trigger an event when the counter reaches a certain value.
+pub struct Timer<'d, T: Instance> {
+    phantom: PhantomData<&'d mut T>,
+}
+
+impl<'d, T: Instance> Timer<'d, T> {
+    pub fn new(
+        timer: impl Unborrow<Target = T> + 'd,
+        irq: impl Unborrow<Target = T::Interrupt> + 'd,
+    ) -> Self {
+        unborrow!(irq);
+
+        irq.set_handler(Self::on_interrupt);
+        irq.unpend();
+        irq.enable();
+
+        Self::new_irqless(timer)
+    }
+
+    /// Create a `Timer` without an interrupt, meaning `Cc::wait` won't work.
+    ///
+    /// This is used by `Uarte` internally.
+    pub(crate) fn new_irqless(_timer: impl Unborrow<Target = T> + 'd) -> Self {
+        let regs = T::regs();
+
+        let mut this = Self {
+            phantom: PhantomData,
+        };
+
+        // Stop the timer before doing anything else,
+        // since changing BITMODE while running can cause 'unpredictable behaviour' according to the specification.
+        this.stop();
+
+        // Set the instance to timer mode.
+        regs.mode.write(|w| w.mode().timer());
+
+        // Make the counter's max value as high as possible.
+        // TODO: is there a reason someone would want to set this lower?
+        regs.bitmode.write(|w| w.bitmode()._32bit());
+
+        // Initialize the counter at 0.
+        this.clear();
+
+        // Default to the max frequency of the lower power clock
+        this.set_frequency(Frequency::F1MHz);
+
+        for n in 0..T::CCS {
+            let cc = this.cc(n);
+            // Initialize all the shorts as disabled.
+            cc.unshort_compare_clear();
+            cc.unshort_compare_stop();
+            // Initialize the CC registers as 0.
+            cc.write(0);
+        }
+
+        this
+    }
+
+    /// Starts the timer.
+    pub fn start(&self) {
+        T::regs().tasks_start.write(|w| unsafe { w.bits(1) })
+    }
+
+    /// Stops the timer.
+    pub fn stop(&self) {
+        T::regs().tasks_stop.write(|w| unsafe { w.bits(1) })
+    }
+
+    /// Reset the timer's counter to 0.
+    pub fn clear(&self) {
+        T::regs().tasks_clear.write(|w| unsafe { w.bits(1) })
+    }
+
+    /// Returns the START task, for use with PPI.
+    ///
+    /// When triggered, this task starts the timer.
+    pub fn task_start(&self) -> Task {
+        Task::from_reg(&T::regs().tasks_start)
+    }
+
+    /// Returns the STOP task, for use with PPI.
+    ///
+    /// When triggered, this task stops the timer.
+    pub fn task_stop(&self) -> Task {
+        Task::from_reg(&T::regs().tasks_stop)
+    }
+
+    /// Returns the CLEAR task, for use with PPI.
+    ///
+    /// When triggered, this task resets the timer's counter to 0.
+    pub fn task_clear(&self) -> Task {
+        Task::from_reg(&T::regs().tasks_clear)
+    }
+
+    /// Change the timer's frequency.
+    ///
+    /// This will stop the timer if it isn't already stopped,
+    /// because the timer may exhibit 'unpredictable behaviour' if it's frequency is changed while it's running.
+    pub fn set_frequency(&self, frequency: Frequency) {
+        self.stop();
+
+        T::regs()
+            .prescaler
+            // SAFETY: `frequency` is a variant of `Frequency`,
+            // whose values are all in the range of 0-9 (the valid range of `prescaler`).
+            .write(|w| unsafe { w.prescaler().bits(frequency as u8) })
+    }
+
+    fn on_interrupt(_: *mut ()) {
+        let regs = T::regs();
+        for n in 0..T::CCS {
+            if regs.events_compare[n].read().bits() != 0 {
+                // Clear the interrupt, otherwise the interrupt will be repeatedly raised as soon as the interrupt handler exits.
+                // We can't clear the event, because it's used to poll whether the future is done or still pending.
+                regs.intenclr
+                    .modify(|r, w| unsafe { w.bits(r.bits() | (1 << (16 + n))) });
+                T::waker(n).wake();
+            }
+        }
+    }
+
+    /// Returns this timer's `n`th CC register.
+    ///
+    /// # Panics
+    /// Panics if `n` >= the number of CC registers this timer has (4 for a normal timer, 6 for an extended timer).
+    pub fn cc(&mut self, n: usize) -> Cc<T> {
+        if n >= T::CCS {
+            panic!(
+                "Cannot get CC register {} of timer with {} CC registers.",
+                n,
+                T::CCS
+            );
+        }
+        Cc {
+            n,
+            phantom: PhantomData,
+        }
+    }
+}
+
+/// A representation of a timer's Capture/Compare (CC) register.
+///
+/// A CC register holds a 32-bit value.
+/// This is used either to store a capture of the timer's current count, or to specify the value for the timer to compare against.
+///
+/// The timer will fire the register's COMPARE event when its counter reaches the value stored in the register.
+/// When the register's CAPTURE task is triggered, the timer will store the current value of its counter in the register
+pub struct Cc<'a, T: Instance> {
+    n: usize,
+    phantom: PhantomData<&'a mut T>,
+}
+
+impl<'a, T: Instance> Cc<'a, T> {
+    /// Get the current value stored in the register.
+    pub fn read(&self) -> u32 {
+        T::regs().cc[self.n].read().cc().bits()
+    }
+
+    /// Set the value stored in the register.
+    ///
+    /// `event_compare` will fire when the timer's counter reaches this value.
+    pub fn write(&self, value: u32) {
+        // SAFETY: there are no invalid values for the CC register.
+        T::regs().cc[self.n].write(|w| unsafe { w.cc().bits(value) })
+    }
+
+    /// Capture the current value of the timer's counter in this register, and return it.
+    pub fn capture(&self) -> u32 {
+        T::regs().tasks_capture[self.n].write(|w| unsafe { w.bits(1) });
+        self.read()
+    }
+
+    /// Returns this CC register's CAPTURE task, for use with PPI.
+    ///
+    /// When triggered, this task will capture the current value of the timer's counter in this register.
+    pub fn task_capture(&self) -> Task {
+        Task::from_reg(&T::regs().tasks_capture[self.n])
+    }
+
+    /// Returns this CC register's COMPARE event, for use with PPI.
+    ///
+    /// This event will fire when the timer's counter reaches the value in this CC register.
+    pub fn event_compare(&self) -> Event {
+        Event::from_reg(&T::regs().events_compare[self.n])
+    }
+
+    /// Enable the shortcut between this CC register's COMPARE event and the timer's CLEAR task.
+    ///
+    /// This means that when the COMPARE event is fired, the CLEAR task will be triggered.
+    ///
+    /// So, when the timer's counter reaches the value stored in this register, the timer's counter will be reset to 0.
+    pub fn short_compare_clear(&self) {
+        T::regs()
+            .shorts
+            .modify(|r, w| unsafe { w.bits(r.bits() | (1 << self.n)) })
+    }
+
+    /// Disable the shortcut between this CC register's COMPARE event and the timer's CLEAR task.
+    pub fn unshort_compare_clear(&self) {
+        T::regs()
+            .shorts
+            .modify(|r, w| unsafe { w.bits(r.bits() & !(1 << self.n)) })
+    }
+
+    /// Enable the shortcut between this CC register's COMPARE event and the timer's STOP task.
+    ///
+    /// This means that when the COMPARE event is fired, the STOP task will be triggered.
+    ///
+    /// So, when the timer's counter reaches the value stored in this register, the timer will stop counting up.
+    pub fn short_compare_stop(&self) {
+        T::regs()
+            .shorts
+            .modify(|r, w| unsafe { w.bits(r.bits() | (1 << (8 + self.n))) })
+    }
+
+    /// Disable the shortcut between this CC register's COMPARE event and the timer's STOP task.
+    pub fn unshort_compare_stop(&self) {
+        T::regs()
+            .shorts
+            .modify(|r, w| unsafe { w.bits(r.bits() & !(1 << (8 + self.n))) })
+    }
+
+    /// Wait until the timer's counter reaches the value stored in this register.
+    ///
+    /// This requires a mutable reference so that this task's waker cannot be overwritten by a second call to `wait`.
+    pub async fn wait(&mut self) {
+        let regs = T::regs();
+
+        // Enable the interrupt for this CC's COMPARE event.
+        regs.intenset
+            .modify(|r, w| unsafe { w.bits(r.bits() | (1 << (16 + self.n))) });
+
+        // Disable the interrupt if the future is dropped.
+        let on_drop = OnDrop::new(|| {
+            regs.intenclr
+                .modify(|r, w| unsafe { w.bits(r.bits() | (1 << (16 + self.n))) });
+        });
+
+        poll_fn(|cx| {
+            T::waker(self.n).register(cx.waker());
+
+            if regs.events_compare[self.n].read().bits() != 0 {
+                // Reset the register for next time
+                regs.events_compare[self.n].reset();
+                Poll::Ready(())
+            } else {
+                Poll::Pending
+            }
+        })
+        .await;
+
+        // The interrupt was already disabled in the interrupt handler, so there's no need to disable it again.
+        on_drop.defuse();
+    }
+}
diff --git a/embassy-nrf/src/uarte.rs b/embassy-nrf/src/uarte.rs
index de093c7f0..67ec5d73f 100644
--- a/embassy-nrf/src/uarte.rs
+++ b/embassy-nrf/src/uarte.rs
@@ -18,7 +18,9 @@ use crate::gpio::{self, OptionalPin as GpioOptionalPin, Pin as GpioPin};
 use crate::interrupt::Interrupt;
 use crate::pac;
 use crate::ppi::{AnyConfigurableChannel, ConfigurableChannel, Event, Ppi, Task};
+use crate::timer::Frequency;
 use crate::timer::Instance as TimerInstance;
+use crate::timer::Timer;
 
 // Re-export SVD variants to allow user to directly set values.
 pub use pac::uarte0::{baudrate::BAUDRATE_A as Baudrate, config::PARITY_A as Parity};
@@ -287,7 +289,7 @@ impl<'d, T: Instance> Write for Uarte<'d, T> {
 /// allowing it to implement the ReadUntilIdle trait.
 pub struct UarteWithIdle<'d, U: Instance, T: TimerInstance> {
     uarte: Uarte<'d, U>,
-    timer: T,
+    timer: Timer<'d, T>,
     ppi_ch1: Ppi<'d, AnyConfigurableChannel>,
     _ppi_ch2: Ppi<'d, AnyConfigurableChannel>,
 }
@@ -316,11 +318,11 @@ impl<'d, U: Instance, T: TimerInstance> UarteWithIdle<'d, U, T> {
     ) -> Self {
         let baudrate = config.baudrate;
         let uarte = Uarte::new(uarte, irq, rxd, txd, cts, rts, config);
+        let mut timer = Timer::new_irqless(timer);
 
-        unborrow!(timer, ppi_ch1, ppi_ch2);
+        unborrow!(ppi_ch1, ppi_ch2);
 
         let r = U::regs();
-        let rt = timer.regs();
 
         // BAUDRATE register values are `baudrate * 2^32 / 16000000`
         // source: https://devzone.nordicsemi.com/f/nordic-q-a/391/uart-baudrate-register-values
@@ -330,25 +332,19 @@ impl<'d, U: Instance, T: TimerInstance> UarteWithIdle<'d, U, T> {
         // This gives us the amount of 16M ticks for 20 bits.
         let timeout = 0x8000_0000 / (baudrate as u32 / 40);
 
-        rt.tasks_stop.write(|w| unsafe { w.bits(1) });
-        rt.bitmode.write(|w| w.bitmode()._32bit());
-        rt.prescaler.write(|w| unsafe { w.prescaler().bits(0) });
-        rt.cc[0].write(|w| unsafe { w.bits(timeout) });
-        rt.mode.write(|w| w.mode().timer());
-        rt.shorts.write(|w| {
-            w.compare0_clear().set_bit();
-            w.compare0_stop().set_bit();
-            w
-        });
+        timer.set_frequency(Frequency::F16MHz);
+        timer.cc(0).write(timeout);
+        timer.cc(0).short_compare_clear();
+        timer.cc(0).short_compare_stop();
 
         let mut ppi_ch1 = Ppi::new(ppi_ch1.degrade_configurable());
         ppi_ch1.set_event(Event::from_reg(&r.events_rxdrdy));
-        ppi_ch1.set_task(Task::from_reg(&rt.tasks_clear));
-        ppi_ch1.set_fork_task(Task::from_reg(&rt.tasks_start));
+        ppi_ch1.set_task(timer.task_clear());
+        ppi_ch1.set_fork_task(timer.task_start());
         ppi_ch1.enable();
 
         let mut ppi_ch2 = Ppi::new(ppi_ch2.degrade_configurable());
-        ppi_ch2.set_event(Event::from_reg(&rt.events_compare[0]));
+        ppi_ch2.set_event(timer.cc(0).event_compare());
         ppi_ch2.set_task(Task::from_reg(&r.tasks_stoprx));
         ppi_ch2.enable();
 
@@ -373,12 +369,10 @@ impl<'d, U: Instance, T: TimerInstance> ReadUntilIdle for UarteWithIdle<'d, U, T
             let r = U::regs();
             let s = U::state();
 
-            let rt = self.timer.regs();
-
-            let drop = OnDrop::new(move || {
+            let drop = OnDrop::new(|| {
                 info!("read drop: stopping");
 
-                rt.tasks_stop.write(|w| unsafe { w.bits(1) });
+                self.timer.stop();
 
                 r.intenclr.write(|w| w.endrx().clear());
                 r.events_rxto.reset();
@@ -413,7 +407,7 @@ impl<'d, U: Instance, T: TimerInstance> ReadUntilIdle for UarteWithIdle<'d, U, T
             let n = r.rxd.amount.read().amount().bits() as usize;
 
             // Stop timer
-            rt.tasks_stop.write(|w| unsafe { w.bits(1) });
+            self.timer.stop();
             r.events_rxstarted.reset();
 
             drop.defuse();