1use bincode::{Decode, Encode};
2use jgenesis_proc_macros::EnumAll;
3use std::array;
4use std::cmp::Ordering;
5
6#[derive(Debug, Clone, Copy, PartialEq, Eq, Encode, Decode, EnumAll)]
7pub enum SchedulerEvent {
8    VBlankIrq = 0,
9    HBlankIrq,
10    VCounterIrq,
11    PpuEvent,
12    TimerOverflow,
13    Dummy,
14}
15
16impl SchedulerEvent {
17    fn as_bit(self) -> u32 {
18        1 << (self as u8)
19    }
20}
21
22#[derive(Debug, Clone, Copy, PartialEq, Eq, Encode, Decode)]
23struct HeapEntry {
24    event: SchedulerEvent,
25    cycles: u64,
26}
27
28impl PartialOrd for HeapEntry {
29    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
30        Some(self.cmp(other))
31    }
32}
33
34impl Ord for HeapEntry {
35    fn cmp(&self, other: &Self) -> Ordering {
36        self.cycles.cmp(&other.cycles)
37    }
38}
39
40#[derive(Debug, Clone, Encode, Decode)]
41pub struct Scheduler {
42    heap: [HeapEntry; SchedulerEvent::ALL.len()],
43    len: usize,
44    scheduled_bits: u32,
45}
46
47impl Scheduler {
48    pub fn new() -> Self {
49        // Initialize with a dummy event to avoid ever needing to check if the heap is empty
50        Self {
51            heap: array::from_fn(|_| HeapEntry { event: SchedulerEvent::Dummy, cycles: u64::MAX }),
52            len: 1,
53            scheduled_bits: SchedulerEvent::Dummy.as_bit(),
54        }
55    }
56
57    // Insert if event is not present, update cycles if it is present
58    pub fn insert_or_update(&mut self, event: SchedulerEvent, cycles: u64) {
59        log::trace!("Inserting event {event:?} at cycles {cycles}");
60
61        if self.scheduled_bits & event.as_bit() != 0 {
62            for i in 0..self.len {
63                if self.heap[i].event != event {
64                    continue;
65                }
66
67                let old_cycles = self.heap[i].cycles;
68                self.heap[i].cycles = cycles;
69
70                match cycles.cmp(&old_cycles) {
71                    Ordering::Less => self.heap_up(i),
72                    Ordering::Greater => self.heap_down(i),
73                    Ordering::Equal => {}
74                }
75
76                return;
77            }
78        }
79        self.scheduled_bits |= event.as_bit();
80
81        self.heap[self.len] = HeapEntry { event, cycles };
82        self.len += 1;
83        self.heap_up(self.len - 1);
84    }
85
86    pub fn remove(&mut self, event: SchedulerEvent) {
87        log::trace!("Removing event {event:?}");
88
89        if self.scheduled_bits & event.as_bit() == 0 {
90            return;
91        }
92        self.scheduled_bits &= !event.as_bit();
93
94        for i in 0..self.len {
95            if self.heap[i].event == event {
96                let old_cycles = self.heap[i].cycles;
97                self.heap.swap(i, self.len - 1);
98                self.len -= 1;
99
100                match self.heap[i].cycles.cmp(&old_cycles) {
101                    Ordering::Less => self.heap_up(i),
102                    Ordering::Greater => self.heap_down(i),
103                    Ordering::Equal => {}
104                }
105
106                return;
107            }
108        }
109    }
110
111    pub fn is_event_ready(&self, cycles: u64) -> bool {
112        cycles >= self.heap[0].cycles
113    }
114
115    pub fn pop(&mut self, cycles: u64) -> Option<(SchedulerEvent, u64)> {
116        if cycles < self.heap[0].cycles {
117            return None;
118        }
119
120        let HeapEntry { event, cycles } = self.heap[0];
121        self.heap.swap(0, self.len - 1);
122        self.len -= 1;
123        self.heap_down(0);
124        self.scheduled_bits &= !event.as_bit();
125
126        log::trace!("Popped event {event:?} at cycles {cycles}");
127
128        Some((event, cycles))
129    }
130
131    fn heap_up(&mut self, mut i: usize) {
132        while i != 0 {
133            let parent = i / 2;
134            if self.heap[parent] <= self.heap[i] {
135                return;
136            }
137
138            self.heap.swap(i, parent);
139            i = parent;
140        }
141    }
142
143    fn heap_down(&mut self, mut i: usize) {
144        loop {
145            let left = 2 * i + 1;
146            if left >= self.len {
147                return;
148            }
149            let right = left + 1;
150
151            if right < self.len
152                && self.heap[right] < self.heap[left]
153                && self.heap[right] < self.heap[i]
154            {
155                self.heap.swap(i, right);
156                i = right;
157            } else if self.heap[left] < self.heap[i] {
158                self.heap.swap(i, left);
159                i = left;
160            } else {
161                return;
162            }
163        }
164    }
165}