Fix label references on control statements

Signed-off-by: Miquel Sabaté Solà <mikisabate@gmail.com>

1 files changed, 47 insertions, 1 deletions

diff --git a/lib/xixanta/src/assembler.rs b/lib/xixanta/src/assembler.rs
index f2cd507..7b5fb65 100644
--- a/lib/xixanta/src/assembler.rs
+++ b/lib/xixanta/src/assembler.rs
@@ -878,6 +878,7 @@ impl Assembler {
             Some(args) => {
                 for arg in args {
                     // Evaluate the argument as a node.
+                    self.literal_mode = None;
                     let mut bundle = self.evaluate_node(arg)?;
 
                     // If there is a missmatch between the expected number of
@@ -904,7 +905,7 @@ impl Assembler {
                             _ => panic!("bad argument when evaluating arguments"),
                         }
                     }
-                    self.push_bundle(bundle, node)?;
+                    self.push_bundle(bundle, arg)?;
                 }
             }
             None => {
@@ -2521,6 +2522,51 @@ code:
     }
 
     #[test]
+    fn proc_reference_another_segment() {
+        let mut asm = Assembler::new(one_two().to_vec());
+        asm.mappings[0].segments[0].bundles = minimal_header();
+        asm.mappings[0].offset = 6;
+        let bundles = &asm
+            .assemble(
+                r#"
+.segment "ONE"
+.addr code
+
+.segment "TWO"
+nop
+code:
+    rts
+"#
+                .as_bytes(),
+            )
+            .unwrap()[0x10..]; // Ignoring HEADER
+
+        assert_eq!(bundles[0].size, 2);
+        assert_eq!(bundles[0].bytes[0], 0x03);
+        assert_eq!(bundles[0].bytes[1], 0x80);
+    }
+
+    #[test]
+    fn cannot_fit_address_in_one_byte() {
+        let mut asm = Assembler::new(one_two().to_vec());
+        assert_error_with_assembler(
+            &mut asm,
+            r#".segment "ONE"
+.byte code
+
+.segment "TWO"
+nop
+code:
+    rts
+"#,
+            "Evaluation",
+            2,
+            false,
+            "expecting an argument that fits into a byte",
+        );
+    }
+
+    #[test]
     fn cannot_switch_to_segment_inside_of_scope() {
         let mut asm = Assembler::new(EMPTY.to_vec());
         asm.mappings[0].segments[0].bundles = minimal_header();