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