Implement .if/.elsif/.else statements

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

4 files changed, 258 insertions, 43 deletions

diff --git a/lib/xixanta/src/assembler.rs b/lib/xixanta/src/assembler.rs
index 2621982..ce34983 100644
--- a/lib/xixanta/src/assembler.rs
+++ b/lib/xixanta/src/assembler.rs
@@ -508,8 +508,8 @@ impl<'a> Assembler<'a> {
                     }
                 }
                 NodeType::Control(control_type) => {
-                    if let Err(e) = self.evaluate_control_statement(node) {
-                        errors.push(e);
+                    if let Err(mut e) = self.evaluate_control_statement(node) {
+                        errors.append(&mut e);
                     }
 
                     // On control statements which modify the context, there are
@@ -1277,7 +1277,7 @@ impl<'a> Assembler<'a> {
         }
     }
 
-    fn evaluate_control_statement(&mut self, node: &'a PNode) -> Result<(), Error> {
+    fn evaluate_control_statement(&mut self, node: &'a PNode) -> Result<(), Vec<Error>> {
         // This might just be a statement that changes the context (e.g.
         // ".macro", ".proc", etc.). In this case change the context and leave
         // early.
@@ -1293,23 +1293,26 @@ impl<'a> Assembler<'a> {
                     line: node.value.line,
                     source: self.source_for(node),
                     global: false,
-                });
+                }
+                .into());
             }
         }
 
         // Otherwise, check the function that could act as a statement that
         // produces bundles.
         match node.node_type {
-            NodeType::Control(ControlType::Byte) => self.push_evaluated_arguments(node, 1),
+            NodeType::Control(ControlType::Byte) => Ok(self.push_evaluated_arguments(node, 1)?),
             NodeType::Control(ControlType::Addr) | NodeType::Control(ControlType::Word) => {
-                self.push_evaluated_arguments(node, 2)
+                Ok(self.push_evaluated_arguments(node, 2)?)
             }
-            NodeType::Control(ControlType::ReserveMemory) => self.reserve_memory(node),
-            NodeType::Control(ControlType::Asciiz) => self.push_ascii_string(node),
-            NodeType::Control(ControlType::Segment) => self.switch_to_segment(node),
+            NodeType::Control(ControlType::ReserveMemory) => Ok(self.reserve_memory(node)?),
+            NodeType::Control(ControlType::Asciiz) => Ok(self.push_ascii_string(node)?),
+            NodeType::Control(ControlType::Segment) => Ok(self.switch_to_segment(node)?),
             NodeType::Control(ControlType::IncBin) => {
-                self.incbin(node.args.as_ref().unwrap().first().unwrap())
+                Ok(self.incbin(node.args.as_ref().unwrap().first().unwrap())?)
             }
+            NodeType::Control(ControlType::If) => self.evaluate_if_block(node),
+            NodeType::Control(ControlType::EndIf) => Ok(()),
             NodeType::Control(ControlType::IncludeSource) => Ok(()),
             _ => Err(Error {
                 line: node.value.line,
@@ -1319,7 +1322,8 @@ impl<'a> Assembler<'a> {
                 ),
                 source: self.source_for(node),
                 global: false,
-            }),
+            }
+            .into()),
         }
     }
 
@@ -1503,6 +1507,45 @@ impl<'a> Assembler<'a> {
         Ok(())
     }
 
