feat(assembler): Add support for two pass collapse instructions and implement subroutine instructions

4 files changed, 84 insertions, 18 deletions

diff --git a/assembler/src/parser/code_gen.rs b/assembler/src/parser/code_gen.rs
index ec4a2d2..2eaa34c 100644
--- a/assembler/src/parser/code_gen.rs
+++ b/assembler/src/parser/code_gen.rs
@@ -3,7 +3,6 @@ use super::{ Program, Address };
 
 impl Program {
     /* Intermediate code generation */
-
     pub(crate) fn collapse_instructions(&mut self) {
         // Fix instruction addresses
         let mut offsets = Vec::new();
@@ -14,8 +13,14 @@ impl Program {
 
         for instr in &mut self.instructions {
             for arg in &mut instr.args {
-                if let Address::Instruction(num) = arg {
-                    *arg = Address::Instruction(offsets[*num]);
+                match arg {
+                    Address::Instruction(num)
+                        => *arg = Address::Instruction(offsets[*num]),
+                    Address::InstructionOffset(num, offset)
+                        => *arg = Address::InstructionOffset(offsets[*num], *offset),
+                    Address::InstructionConstant(num)
+                        => *arg = Address::InstructionConstant(offsets[*num]),
+                    _ => {},
                 }
             }
         }
@@ -32,9 +37,27 @@ impl Program {
 
     /* Machine code generation */
 
-    pub(crate) fn instructions_to_constants(&mut self, mut absolute_start: usize) {
-        absolute_start += 1;
-        let constlen = self.constants.len();
+    pub(crate) fn instructions_to_constants(&mut self) {
+        let absolute_start = self.absolute_start + 1;
+
+        let mut constlen = self.constants.len();
+        for instr in &self.instructions {
+            for arg in &instr.args {
+                if let Address::InstructionConstant(_) = arg {
+                    constlen += 1;
+                }
+            }
+        }
+
+        for instr in &mut self.instructions {
+            for arg in &mut instr.args {
+                if let Address::InstructionConstant(num) = arg {
+                    self.constants.push((constlen + absolute_start + *num * 3) as Word);
+                    *arg = Address::Constants(self.constants.len() - 1);
+                }
+            }
+        }
+
         self.constants.insert(0, (absolute_start + constlen) as Word);
 
         for instr in &mut self.instructions {
@@ -42,7 +65,10 @@ impl Program {
                 self.constants.push(match arg {
                     Address::Memory(num) => *num,
                     Address::Constants(num) => absolute_start + *num,
-                    Address::Instruction(num) => constlen + absolute_start + *num * 3,
+                    Address::Instruction(num)
+                        => constlen + absolute_start + *num * 3,
+                    Address::InstructionOffset(num, off)
+                        => *off + constlen + absolute_start + *num * 3,
                     _ => unreachable!("Definition addresses should've already been removed!"),
                 } as Word);
             }
@@ -50,9 +76,9 @@ impl Program {
         self.instructions.clear();
     }
 
-    pub(crate) fn to_bytes(self, absolute_start: usize, pc: usize) -> Vec<u8> {
+    pub(crate) fn to_bytes(self, pc: usize) -> Vec<u8> {
         let mut res = Vec::new();
-        res.resize(absolute_start, Default::default());
+        res.resize(self.absolute_start, Default::default());
         res[pc] = self.constants[0];
         res.extend(self.constants);
         res
diff --git a/assembler/src/parser/instruction.rs b/assembler/src/parser/instruction.rs
index 5e19b7e..3ccd309 100644
--- a/assembler/src/parser/instruction.rs
+++ b/assembler/src/parser/instruction.rs
@@ -1,8 +1,9 @@
-use super::{ Instruction, Address };
+use super::{ Instruction, Address, Word };
 
 static Z: Address = Address::Memory(0);
 static E: Address = Address::Memory(1);
 static N: Address = Address::Memory(2);
+static STACK_POINTER: Address = Address::Memory(4);
 
 #[derive(Debug)]
 pub(crate) enum ArgumentError {
@@ -15,14 +16,21 @@ pub(crate) enum ArgumentError {
 }
 
 impl Instruction {
+    fn new1(name: &str, a: &Address) -> Instruction {
+        Instruction { name: name.to_string(), args: vec![a.clone()] }
+    }
+    fn new2(name: &str, a: &Address, b: &Address) -> Instruction {
+        Instruction { name: name.to_string(), args: vec![a.clone(), b.clone()] }
+    }
+    fn new3(name: &str, a: &Address, b: &Address, c: &Address) -> Instruction {
+        Instruction { name: name.to_string(), args: vec![a.clone(), b.clone(), c.clone()] }
+    }
+
     fn new_subleq(a: &Address, b: &Address, c: &Address) -> Instruction {
         if let Address::Constants(_) = c {
             panic!("Constants cannot be the third subleq argument!");
         }
-        Instruction {
-            name: "subleq".to_string(),
-            args: vec![a.clone(), b.clone(), c.clone()],
-        }
+        Self::new3("subleq", a, b, c)
     }
 
     fn check_args_len(&self, expected: usize) -> Result<(), ArgumentError> {
@@ -46,7 +54,7 @@ impl Instruction {
                     }
                 }
             },
-            "jmp" => {
+            "jmp" | "jsr" => {
                 self.check_args_len(1)?;
 
                 if ! self.args[0].is_location() {
@@ -69,6 +77,9 @@ impl Instruction {
                     return Err(ArgumentError::NotLocation);
                 }
             },
+            "rts" => {
+                self.check_args_len(0)?;
+            },
             _ => return Err(ArgumentError::UnknownInstruction),
         }
         Ok(())
@@ -78,6 +89,7 @@ impl Instruction {
         self.name = self.name.to_lowercase();
 
         match &*self.name {
+            /* Primitive instrutions (need only one collapse) */
             "subleq" => vec![self.clone()],
             "jmp" => vec![
                 Instruction::new_subleq(&Z, &Z, &self.args[0]),
@@ -105,6 +117,29 @@ impl Instruction {
                 Instruction::new_subleq(&Z,            &Z,            &Address::Instruction(next + 3)),
                 Instruction::new_subleq(&Z,            &self.args[0], &self.args[1]                  ),
             ],
+
+            /* Composite instructions (require two collapses) */
+            "jsr" => {
+                vec![
+                    Instruction::new1("inc", &STACK_POINTER),
+
+                    // We're modifying the generated mov instruction
+                    // Every &self.args[1] value gets replaced with the value which SP points to
+                    Instruction::new2("mov", &STACK_POINTER, &Address::InstructionOffset(next + 4, 0)),
+                    Instruction::new2("mov", &STACK_POINTER, &Address::InstructionOffset(next + 4, 1)),
+                    Instruction::new2("mov", &STACK_POINTER, &Address::InstructionOffset(next + 4, 7)),
+
+                    Instruction::new2("mov", &Address::InstructionConstant(next + 6), &Z),
+                    Instruction::new1("jmp", &self.args[0]),
+                ]
+            },
+            "rts" => vec![
+                Instruction::new2("mov", &STACK_POINTER, &Address::InstructionOffset(next + 1, 3)),
+                Instruction::new2("mov", &Z,             &Address::InstructionOffset(next + 3, 2)),
+                Instruction::new1("dec", &STACK_POINTER),
+                Instruction::new1("jmp", &Z),
+            ],
+
             _ => panic!("Unknown instruction {}!", self.name),
         }
     }
@@ -115,7 +150,7 @@ impl Instruction {
                 => 1,
             "add"
                 => 3,
-            "mov" | "beq"
+            "mov" | "beq" | "jsr" | "rts"
                 => 4,
             _ => panic!("Unknown instruction {}!", self.name),
         }
diff --git a/assembler/src/parser/mod.rs b/assembler/src/parser/mod.rs
index f93875b..f22fc36 100644
--- a/assembler/src/parser/mod.rs
+++ b/assembler/src/parser/mod.rs
@@ -14,6 +14,8 @@ pub(crate) struct Program {
     definitions: HashMap<String, Address>,
     constants: Vec<Word>,
     instructions: Vec<Instruction>,
+
+    absolute_start: usize,
 }
 
 impl Program {
@@ -43,6 +45,9 @@ pub(crate) enum Address {
     Constants(usize),
     Instruction(usize),
     Definitions(String),
+
+    InstructionOffset(usize, usize), // Internal
+    InstructionConstant(usize), // Internal
     #[default]
     None,
 }
diff --git a/assembler/src/parser/syntax.rs b/assembler/src/parser/syntax.rs
index e6eb810..563a67e 100644
--- a/assembler/src/parser/syntax.rs
+++ b/assembler/src/parser/syntax.rs
@@ -2,8 +2,8 @@ use crate::{ Word, lexer::Token };
 use super::{ Program, Address, Instruction };
 
 impl Program {
-    pub(crate) fn new(tokens: Vec<Token>) -> Program {
-        let mut prog = Program {..Default::default() };
+    pub(crate) fn new(tokens: Vec<Token>, absolute_start: usize) -> Program {
+        let mut prog = Program {absolute_start, ..Default::default() };
 
         let mut op_name = String::new();
         let mut op_args = Vec::new();