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     events_count: u32,

    27 }

    28 

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

    30     pub fn new() -> Self {

    31         Self {

    32             events: [0; MAX_TIMEOUT].map(|_| Slab::new()),

    33             current_time: Instant::now(),

    34             current_tick_index: 0,

    35             next_event_id: 0,

    36             events_count: 0,

    37         }

    38     }

    39 

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

    41         let tick_index = (self.current_tick_index

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

    43             % MAX_TIMEOUT as u32;

    44         let event_id = self.next_event_id;

    45         self.next_event_id += 1;

    46         let event = Event { event_id, data };

    47 

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

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

    50         entry.insert(event);

    51 

    52         self.events_count += 1;

    53 

    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             self.events_count -= 1;

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

    67         } else {

    68             None

    69         }

    70     }

    71 

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

    73         let mut result = vec![];

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

    75         while time - self.current_time > second {

    76             self.current_time += second;

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

    78             result.extend(

    79                 self.events[self.current_tick_index as usize]

    80                     .drain()

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

    82             );

    83         }

    84         self.events_count -= result.len() as u32;

    85         result

    86     }

    87 

    88     pub fn is_empty(&self) -> bool {

    89         self.events_count == 0

    90     }

    91 }

    92 

    93 mod test {

    94     use super::TimedEvents;

    95     use std::{

    96         num::NonZeroU32,

    97         time::{Duration, Instant},

    98     };

    99 

   100     #[test]

   101     fn events_test() {

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

   103         let now = Instant::now();

   104 

   105         let timeouts = (1..=3)

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

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

   108 

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

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

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

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

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

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

   115     }

   116 }