+    // Evaluate the given `node` by assuming that it's an .if/.elsif/.else
+    // block.
+    fn evaluate_if_block(&mut self, node: &'a PNode) -> Result<(), Vec<Error>> {
+        // Evaluate the condition if possible. If that's not the case (i.e.
+        // '.else'), then we just assume it to be true.
+        let cond = match &node.args {
+            Some(args) => {
+                self.literal_mode = None;
+                let cond_node = self.evaluate_node(args.first().unwrap())?;
+                cond_node.value() == 1
+            }
+            None => true,
+        };
+
+        // If condition was false, try to go into the .elsif/.else statement. If
+        // that doesn't exist, then we are done.
+        if !cond {
+            match &node.left {
+                Some(else_node) => return self.evaluate_if_block(else_node),
+                None => return Ok(()),
+            }
+        }
+
+        // And evaluate the inner block if the condition was true.
+        let inner = &node.right.as_ref().unwrap().args.as_ref().unwrap();
+        if inner.is_empty() {
+            self.warnings.push(Error {
+                line: node.value.line,
+                message: "empty body".to_string(),
+                source: self.source_for(node),
+                global: false,
+            });
+        } else {
+            self.bundle(inner)?;
+        }
+
+        Ok(())
+    }
+
     fn evaluate_control_expression(&mut self, node: &PNode) -> Result<Bundle, Error> {
         match node.node_type {
             NodeType::Control(ControlType::Hibyte) => self.evaluate_byte(node, true),
@@ -3250,6 +3293,85 @@ jsr Movement::update
         assert_eq!(warnings[3].to_string(), "segment 'CODE' is empty");
     }
 
+    #[test]
+    fn if_block() {
+        let res = just_bundles(
+            r#"Var = 0
+.if Var == 0
+  lda #1
+.endif
+
+.if Var == 1
+  lda #1
+.elsif Var == 0
+  lda #2
+.elsif Var == 2
+  lda #3
+.endif
+
+.if Var == 1
+  lda #1
+.elsif Var == 2
+  lda #2
+.else
+  lda #3
+.endif
+"#,
+        );
+
+        assert_eq!(res.len(), 3);
+
+        let instrs: Vec<[u8; 2]> = vec![[0xA9, 0x01], [0xA9, 0x02], [0xA9, 0x03]];
+        for i in 0..3 {
+            assert_eq!(res[i].size, 2);
+            assert_eq!(res[i].bytes[0], instrs[i][0]);
+            assert_eq!(res[i].bytes[1], instrs[i][1]);
+        }
+    }
+
+    #[test]
+    fn bad_if() {
+        assert_error(
+            r#"Var = 1
+.if Var == 1
+  lda #1
+    "#,
+            5,
+            false,
+            "expecting a 'control function (.endif)' but there are no more statements",
+        );
+
+        let res = just_assemble(
+            r#"Var = 1
+;;;.if Var1 == 1
+  lda #1
+.elsif Var1 == 2
+  lda #2
+.else
+  lda #3
+.endif
+    "#,
+        );
+
+        assert_eq!(res.errors.len(), 3);
+
+        assert_eq!(res.errors[0].line, 6);
+        assert_eq!(
+            res.errors[0].message,
+            "unexpected 'control function (.elsif)'"
+        );
+        assert_eq!(res.errors[1].line, 8);
+        assert_eq!(
+            res.errors[1].message,
+            "unexpected 'control function (.else)'"
+        );
+        assert_eq!(res.errors[2].line, 10);
+        assert_eq!(
+            res.errors[2].message,
+            "unexpected 'control function (.endif)'"
+        );
+    }
+
     // Macros
 
     #[test]
diff --git a/lib/xixanta/src/node.rs b/lib/xixanta/src/node.rs
index a31c578..ae79f1e 100644
--- a/lib/xixanta/src/node.rs
+++ b/lib/xixanta/src/node.rs
@@ -143,6 +143,10 @@ pub enum ControlType {
     IncludeSource,
     ReserveMemory,
     Asciiz,
+    If,
+    Elsif,
+    Else,
+    EndIf,
 }
 
 impl fmt::Display for ControlType {
@@ -166,6 +170,10 @@ impl fmt::Display for ControlType {
             ControlType::IncludeSource => write!(f, ".include"),
             ControlType::ReserveMemory => write!(f, ".res"),
             ControlType::Asciiz => write!(f, ".asciiz"),
+            ControlType::If => write!(f, ".if"),
+            ControlType::Elsif => write!(f, ".elsif"),
+            ControlType::Else => write!(f, ".else"),
+            ControlType::EndIf => write!(f, ".endif"),
         }
     }
 }
