Add the .fallthrough control statement

This is exclusive to 'nasm' and it allows the developer to explicitly
tell the assembler than a "fall through" condition is actually desired:
it's not a mistake.

This comes in two flavors. The first, without arguments, just makes this
explicit without much enforcement. The second allows you to pass an
argument which is the name of the function or label you are expecting to
fall through. The assembler will error out if the fall through address
is not the expected one, hence telling the programmer whenever the fall
through condition they thought in the past is no longer true (e.g. the
function has moved somewhere else in the code).

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

9 files changed, 164 insertions, 1 deletions

diff --git a/lib/xixanta/src/assembler.rs b/lib/xixanta/src/assembler.rs
index d7e036f..5c43387 100644
--- a/lib/xixanta/src/assembler.rs
+++ b/lib/xixanta/src/assembler.rs
@@ -30,7 +30,7 @@ enum LiteralMode {
 
 /// The different stages that the assembler goes through and which are relevant
 /// for the process.
-#[derive(PartialEq)]
+#[derive(Debug, PartialEq)]
 enum Stage {
     /// The context is still building up (i.e. we don't have all the variable
     /// values, labels and their addresses yet).
@@ -883,6 +883,23 @@ impl<'a> Assembler<'a> {
             self.labels_seen = pn.labels_seen;
             self.context.force_context_switch(&pn.context);
 
+            // .fallthrough is handled here, whenever we already know addresses,
+            // sizes, etc. If this is the case, this is not a real node that can
+            // be bundled, but perform its check and move into the next
+            // iteration.
+            //
+            // NOTE: this has to happen with a context switch, otherwise the
+            // name resolution won't be accurate.
+            if matches!(
+                pn.node.node_type,
+                NodeType::Control(ControlType::Fallthrough)
+            ) {
+                if let Err(e) = self.fallthrough(&pn) {
+                    errors.push(e);
+                }
+                continue;
+            }
+
             self.literal_mode = None;
             match self.evaluate_node(&pn.node) {
                 Ok(mut bundle) => {
@@ -949,6 +966,59 @@ impl<'a> Assembler<'a> {
         }
     }
 
+    fn fallthrough(&mut self, pn: &PendingNode) -> Result<(), Error> {
+        // If we are not in 'crunch' mode, then it's a bug.
+        assert_eq!(self.stage, Stage::Crunching);
+
+        let node = pn.node.left.as_ref().unwrap();
+
+        // If there was no "argument", then skip things altogether. The
+        // programmer opted for an explicit fallthrough without further checks.
+        if node.value.is_empty() {
+            return Ok(());
+        }
+
+        // If there is an argument, it has to be a valid address identifier.
+        if let Err(message) = node.value.is_valid_identifier(false) {
+            return Err(Error {
+                line: pn.node.value.line,
+                message,
+                source: self.source_for(&pn.node),
+                expanded_from: pn.macro_context.clone(),
+                global: false,
+            });
+        }
+
+        // The check looks scarier than it is. We first grab the target address
+        // by evaluating the identifier as a variable. This should just gives as
+        // the address as is by calling .value(). The effective address is taken
+        // from the 'bundle_index' from the PendingNode, as it was pushed while
+        // pointing to the "next" bundle. Hence, we just fetch that bundle and
+        // get its address.
+        let target_address = self.evaluate_variable(node)?.value() as usize;
+        let current = &self.mappings[pn.mapping].segments[pn.segment];
+        let Some(effective) = current.bundles.get(pn.bundle_index) else {
+            return Err(Error {
+                line: pn.node.value.line,
+                message: String::from("statement does not fall through"),
+                source: self.source_for(&pn.node),
+                expanded_from: pn.macro_context.clone(),
+                global: false,
+            });
+        };
+
+        if effective.address != target_address {
+            return Err(Error {
+                line: pn.node.value.line,
+                message: String::from("statement does not fall through"),
+                source: self.source_for(&pn.node),
+                expanded_from: pn.macro_context.clone(),
+                global: false,
+            });
+        }
+        Ok(())
+    }
+
     fn asan(&mut self, memory: &mut MemoryResult) -> Result<(), Vec<Error>> {
         let mut errors = vec![];
 
@@ -1865,6 +1935,21 @@ impl<'a> Assembler<'a> {
             NodeType::Control(ControlType::EndIf) => Ok(()),
             NodeType::Control(ControlType::IncludeSource) => Ok(()),
             NodeType::Control(ControlType::Echo(_)) => Ok(()),
+            NodeType::Control(ControlType::Fallthrough) => {
+                // We don't do much other than pushing a fake PendingNode with a
+                // 'bundle_index' value that will point to the next node.
+                let current = &mut self.mappings[self.current_mapping];
+                self.pending.push(PendingNode {
+                    mapping: self.current_mapping,
+                    segment: self.current_segment,
+                    context: self.context.name().to_string(),
+                    bundle_index: current.segments[self.current_segment].bundles.len(),
+                    node: node.to_owned(),
+                    labels_seen: self.context.labels_seen(),
+                    macro_context: self.macro_context.clone(),
+                });
+                Ok(())
+            }
             _ => Err(Error {
                 line: node.value.line,
                 message: format!(
diff --git a/lib/xixanta/src/node.rs b/lib/xixanta/src/node.rs
index 3f1778a..e13dd3c 100644
--- a/lib/xixanta/src/node.rs
+++ b/lib/xixanta/src/node.rs
@@ -166,6 +166,7 @@ pub enum ControlType {
     EndIf,
     Defined,
     Echo(EchoKind),
+    Fallthrough,
 }
 
 impl fmt::Display for ControlType {
@@ -201,6 +202,7 @@ impl fmt::Display for ControlType {
                 EchoKind::Warning => write!(f, ".warning"),
                 EchoKind::Error => write!(f, ".error"),
             },
+            ControlType::Fallthrough => write!(f, ".fallthrough"),
         }
     }
 }
diff --git a/lib/xixanta/src/opcodes.rs b/lib/xixanta/src/opcodes.rs
index 48e4fca..bd3c207 100644
--- a/lib/xixanta/src/opcodes.rs
+++ b/lib/xixanta/src/opcodes.rs
@@ -1999,6 +1999,16 @@ pub static CONTROL_FUNCTIONS: LazyLock<HashMap<String, Control>> = LazyLock::new
             only_string: true,
         },
     );
+    functions.insert(
+        String::from(".fallthrough"),
+        Control {
+            control_type: ControlType::Fallthrough,
+            has_identifier: Some(false),
+            required_args: Some((0, 0)),
+            touches_context: false,
+            only_string: false,
+        },
+    );
 
     functions
 });
diff --git a/scripts/test-e2e.sh b/scripts/test-e2e.sh
index 2f7ff5f..18dbeff 100755
--- a/scripts/test-e2e.sh
+++ b/scripts/test-e2e.sh
@@ -76,6 +76,18 @@ echo "test: custom => bad_jal.nes"
 diff tests/out/bad_jal.txt tests/expected/bad_jal.txt
 exit_code=$((exit_code + $?))
 
+echo "test: custom => bad_fallthrough.nes"
+./target/debug/nasm -c empty --asan tests/bad_fallthrough.s -o /dev/null 2>tests/out/bad_fallthrough.txt
+diff tests/out/bad_fallthrough.txt tests/expected/bad_fallthrough.txt
+exit_code=$((exit_code + $?))
+
+echo "test: custom => fallthrough.nes"
+./target/debug/nasm -c empty --asan tests/fallthrough.s -o tests/out/fallthrough.nes 2>tests/out/fallthrough.txt
+diff tests/out/fallthrough.txt tests/expected/fallthrough.txt
+exit_code=$((exit_code + $?))
+diff tests/out/fallthrough.nes tests/expected/fallthrough.nes
+exit_code=$((exit_code + $?))
+
 ##
 # code.nes
 
diff --git a/tests/bad_fallthrough.s b/tests/bad_fallthrough.s
new file mode 100644
index 0000000..ac17ade
--- /dev/null
+++ b/tests/bad_fallthrough.s
@@ -0,0 +1,27 @@
+.segment "HEADER"
+    .byte 'N', 'E', 'S', $1A
+    .byte $02, $01
+    .byte $00
+    .byte $00
+
+.segment "CODE"
+
+.proc foo
+    lda #0
+    .fallthrough bar
+.endproc
+
+.fallthrough bar
+.fallthrough other
+ lda #0
+
+.proc bar
+    lda #0
+    .fallthrough other
+.endproc
+
+.proc other
+    rts
+.endproc
+
+.fallthrough other
diff --git a/tests/expected/bad_fallthrough.txt b/tests/expected/bad_fallthrough.txt
new file mode 100644
index 0000000..f105880
--- /dev/null
+++ b/tests/expected/bad_fallthrough.txt
@@ -0,0 +1,4 @@
+error: statement does not fall through (bad_fallthrough.s: line 11)
+error: statement does not fall through (bad_fallthrough.s: line 14)
+error: statement does not fall through (bad_fallthrough.s: line 15)
+error: statement does not fall through (bad_fallthrough.s: line 27)
diff --git a/tests/expected/fallthrough.nes b/tests/expected/fallthrough.nes
new file mode 100644
index 0000000..976f8fc
--- /dev/null
+++ b/tests/expected/fallthrough.nes
diff --git a/tests/expected/fallthrough.txt b/tests/expected/fallthrough.txt
new file mode 100644
index 0000000..e69de29
--- /dev/null
+++ b/tests/expected/fallthrough.txt
diff --git a/tests/fallthrough.s b/tests/fallthrough.s
new file mode 100644
index 0000000..2340ee8
--- /dev/null
+++ b/tests/fallthrough.s
@@ -0,0 +1,23 @@
+.segment "HEADER"
+    .byte 'N', 'E', 'S', $1A
+    .byte $02, $01
+    .byte $00
+    .byte $00
+
+.segment "CODE"
+
+.proc foo
+    lda #0
+    .fallthrough bar
+.endproc
+
+.fallthrough bar
+
+.proc bar
+    lda #0
+    .fallthrough other
+.endproc
+
+.proc other
+    rts
+.endproc