Find variable values on parent scopes

If a given variable cannot be found on the current scope, attempt to go
up the context hierarchy to find it.

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

2 files changed, 99 insertions, 12 deletions

diff --git a/lib/xixanta/src/assembler.rs b/lib/xixanta/src/assembler.rs
index 4e33d66..888e122 100644
--- a/lib/xixanta/src/assembler.rs
+++ b/lib/xixanta/src/assembler.rs
@@ -1775,6 +1775,40 @@ adc Variable
     }
 
     #[test]
+    fn reference_outer_variables() {
+        let mut asm = Assembler::new(EMPTY.to_vec());
+        asm.mappings[0].segments[0].bundles = minimal_header();
+        asm.mappings[0].offset = 6;
+        asm.current_mapping = 1;
+        let res = &asm
+            .assemble(
+                std::env::current_dir().unwrap().to_path_buf(),
+                r#"
+foo:
+  rts
+
+.proc inner
+    jsr foo
+.endproc
+"#
+                .as_bytes(),
+            )
+            .unwrap()[0x10..];
+
+        assert_eq!(res.len(), 2);
+
+        // foo: rts
+        assert_eq!(res[0].size, 1);
+        assert_eq!(res[0].bytes[0], 0x60);
+
+        // inner: jsr foo
+        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);
+    }
+
+    #[test]
     fn bad_variable_but_valid_identifier_in_instruction() {
         assert_eval_error(
             "adc Variable",
diff --git a/lib/xixanta/src/object.rs b/lib/xixanta/src/object.rs
index 2708647..41ec67e 100644
--- a/lib/xixanta/src/object.rs
+++ b/lib/xixanta/src/object.rs
@@ -177,12 +177,25 @@ impl Context {
     pub fn get_variable(&self, id: &PString, mappings: &[Mapping]) -> Result<Object, ContextError> {
         // First of all, figure out the name of the scope and the real name of
         // the variable. If this was not scoped at all (None case when trying to
-        // rsplit by the "::" operator), then we assume it's a global variable.
+        // rsplit by the "::" operator), then we assume on the current scope.
         let (scope_name, var_name) = match id.value.rsplit_once("::") {
             Some((scope, name)) => (scope, name),
             None => (self.name(), id.value.as_str()),
         };
 
+        self.get_variable_in_scope(id.line, scope_name, var_name, mappings)
+    }
+
+    // Get the `var_name` variable on the `scope_name` scope (or parents). For
+    // further context, take the `mappings` into consideration when resolving
+    // labels, and `line` when producing context errors.
+    fn get_variable_in_scope(
+        &self,
+        line: usize,
+        scope_name: &str,
+        var_name: &str,
+        mappings: &[Mapping],
+    ) -> Result<Object, ContextError> {
         // And with that, the only thing left is to find the scope and the
         // variable in it.
         match self.map.get(scope_name) {
@@ -191,20 +204,46 @@ impl Context {
                     ObjectType::Value => Ok(var.clone()),
                     ObjectType::Address => Ok(self.resolve_label(mappings, var)?),
                 },
-                None => Err(ContextError {
-                    message: format!(
-                        "could not find variable '{}' in {}",
-                        var_name,
-                        self.to_human_with(scope_name)
-                    ),
-                    line: id.line,
-                    reason: ContextErrorReason::UnknownVariable,
-                    global: false,
-                }),
+                None => {
+                    // If it cannot be found, then we have to move up through
+                    // the scope hierarchy to see if we can fetch it there. For
+                    // that, though, we first prepare an error so we return the
+                    // original one, not the propagated one (see below).
+                    let err = Err(ContextError {
+                        message: format!(
+                            "could not find variable '{}' in {}",
+                            var_name,
+                            self.to_human_with(scope_name)
+                        ),
+                        line,
+                        reason: ContextErrorReason::UnknownVariable,
+                        global: false,
+                    });
+
+                    // If we are already in the global context and the object
+                    // was not found, just leave with an error.
+                    if scope_name == GLOBAL_CONTEXT {
+                        err
+                    } else {
+                        // Recursive call to find the object on the parent
+                        // scope. If this cannot be found, instead of using the
+                        // error from the recursive call, preserve the original
+                        // error so it better reflects the original scope where
+                        // this was first attempted.
+                        let parent = self.parent(scope_name);
+                        if let Ok(object) =
+                            self.get_variable_in_scope(line, parent, var_name, mappings)
+                        {
+                            Ok(object)
+                        } else {
+                            err
+                        }
+                    }
+                }
             },
             None => Err(ContextError {
                 message: format!("did not find scope '{}'", scope_name),
-                line: id.line,
+                line,
                 reason: ContextErrorReason::BadScope,
                 global: false,
             }),
@@ -443,6 +482,20 @@ impl Context {
         Ok(())
     }
 
+    // Returns the name context that is directly above the one named `name`.
+    fn parent(&self, name: &str) -> &str {
+        let index = self
+            .stack
+            .iter()
+            .position(|n| n.as_str() == name)
+            .unwrap_or(0);
+        if index < 2 {
+            GLOBAL_CONTEXT
+        } else {
+            self.stack.get(index - 1).unwrap()
+        }
+    }
+
     /// Returns the name of the current context.
     pub fn name(&self) -> &str {
         match self.stack.last() {