Implement .ifdef/.ifndef statements

They are just synonyms for ".if .defined" and ".if !.defined"
respectively.

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

4 files changed, 78 insertions, 4 deletions

diff --git a/lib/xixanta/src/assembler.rs b/lib/xixanta/src/assembler.rs
index d0a0e6e..8699529 100644
--- a/lib/xixanta/src/assembler.rs
+++ b/lib/xixanta/src/assembler.rs
@@ -1312,6 +1312,9 @@ impl<'a> Assembler<'a> {
                 Ok(self.incbin(node.args.as_ref().unwrap().first().unwrap())?)
             }
             NodeType::Control(ControlType::If) => self.evaluate_if_block(node),
+            NodeType::Control(ControlType::IfDef) | NodeType::Control(ControlType::IfNDef) => {
+                self.evaluate_ifdef_block(node)
+            }
             NodeType::Control(ControlType::EndIf) => Ok(()),
             NodeType::Control(ControlType::IncludeSource) => Ok(()),
             _ => Err(Error {
@@ -1521,6 +1524,24 @@ impl<'a> Assembler<'a> {
             None => true,
         };
 
+        self.evaluate_if_cond(node, cond)
+    }
+
+    // Evaluate the given `node` by assuming that it's an .ifdef/.ifndef block.
+    fn evaluate_ifdef_block(&mut self, node: &'a PNode) -> Result<(), Vec<Error>> {
+        let defined = self.evaluate_defined(node)?.value();
+        let cond = match &node.node_type {
+            NodeType::Control(ControlType::IfDef) => defined == 1,
+            NodeType::Control(ControlType::IfNDef) => defined == 0,
+            _ => panic!("unexpected .ifdef block"),
+        };
+
+        self.evaluate_if_cond(node, cond)
+    }
+
+    // Evaluate the `node`'s body if `cond` is true, otherwise try to evaluate
+    // the "else" branch for this .if-looking statement.
+    fn evaluate_if_cond(&mut self, node: &'a PNode, cond: bool) -> Result<(), Vec<Error>> {
         // If condition was false, try to go into the .elsif/.else statement. If
         // that doesn't exist, then we are done.
         if !cond {
@@ -3425,6 +3446,34 @@ jsr Movement::update
         }
     }
 
+    #[test]
+    fn ifdef_block() {
+        let res = just_bundles(
+            r#"Var = 0
+.ifdef Var
+  lda #1
+.endif
+
+.ifndef Var
+  lda #1
+.elsif Var == 0
+  lda #2
+.elsif Var == 2
+  lda #3
+.endif
+"#,
+        );
+
+        assert_eq!(res.len(), 2);
+
+        let instrs: Vec<[u8; 2]> = vec![[0xA9, 0x01], [0xA9, 0x02]];
+        for i in 0..instrs.len() {
+            assert_eq!(res[i].size, 2);
+            assert_eq!(res[i].bytes[0], instrs[i][0]);
+            assert_eq!(res[i].bytes[1], instrs[i][1]);
+        }
+    }
+
     // Macros
 
     #[test]
diff --git a/lib/xixanta/src/node.rs b/lib/xixanta/src/node.rs
index f1eb558..243b478 100644
--- a/lib/xixanta/src/node.rs
+++ b/lib/xixanta/src/node.rs
@@ -144,6 +144,8 @@ pub enum ControlType {
     ReserveMemory,
     Asciiz,
     If,
+    IfDef,
+    IfNDef,
     Elsif,
     Else,
     EndIf,
@@ -172,6 +174,8 @@ impl fmt::Display for ControlType {
             ControlType::ReserveMemory => write!(f, ".res"),
             ControlType::Asciiz => write!(f, ".asciiz"),
             ControlType::If => write!(f, ".if"),
+            ControlType::IfDef => write!(f, ".ifdef"),
+            ControlType::IfNDef => write!(f, ".ifndef"),
             ControlType::Elsif => write!(f, ".elsif"),
             ControlType::Else => write!(f, ".else"),
             ControlType::EndIf => write!(f, ".endif"),
@@ -192,7 +196,14 @@ impl ControlType {
     pub fn has_body(&self) -> bool {
         matches!(
             self,
-            ControlType::StartMacro | Self::StartProc | Self::StartScope
+            ControlType::StartMacro
+                | Self::StartProc
+                | Self::StartScope
+                | Self::If
+                | Self::Elsif
+                | Self::Else
+                | Self::IfDef
+                | Self::IfNDef
         )
     }
 }
@@ -335,6 +346,8 @@ impl NodeType {
                 Some(NodeType::Control(ControlType::EndRepeat))
             }
             NodeType::Control(ControlType::If) => Some(NodeType::Control(ControlType::EndIf)),
+            NodeType::Control(ControlType::IfDef) => Some(NodeType::Control(ControlType::EndIf)),
+            NodeType::Control(ControlType::IfNDef) => Some(NodeType::Control(ControlType::EndIf)),
             NodeType::Control(ControlType::Elsif) => Some(NodeType::Control(ControlType::EndIf)),
             NodeType::Control(ControlType::Else) => Some(NodeType::Control(ControlType::EndIf)),
             _ => None,
@@ -404,6 +417,8 @@ impl PNode {
             | NodeType::Control(ControlType::StartScope)
             | NodeType::Control(ControlType::StartRepeat)
             | NodeType::Control(ControlType::If)
+            | NodeType::Control(ControlType::IfDef)
+            | NodeType::Control(ControlType::IfNDef)
             | NodeType::Control(ControlType::Elsif)
             | NodeType::Control(ControlType::Else) => NodeBodyType::Starts,
             NodeType::Control(ControlType::EndMacro)
diff --git a/lib/xixanta/src/opcodes.rs b/lib/xixanta/src/opcodes.rs
index 80878ec..865b5cd 100644
--- a/lib/xixanta/src/opcodes.rs
+++ b/lib/xixanta/src/opcodes.rs
@@ -758,6 +758,8 @@ lazy_static! {
         functions.insert(String::from(".res"), Control { control_type: ControlType::ReserveMemory, has_identifier: None, required_args: Some((1, 2)), touches_context: false, only_string: false });
         functions.insert(String::from(".asciiz"), Control { control_type: ControlType::Asciiz, has_identifier: None, required_args: Some((1, 1)), touches_context: false, only_string: true });
         functions.insert(String::from(".if"), Control { control_type: ControlType::If, has_identifier: None, required_args: Some((1, 1)), touches_context: false, only_string: false });
+        functions.insert(String::from(".ifdef"), Control { control_type: ControlType::IfDef, has_identifier: None, required_args: Some((1, 1)), touches_context: false, only_string: false });
+        functions.insert(String::from(".ifndef"), Control { control_type: ControlType::IfNDef, has_identifier: None, required_args: Some((1, 1)), touches_context: false, only_string: false });
         functions.insert(String::from(".elsif"), Control { control_type: ControlType::Elsif, has_identifier: None, required_args: Some((1, 1)), touches_context: false, only_string: false });
         functions.insert(String::from(".else"), Control { control_type: ControlType::Else, has_identifier: None, required_args: Some((0, 0)), touches_context: false, only_string: false });
         functions.insert(String::from(".endif"), Control { control_type: ControlType::EndIf, has_identifier: None, required_args: Some((0, 0)), touches_context: false, only_string: false });
diff --git a/lib/xixanta/src/parser.rs b/lib/xixanta/src/parser.rs
index 13594e7..50f2af1 100644
--- a/lib/xixanta/src/parser.rs
+++ b/lib/xixanta/src/parser.rs
@@ -651,14 +651,19 @@ impl Parser {
             // If this is an .if already, then we can quit, otherwise we must
             // ensure that the next node is an .if/.elsif and continue the loop.
             match node.node_type {
-                NodeType::Control(ControlType::If) => {
+                NodeType::Control(ControlType::If)
+                | NodeType::Control(ControlType::IfDef)
+                | NodeType::Control(ControlType::IfNDef) => {
                     nodes.truncate(nodes.len() - count);
                     return Ok(());
                 }
                 NodeType::Control(ControlType::Elsif) => {
                     if !matches!(
                         next.node_type,
-                        NodeType::Control(ControlType::If) | NodeType::Control(ControlType::Elsif)
+                        NodeType::Control(ControlType::If)
+                            | NodeType::Control(ControlType::IfDef)
+                            | NodeType::Control(ControlType::IfNDef)
+                            | NodeType::Control(ControlType::Elsif)
                     ) {
                         return Err(self.parser_error("expecting an .if statement").into());
                     }
@@ -666,7 +671,10 @@ impl Parser {
                 NodeType::Control(ControlType::Else) => {
                     if !matches!(
                         next.node_type,
-                        NodeType::Control(ControlType::If) | NodeType::Control(ControlType::Elsif)
+                        NodeType::Control(ControlType::If)
+                            | NodeType::Control(ControlType::IfDef)
+                            | NodeType::Control(ControlType::IfNDef)
+                            | NodeType::Control(ControlType::Elsif)
                     ) {
                         return Err(self.parser_error("expecting an .if statement").into());
                     }