1 files changed, 41 insertions, 167 deletions

diff --git a/embassy-time/src/driver_std.rs b/embassy-time/src/driver_std.rs
index cbef7aae1..35888fddd 100644
--- a/embassy-time/src/driver_std.rs
+++ b/embassy-time/src/driver_std.rs
@@ -1,160 +1,66 @@
-use core::sync::atomic::{AtomicU8, Ordering};
-use std::cell::{RefCell, UnsafeCell};
-use std::mem::MaybeUninit;
-use std::sync::{Condvar, Mutex, Once};
+use std::sync::{Condvar, Mutex};
+use std::thread;
 use std::time::{Duration as StdDuration, Instant as StdInstant};
-use std::{mem, ptr, thread};
 
-use critical_section::Mutex as CsMutex;
-use embassy_time_driver::{AlarmHandle, Driver};
+use embassy_time_driver::Driver;
+use embassy_time_queue_driver::Queue;
 
-const ALARM_COUNT: usize = 4;
-
-struct AlarmState {
-    timestamp: u64,
-
-    // This is really a Option<(fn(*mut ()), *mut ())>
-    // but fn pointers aren't allowed in const yet
-    callback: *const (),
-    ctx: *mut (),
+struct TimeDriver {
+    signaler: Signaler,
+    inner: Mutex<Inner>,
 }
 
-unsafe impl Send for AlarmState {}
-
-impl AlarmState {
-    const fn new() -> Self {
-        Self {
-            timestamp: u64::MAX,
-            callback: ptr::null(),
-            ctx: ptr::null_mut(),
-        }
-    }
-}
-
-struct TimeDriver {
-    alarm_count: AtomicU8,
-
-    once: Once,
-    // The STD Driver implementation requires the alarms' mutex to be reentrant, which the STD Mutex isn't
-    // Fortunately, mutexes based on the `critical-section` crate are reentrant, because the critical sections
-    // themselves are reentrant
-    alarms: UninitCell<CsMutex<RefCell<[AlarmState; ALARM_COUNT]>>>,
-    zero_instant: UninitCell<StdInstant>,
-    signaler: UninitCell<Signaler>,
+struct Inner {
+    zero_instant: Option<StdInstant>,
+    queue: Queue,
 }
 
 embassy_time_driver::time_driver_impl!(static DRIVER: TimeDriver = TimeDriver {
-    alarm_count: AtomicU8::new(0),
-
-    once: Once::new(),
-    alarms: UninitCell::uninit(),
-    zero_instant: UninitCell::uninit(),
-    signaler: UninitCell::uninit(),
+    inner: Mutex::new(Inner{
+        zero_instant: None,
+        queue: Queue::new(),
+    }),
+    signaler: Signaler::new(),
 });
 
-impl TimeDriver {
-    fn init(&self) {
-        self.once.call_once(|| unsafe {
-            self.alarms
-                .write(CsMutex::new(RefCell::new([const { AlarmState::new() }; ALARM_COUNT])));
-            self.zero_instant.write(StdInstant::now());
-            self.signaler.write(Signaler::new());
-
-            thread::spawn(Self::alarm_thread);
-        });
-    }
-
-    fn alarm_thread() {
-        let zero = unsafe { DRIVER.zero_instant.read() };
-        loop {
-            let now = DRIVER.now();
-
-            let next_alarm = critical_section::with(|cs| {
-                let alarms = unsafe { DRIVER.alarms.as_ref() }.borrow(cs);
-                loop {
-                    let pending = alarms
-                        .borrow_mut()
-                        .iter_mut()
-                        .find(|alarm| alarm.timestamp <= now)
-                        .map(|alarm| {
-                            alarm.timestamp = u64::MAX;
-
-                            (alarm.callback, alarm.ctx)
-                        });
-
-                    if let Some((callback, ctx)) = pending {
-                        // safety:
-                        // - we can ignore the possiblity of `f` being unset (null) because of the safety contract of `allocate_alarm`.
-                        // - other than that we only store valid function pointers into alarm.callback
-                        let f: fn(*mut ()) = unsafe { mem::transmute(callback) };
-                        f(ctx);
-                    } else {
-                        // No alarm due
-                        break;
-                    }
-                }
-
-                alarms
-                    .borrow()
-                    .iter()
-                    .map(|alarm| alarm.timestamp)
-                    .min()
-                    .unwrap_or(u64::MAX)
-            });
-
-            // Ensure we don't overflow
-            let until = zero
-                .checked_add(StdDuration::from_micros(next_alarm))
-                .unwrap_or_else(|| StdInstant::now() + StdDuration::from_secs(1));
-
-            unsafe { DRIVER.signaler.as_ref() }.wait_until(until);
-        }
+impl Inner {
+    fn init(&mut self) -> StdInstant {
+        *self.zero_instant.get_or_insert_with(|| {
+            thread::spawn(alarm_thread);
+            StdInstant::now()
+        })
     }
 }
 
 impl Driver for TimeDriver {
     fn now(&self) -> u64 {
-        self.init();
-
-        let zero = unsafe { self.zero_instant.read() };
+        let mut inner = self.inner.lock().unwrap();
+        let zero = inner.init();
         StdInstant::now().duration_since(zero).as_micros() as u64
     }
 
-    unsafe fn allocate_alarm(&self) -> Option<AlarmHandle> {
-        let id = self.alarm_count.fetch_update(Ordering::AcqRel, Ordering::Acquire, |x| {
-            if x < ALARM_COUNT as u8 {
-                Some(x + 1)
-            } else {
-                None
-            }
-        });
-
-        match id {
-            Ok(id) => Some(AlarmHandle::new(id)),
-            Err(_) => None,
+    fn schedule_wake(&self, at: u64, waker: &core::task::Waker) {
+        let mut inner = self.inner.lock().unwrap();
+        inner.init();
+        if inner.queue.schedule_wake(at, waker) {
+            self.signaler.signal();
         }
     }
+}
 
-    fn set_alarm_callback(&self, alarm: AlarmHandle, callback: fn(*mut ()), ctx: *mut ()) {
-        self.init();
-        critical_section::with(|cs| {
-            let mut alarms = unsafe { self.alarms.as_ref() }.borrow_ref_mut(cs);
-            let alarm = &mut alarms[alarm.id() as usize];
-            alarm.callback = callback as *const ();
-            alarm.ctx = ctx;
-        });
-    }
+fn alarm_thread() {
+    let zero = DRIVER.inner.lock().unwrap().zero_instant.unwrap();
+    loop {
+        let now = DRIVER.now();
+
+        let next_alarm = DRIVER.inner.lock().unwrap().queue.next_expiration(now);
 
-    fn set_alarm(&self, alarm: AlarmHandle, timestamp: u64) -> bool {
-        self.init();
-        critical_section::with(|cs| {
-            let mut alarms = unsafe { self.alarms.as_ref() }.borrow_ref_mut(cs);
-            let alarm = &mut alarms[alarm.id() as usize];
-            alarm.timestamp = timestamp;
-            unsafe { self.signaler.as_ref() }.signal();
-        });
+        // Ensure we don't overflow
+        let until = zero
+            .checked_add(StdDuration::from_micros(next_alarm))
+            .unwrap_or_else(|| StdInstant::now() + StdDuration::from_secs(1));
 
-        true
+        DRIVER.signaler.wait_until(until);
     }
 }
 
@@ -164,7 +70,7 @@ struct Signaler {
 }
 
 impl Signaler {
-    fn new() -> Self {
+    const fn new() -> Self {
         Self {
             mutex: Mutex::new(false),
             condvar: Condvar::new(),
@@ -196,35 +102,3 @@ impl Signaler {
         self.condvar.notify_one();
     }
 }
-
-pub(crate) struct UninitCell<T>(MaybeUninit<UnsafeCell<T>>);
-unsafe impl<T> Send for UninitCell<T> {}
-unsafe impl<T> Sync for UninitCell<T> {}
-
-impl<T> UninitCell<T> {
-    pub const fn uninit() -> Self {
-        Self(MaybeUninit::uninit())
-    }
-
-    pub unsafe fn as_ptr(&self) -> *const T {
-        (*self.0.as_ptr()).get()
-    }
-
-    pub unsafe fn as_mut_ptr(&self) -> *mut T {
-        (*self.0.as_ptr()).get()
-    }
-
-    pub unsafe fn as_ref(&self) -> &T {
-        &*self.as_ptr()
-    }
-
-    pub unsafe fn write(&self, val: T) {
-        ptr::write(self.as_mut_ptr(), val)
-    }
-}
-
-impl<T: Copy> UninitCell<T> {
-    pub unsafe fn read(&self) -> T {
-        ptr::read(self.as_mut_ptr())
-    }
-}