vnf: add support for all stack instructions

This also includes brk and rti, for which we are also storing and
exporting the break mark, just in case it can be useful to library users
in case they want to add error codes in break marks.

This commit also includes a fix for accounting the Z and N flags on
transfer instructions, which couldn't easily be split into a separate
commit, but it's a fairly simple change.

Signed-off-by: Miquel Sabaté Solà <mssola@mssola.com>

8 files changed, 383 insertions, 24 deletions

diff --git a/Cargo.lock b/Cargo.lock
index ea29435..1a180c6 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -39,6 +39,14 @@ dependencies = [
 ]
 
 [[package]]
+name = "vnf-tests"
+version = "0.1.0"
+dependencies = [
+ "vnf",
+ "xixanta",
+]
+
+[[package]]
 name = "xa65"
 version = "0.1.0"
 
diff --git a/Cargo.toml b/Cargo.toml
index 11a627c..cdb3983 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -1,5 +1,5 @@
 [workspace]
-members = ["lib/*", "crates/*"]
+members = ["lib/*", "crates/*", "tests/vnf-tests"]
 resolver = "2"
 
 [workspace.package]
diff --git a/lib/vnf/src/lib.rs b/lib/vnf/src/lib.rs
index b1b596b..e9c7539 100644
--- a/lib/vnf/src/lib.rs
+++ b/lib/vnf/src/lib.rs
@@ -17,6 +17,9 @@ pub struct StatusRegister {
     pub interrupt: bool,
     pub zero: bool,
     pub carry: bool,
+
+    /// The byte after a 'brk' instruction.
+    pub break_mark: u8,
 }
 
 impl Default for StatusRegister {
@@ -29,11 +32,24 @@ impl Default for StatusRegister {
             interrupt: true,
             zero: false,
             carry: false,
+            break_mark: 0,
         }
     }
 }
 
 impl StatusRegister {
+    // Bitmap constants for how the different bools are laid out on the Status
+    // Register. These are defined like so: (from most to least significant):
+    //   N V - B D I Z C
+    pub const CARRY: u8 = 1 << 0;
+    pub const ZERO: u8 = 1 << 1;
+    pub const INTERRUPT: u8 = 1 << 2;
+    pub const DECIMAL: u8 = 1 << 3;
+    pub const BRK: u8 = 1 << 4;
+    // NOTE: bit 5 unused.
+    pub const OVERFLOW: u8 = 1 << 6;
+    pub const NEGATIVE: u8 = 1 << 7;
+
     /// Returns a string with the initial letter for each status bit that is
     /// set. Otherwise, for unset bits, a '-' is given.
     fn humanize(&self) -> String {
@@ -145,6 +161,7 @@ pub struct MemoryPolicy {
 
 impl Default for MemoryPolicy {
     fn default() -> Self {
+        #[allow(clippy::single_range_in_vec_init)]
         Self {
             initial_value: MemoryInitialValue::Fixed(0),
             allowed_reads: vec![(0..0x800)],
@@ -298,6 +315,11 @@ fn init_memory(policy: &MemoryPolicy) -> Vec<MemoryCell> {
         });
     }
 
+    // Force these initial stack values as these should be the values regardless
+    // of the memory policy, as it's per 6502 specification..
+    vec[0x2FF].value = 0x00;
+    vec[0x2FE].value = 0x00;
+
     vec
 }
 
@@ -661,6 +683,67 @@ impl Machine {
         self.read_memory(address, account_read)
     }
 
+    // Push the given address onto the stack.
+    fn push_address(&mut self, address: usize) -> Result<(), String> {
+        let low = (address as u16 & 0x00FF) as u8;
+        let high = ((address as u16 & 0xFF00) >> 8) as u8;
+
+        self.push_stack(high, false)?;
+        self.push_stack(low, true)?;
+
+        Ok(())
+    }
+
+    // Push the status register onto the stack.
+    fn push_status_register(&mut self) -> Result<(), String> {
+        // brk flag and bit 5 are always set on 'php'/'brk' as per 6502
+        // specification.
+        let mut val = 0b00110000;
+        self.status_register.brk = true;
+
+        if self.status_register.carry {
+            val |= StatusRegister::CARRY;
+        }
+        if self.status_register.zero {
+            val |= StatusRegister::ZERO;
+        }
+        if self.status_register.interrupt {
+            val |= StatusRegister::INTERRUPT;
+        }
+        if self.status_register.decimal {
+            val |= StatusRegister::DECIMAL;
+        }
+        if self.status_register.overflow {
+            val |= StatusRegister::OVERFLOW;
+        }
+        if self.status_register.negative {
+            val |= StatusRegister::NEGATIVE;
+        }
+        self.push_stack(val, true)?;
+        Ok(())
+    }
+
+    // Pull the status register from the stack.
+    fn pop_status_register(&mut self) -> Result<(), String> {
+        let val = self.pop_stack(true, true)?;
+
+        self.status_register.carry = (val & StatusRegister::CARRY) == StatusRegister::CARRY;
+        self.status_register.zero = (val & StatusRegister::ZERO) == StatusRegister::ZERO;
+        self.status_register.interrupt =
+            (val & StatusRegister::INTERRUPT) == StatusRegister::INTERRUPT;
+        self.status_register.decimal = (val & StatusRegister::DECIMAL) == StatusRegister::DECIMAL;
+        self.status_register.overflow =
+            (val & StatusRegister::OVERFLOW) == StatusRegister::OVERFLOW;
+        self.status_register.negative =
+            (val & StatusRegister::NEGATIVE) == StatusRegister::NEGATIVE;
+
+        // BRK is always cleared. We also clear the break mark now.
+        self.status_register.brk = false;
+        self.status_register.break_mark = 0;
+
+        Ok(())
+    }
+
     // Returns true of the stack is empty, false otherwise. Note that this
     // just means that the value of the 's' register is the one set as its
     // initial value.
@@ -686,14 +769,8 @@ impl Machine {
 
         match self.current_instruction.identifier {
             // TODO
-            InstructionIdentifier::Brk => todo!(),
             InstructionIdentifier::Bvc => todo!(),
             InstructionIdentifier::Bvs => todo!(),
-            InstructionIdentifier::Pha => todo!(),
-            InstructionIdentifier::Pla => todo!(),
-            InstructionIdentifier::Php => todo!(),
-            InstructionIdentifier::Plp => todo!(),
-            InstructionIdentifier::Rti => todo!(),
 
             // Flag instructions.
             InstructionIdentifier::Sec => self.status_register.carry = true,
@@ -913,16 +990,12 @@ impl Machine {
                 }
 
                 // NOTE: as per 6502 specification, the address should be - 1
-                // because the 'rts'/'rti' instruction will be the one in charge
-                // of adding its size to the end PC upon execution. This is kind
-                // of pedantic but in the end we record the address being pushed
+                // because the 'rts' instruction will be the one in charge of
+                // adding its size to the end PC upon execution. This is kind of
+                // pedantic but in the end we record the address being pushed
                 // onto the stack and that should be precise.
                 let next_address = self.pc + self.current_instruction.size as usize - 1;
-                let low = (next_address as u16 & 0x00FF) as u8;
-                let high = ((next_address as u16 & 0xFF00) >> 8) as u8;
-
-                self.push_stack(high, false)?;
-                self.push_stack(low, true)?;
+                self.push_address(next_address)?;
 
                 self.pc = address;
                 self.skip_pc = true;
@@ -980,24 +1053,84 @@ impl Machine {
                 //
                 // NOTE: as per 6502 specification, the address saved onto the
                 // stack was the next instruction before the call, but the size
-                // of the 'rts/rti' should also be accounted. Hence the + 1 to
-                // the resulting PC.
-                let low = self.pop_stack(false, false)? as u16;
-                let high = (self.pop_stack(false, true)? as u16) << 8;
+                // of the 'rts' should also be accounted. Hence the + 1 to the
+                // resulting PC.
+                let low = self.pop_stack(true, false)? as u16;
+                let high = (self.pop_stack(true, true)? as u16) << 8;
                 self.pc = (high + low) as usize + 1;
                 self.skip_pc = true;
             }
 
+            // Stack
+            InstructionIdentifier::Pha => self.push_stack(self.a, true)?,
+            InstructionIdentifier::Pla => self.a = self.pop_stack(true, true)?,
+            InstructionIdentifier::Php => self.push_status_register()?,
+            InstructionIdentifier::Plp => self.pop_status_register()?,
+
+            InstructionIdentifier::Brk => {
+                // The next address after a 'brk' is: pc + size of 'brk' (1) +
+                // break mark (1).
+                let next_address = self.pc + 2;
+                self.push_address(next_address)?;
+                self.push_status_register()?;
+
+                // The break mark is basically the byte which is in current PC +
+                // 1. If that's not possible, then we have a 'brk' as the last
+                // instruction with no break mark or something like that, which
+                // is just nonsense.
+                self.status_register.break_mark = *self
+                    .prg_rom
+                    .get(self.pc - 0x8000 + 1)
+                    .expect("you need to reserve a byte for the break mark");
+
+                // TODO: the next PC is the advertized IRQ one. This one is
+                // picked as the last byte from PRG-ROM, but depending on how
+                // bank mapping is done, this is not necessarily true.
+                let len = self.prg_rom.len();
+                let high = &(self.prg_rom[len - 1] as usize) << 8;
+                let low = &(self.prg_rom[len - 2] as usize);
+                self.pc = high + low;
+                self.skip_pc = true;
+            }
+            InstructionIdentifier::Rti => {
+                self.pop_status_register()?;
+
+                let low = self.pop_stack(true, false)? as usize;
+                let high = (self.pop_stack(true, true)? as usize) << 8;
+                self.pc = high + low;
+                self.skip_pc = true;
+            }
+
             // transfer
-            InstructionIdentifier::Tax => self.x = self.a,
-            InstructionIdentifier::Tay => self.y = self.a,
-            InstructionIdentifier::Tsx => self.x = self.s,
-            InstructionIdentifier::Txa => self.a = self.x,
+            InstructionIdentifier::Tax => {
+                self.x = self.a;
+                self.status_register.zero = self.x == 0;
+                self.status_register.negative = (self.x & 0x80) == 0x80;
+            }
+            InstructionIdentifier::Tay => {
+                self.y = self.a;
+                self.status_register.zero = self.y == 0;
+                self.status_register.negative = (self.y & 0x80) == 0x80;
+            }
+            InstructionIdentifier::Tsx => {
+                self.x = self.s;
+                self.status_register.zero = self.x == 0;
+                self.status_register.negative = (self.x & 0x80) == 0x80;
+            }
+            InstructionIdentifier::Txa => {
+                self.a = self.x;
+                self.status_register.zero = self.a == 0;
+                self.status_register.negative = (self.a & 0x80) == 0x80;
+            }
             InstructionIdentifier::Txs => {
                 self.s = self.x;
                 self.initial_stack_value = self.x;
             }
-            InstructionIdentifier::Tya => self.a = self.y,
+            InstructionIdentifier::Tya => {
+                self.a = self.y;
+                self.status_register.zero = self.a == 0;
+                self.status_register.negative = (self.a & 0x80) == 0x80;
+            }
 
             // other
             InstructionIdentifier::Bit => {
diff --git a/scripts/test-e2e.sh b/scripts/test-e2e.sh
index 4496d85..7893274 100755
--- a/scripts/test-e2e.sh
+++ b/scripts/test-e2e.sh
@@ -282,6 +282,12 @@ echo "test: runrom => arithlog.nes"
 diff tests/out/arithlog.txt tests/expected/runrom/arithlog.txt
 exit_code=$((exit_code + $?))
 
+echo "test: runrom => stack.nes"
+./target/debug/nasm -Werror -o tests/out/stack.nes tests/runrom/stack.s
+./target/debug/runrom --function --dump-memory tests/out/stack.nes > tests/out/stack.txt
+diff tests/out/stack.txt tests/expected/runrom/stack.txt
+exit_code=$((exit_code + $?))
+
 echo "test: runrom => misc.nes"
 ./target/debug/nasm -Werror --asan --write-info --allow-unused --out tests/out/misc.nes tests/runrom/misc.s
 ./target/debug/runrom --nasm .nasm --dump-memory --function tests/out/misc.nes > tests/out/misc-all.txt
@@ -297,6 +303,12 @@ diff tests/out/misc-foo.txt tests/expected/runrom/misc-foo.txt
 exit_code=$((exit_code + $?))
 
 ##
+# vnf-tests
+
+echo "test: vnf-tests"
+cargo run --bin vnf-tests tests/
+
+##
 # Done!
 
 exit $exit_code
diff --git a/tests/expected/runrom/stack.txt b/tests/expected/runrom/stack.txt
new file mode 100644
index 0000000..f35ea16
--- /dev/null
+++ b/tests/expected/runrom/stack.txt
@@ -0,0 +1,44 @@
+<start>		PC: $8000, cycles: 7, registers: [a: $00, x: $00, y: $00, sp: $FD], status: ----I--
+		[STACK]: 02 80 00 00 
+		[STACK]: 34 02 80 00 00 
+brk		PC: $8003, cycles: 14, registers: [a: $00, x: $00, y: $00, sp: $FA], status: --B-I--
+		[STACK]: 00 34 02 80 00 00 
+pha		PC: $8004, cycles: 17, registers: [a: $00, x: $00, y: $00, sp: $F9], status: --B-I--
+txa		PC: $8005, cycles: 19, registers: [a: $00, x: $00, y: $00, sp: $F9], status: --B-IZ-
+		[STACK]: 00 00 34 02 80 00 00 
+pha		PC: $8006, cycles: 22, registers: [a: $00, x: $00, y: $00, sp: $F8], status: --B-IZ-
+tya		PC: $8007, cycles: 24, registers: [a: $00, x: $00, y: $00, sp: $F8], status: --B-IZ-
+		[STACK]: 00 00 00 34 02 80 00 00 
+pha		PC: $8008, cycles: 27, registers: [a: $00, x: $00, y: $00, sp: $F7], status: --B-IZ-
+		[STACK]: 36 00 00 00 34 02 80 00 00 
+php		PC: $8009, cycles: 30, registers: [a: $00, x: $00, y: $00, sp: $F6], status: --B-IZ-
+lda #$01	PC: $800B, cycles: 32, registers: [a: $01, x: $00, y: $00, sp: $F6], status: --B-I--
+ldy #$02	PC: $800D, cycles: 34, registers: [a: $01, x: $00, y: $02, sp: $F6], status: --B-I--
+ldx #$FF	PC: $800F, cycles: 36, registers: [a: $01, x: $FF, y: $02, sp: $F6], status: N-B-I--
+inx		PC: $8010, cycles: 38, registers: [a: $01, x: $00, y: $02, sp: $F6], status: --B-IZ-
+inx		PC: $8011, cycles: 40, registers: [a: $01, x: $01, y: $02, sp: $F6], status: --B-I--
+		[STACK]: 00 00 00 34 02 80 00 00 
+plp		PC: $8012, cycles: 44, registers: [a: $01, x: $01, y: $02, sp: $F7], status: ----IZ-
+		[STACK]: 00 00 34 02 80 00 00 
+pla		PC: $8013, cycles: 48, registers: [a: $00, x: $01, y: $02, sp: $F8], status: ----IZ-
+tay		PC: $8014, cycles: 50, registers: [a: $00, x: $01, y: $00, sp: $F8], status: ----IZ-
+		[STACK]: 00 34 02 80 00 00 
+pla		PC: $8015, cycles: 54, registers: [a: $00, x: $01, y: $00, sp: $F9], status: ----IZ-
+tax		PC: $8016, cycles: 56, registers: [a: $00, x: $00, y: $00, sp: $F9], status: ----IZ-
+		[STACK]: 34 02 80 00 00 
+pla		PC: $8017, cycles: 60, registers: [a: $00, x: $00, y: $00, sp: $FA], status: ----IZ-
+		[STACK]: 02 80 00 00 
+		[STACK]: 00 00 
+rti		PC: $8002, cycles: 66, registers: [a: $00, x: $00, y: $00, sp: $FD], status: ----I--
+rts		PC: $8003, cycles: 72, registers: [a: $00, x: $00, y: $00, sp: $FD], status: ----I--
+<end>
+
+== Memory dump ==
+
+[$2F7] = $36 [reads=1, writes=1]
+[$2F8] = $00 [reads=1, writes=1]
+[$2F9] = $00 [reads=1, writes=1]
+[$2FA] = $00 [reads=1, writes=1]
+[$2FB] = $34 [reads=1, writes=1]
+[$2FC] = $02 [reads=1, writes=1]
+[$2FD] = $80 [reads=1, writes=1]
diff --git a/tests/runrom/stack.s b/tests/runrom/stack.s
new file mode 100644
index 0000000..ec15718
--- /dev/null
+++ b/tests/runrom/stack.s
@@ -0,0 +1,44 @@
+.segment "HEADER"
+    .byte 'N', 'E', 'S', $1A
+    .byte $02, $01
+    .byte $00
+    .byte $00
+
+.segment "CHARS"
+.byte 0
+
+.segment "CODE"
+;; asan:stack full
+
+.proc reset
+    brk
+    .byte $42
+    rts
+.endproc
+
+.proc irq
+    pha
+    txa
+    pha
+    tya
+    pha
+    php
+
+    lda #1
+    ldy #2
+    ldx #$FF
+    inx
+    inx
+
+    plp
+    pla
+    tay
+    pla
+    tax
+    pla
+
+    rti
+.endproc
+
+.segment "VECTORS"
+    .addr reset, reset, irq
diff --git a/tests/vnf-tests/Cargo.toml b/tests/vnf-tests/Cargo.toml
new file mode 100644
index 0000000..8ad6b70
--- /dev/null
+++ b/tests/vnf-tests/Cargo.toml
@@ -0,0 +1,10 @@
+[package]
+name = "vnf-tests"
+version = "0.1.0"
+edition = "2024"
+license.workspace = true
+authors.workspace = true
+
+[dependencies]
+vnf.workspace = true
+xixanta.workspace = true
diff --git a/tests/vnf-tests/src/main.rs b/tests/vnf-tests/src/main.rs
new file mode 100644
index 0000000..6c62dbd
--- /dev/null
+++ b/tests/vnf-tests/src/main.rs
@@ -0,0 +1,108 @@
+use std::path::PathBuf;
+
+use vnf::{Machine, MemoryPolicy};
+use xixanta::opcodes::InstructionIdentifier;
+
+#[derive(Default)]
+struct Args {
+    file: String,
+}
+
+fn print_help() {
+    println!("End-to-end tests for the vnf library\n");
+    println!("usage: vnf-tests [OPTIONS] <tests/ directory>\n");
+    println!("Options:");
+    println!("  -h, --help\t\t\tPrint this message and quit.");
+    std::process::exit(0);
+}
+
+// Print the given `message` and exit(1).
+fn die(message: String) -> ! {
+    eprintln!("error: {message}");
+    std::process::exit(1);
+}
+
+fn parse_arguments() -> Args {
+    let mut args = std::env::args();
+    let mut res = Args::default();
+
+    // Skip command name.
+    args.next();
+
+    for arg in args {
+        match arg.as_str() {
+            "-h" | "--help" => print_help(),
+            _ => {
+                if arg.starts_with('-') {
+                    die(format!("don't know how to handle the '{arg}' flag"));
+                }
+                if !res.file.is_empty() {
+                    die("cannot have multiple paths".to_string());
+                }
+                res.file = arg;
+            }
+        }
+    }
+
+    if res.file.is_empty() {
+        die("you have to provide a path to the tests/ directory".to_string());
+    }
+
+    res
+}
+
+fn run_break_mark_test(path: &String) -> Result<(), String> {
+    let rom = PathBuf::from(path).join("out/stack.nes");
+    let roms: String = rom.to_string_lossy().into_owned().to_string();
+
+    let mut machine = Machine::from(&roms, 0x8000, MemoryPolicy::default())?;
+
+    // Get out of <start>
+    let _ = machine.next_iteration();
+
+    // Assert that we are at the expected location.
+    assert!(matches!(
+        machine.current_instruction.identifier,
+        InstructionIdentifier::Brk
+    ));
+
+    // Before vs after running 'brk'.
+    assert_eq!(machine.status_register.break_mark, 0);
+    assert!(!machine.status_register.brk);
+    let _ = machine.next_iteration();
+    assert_eq!(machine.status_register.break_mark, 0x42);
+    assert!(machine.status_register.brk);
+
+    // Skip instructions we don't care about here.
+    for _ in 0..11 {
+        let _ = machine.next_iteration();
+    }
+
+    // Before running 'plp', everything is still as expected.
+    assert!(matches!(
+        machine.current_instruction.identifier,
+        InstructionIdentifier::Plp
+    ));
+    assert_eq!(machine.status_register.break_mark, 0x42);
+    assert!(machine.status_register.brk);
+
+    // After running 'plp'.
+    let _ = machine.next_iteration();
+    assert_eq!(machine.status_register.break_mark, 0x00);
+    assert!(!machine.status_register.brk);
+
+    Ok(())
+}
+
+fn main() {
+    let args = parse_arguments();
+    let file = PathBuf::from(args.file.clone());
+
+    if !file.is_dir() {
+        die("you have to provide a path to the tests/ directory".to_string());
+    }
+
+    if let Err(e) = run_break_mark_test(&args.file) {
+        die(e)
+    }
+}