add add_ri, sub_ri

5 files changed, 101 insertions, 29 deletions

diff --git a/README.md b/README.md
index 18a02e8..d5fc644 100644
--- a/README.md
+++ b/README.md
@@ -4,15 +4,15 @@
 - Word Size
     - **Data Width:** 8 bits
     - **Address width:** 16 bits
-    - **Address space:** 64 KB (0x0000- 0xFFFF)
+    - **Address space:** 64 KB (0x0000-0xFFFF)
 
 ## Supported Instructions
 
 | Instruction | Syntax       |
 | ----------- | ------------ |
 | MOV         | mov reg, imm OR mov reg, reg |
-| ADD         | add r1, r2   |
-| SUB         | sub r1, r2   |
+| ADD         | add reg, imm OR add r1, r2   |
+| SUB         | sub reg, imm OR add r1, r2   |
 | JMP         | jmp addr     |
 | JZ          | jz addr      |
 | JNZ         | jnz addr     |
diff --git a/examples/addition.asm b/examples/addition.asm
index 9914dee..ea2d317 100644
--- a/examples/addition.asm
+++ b/examples/addition.asm
@@ -3,6 +3,7 @@ mov a, 10
 ; move imm 5 to register B
 mov b, 5
 ; Add a & B
-add a, b
+add a, 2
+sub a, 2
 ; Halt
 hlt
diff --git a/src/assembler.rs b/src/assembler.rs
index c4751c1..e7aa305 100644
--- a/src/assembler.rs
+++ b/src/assembler.rs
@@ -2,12 +2,13 @@ use crate::instructions::Instruction;
 use std::collections::HashMap;
 
 fn tokenize(line: &str) -> Vec<String> {
-    line.split(|c| c == ' ' || c == ',' || c == '\t')
+    line.split([' ', ',', '\t'])
         .filter(|s| !s.is_empty())
         .map(|s| s.to_lowercase())
         .collect()
 }
 
