From a15a6a25b6f5fc4119ffe24163d6630721cdd38f Mon Sep 17 00:00:00 2001
From: Rekai Musuka <rekai@musuka.dev>
Date: Mon, 5 Apr 2021 00:52:12 -0500
Subject: [PATCH] feat(cpu): properly implement EI instruction

---
 src/cpu.rs         | 62 ++++++++++++++++++++++++++++++++++++++--------
 src/instruction.rs | 10 ++++----
 2 files changed, 57 insertions(+), 15 deletions(-)

diff --git a/src/cpu.rs b/src/cpu.rs
index a4bd8ba..a0143e2 100644
--- a/src/cpu.rs
+++ b/src/cpu.rs
@@ -4,13 +4,14 @@ 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 {
     pub bus: Bus,
     reg: Registers,
     flags: Flags,
-    ime: bool,
+    ime: ImeState,
     halted: Option<HaltState>,
     state: State,
 }
@@ -41,12 +42,12 @@ impl Cpu {
         })
     }
 
-    pub fn ime(&self) -> bool {
+    pub fn ime(&self) -> ImeState {
         self.ime
     }
 
-    pub fn set_ime(&mut self, enabled: bool) {
-        self.ime = enabled;
+    pub fn set_ime(&mut self, state: ImeState) {
+        self.ime = state;
     }
 
     pub fn halt(&mut self, state: HaltState) {
@@ -86,9 +87,9 @@ impl Cpu {
     }
 
     pub fn step(&mut self) -> Cycle {
-        if !self.bus.boot_enabled() {
-            self.log_state().unwrap();
-        }
+        // if !self.bus.boot_enabled() {
+        //     self.log_state().unwrap();
+        // }
 
         let cycles = match self.halted() {
             Some(state) => {
@@ -106,7 +107,9 @@ impl Cpu {
                 let instr = self.decode(opcode);
                 let cycles = self.execute(instr);
 
+                self.check_ime();
                 self.bus.step(cycles);
+
                 cycles
             }
         };
@@ -151,7 +154,17 @@ impl Cpu {
         &mut self.bus.ppu
     }
 
-    pub fn handle_interrupts(&mut self) {
+    fn check_ime(&mut self) {
+        if let ImeState::EnablePending(count) = self.ime {
+            self.ime = if count < 2 {
+                self.ime.wait()
+            } else {
+                ImeState::Enabled
+            }
+        };
+    }
+
+    fn handle_interrupts(&mut self) {
         let req = self.read_byte(0xFF0F);
         let enabled = self.read_byte(0xFFFF);
 
@@ -168,7 +181,14 @@ impl Cpu {
             }
         }
 
-        if self.ime() {
+        if let ImeState::Enabled = self.ime() {
+            println!(
+                "req & enabled: {:#010b} | Halted: {:?} | PC: {:#06X}",
+                req & enabled,
+                self.halted,
+                self.reg.pc,
+            );
+
             let mut req: InterruptFlag = req.into();
             let enabled: InterruptEnable = enabled.into();
 
@@ -212,7 +232,7 @@ impl Cpu {
                     self.write_byte(0xFF0F, req.into());
 
                     // Disable all future interrupts
-                    self.set_ime(false);
+                    self.set_ime(ImeState::Disabled);
                     self.execute(Instruction::RST(register))
                 }
                 None => Cycle::new(0), // NO Interrupts were enabled and / or requested
@@ -454,3 +474,25 @@ pub enum HaltState {
     NonePending,
     SomePending,
 }
+
+#[derive(Debug, Clone, Copy)]
+pub enum ImeState {
+    Disabled,
+    Enabled,
+    EnablePending(u8),
+}
+
+impl Default for ImeState {
+    fn default() -> Self {
+        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 391f581..a20d901 100644
--- a/src/instruction.rs
+++ b/src/instruction.rs
@@ -1,4 +1,4 @@
-use super::cpu::{Cpu, Flags, HaltState, Register, RegisterPair};
+use super::cpu::{Cpu, Flags, HaltState, ImeState, Register, RegisterPair};
 use std::{convert::TryFrom, fmt::Debug};
 
 #[derive(Debug, Copy, Clone)]
@@ -611,7 +611,7 @@ impl Instruction {
                 let req = cpu.read_byte(0xFF0F);
                 let enabled = cpu.read_byte(0xFFFF);
 
-                let halt_state = if cpu.ime() {
+                let halt_state = if let ImeState::Enabled = cpu.ime() {
                     ImeSet
                 } else if req & enabled != 0 {
                     SomePending
@@ -957,7 +957,7 @@ impl Instruction {
                 // Same as RET, after which interrupts are enabled.
                 let addr = Self::pop(cpu);
                 cpu.set_register_pair(RegisterPair::PC, addr);
-                cpu.set_ime(true);
+                cpu.set_ime(ImeState::Enabled);
                 Cycle::new(16)
             }
             Instruction::JP(cond, target) => match target {
@@ -1016,13 +1016,13 @@ impl Instruction {
             },
             Instruction::DI => {
                 // Disable IME
-                cpu.set_ime(false);
+                cpu.set_ime(ImeState::Disabled);
                 Cycle::new(4)
             }
             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(true);
+                cpu.set_ime(ImeState::EnablePending(0));
                 Cycle::new(4)
             }
             Instruction::CALL(cond, nn) => {