add labels, improve documentation, add step debug

4 files changed, 94 insertions, 2 deletions

diff --git a/src/assembler.rs b/src/assembler.rs
index 3cd2bf8..3eeec64 100644
--- a/src/assembler.rs
+++ b/src/assembler.rs
@@ -1,4 +1,5 @@
 use crate::instructions::Instruction;
+use std::collections::HashMap;
 
 fn tokenize(line: &str) -> Vec<String> {
     line.split(|c| c == ' ' || c == ',' || c == '\t')
@@ -17,14 +18,51 @@ fn parse_reg(s: &str) -> u8 {
     }
 }
 
+fn instr_size(tokens: &[String]) ->  u16 {
+    match tokens[0].as_str() {
+        "mov" | "add" | "sub" | "jmp" | "jz" | "jnz" => 3,
+        "hlt" => 1,
+        _ => panic!("Unknown instruction {}", tokens[0])
+
+    }
+}
+
+
+fn first_pass(source: &str) -> HashMap<String, u16> {
+    let mut symbols = HashMap::new();
+    let mut addr: u16 = 0;
+
+    for line in source.lines() {
+        let line = line.trim();
+
+        // Ignoring comments and empty lines
+        if line.is_empty() || line.starts_with(";") {
+            continue;
+        }
+
+        // Labels (ends with ":")
+        if line.ends_with(":") {
+            let label = line.trim_end_matches(":").to_string();
+            symbols.insert(label, addr);
+            continue;
+        }
+
+        let tokens = tokenize(line);
+        addr += instr_size(&tokens);
+    }
+
+    symbols
+}
+
 pub fn assembler(source: &str) -> Vec<u8> {
+    let symbols = first_pass(source);
     let mut bytes = Vec::new();
 
     for (line_no, line) in source.lines().enumerate() {
         let line = line.trim();
 
         // Comments in assembly start with ";"
-        if line.is_empty() || line.starts_with(';') {
+        if line.is_empty() || line.starts_with(';') || line.ends_with(":") {
             continue;
         }
 
@@ -61,6 +99,22 @@ pub fn assembler(source: &str) -> Vec<u8> {
                 bytes.push(r2);
             }
 
+            "jmp" | "jz" | "jnz" => {
+                let opcode = match tokens[0].as_str() {
+                    "jmp" => Instruction::JMP,
+                    "jz" => Instruction::JZ,
+                    "jnz" => Instruction::JNZ,
+                    _ => unreachable!()
+                };
+
+                let label = &tokens[1];
+                let addr = *symbols.get(label).expect("Uknown label");
+
+                bytes.push(opcode as u8);
+                bytes.push((addr & 0xFF) as u8); // low
+                bytes.push((addr >> 8) as u8); // high
+            }
+
             "hlt" => {
                 bytes.push(Instruction::HLT as u8);
             }
diff --git a/src/cpu.rs b/src/cpu.rs
index 13d2fc8..df8b3a2 100644
--- a/src/cpu.rs
+++ b/src/cpu.rs
@@ -2,6 +2,7 @@ use crate::instructions::Instruction;
 use crate::memory::Memory;
 
 #[derive(Default, Debug)]
+#[allow(dead_code)]
 pub struct CPU {
     pub a: u8,
     pub b: u8,
@@ -18,6 +19,23 @@ pub struct CPU {
 }
 
 impl CPU {
+
+    pub fn debug_instr(&self, mem: &Memory) {
+        let opcode = mem.read(self.pc);
+
+        println!(
+            "PC={:04X} {:<3} | A={:02X} B={:02X} C={:02X} D={:02X} | Z={} C={}",
+            self.pc,
+            Instruction::opcode_name(opcode),
+            self.a,
+            self.b,
+            self.c,
+            self.d,
+            self.zero as u8,
+            self.carry as u8
+        );
+    }
+
     pub fn step(&mut self, mem: &mut Memory) {
         let opcode = mem.read(self.pc);
         self.inc_pc();
diff --git a/src/instructions.rs b/src/instructions.rs
index acfb8bd..ac88e94 100644
--- a/src/instructions.rs
+++ b/src/instructions.rs
@@ -8,3 +8,18 @@ pub enum Instruction {
     JNZ = 0x06,
     HLT = 0xFF,
 }
+
+impl Instruction {
+    pub  fn opcode_name(op: u8) -> &'static str{
+        match op {
+            0x01 => "MOV",
+            0x02 => "ADD",
+            0x03 => "SUB",
+            0x04 => "JMP",
+            0x05 => "JZ",
+            0x06 => "JNZ",
+            0xFF => "HLT",
+            _ => "???",
+        }
+    }
+}
diff --git a/src/main.rs b/src/main.rs
index 607fb34..3dc0121 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -3,6 +3,8 @@ mod cpu;
 mod instructions;
 mod memory;
 
+use std::io;
+
 use cpu::CPU;
 use memory::Memory;
 use clap::Parser;
@@ -30,7 +32,10 @@ fn main() {
     }
 
     while !cpu.halted {
+        cpu.debug_instr(&mem);
         cpu.step(&mut mem);
-        println!("{:?}", cpu);
+
+        let mut buf = String::new();
+        io::stdin().read_line(&mut buf).unwrap();
     }
 }