Properly fill the context stack when forcing it

When calling `force_context_switch` the stack was mindlessly pushing the
given name without taking into consideration how scopes are to be laid
out. This made some variables/addresses that were previously preserved
no longer reachable when crunching pending nodes.

This patch also makes `force_context_switch` reset the stack before
doing anything at all, which means that `force_context_pop` was no
longer relevant.

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

2 files changed, 67 insertions, 16 deletions

diff --git a/lib/xixanta/src/assembler.rs b/lib/xixanta/src/assembler.rs
index f190d33..06f0b24 100644
--- a/lib/xixanta/src/assembler.rs
+++ b/lib/xixanta/src/assembler.rs
@@ -587,8 +587,6 @@ impl<'a> Assembler<'a> {
                 }
                 Err(e) => errors.push(e),
             }
-
-            self.context.force_context_pop();
         }
 
         if errors.is_empty() {
@@ -2936,6 +2934,54 @@ cpx #(4 * var2)"#,
         assert_instruction("stx $020, y", &[0x96, 0x20]);
     }
 
+    #[test]
+    fn jsr_inside_multiple_contexts() {
+        let res = just_bundles(
+            r#".scope Movement
+    .proc accelerate
+        rts
+    .endproc
+
+    .proc update
+        jsr accelerate
+        rts
+    .endproc
+.endscope
+
+jsr Movement::accelerate
+jsr Movement::update
+    "#,
+        );
+
+        assert_eq!(res.len(), 5);
+
+        // accelerate -> rts
+        assert_eq!(res[0].size, 1);
+        assert_eq!(res[0].bytes[0], 0x60);
+
+        // jsr accelerate
+        assert_eq!(res[1].size, 3);
+        assert_eq!(res[1].bytes[0], 0x20);
+        assert_eq!(res[1].bytes[1], 0x00);
+        assert_eq!(res[1].bytes[2], 0x80);
+
+        // update -> rts
+        assert_eq!(res[2].size, 1);
+        assert_eq!(res[2].bytes[0], 0x60);
+
+        // jsr Movement::accelerate
+        assert_eq!(res[3].size, 3);
+        assert_eq!(res[3].bytes[0], 0x20);
+        assert_eq!(res[3].bytes[1], 0x00);
+        assert_eq!(res[3].bytes[2], 0x80);
+
+        // jsr Movement::update
+        assert_eq!(res[4].size, 3);
+        assert_eq!(res[4].bytes[0], 0x20);
+        assert_eq!(res[4].bytes[1], 0x01);
+        assert_eq!(res[4].bytes[2], 0x80);
+    }
+
     // Control statements
 
     #[test]
diff --git a/lib/xixanta/src/object.rs b/lib/xixanta/src/object.rs
index ad02747..a92f9fb 100644
--- a/lib/xixanta/src/object.rs
+++ b/lib/xixanta/src/object.rs
@@ -339,20 +339,25 @@ impl Context {
     }
 
     /// Change the current context to the given one identified by `name`,
-    /// disregarding any check. This is to be used when switching a context to
-    /// set a very specific value for that context. You should call
-    /// `force_context_pop` immediately.
-    pub fn force_context_switch(&mut self, name: &String) {
-        self.stack.push(name.to_owned());
-    }
-
-    /// Remove the last context being used if any. In contrast with
-    /// `context_pop`, this one does not error out, but does nothing in case we
-    /// are in the global context. This is to be used in conjunction with
-    /// `force_context_switch`.
-    pub fn force_context_pop(&mut self) {
-        if !self.stack.is_empty() {
-            self.stack.truncate(self.stack.len() - 1);
+    /// disregarding any check. This is a destructive operation and will clear
+    /// out any previous context. Hence, only call this when you are sure that
+    /// there's no reliance on a specific context in the future.
+    pub fn force_context_switch(&mut self, name: &str) {
+        let mut ax = String::with_capacity(name.len());
+
+        // Clear the context stack as we will set it manually.
+        self.stack.truncate(0);
+
+        // Forcing a context switch is not as simple as pushing the given name,
+        // but we have to make sure that the whole scope hierarchy is laid out
+        // so the code that fetches variables continues to work (e.g. if
+        // "A::B::C", then the stack must look like ["A", "A::B", "A::B::C"]).
+        for n in name.split("::") {
+            if !ax.is_empty() {
+                ax += "::";
+            }
+            ax += n;
+            self.stack.push(ax.clone());
         }
     }