@@ -324,6 +332,9 @@ impl NodeType {
             NodeType::Control(ControlType::StartRepeat) => {
                 Some(NodeType::Control(ControlType::EndRepeat))
             }
+            NodeType::Control(ControlType::If) => Some(NodeType::Control(ControlType::EndIf)),
+            NodeType::Control(ControlType::Elsif) => Some(NodeType::Control(ControlType::EndIf)),
+            NodeType::Control(ControlType::Else) => Some(NodeType::Control(ControlType::EndIf)),
             _ => None,
         }
     }
@@ -389,11 +400,15 @@ impl PNode {
             NodeType::Control(ControlType::StartMacro)
             | NodeType::Control(ControlType::StartProc)
             | NodeType::Control(ControlType::StartScope)
-            | NodeType::Control(ControlType::StartRepeat) => NodeBodyType::Starts,
+            | NodeType::Control(ControlType::StartRepeat)
+            | NodeType::Control(ControlType::If)
+            | NodeType::Control(ControlType::Elsif)
+            | NodeType::Control(ControlType::Else) => NodeBodyType::Starts,
             NodeType::Control(ControlType::EndMacro)
             | NodeType::Control(ControlType::EndProc)
             | NodeType::Control(ControlType::EndScope)
-            | NodeType::Control(ControlType::EndRepeat) => NodeBodyType::Ends,
+            | NodeType::Control(ControlType::EndRepeat)
+            | NodeType::Control(ControlType::EndIf) => NodeBodyType::Ends,
             _ => NodeBodyType::None,
         }
     }
