2 files changed, 63 insertions, 0 deletions
diff --git a/assembler/src/parser/code_gen.rs b/assembler/src/parser/code_gen.rs
index 2eaa34c..d183145 100644
--- a/assembler/src/parser/code_gen.rs
+++ b/assembler/src/parser/code_gen.rs
@@ -88,3 +88,32 @@ impl Program {
             .collect()
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::lexer::Token;
+
+    #[test]
+    fn after_two_collapses_only_subleq_remains() {
+        let mut prog = Program::new(vec![
+            Token::Operator("rts".to_string()),
+            Token::Operator("jsr".to_string()),
+            Token::Address(1000),
+            Token::Operator("mov".to_string()),
+            Token::Address(1001),
+            Token::Address(1002),
+            Token::Operator("subleq".to_string()),
+            Token::Address(1001),
+            Token::Address(1002),
+            Token::Address(1001),
+
+        ], 10);
+        prog.collapse_instructions();
+        prog.collapse_instructions();
+
+        for i in prog.instructions {
+            assert_eq!(i.name, "subleq");
+        }
+    }
+}
diff --git a/assembler/src/parser/semantic.rs b/assembler/src/parser/semantic.rs
index 47ba201..438b2f1 100644
--- a/assembler/src/parser/semantic.rs
+++ b/assembler/src/parser/semantic.rs
@@ -21,3 +21,37 @@ impl Program {
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::lexer::Token;
+
+    #[test]
+    fn remove_definition_addresses_removes_symbols() {
+        let mut prog = Program::new(vec![
+            Token::SymbolDef("hello".to_string()),
+            Token::Number(1000),
+            Token::SymbolDef("world".to_string()),
+            Token::Number(64),
+            Token::Operator("jsr".to_string()),
+            Token::Symbol("hello".to_string()),
+            Token::Operator("mov".to_string()),
+            Token::Address(1001),
+            Token::Symbol("world".to_string()),
+            Token::Operator("subleq".to_string()),
+            Token::Symbol("hello".to_string()),
+            Token::Address(1002),
+            Token::Symbol("world".to_string()),
+
+        ], 10);
+        prog.remove_definition_addresses();
+
+        assert_eq!(prog.definitions.len(), 0);
+        for i in prog.instructions {
+            for arg in i.args {
+                assert!(! matches!(arg, Address::Definitions(_)));
+            }
+        }
+    }
+}