+// Returns corrensponding u8 value of the register
 fn parse_reg(s: &str) -> u8 {
     match s {
         "a" => 0,
@@ -19,10 +20,7 @@ fn parse_reg(s: &str) -> u8 {
 }
 
 fn is_reg(s: &str) -> bool {
-    match s {
-        "a" | "b" | "c" | "d" => true,
-        _ => false,
-    }
+    matches!(s, "a" | "b" | "c" | "d")
 }
 
 fn instr_size(tokens: &[String]) -> u16 {
@@ -76,16 +74,16 @@ pub fn assembler(source: &str) -> Vec<u8> {
         match tokens[0].as_str() {
             "mov" => {
                 // mov reg, imm
-                let reg = parse_reg(&tokens[1]);
+                let r1= parse_reg(&tokens[1]);
                 if is_reg(&tokens[2]) {
                     let r2 = parse_reg(&tokens[2]);
                     bytes.push(Instruction::MOV_RR as u8);
-                    bytes.push(reg);
+                    bytes.push(r1);
                     bytes.push(r2);
                 } else {
                     let imm: u8 = tokens[2].parse().unwrap();
                     bytes.push(Instruction::MOV_RI as u8);
-                    bytes.push(reg);
+                    bytes.push(r1);
                     bytes.push(imm);
                 }
             }
@@ -93,21 +91,35 @@ pub fn assembler(source: &str) -> Vec<u8> {
             "add" => {
                 // add a, b
                 let r1 = parse_reg(&tokens[1]);
-                let r2 = parse_reg(&tokens[2]);
 
-                bytes.push(Instruction::ADD as u8);
-                bytes.push(r1);
-                bytes.push(r2);
+                if is_reg(&tokens[2]) {
+                    let r2 = parse_reg(&tokens[2]);
+                    bytes.push(Instruction::ADD_RR as u8);
+                    bytes.push(r1);
+                    bytes.push(r2);
+                } else {
+                    let imm: u8 = tokens[2].parse().unwrap();
+                    bytes.push(Instruction::ADD_RI as u8);
+                    bytes.push(r1);
+                    bytes.push(imm);
+                }
+
             }
 
             "sub" => {
                 // sub a, b
-                let r2 = parse_reg(&tokens[2]);
                 let r1 = parse_reg(&tokens[1]);
-
-                bytes.push(Instruction::SUB as u8);
-                bytes.push(r1);
-                bytes.push(r2);
+                if is_reg(&tokens[2]) {
+                    let r2 = parse_reg(&tokens[2]);
+                    bytes.push(Instruction::SUB_RR as u8);
+                    bytes.push(r1);
+                    bytes.push(r2);
+                } else {
+                    let imm: u8 = tokens[2].parse().unwrap();
+                    bytes.push(Instruction::SUB_RI as u8);
+                    bytes.push(r1);
+                    bytes.push(imm);
+                }
             }
 
             "jmp" | "jz" | "jnz" => {
diff --git a/src/cpu.rs b/src/cpu.rs
index 116d20d..d0365ba 100644
--- a/src/cpu.rs
+++ b/src/cpu.rs
@@ -42,8 +42,10 @@ impl CPU {
         match opcode {
             x if x == Instruction::MOV_RI as u8 => self.mov_ri(mem),
             x if x == Instruction::MOV_RR as u8 => self.mov_rr(mem),
-            x if x == Instruction::ADD as u8 => self.add(mem),
-            x if x == Instruction::SUB as u8 => self.sub(mem),
+            x if x == Instruction::ADD_RR as u8 => self.add_rr(mem),
+            x if x == Instruction::ADD_RI as u8 => self.add_ri(mem),
+            x if x == Instruction::SUB_RR as u8 => self.sub_rr(mem),
+            x if x == Instruction::SUB_RI as u8 => self.sub_ri(mem),
             x if x == Instruction::JMP as u8 => self.jmp(mem),
             x if x == Instruction::JZ as u8 => self.jz(mem),
             x if x == Instruction::JNZ as u8 => self.jnz(mem),
@@ -92,7 +94,7 @@ impl CPU {
     }
 
 
-    pub fn add(&mut self, mem: &Memory) {
+    pub fn add_rr(&mut self, mem: &Memory) {
         let dest = mem.read(self.pc);
         self.pc += 1;
         let src = mem.read(self.pc);
@@ -133,8 +135,35 @@ impl CPU {
         self.zero = result == 0;
         self.carry = carry;
     }
+    pub fn add_ri(&mut self, mem: &Memory) {
 
-    pub fn sub(&mut self, mem: &Memory) {
+        let dest = mem.read(self.pc);
+        self.pc += 1;
+        let imm = mem.read(self.pc);
+        self.pc += 1;
+
+        let (result, carry) = match dest {
+            0 => self.a.overflowing_add(imm),
+            1 => self.b.overflowing_add(imm),
+            2 => self.c.overflowing_add(imm),
+            3 => self.d.overflowing_add(imm),
+            _ => (0, false),
+        };
+
+        match dest {
+            0 => self.a = result,
+            1 => self.b = result,
+            2 => self.c = result,
+            3 => self.d = result,
+            _ => {}
+        }
+
+        self.zero = result == 0;
+        self.carry = carry;
+    }
+
+
+    pub fn sub_rr(&mut self, mem: &Memory) {
         let dest = mem.read(self.pc);
         self.pc += 1;
         let src = mem.read(self.pc);
@@ -176,6 +205,33 @@ impl CPU {
         self.carry = borrow;
     }
 
+ pub fn sub_ri(&mut self, mem: &Memory) {
+        let dest = mem.read(self.pc);
+        self.pc += 1;
+        let imm = mem.read(self.pc);
+        self.pc += 1;
+
+        let (result, borrow) = match dest {
+            0 => self.a.overflowing_sub(imm),
+            1 => self.b.overflowing_sub(imm),
+            2 => self.c.overflowing_sub(imm),
+            3 => self.d.overflowing_sub(imm),
+            _ => (0, false),
+        };
+
+        match dest {
+            0 => self.a = result,
+            1 => self.b = result,
+            2 => self.c = result,
+            3 => self.d = result,
+            _ => {}
+        }
+
+        self.zero = result == 0;
+        self.carry = borrow;
+    }
+
+
     pub fn jmp(&mut self, mem: &Memory) {
         let low = mem.read(self.pc) as u16;
         self.inc_pc();
diff --git a/src/instructions.rs b/src/instructions.rs
index 77de7f4..bfc4fc6 100644
--- a/src/instructions.rs
+++ b/src/instructions.rs
@@ -1,9 +1,12 @@
 #[repr(u8)]
+#[allow(non_camel_case_types)]
 pub enum Instruction {
     MOV_RI = 0x01,
     MOV_RR = 0x08,
-    ADD = 0x02,
-    SUB = 0x03,
+    ADD_RR = 0x02,
+    ADD_RI = 0x0A,
+    SUB_RR = 0x03,
+    SUB_RI = 0x0B,
     JMP = 0x04,
     JZ = 0x05,
     JNZ = 0x06,
@@ -16,8 +19,8 @@ impl Instruction {
     pub  fn opcode_name(op: u8) -> &'static str{
         match op {
             0x01 | 0x08 => "MOV",
-            0x02 => "ADD",
-            0x03 => "SUB",
+            0x02 | 0x0A => "ADD",
+            0x03 | 0x0B => "SUB",
             0x04 => "JMP",
             0x05 => "JZ",
             0x06 => "JNZ",