diff --git a/lib/xixanta/src/opcodes.rs b/lib/xixanta/src/opcodes.rs
index 0e23594..8488b77 100644
--- a/lib/xixanta/src/opcodes.rs
+++ b/lib/xixanta/src/opcodes.rs
@@ -757,6 +757,10 @@ lazy_static! {
         functions.insert(String::from(".include"), Control { control_type: ControlType::IncludeSource, has_identifier: None, required_args: Some((1, 1)), touches_context: false, only_string: true });
         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(".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 });
 
         functions
     };
diff --git a/lib/xixanta/src/parser.rs b/lib/xixanta/src/parser.rs
index 5bac127..13594e7 100644
--- a/lib/xixanta/src/parser.rs
+++ b/lib/xixanta/src/parser.rs
@@ -606,42 +606,27 @@ impl Parser {
         let node_type = node.node_type.clone();
         match node.body_type() {
             NodeBodyType::Starts => {
+                // .elsif/.else statements close the previous block by mocking
+                // an .endif, and then they start a new one.
+                if matches!(
+                    node_type,
+                    NodeType::Control(ControlType::Elsif) | NodeType::Control(ControlType::Else)
+                ) {
+                    self.close_body(&node_type, &NodeType::Control(ControlType::EndIf))?;
+                }
+
                 self.bodies.push(node_type.closing_type().unwrap());
                 self.nodes.last_mut().unwrap().push(node);
                 self.nodes.push(vec![]);
             }
             NodeBodyType::Ends => {
-                // Pop out which start statement was last seen and check that it
-                // makes sense to the end statement we are parsing now.
-                let expected_close = match self.bodies.pop() {
-                    Some(ec) => ec,
-                    None => {
-                        return Err(self
-                            .parser_error(format!("unexpected '{}'", node_type).as_str())
-                            .into())
-                    }
-                };
-                if node_type != expected_close {
-                    return Err(self
-                        .parser_error(
-                            format!("expecting '{}', found '{}'", expected_close, node_type)
-                                .as_str(),
-                        )
-                        .into());
-                }
+                self.close_body(&node_type, &node_type)?;
 
-                // Note that empty bodies are possible. This is left
-                // to the caller (e.g. assembler) to decide whether
-                // it makes sense or not.
-                let nodes = self.nodes.pop().unwrap();
-                self.nodes.last_mut().unwrap().last_mut().unwrap().right = Some(Box::new(PNode {
-                    node_type: NodeType::ControlBody,
-                    value: PString::default(),
-                    left: None,
-                    right: None,
-                    args: Some(nodes),
-                    source: self.current_source,
-                }));
+                // If this is an .endif statement, then we need to fold
+                // .if/.elsif/.else statements that have accumulated.
+                if matches!(node_type, NodeType::Control(ControlType::EndIf)) {
+                    self.fold_if_branches()?;
+                }
                 self.nodes.last_mut().unwrap().push(node);
             }
             NodeBodyType::None => self.nodes.last_mut().unwrap().push(node),
@@ -650,6 +635,95 @@ impl Parser {
         Ok(())
     }
 
+    // Fold .if/.elsif/.else statements so each branch of the block is on the
+    // `left` node of the previous one. The last layer of nodes will also be
+    // truncated accordingly.
+    fn fold_if_branches(&mut self) -> Result<(), Vec<Error>> {
+        let nodes = self.nodes.last_mut().unwrap();
+        let mut count = 0;
+
+        // Just iterate in reverse order from the last node until the .if
+        // statement that started the whole .if/.elsif/.else block.
+        for idx in (1..nodes.len()).rev() {
+            let node = &nodes[idx];
+            let next = &nodes[idx - 1];
+
+            // 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) => {
+                    nodes.truncate(nodes.len() - count);
+                    return Ok(());
+                }
+                NodeType::Control(ControlType::Elsif) => {
+                    if !matches!(
+                        next.node_type,
+                        NodeType::Control(ControlType::If) | NodeType::Control(ControlType::Elsif)
+                    ) {
+                        return Err(self.parser_error("expecting an .if statement").into());
+                    }
+                }
+                NodeType::Control(ControlType::Else) => {
+                    if !matches!(
+                        next.node_type,
+                        NodeType::Control(ControlType::If) | NodeType::Control(ControlType::Elsif)
+                    ) {
+                        return Err(self.parser_error("expecting an .if statement").into());
+                    }
+                }
+                _ => return Err(self.parser_error("expecting an .if statement").into()),
+            }
+
+            // After all the checks are done, folding simply means to assign the
+            // left node to the current one.
+            count += 1;
+            nodes[idx - 1].left = Some(Box::new(nodes[idx].clone()));
+        }
+
+        nodes.truncate(nodes.len() - count);
+        Ok(())
+    }
+
+    // Close the currently open block body with the given `node_type`. The given
+    // `node_type` is one that ideally closes the current block, but it might
+    // not be necessarily what it was really found, which should be passed as
+    // `real_type`. This is so statements such as '.elsif' can also close
+    // previous `.if/.elsif` blocks.
+    fn close_body(&mut self, real_type: &NodeType, node_type: &NodeType) -> Result<(), Vec<Error>> {
+        // Pop out which start statement was last seen and check that it
+        // makes sense to the end statement we are parsing now.
+        let expected_close = match self.bodies.pop() {
+            Some(ec) => ec,
+            None => {
+                return Err(self
+                    .parser_error(format!("unexpected '{}'", real_type).as_str())
+                    .into())
+            }
+        };
+        if *node_type != expected_close {
+            return Err(self
+                .parser_error(
+                    format!("expecting '{}', found '{}'", expected_close, node_type).as_str(),
+                )
+                .into());
+        }
+
+        // Note that empty bodies are possible. This is left
+        // to the caller (e.g. assembler) to decide whether
+        // it makes sense or not.
+        let nodes = self.nodes.pop().unwrap();
+        self.nodes.last_mut().unwrap().last_mut().unwrap().right = Some(Box::new(PNode {
+            node_type: NodeType::ControlBody,
+            value: PString::default(),
+            left: None,
+            right: None,
+            args: Some(nodes),
+            source: self.current_source,
+        }));
+
+        Ok(())
+    }
+
     // Consume the given `node` by assuming it's an `.include` statement. This
     // will in turn produce a new parsing session for the given file if possible
     // and push the parsed nodes from it to our current list.