From a9f50efdef4008822287911025e0f1783dcc9638 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Miquel=20Sabat=C3=A9=20Sol=C3=A0?= <mikisabate@gmail.com>
Date: Wed, 18 Dec 2024 15:33:35 +0100
Subject: [PATCH] Fix label references on control statements
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: Miquel Sabaté Solà <mikisabate@gmail.com>
---
 lib/xixanta/src/assembler.rs | 48 +++++++++++++++++++++++++++++++++++-
 1 file changed, 47 insertions(+), 1 deletion(-)

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 => {
@@ -2520,6 +2521,51 @@ code:
         assert_eq!(bundles[2].bytes[2], 0x80);
     }
 
+    #[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());