1 use slab::Slab;

     2 use std::{

     3     convert::TryInto,

     4     iter,

     5     num::NonZeroU32,

     6     time::{Duration, Instant},

     7 };

     8 

     9 struct Event<Data> {

    10     event_id: u32,

    11     data: Data,

    12 }

    13 

    14 #[derive(Clone)]

    15 pub struct Timeout {

    16     tick_index: u32,

    17     event_index: u32,

    18     event_id: u32,

    19 }

    20 

    21 pub struct TimedEvents<Data, const MAX_TIMEOUT: usize> {

    22     events: [Slab<Event<Data>>; MAX_TIMEOUT],

    23     current_time: Instant,

    24     current_tick_index: u32,

    25     next_event_id: u32,

    26 }

    27 

    28 impl<Data, const MAX_TIMEOUT: usize> TimedEvents<Data, MAX_TIMEOUT> {

    29     pub fn new() -> Self {

    30         Self {

    31             events: iter::repeat_with(|| Slab::new())

    32                 .take(MAX_TIMEOUT)

    33                 .collect::<Vec<_>>()

    34                 .try_into()

    35                 .ok()

    36                 .unwrap(),

    37             current_time: Instant::now(),

    38             current_tick_index: 0,

    39             next_event_id: 0,

    40         }

    41     }

    42 

    43     pub fn set_timeout(&mut self, seconds_delay: NonZeroU32, data: Data) -> Timeout {

    44         let tick_index = (self.current_tick_index

    45             + std::cmp::min(seconds_delay.get(), MAX_TIMEOUT as u32))

    46             % MAX_TIMEOUT as u32;

    47         let event_id = self.next_event_id;

    48         self.next_event_id += 1;

    49         let event = Event { event_id, data };

    50 

    51         let entry = self.events[tick_index as usize].vacant_entry();

    52         let event_index = entry.key() as u32;

    53         entry.insert(event);

    54         Timeout {

    55             tick_index,

    56             event_index,

    57             event_id,

    58         }

    59     }

    60 

    61     pub fn cancel_timeout(&mut self, timeout: Timeout) -> Option<Data> {

    62         let events = &mut self.events[timeout.tick_index as usize];

    63         if matches!(events.get(timeout.event_index as usize), Some(Event { event_id: id, ..}) if *id == timeout.event_id)

    64         {

    65             Some(events.remove(timeout.event_index as usize).data)

    66         } else {

    67             None

    68         }

    69     }

    70 

    71     pub fn poll(&mut self, time: Instant) -> Vec<Data> {

    72         let mut result = vec![];

    73         let second = Duration::from_secs(1);

    74         while time - self.current_time > second {

    75             self.current_time += second;

    76             self.current_tick_index = (self.current_tick_index + 1) % MAX_TIMEOUT as u32;

    77             result.extend(

    78                 self.events[self.current_tick_index as usize]

    79                     .drain()

    80                     .map(|e| e.data),

    81             );

    82         }

    83         result

    84     }

    85 }

    86 

    87 mod test {

    88     use super::TimedEvents;

    89     use std::{

    90         num::NonZeroU32,

    91         time::{Duration, Instant},

    92     };

    93 

    94     #[test]

    95     fn events_test() {

    96         let mut events = TimedEvents::<u32, 30>::new();

    97         let now = Instant::now();

    98 

    99         let timeouts = (1..=3)

   100             .map(|n| events.set_timeout(NonZeroU32::new(n).unwrap(), n))

   101             .collect::<Vec<_>>();

   102 

   103         let second = Duration::from_secs(1);

   104         assert_eq!(events.cancel_timeout(timeouts[1].clone()), Some(2));

   105         assert_eq!(events.poll(now + second), vec![1]);

   106         assert!(events.poll(now + second).is_empty());

   107         assert!(events.poll(now + 2 * second).is_empty());

   108         assert_eq!(events.poll(now + 3 * second), vec![3]);

   109     }

   110 }