Re-create the values for bundle call arguments

Bundle call arguments were fine most of the times, when arguments could
be processed as-is and there were no issues with arguments being
overwritten by successive calls.

This was not the case, though, whenever a given instruction was delayed
into a PendingNode status. In this case, the argument would get the
last value, and in some extreme cases that definition might not have
been there any more.

Prevent all of this by providing a list of PendingDefine's, which are a
way to re-create these call-only arguments as they were initially
found. These PendingDefine's are then created on "crunch" on each
PendingNode.

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

5 files changed, 94 insertions, 4 deletions

diff --git a/lib/xixanta/src/assembler.rs b/lib/xixanta/src/assembler.rs
index d624c7e..1b963f2 100644
--- a/lib/xixanta/src/assembler.rs
+++ b/lib/xixanta/src/assembler.rs
@@ -58,6 +58,15 @@ struct PendingNode {
     node: PNode,
     labels_seen: usize,
     macro_context: Vec<ExpandedFrom>,
+    pending_defines: Option<Vec<PendingDefine>>,
+}
+
+/// A definition that is pending to be re-created whenever we crunch for
+/// PendingNode's.
+#[derive(Clone, Debug)]
+struct PendingDefine {
+    id: PString,
+    obj: Object,
 }
 
 /// Memory range that can be identified by a name.
@@ -143,6 +152,10 @@ struct Assembler<'a> {
     // 'Error' so the programmer gets information on all the macro expansions
     // that happened before reaching a given error.
     macro_context: Vec<ExpandedFrom>,
+
+    // Stack of pending definitions. The last element from this list is the list
+    // to be set to the next PendingNode push.
+    pending_defines: Vec<Vec<PendingDefine>>,
 }
 
 /// The result to be given at the end of `assembler::assemble` and
@@ -382,6 +395,7 @@ impl<'a> Assembler<'a> {
             asan_next_ignore: false,
             asan_next_reserve: 1,
             macro_context: vec![],
+            pending_defines: vec![],
         }
     }
 
@@ -906,6 +920,7 @@ impl<'a> Assembler<'a> {
                         node: node.to_owned(),
                         labels_seen: self.context.labels_seen(),
                         macro_context: self.macro_context.clone(),
+                        pending_defines: self.pending_defines.last().cloned(),
                     });
                 }
                 _ => {}
@@ -949,6 +964,14 @@ impl<'a> Assembler<'a> {
                 continue;
             }
 
+            // Re-create the definitions for this context which only made sense
+            // locally. These are basically arguments from a macro call.
+            if let Some(pd) = &pn.pending_defines {
+                for d in pd {
+                    let _ = self.context.set_variable(&d.id, &d.obj, true);
+                }
+            }
+
             self.literal_mode = None;
             match self.evaluate_node(&pn.node) {
                 Ok(mut bundle) => {
@@ -1302,6 +1325,12 @@ impl<'a> Assembler<'a> {
             .into());
         }
 
+        // All arguments are going to be added here as well so the inner block
+        // can take it as a list of PendingDefine's. This way we can re-create
+        // the values for the arguments used for this bundle call without having
+        // to go over the full bundle call.
+        let mut defines = vec![];
+
         // If there are arguments defined by the macro, set their values now.
         if given_args > 0 {
             let mut margs = mcr.args.as_ref().unwrap().iter();
@@ -1322,10 +1351,8 @@ impl<'a> Assembler<'a> {
                 // Note that we overwrite the variable value from previous
                 // calls, just in case a macro is applied multiple times and we
                 // need to get the latest value.
-                if let Err(message) =
-                    self.context
-                        .set_variable(&margs.next().unwrap().value, &obj, true)
-                {
+                let id = &margs.next().unwrap().value;
+                if let Err(message) = self.context.set_variable(id, &obj, true) {
                     return Err(Error {
                         line: node.value.line,
                         message,
@@ -1335,6 +1362,10 @@ impl<'a> Assembler<'a> {
                     }
                     .into());
                 }
+                defines.push(PendingDefine {
+                    id: id.clone(),
+                    obj,
+                });
             }
         }
 
@@ -1358,7 +1389,9 @@ impl<'a> Assembler<'a> {
                 line: node.value.line,
                 source: self.source_for(node),
             });
+            self.pending_defines.push(defines);
             self.bundle(inner)?;
+            self.pending_defines.pop();
             let _ = self.macro_context.pop();
         }
         Ok(())
@@ -1379,6 +1412,7 @@ impl<'a> Assembler<'a> {
                 node: node.to_owned(),
                 labels_seen: self.context.labels_seen(),
                 macro_context: self.macro_context.clone(),
+                pending_defines: self.pending_defines.last().cloned(),
             });
         }
         current.segments[self.current_segment].bundles.push(bundle);
diff --git a/scripts/test-e2e.sh b/scripts/test-e2e.sh
index 02d1ca3..ddb7b0a 100755
--- a/scripts/test-e2e.sh
+++ b/scripts/test-e2e.sh
@@ -102,6 +102,13 @@ exit_code=$((exit_code + $?))
 diff tests/out/indirect.nes tests/expected/indirect.nes
 exit_code=$((exit_code + $?))
 
+echo "test: custom => delayed_macro_arguments.nes"
+./target/debug/nasm -c empty --asan tests/delayed_macro_arguments.s -o tests/out/delayed_macro_arguments.nes 2>tests/out/delayed_macro_arguments.txt
+diff tests/out/delayed_macro_arguments.txt tests/expected/delayed_macro_arguments.txt
+exit_code=$((exit_code + $?))
+diff tests/out/delayed_macro_arguments.nes tests/expected/delayed_macro_arguments.nes
+exit_code=$((exit_code + $?))
+
 ##
 # code.nes
 
diff --git a/tests/delayed_macro_arguments.s b/tests/delayed_macro_arguments.s
new file mode 100644
index 0000000..b91ec92
--- /dev/null
+++ b/tests/delayed_macro_arguments.s
@@ -0,0 +1,49 @@
+.segment "HEADER"
+    .byte 'N', 'E', 'S', $1A
+    .byte $02
+    .byte $01
+    .byte $00
+    .byte $00
+
+.segment "CODE"
+
+.macro JAL ADDR
+    jmp ADDR
+.endmacro
+
+.macro DEC_MOVEMENT_X ADDR
+    lda ADDR
+    rts
+.endmacro
+
+.proc bounce
+    rts
+.endproc
+
+.proc chase
+    JAL bounce
+@wait:
+    DEC_MOVEMENT_X #1
+    DEC_MOVEMENT_X #2
+    JAL bounce
+    DEC_MOVEMENT_X #3
+    JAL chase
+    bne @wait
+    rts
+.endproc
+
+.macro NEXT_EXPLOSION_INDEX_X
+    inx
+    inx
+    inx
+.endmacro
+
+.scope Explosions
+    .proc init
+    @loop:
+        NEXT_EXPLOSION_INDEX_X
+        bne @loop
+
+        rts
+    .endproc
+.endscope
diff --git a/tests/expected/delayed_macro_arguments.nes b/tests/expected/delayed_macro_arguments.nes
new file mode 100644
index 0000000..20c81cd
--- /dev/null
+++ b/tests/expected/delayed_macro_arguments.nes
diff --git a/tests/expected/delayed_macro_arguments.txt b/tests/expected/delayed_macro_arguments.txt
new file mode 100644
index 0000000..e69de29
--- /dev/null
+++ b/tests/expected/delayed_macro_arguments.txt