4 files changed, 64 insertions, 3 deletions

diff --git a/src/assembler.rs b/src/assembler.rs
index 1b9c85b..73dec00 100644
--- a/src/assembler.rs
+++ b/src/assembler.rs
@@ -25,8 +25,8 @@ fn is_reg(s: &str) -> bool {
 
 fn instr_size(tokens: &[String]) -> u16 {
     match tokens[0].as_str() {
-        "mov" | "add" | "sub" | "jmp" | "jz" | "jnz" | "cmp" | "mul" | "div" => 3,
-        "hlt" => 1,
+        "mov" | "add" | "sub" | "jmp" | "jz" | "jnz" | "cmp" | "mul" | "div" | "call" => 3,
+        "hlt" | "ret" => 1,
         _ => panic!("Unknown instruction {}", tokens[0]),
     }
 }
@@ -172,6 +172,20 @@ pub fn assembler(source: &str) -> Vec<u8> {
                 bytes.push(r2);
             }
 
+            "call" => {
+                let addr = *symbols
+                    .get(&tokens[1])
+                    .expect("Unknown label");
+
+                bytes.push(Instruction::CALL as u8);
+                bytes.push((addr & 0xFF) as u8);   // low
+                bytes.push((addr >> 8) as u8);     // high
+            }
+
+            "ret" => {
+                bytes.push(Instruction::RET as u8);
+            }
+
             "hlt" => {
                 bytes.push(Instruction::HLT as u8);
             }
diff --git a/src/cpu.rs b/src/cpu.rs
index 9416783..1393e95 100644
--- a/src/cpu.rs
+++ b/src/cpu.rs
@@ -19,6 +19,22 @@ pub struct CPU {
 }
 
 impl CPU {
+    fn push16(&mut self, mem: &mut Memory, val: u16) {
+        self.dec_sp();
+        mem.write(self.sp, (val >> 8) as u8); // high
+        self.dec_sp();
+        mem.write(self.sp, (val & 0xFF) as u8); // low
+    }
+
+    fn pop16(&mut self, mem: &mut Memory) -> u16 {
+        let low = mem.read(self.sp) as u16;
+        self.inc_sp();
+        let high = mem.read(self.sp) as u16;
+        self.inc_sp();
+
+        (high << 8) | low
+    }
+
     pub fn debug_instr(&self, mem: &Memory) {
         let opcode = mem.read(self.pc);
 
@@ -53,6 +69,21 @@ impl CPU {
             x if x == Instruction::CMP_RR as u8 => self.cmp_rr(mem),
             x if x == Instruction::MUL as u8 => self.mul(mem),
             x if x == Instruction::DIV as u8 => self.div(mem),
+            x if x == Instruction::CALL as u8 => {
+                let low  = mem.read(self.pc) as u16; self.inc_pc();
+                let high = mem.read(self.pc) as u16; self.inc_pc();
+                let addr = (high << 8) | low;
+
+                let return_addr = self.pc;
+
+                self.push16(mem, return_addr);
+                self.pc = addr;
+
+            },
+            x if x == Instruction::RET as u8 => {
+                let addr = self.pop16(mem);
+                self.pc = addr;
+            },
             x if x == Instruction::HLT as u8 => self.halt(),
             _ => panic!("Unknown opcode {:02X}", opcode),
         }
@@ -62,6 +93,17 @@ impl CPU {
         self.pc += 1;
     }
 
+    pub fn inc_sp(&mut self) {
+        // self.sp += 1;
+        self.sp = self.sp.wrapping_add(1);
+    }
+
+    pub fn dec_sp(&mut self) {
+        // self.sp -= 1;
+        self.sp = self.sp.wrapping_sub(1);
+    }
+
+
     pub fn halt(&mut self) {
         self.halted = true;
     }
diff --git a/src/instructions.rs b/src/instructions.rs
index 9f81b9a..620b0ad 100644
--- a/src/instructions.rs
+++ b/src/instructions.rs
@@ -14,6 +14,8 @@ pub enum Instruction {
     CMP_RR = 0x09,
     MUL = 0x0C,
     DIV = 0x0D,
+    CALL = 0x0E,
+    RET = 0x0F,
     HLT = 0xFF,
 }
 
@@ -29,7 +31,10 @@ impl Instruction {
             0x07 | 0x09 => "CMP",
             0x0C => "MUL",
             0x0D => "DIV",
+            0x0E => "CALL",
+            0x0F => "RET",
             0xFF => "HLT",
+
             _ => "???",
         }
     }
diff --git a/src/main.rs b/src/main.rs
index 3dc0121..1c7fdb0 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -21,7 +21,7 @@ struct Args {
 fn main() {
     let args = Args::parse();
 
-    let mut cpu = CPU::default();
+    let mut cpu = CPU {sp: 0xFFFE, ..Default::default( )};
     let mut mem = Memory::new();
 
     let asm = std::fs::read_to_string(args.filename).unwrap();