From 748c32c44612156b6c56d10f65fd590c8ea06c0e Mon Sep 17 00:00:00 2001
From: Rekai Musuka <rekai@musuka.dev>
Date: Mon, 5 Apr 2021 01:10:03 -0500
Subject: [PATCH] fix(cpu): use enums only of maintaining IME register state

---
 src/cpu.rs         | 30 +++++++++++++-----------------
 src/instruction.rs |  2 +-
 2 files changed, 14 insertions(+), 18 deletions(-)

diff --git a/src/cpu.rs b/src/cpu.rs
index dc957dc..3407858 100644
--- a/src/cpu.rs
+++ b/src/cpu.rs
@@ -4,7 +4,6 @@ use super::interrupt::{InterruptEnable, InterruptFlag};
 use super::ppu::Ppu;
 use bitfield::bitfield;
 use std::fmt::{Display, Formatter, Result as FmtResult};
-use std::ops::Add;
 
 #[derive(Debug, Clone, Default)]
 pub struct Cpu {
@@ -155,13 +154,18 @@ impl Cpu {
     }
 
     fn check_ime(&mut self) {
-        if let ImeState::EnablePending(count) = self.ime {
-            self.ime = if count < 2 {
-                self.ime.wait()
-            } else {
-                ImeState::Enabled
+        match self.ime {
+            ImeState::Pending => {
+                // This is within the context of the EI instruction, we need to not update EI until the end of the
+                // next executed Instruction
+                self.ime = ImeState::PendingEnd;
             }
-        };
+            ImeState::PendingEnd => {
+                // The Instruction after EI has now been executed, so we want to enable the IME flag here
+                self.ime = ImeState::Enabled;
+            }
+            ImeState::Disabled | ImeState::Enabled => {} // Do Nothing
+        }
     }
 
     fn handle_interrupts(&mut self) {
@@ -478,8 +482,9 @@ pub enum HaltState {
 #[derive(Debug, Clone, Copy)]
 pub enum ImeState {
     Disabled,
+    Pending,
+    PendingEnd,
     Enabled,
-    EnablePending(u8),
 }
 
 impl Default for ImeState {
@@ -487,12 +492,3 @@ impl Default for ImeState {
         Self::Disabled
     }
 }
-
-impl ImeState {
-    pub fn wait(self) -> Self {
-        match self {
-            Self::EnablePending(count) => Self::EnablePending(count + 1),
-            _ => panic!("IME is {:?}, however wait() was called", self),
-        }
-    }
-}
diff --git a/src/instruction.rs b/src/instruction.rs
index 2e7897c..331c7af 100644
--- a/src/instruction.rs
+++ b/src/instruction.rs
@@ -1022,7 +1022,7 @@ impl Instruction {
             Instruction::EI => {
                 // Enable IME (After the next instruction)
                 // FIXME: IME is set after the next instruction, this currently is not represented in this emulator.
-                cpu.set_ime(ImeState::EnablePending(0));
+                cpu.set_ime(ImeState::Pending);
                 Cycle::new(4)
             }
             Instruction::CALL(cond, nn) => {