Add support for branching

This is a bare bones implementation of it, and there's a lot of nuance
that it's being missed, but at least the fundamentals have been
implemented already.

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

5 files changed, 295 insertions, 75 deletions

diff --git a/lib/xixanta/src/assembler.rs b/lib/xixanta/src/assembler.rs
index 783cccd..7546e1d 100644
--- a/lib/xixanta/src/assembler.rs
+++ b/lib/xixanta/src/assembler.rs
@@ -1,11 +1,12 @@
-use crate::context::Context;
+use crate::context::{Context, PValue};
 use crate::errors::ParseError;
 use crate::instruction::{
-    AddressingMode, Encodable, Generic, Instruction, Literal, Node, PString, Scoped,
+    AddressingMode, Encodable, Fill, Generic, Instruction, Label, Literal, Node, PString, Scoped,
 };
 use crate::mapping::{Mapping, Segment};
 use crate::opcodes::{INSTRUCTIONS, OPCODES};
 use std::collections::hash_map::Entry;
+use std::collections::HashMap;
 use std::io::{self, BufRead, Read};
 use std::ops::Range;
 
@@ -14,9 +15,9 @@ type Result<T> = std::result::Result<T, ParseError>;
 pub struct Assembler {
     line: usize,
     column: usize,
-    offset_address: usize,
     context: Context,
     mapping: Mapping,
+    offsets: HashMap<String, usize>,
 }
 
 impl Assembler {
@@ -24,16 +25,15 @@ impl Assembler {
         Self {
             line: 0,
             column: 0,
-            offset_address: 0,
             context: Context::new(),
             mapping: Mapping::new(segments),
+            offsets: HashMap::new(),
         }
     }
 
     pub fn reset(&mut self) {
         self.line = 0;
         self.column = 0;
-        self.offset_address = 0;
         self.context = Context::new();
         self.mapping.reset();
     }
@@ -42,6 +42,7 @@ impl Assembler {
         self.from_reader(reader)?;
         self.context.reset();
         self.evaluate()?;
+        self.resolve_labels()?;
 
         Ok(())
     }
@@ -51,9 +52,18 @@ impl Assembler {
 
         self.assemble_nodes(reader)?;
 
+        let mut idx: usize = 0;
         for segment in &self.mapping.segments {
             let mut size: usize = 0;
 
+            while idx < segment.start.into() {
+                match &segment.fill_value {
+                    Some(fill) => instructions.push(fill),
+                    None => instructions.push(&Fill { value: 0x00 }),
+                }
+                idx += 1;
+            }
+
             for node in &self.mapping.nodes[&segment.name] {
                 match node {
                     Node::Instruction(instr) => {
@@ -77,6 +87,7 @@ impl Assembler {
                     ),
                 });
             }
+            idx += size;
             if segment.fill_value.is_none() {
                 continue;
             }
@@ -84,6 +95,7 @@ impl Assembler {
             while size < segment.size {
                 instructions.push(segment.fill_value.as_ref().unwrap());
                 size += 1;
+                idx += 1;
             }
         }
 
@@ -129,22 +141,99 @@ impl Assembler {
     }
 
     pub fn evaluate(&mut self) -> Result<()> {
-        for node in self.mapping.current_mut() {
-            match node {
-                Node::Instruction(instr) => {
-                    Self::update_instruction_with_context(instr, &self.context)?
-                }
-                Node::Scoped(scope) => {
-                    if scope.start {
-                        self.context.push(&scope.identifier.value);
-                    } else {
-                        _ = self.context.pop();
+        for segment in &self.mapping.segments {
+            for node in self.mapping.nodes.get_mut(&segment.name).unwrap() {
+                match node {
+                    Node::Instruction(instr) => {
+                        Self::update_instruction_with_context(instr, &self.context)?;
+                        instr.address = segment.start;
+
+                        self.offsets
+                            .entry(segment.name.clone())
+                            .and_modify(|value| {
+                                instr.address += *value as u16;
+                                *value += usize::from(instr.size())
+                            })
+                            .or_insert(instr.size().into());
                     }
+                    Node::Scoped(scope) => {
+                        if scope.start {
+                            self.context.push(&scope.identifier.value);
+                        } else {
+                            _ = self.context.pop();
+                        }
+                    }
+                    Node::Literal(literal) => {
+                        Self::update_literal_with_context(literal, &self.context)?;
+                        self.offsets
+                            .entry(segment.name.clone())
+                            .and_modify(|value| *value += usize::from(literal.size()))
+                            .or_insert(literal.size().into());
+                    }
+                    Node::Label(label) => {
+                        let address =
+                            usize::from(segment.start) + self.offsets.get(&segment.name).unwrap();
+
+                        self.context
+                            .current_mut()
+                            .unwrap()
+                            .entry(label.value.clone())
+                            .and_modify(|e| e.value = address);
+                    }
+                    _ => {}
                 }
-                Node::Literal(literal) => {
-                    Self::update_literal_with_context(literal, &self.context)?
+            }
+        }
+
+        Ok(())
+    }
+
+    // TODO: oh boy...
+    pub fn resolve_labels(&mut self) -> Result<()> {
+        for segment in &self.mapping.segments {
+            for node in self.mapping.nodes.get_mut(&segment.name).unwrap() {
+                match node {
+                    Node::Instruction(instr) => {
+                        if !instr.resolved {
+                            match &instr.left {
+                                Some(pstring) => {
+                                    match self.context.current().unwrap().get(&pstring.value) {
+                                        Some(entry) => {
+                                            if instr.mode == AddressingMode::Absolute {
+                                                let bytes = entry.value.to_le_bytes();
+                                                instr.bytes = [bytes[0], bytes[1]];
+                                            } else {
+                                                let diff: isize = entry.value as isize
+                                                    - (instr.address as isize + 2);
+                                                if diff < -128 || diff > 127 {
+                                                    return Err(instr.mnemonic.parser_error(
+                                                        format!("relative addressing out of range")
+                                                            .as_str(),
+                                                    ));
+                                                }
+                                                let bytes = diff.to_le_bytes();
+                                                instr.bytes = [bytes[0], 0];
+                                            }
+                                        }
+                                        None => {
+                                            return Err(instr.mnemonic.parser_error(
+                                                format!("label '{}' not found", pstring.value)
+                                                    .as_str(),
+                                            ))
+                                        }
+                                    }
+                                }
+                                None => {
+                                    return Err(instr.mnemonic.parser_error(
+                                        format!("there is no label for the given jump instruction")
+                                            .as_str(),
+                                    ))
+                                }
+                            }
+                        }
+                    }
+                    _ => {}
                 }
-                _ => {}
             }
         }
 
@@ -204,7 +293,16 @@ impl Assembler {
         if left.value.to_lowercase() == "a" {
             instr.mode = AddressingMode::Implied;
         } else {
-            let nleft = Self::replace_variable(left, context)?;
+            let (nleft, resolved) = Self::replace_variable(left, context)?;
+            // TODO
+            instr.resolved = resolved;
+            if !resolved {
+                if instr.mnemonic.value == "jmp" {
+                    instr.mode = AddressingMode::Absolute;
+                } else {
+                    instr.mode = AddressingMode::RelativeOrZeropage;
+                }
+            }
 
             if nleft.value.starts_with('$') {
                 // This is an address. At this point we should assume that the
@@ -327,13 +425,15 @@ impl Assembler {
                         ).as_str(),
                     ));
                 }
-                return Err(instr.mnemonic.parser_error(
-                    format!(
-                        "unknown addressing mode for instruction '{}'",
-                        instr.mnemonic.value
-                    )
-                    .as_str(),
-                ));
+                // TODO:
+                // instr.mode = AddressingMode::Absolute;
+                // return Err(instr.mnemonic.parser_error(
+                //     format!(
+                //         "unknown addressing mode for instruction '{}'",
+                //         instr.mnemonic.value
+                //     )
+                //     .as_str(),
+                // ));
             }
         }
 
@@ -342,7 +442,7 @@ impl Assembler {
 
     fn update_literal_with_context(literal: &mut Literal, context: &Context) -> Result<()> {
         // Evaluate any possible variable being used inside of this literal.
-        let evaled = Self::replace_variable(&literal.identifier, context)?;
+        let (evaled, _) = Self::replace_variable(&literal.identifier, context)?;
 
         // Parse the numeric value after a possible variable has been replaced.
         let two_bytes_allowed = literal.size == 2;
@@ -450,8 +550,12 @@ impl Assembler {
         }
     }
 
-    fn replace_variable(node: &PString, context: &Context) -> Result<PString> {
-        match node.value.find(|c: char| c.is_alphabetic() || c == '_') {
+    // TODO: returns if resolved
+    fn replace_variable(node: &PString, context: &Context) -> Result<(PString, bool)> {
+        match node
+            .value
+            .find(|c: char| c.is_alphabetic() || c == '_' || c == '@')
+        {
             Some(idx) => {
                 // Before doing any replacement, let's check the character
                 // before the one that was found. In this case, if it was a
@@ -461,7 +565,7 @@ impl Assembler {
                 if idx > 0 {
                     let prev = node.value.chars().nth(idx - 1).unwrap_or(' ');
                     if prev.is_ascii_digit() {
-                        return Ok(node.clone());
+                        return Ok((node.clone(), true));
                     }
                 }
 
@@ -492,7 +596,7 @@ impl Assembler {
                         // end and hence string == ""), or this is just the regular
                         // X or Y index, just return early.
                         match string.to_lowercase().as_str() {
-                            "x" | "y" | "" => return Ok(node.clone()),
+                            "x" | "y" | "" => return Ok((node.clone(), true)),
                             _ => {}
                         }
 
@@ -500,16 +604,25 @@ impl Assembler {
                         // the current scope.
                         match hash.get(string) {
                             Some(var) => {
+                                // If this is just a memory address (e.g.
+                                // label), then just return it as is.
+                                if var.label {
+                                    return Ok((node.clone(), false));
+                                }
+
                                 let value = String::from(node.value.get(..idx).unwrap_or(""))
-                                    + var.value.as_str();
-                                Ok(PString {
-                                    value: value.clone() + tail,
-                                    line: node.line,
-                                    range: Range {
-                                        start: node.range.start,
-                                        end: node.range.start + value.len(),
+                                    + var.node.value.as_str();
+                                Ok((
+                                    PString {
+                                        value: value.clone() + tail,
+                                        line: node.line,
+                                        range: Range {
+                                            start: node.range.start,
+                                            end: node.range.start + value.len(),
+                                        },
                                     },
-                                })
+                                    true,
+                                ))
                             }
                             None => {
                                 // If a variable could not be found, check that
@@ -517,7 +630,7 @@ impl Assembler {
                                 // 'AA'). If that's the case, then just return
                                 // its value.
                                 if Self::parse_hex_from(string, node, true, false, false).is_ok() {
-                                    return Ok(node.clone());
+                                    return Ok((node.clone(), true));
                                 }
 
                                 // We've tried hard to not assume the programmer
@@ -534,7 +647,7 @@ impl Assembler {
                     }
                 }
             }
-            None => Ok(node.clone()),
+            None => Ok((node.clone(), true)),
         }
     }
 
@@ -918,16 +1031,57 @@ impl Assembler {
     }
 
     fn parse_statement(&mut self, id: PString, line: &str) -> Result<()> {
-        if line.chars().nth(self.column).unwrap_or(' ') == ':' {
+        if id.value.chars().nth(self.column - 1).unwrap_or(' ') == ':' {
             self.parse_label(id, line)
         } else {
             self.parse_assignment(id, line)
         }
     }
 
-    fn parse_label(&mut self, _id: PString, _line: &str) -> Result<()> {
-        // TODO:
+    fn parse_label(&mut self, id: PString, _line: &str) -> Result<()> {
+        let name = &id.value.as_str()[..id.value.len() - 1].to_string();
+
+        // Forbid weird scenarios.
+        if name.contains("::") {
+            return Err(id.parser_error(
+                format!(
+                    "the label '{}' is scoped: do not declare variables this way",
+                    id.value
+                )
+                .as_str(),
+            ));
+        }
 
+        // Insert the given label into the context.
+        if let Some(entry) = self.context.current_mut() {
+            match entry.entry(name.clone()) {
+                Entry::Occupied(e) => {
+                    return Err(ParseError {
+                        line: self.line,
+                        message: format!(
+                        "label '{}' already exists for this context: it was previously defined in line {}",
+                        id.value, e.get().node.line),
+                    })
+                }
+                Entry::Vacant(e) => e.insert(PValue {
+                    node: PString {
+                    value: name.clone(),
+                    line: self.line,
+                    range: Range {
+                        start: id.range.start,
+                        end: id.range.end,
+                    },
+                    },
+                    value: 0,
+                    label: true,
+                }),
+            };
+        }
+
+        // And add the node so it's picked up later.
+        self.mapping.push(Node::Label(Label {
+            value: name.to_string(),
+        }));
         Ok(())
     }
 
@@ -990,16 +1144,20 @@ impl Assembler {
                         line: self.line,
                         message: format!(
                         "variable '{}' is being re-assigned: it was previously defined in line {}",
-                        id.value, e.get().clone().line),
+                        id.value, e.get().node.line),
                     })
                 }
-                Entry::Vacant(e) => e.insert(PString {
-                    value: l,
-                    line: self.line,
-                    range: Range {
-                        start: id.range.start,
-                        end: line.len(),
+                Entry::Vacant(e) => e.insert(PValue {
+                    node: PString {
+                        value: l,
+                        line: self.line,
+                        range: Range {
+                            start: id.range.start,
+                            end: line.len(),
+                        },
                     },
+                    value: 0,
+                    label: false,
                 }),
             };
         }
@@ -1184,6 +1342,8 @@ impl Assembler {
                         mode: v.mode.to_owned(),
                         cycles: v.cycles,
                         affected_on_page: v.affected_on_page,
+                        address: 0, // TODO
+                        resolved: true,
                     }))
                 }
 
diff --git a/lib/xixanta/src/context.rs b/lib/xixanta/src/context.rs
index 62ec7e0..aade200 100644
--- a/lib/xixanta/src/context.rs
+++ b/lib/xixanta/src/context.rs
@@ -4,9 +4,16 @@ use std::collections::HashMap;
 const GLOBAL_CONTEXT: &str = "Global";
 
 #[derive(Debug)]
+pub struct PValue {
+    pub node: PString,
+    pub value: usize,
+    pub label: bool,
+}
+
+#[derive(Debug)]
 pub struct Context {
     stack: Vec<String>,
-    map: HashMap<String, HashMap<String, PString>>,
+    map: HashMap<String, HashMap<String, PValue>>,
 }
 
 impl Default for Context {
@@ -27,18 +34,18 @@ impl Context {
         self.stack = vec![];
     }
 
-    pub fn find(&self, name: &str) -> Option<&HashMap<String, PString>> {
+    pub fn find(&self, name: &str) -> Option<&HashMap<String, PValue>> {
         self.map.get(name)
     }
 
-    pub fn current(&self) -> Option<&HashMap<String, PString>> {
+    pub fn current(&self) -> Option<&HashMap<String, PValue>> {
         match self.stack.last() {
             Some(name) => self.map.get(name),
             None => self.map.get(GLOBAL_CONTEXT),
         }
     }
 
-    pub fn current_mut(&mut self) -> Option<&mut HashMap<String, PString>> {
+    pub fn current_mut(&mut self) -> Option<&mut HashMap<String, PValue>> {
         match self.stack.last() {
             Some(name) => self.map.get_mut(name),
             None => self.map.get_mut(GLOBAL_CONTEXT),
diff --git a/lib/xixanta/src/instruction.rs b/lib/xixanta/src/instruction.rs
index 227a3e1..c635644 100644
--- a/lib/xixanta/src/instruction.rs
+++ b/lib/xixanta/src/instruction.rs
@@ -112,8 +112,10 @@ pub struct Instruction {
     pub left: Option<PString>,
     pub right: Option<PString>,
     pub mode: AddressingMode,
-    pub cycles: u8,
-    pub affected_on_page: bool,
+    pub cycles: u8, // NOTE: relative addressing makes this runtime-dependant (if branch is taken, then +1 cycle to the base cycle here).
+    pub affected_on_page: bool, // TODO: needed?
+    pub address: u16,
+    pub resolved: bool,
 }
 
 impl Instruction {
@@ -128,6 +130,8 @@ impl Instruction {
             mode: AddressingMode::Unknown,
             cycles: 0,
             affected_on_page: false,
+            address: 0,
+            resolved: true,
         }
     }
 
@@ -142,6 +146,8 @@ impl Instruction {
             mode: AddressingMode::Unknown,
             cycles: 0,
             affected_on_page: false,
+            address: 0,
+            resolved: true,
         }
     }
 }
@@ -299,12 +305,18 @@ impl Encodable for Fill {
 }
 
 #[derive(Debug, Clone, PartialEq)]
+pub struct Label {
+    pub value: String,
+}
+
+#[derive(Debug, Clone, PartialEq)]
 pub enum Node {
     Generic(Generic),
     Instruction(Instruction),
     Scoped(Scoped),
     Literal(Literal),
     Fill(Fill),
+    Label(Label),
 }
 
 impl Node {
diff --git a/lib/xixanta/src/mapping.rs b/lib/xixanta/src/mapping.rs
index 56acf18..c097fd8 100644
--- a/lib/xixanta/src/mapping.rs
+++ b/lib/xixanta/src/mapping.rs
@@ -16,7 +16,7 @@ pub struct Segment {
 pub struct Mapping {
     pub segments: Vec<Segment>,
     pub nodes: HashMap<String, Vec<Node>>,
-    current: String,
+    pub current: String,
 }
 
 impl Mapping {
diff --git a/lib/xixanta/src/opcodes.rs b/lib/xixanta/src/opcodes.rs
index bca8104..a716545 100644
--- a/lib/xixanta/src/opcodes.rs
+++ b/lib/xixanta/src/opcodes.rs
@@ -20,7 +20,6 @@ pub struct Entry {
 }
 
 lazy_static! {
-    // TODO
     pub static ref INSTRUCTIONS: HashMap<String, HashMap<AddressingMode, ShortEntry>> = {
         let mut instrs = HashMap::new();
 
@@ -57,9 +56,20 @@ lazy_static! {
         asl.insert(AddressingMode::IndexedX, ShortEntry{ cycles: 7, size: 3, opcode: 0x1E, affected_on_page: false });
         instrs.insert(String::from("asl"), asl);
 
-        // BCC: TODO
-        // BCS: TODO
-        // BEQ: TODO
+        // bcc
+        let mut bcc = HashMap::new();
+        bcc.insert(AddressingMode::RelativeOrZeropage, ShortEntry{ cycles: 2, size: 2, opcode: 0x90, affected_on_page: true });
+        instrs.insert(String::from("bcc"), bcc);
+
+        // bcs
+        let mut bcs = HashMap::new();
+        bcs.insert(AddressingMode::RelativeOrZeropage, ShortEntry{ cycles: 2, size: 2, opcode: 0xB0, affected_on_page: true });
+        instrs.insert(String::from("bcs"), bcs);
+
+        // beq
+        let mut beq = HashMap::new();
+        beq.insert(AddressingMode::RelativeOrZeropage, ShortEntry{ cycles: 2, size: 2, opcode: 0xF0, affected_on_page: true });
+        instrs.insert(String::from("beq"), beq);
 
         // bit
         let mut bit = HashMap::new();
@@ -67,17 +77,35 @@ lazy_static! {
         bit.insert(AddressingMode::Absolute, ShortEntry{ cycles: 4, size: 3, opcode: 0x2C, affected_on_page: false });
         instrs.insert(String::from("bit"), bit);
 
-        // BMI: TODO
-        // BNE: TODO
-        // BPL: TODO
+        // bmi
+        let mut bmi = HashMap::new();
+        bmi.insert(AddressingMode::RelativeOrZeropage, ShortEntry{ cycles: 2, size: 2, opcode: 0x30, affected_on_page: true });
+        instrs.insert(String::from("bmi"), bmi);
+
+        // bne
+        let mut bne = HashMap::new();
+        bne.insert(AddressingMode::RelativeOrZeropage, ShortEntry{ cycles: 2, size: 2, opcode: 0xD0, affected_on_page: true });
+        instrs.insert(String::from("bne"), bne);
+
+        // bpl
+        let mut bpl = HashMap::new();
+        bpl.insert(AddressingMode::RelativeOrZeropage, ShortEntry{ cycles: 2, size: 2, opcode: 0x10, affected_on_page: true });
+        instrs.insert(String::from("bpl"), bpl);
 
         // brk
         let mut brk = HashMap::new();
         brk.insert(AddressingMode::Implied, ShortEntry{ cycles: 7, size: 1, opcode: 0x00, affected_on_page: false });
         instrs.insert(String::from("brk"), brk);
 
-        // BVC: TODO
-        // BVS: TODO
+        // bvc
+        let mut bvc = HashMap::new();
+        bvc.insert(AddressingMode::RelativeOrZeropage, ShortEntry{ cycles: 2, size: 2, opcode: 0x50, affected_on_page: true });
+        instrs.insert(String::from("bvc"), bvc);
+
+        // bvs
+        let mut bvs = HashMap::new();
+        bvs.insert(AddressingMode::RelativeOrZeropage, ShortEntry{ cycles: 2, size: 2, opcode: 0x70, affected_on_page: true });
+        instrs.insert(String::from("bvs"), bvs);
 
         // clc
         let mut clc = HashMap::new();
@@ -403,23 +431,36 @@ lazy_static! {
         opcodes.insert(0x0E, Entry { mode: AddressingMode::Absolute, mnemonic: String::from("asl"), cycles: 6, size: 3, opcode: 0x0E, affected_on_page: false });
         opcodes.insert(0x1E, Entry { mode: AddressingMode::IndexedX, mnemonic: String::from("asl"), cycles: 7, size: 3, opcode: 0x1E, affected_on_page: false });
 
-        // BCC
-        // BCS
-        // BEQ
+        // bcc
+        opcodes.insert(0x90, Entry { mode: AddressingMode::RelativeOrZeropage, mnemonic: String::from("bcc"), cycles: 2, size: 2, opcode: 0x90, affected_on_page: false });
+
+        // bcs
+        opcodes.insert(0xB0, Entry { mode: AddressingMode::RelativeOrZeropage, mnemonic: String::from("bcs"), cycles: 2, size: 2, opcode: 0xB0, affected_on_page: false });
+
+        // beq
+        opcodes.insert(0xF0, Entry { mode: AddressingMode::RelativeOrZeropage, mnemonic: String::from("beq"), cycles: 2, size: 2, opcode: 0xF0, affected_on_page: true });
 
         // bit
         opcodes.insert(0x24, Entry { mode: AddressingMode::RelativeOrZeropage, mnemonic: String::from("bit"), cycles: 3, size: 2, opcode: 0x24, affected_on_page: false });
         opcodes.insert(0x2C, Entry { mode: AddressingMode::Absolute, mnemonic: String::from("bit"), cycles: 4, size: 3, opcode: 0x2C, affected_on_page: false });
 
-        // BMI
-        // BNE
-        // BPL
+        // bmi
+        opcodes.insert(0x30, Entry { mode: AddressingMode::RelativeOrZeropage, mnemonic: String::from("bmi"), cycles: 2, size: 2, opcode: 0x30, affected_on_page: false });
+
+        // bne
+        opcodes.insert(0xD0, Entry { mode: AddressingMode::RelativeOrZeropage, mnemonic: String::from("bne"), cycles: 2, size: 2, opcode: 0xD0, affected_on_page: false });
+
+        // bpl
+        opcodes.insert(0x10, Entry { mode: AddressingMode::RelativeOrZeropage, mnemonic: String::from("bpl"), cycles: 2, size: 2, opcode: 0x10, affected_on_page: false });
 
         // brk
         opcodes.insert(0x00, Entry { mode: AddressingMode::Implied, mnemonic: String::from("brk"), cycles: 7, size: 1, opcode: 0x00, affected_on_page: false });
 
-        // BVC
-        // BVS
+        // bvc
+        opcodes.insert(0x50, Entry { mode: AddressingMode::RelativeOrZeropage, mnemonic: String::from("bvc"), cycles: 2, size: 2, opcode: 0x50, affected_on_page: false });
+
+        // bvs
+        opcodes.insert(0x70, Entry { mode: AddressingMode::RelativeOrZeropage, mnemonic: String::from("bvs"), cycles: 2, size: 2, opcode: 0x70, affected_on_page: false });
 
         // clc
         opcodes.insert(0x18, Entry { mode: AddressingMode::Implied, mnemonic: String::from("clc"), cycles: 2, size: 1, opcode: 0x18, affected_on_page: false });