feat(assembler): Add instruction collapse code gen step

1 files changed, 21 insertions, 1 deletions

diff --git a/assembler/src/parser/code_gen.rs b/assembler/src/parser/code_gen.rs
index 375a8cb..dd1b872 100644
--- a/assembler/src/parser/code_gen.rs
+++ b/assembler/src/parser/code_gen.rs
@@ -1,4 +1,4 @@
-use super::{ Program, Address };
+use super::{ Program, Address, Instruction };
 
 impl Program {
     pub(crate) fn remove_definition_addresses(&mut self) {
@@ -10,4 +10,24 @@ impl Program {
             }
         }
     }
+
+    pub(crate) fn collapse_instructions(&mut self) {
+        let mut collapsed = Vec::new();
+        let mut update_addrs = Vec::new();
+
+        for instr in &mut self.instructions {
+            update_addrs.push(collapsed.len());
+            collapsed.extend(instr.collapse());
+        }
+
+        for instr in &mut self.instructions {
+            for arg in &mut instr.args {
+                if let Address::Instruction(num) = arg {
+                    *arg = Address::Instruction(update_addrs[*num]);
+                }
+            }
+        }
+
+        self.instructions = collapsed;
+    }
 }