nasm: Implement the -D flag

This allows users to define variables directly from the command line,
which is useful for testing purposes.

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

8 files changed, 134 insertions, 5 deletions

diff --git a/crates/nasm/src/main.rs b/crates/nasm/src/main.rs
index ef7c7d4..5179397 100644
--- a/crates/nasm/src/main.rs
+++ b/crates/nasm/src/main.rs
@@ -14,6 +14,7 @@ struct Args {
     out: Option<String>,
     werror: bool,
     stdout: bool,
+    defines: Vec<(String, u8)>,
 }
 
 // Print the help message and quit.
@@ -22,12 +23,46 @@ fn print_help() {
     println!("usage: nasm [OPTIONS] <FILE>\n");
     println!("Options:");
     println!("  -c, --config <FILE>\tLinker configuration to be used, whether an identifier or a file path.");
+    println!("  -D <NAME>(=VALUE)\tDefine an 8-bit variable on the global scope (default: 1)");
     println!("  -o, --out <FILE>\tFile path where the output should be located after execution.");
     println!("  --stdout\t\tPrint the output binary to the standard output.");
     println!("  -Werror\t\tWarnings should be treated as errors.");
     std::process::exit(0);
 }
 
+// Parse a value from the '-D' flag which is expected to be 'NAME(=VALUE)'.
+fn parse_define(arg: &str) -> (String, u8) {
+    let mut key_value = arg.split('=');
+
+    let Some(name) = key_value.next() else {
+        die(format!("bad format for define '{}'", arg));
+        return (String::default(), 0);
+    };
+
+    if name
+        .chars()
+        .any(|c| !c.is_ascii_alphanumeric() && c != '_' && c != '@' && c != '.')
+    {
+        die(format!(
+            "trying to define '{}' which has invalid characters",
+            arg
+        ));
+    }
+
+    let value = match key_value.next().unwrap_or("1").parse::<u8>() {
+        Ok(integer) => integer,
+        Err(_) => {
+            die(format!(
+                "value for define '{}' must be a valid 8-bit integer",
+                arg
+            ));
+            return (String::default(), 0);
+        }
+    };
+
+    (name.to_string(), value)
+}
+
 // Parse the arguments given to the program and returns an Args object with the
 // given information.
 fn parse_arguments() -> Args {
@@ -48,6 +83,10 @@ fn parse_arguments() -> Args {
                     }
                 },
             },
+            "-D" => match args.next() {
+                Some(a) => res.defines.push(parse_define(&a)),
+                None => die("you need to provide a value for the '-D' flag".to_string()),
+            },
             "-h" | "--help" => print_help(),
             "-o" | "--out" => match res.out {
                 Some(_) => die("only specify the '-o/--out' flag once".to_string()),
@@ -146,6 +185,7 @@ fn main() {
     let res = assemble(
         input,
         args.config.unwrap_or("nrom".to_string()).as_str(),
+        &args.defines,
         source,
     );
 
diff --git a/lib/xixanta/fuzz/fuzz_targets/fuzz_target_assembler.rs b/lib/xixanta/fuzz/fuzz_targets/fuzz_target_assembler.rs
index a402998..ff389a1 100644
--- a/lib/xixanta/fuzz/fuzz_targets/fuzz_target_assembler.rs
+++ b/lib/xixanta/fuzz/fuzz_targets/fuzz_target_assembler.rs
@@ -4,5 +4,5 @@ use libfuzzer_sys::fuzz_target;
 use xixanta::assembler::assemble;
 
 fuzz_target!(|data: &[u8]| {
-    let _ = assemble(data, "empty", xixanta::SourceInfo::default());
+    let _ = assemble(data, "empty", &[], xixanta::SourceInfo::default());
 });
diff --git a/lib/xixanta/src/assembler.rs b/lib/xixanta/src/assembler.rs
index dfaff19..4ee769b 100644
--- a/lib/xixanta/src/assembler.rs
+++ b/lib/xixanta/src/assembler.rs
@@ -1,5 +1,5 @@
 use crate::mapping::{get_mapping_configuration, Mapping};
-use crate::node::{ControlType, NodeType, OperationType, PNode};
+use crate::node::{ControlType, NodeType, OperationType, PNode, PString};
 use crate::object::{Bundle, Context, Object, ObjectType};
 use crate::opcodes::{AddressingMode, INSTRUCTIONS};
 use crate::parser::Parser;
@@ -96,7 +96,12 @@ pub struct AssemblerResult {
 /// statements like ".import" or ".incbin" wouldn't know how to resolve relative
 /// paths. You can specify the mapper to be used as an identifier in `mapping`,
 /// which will be handled via `get_mapping_configuration`.
-pub fn assemble(reader: impl Read, mapping: &str, source: SourceInfo) -> AssemblerResult {
+pub fn assemble(
+    reader: impl Read,
+    mapping: &str,
+    defines: &[(String, u8)],
+    source: SourceInfo,
+) -> AssemblerResult {
     let config = match get_mapping_configuration(mapping) {
         Ok(config) => config,
         Err(e) => {
@@ -113,7 +118,7 @@ pub fn assemble(reader: impl Read, mapping: &str, source: SourceInfo) -> Assembl
         }
     };
 
-    assemble_with_mapping(reader, config, source)
+    assemble_with_mapping(reader, config, defines, source)
 }
 
 /// Read the contents from the `reader` as a source file and produce a list of
@@ -125,6 +130,7 @@ pub fn assemble(reader: impl Read, mapping: &str, source: SourceInfo) -> Assembl
 pub fn assemble_with_mapping(
     reader: impl Read,
     mapping: Vec<Mapping>,
+    defines: &[(String, u8)],
     source: SourceInfo,
 ) -> AssemblerResult {
     let mut asm = Assembler::new(mapping);
@@ -143,6 +149,18 @@ pub fn assemble_with_mapping(
     let nodes = parser.nodes();
     asm.sources = parser.sources;
 
+    // Before building a context, add variables if they were defined by the
+    // caller.
+    for (name, value) in defines {
+        if let Err(e) = asm.define_variable_value(name, *value) {
+            return AssemblerResult {
+                bundles: vec![],
+                errors: e.into(),
+                warnings: vec![],
+            };
+        }
+    }
+
     // Build the context by iterating over the parsed nodes and checking
     // where scopes start/end, evaluating values for variables, labels, etc.
     if let Err(errors) = asm.eval_context(&nodes) {
@@ -210,6 +228,42 @@ impl<'a> Assembler<'a> {
         }
     }
 
+    /// Define a new variable by using the given `name` and `value`.
+    pub fn define_variable_value(&mut self, name: &String, value: u8) -> Result<(), Error> {
+        if name.is_empty() {
+            return Err(Error {
+                global: true,
+                line: 0,
+                source: self.sources.first().unwrap().clone(),
+                message: "empty variable name".to_string(),
+            });
+        }
+
+        let var_name = PString {
+            value: name.to_string(),
+            line: 0,
+            start: 0,
+            end: name.len(),
+        };
+        let var_value = Object {
+            bundle: Bundle::fill(value),
+            mapping: self.current_mapping,
+            segment: self.current_segment,
+            object_type: ObjectType::Value,
+        };
+
+        if let Err(err) = self.context.set_variable(&var_name, &var_value, false) {
+            Err(Error {
+                global: true,
+                line: 0,
+                source: self.sources.first().unwrap().clone(),
+                message: err,
+            })
+        } else {
+            Ok(())
+        }
+    }
+
     // Define a new variable by taking the given `id`. This variable will only
     // be created if `id` is not empty. The function will error out if the given
     // name is already taken.
@@ -2201,7 +2255,7 @@ mod tests {
         // the assembler will freak out.
         let real_line = minimal_header().to_string() + line;
 
-        assemble_with_mapping(real_line.as_bytes(), empty(), SourceInfo::default())
+        assemble_with_mapping(real_line.as_bytes(), empty(), &[], SourceInfo::default())
     }
 
     // Like `just_assemble` but it only returns bundles passed the header.
@@ -3965,6 +4019,7 @@ lda #Variable
     "#
             .as_bytes(),
             one_two().to_vec(),
+            &[],
             SourceInfo::default(),
         );
 
@@ -4006,6 +4061,7 @@ lda #Variable
     "#
             .as_bytes(),
             one_two().to_vec(),
+            &[],
             SourceInfo::default(),
         );
 
@@ -4079,6 +4135,7 @@ lda #Variable
     "#
             .as_bytes(),
             one_two().to_vec(),
+            &[],
             SourceInfo::default(),
         );
 
@@ -4126,6 +4183,7 @@ lda #Variable
     "#
             .as_bytes(),
             one_two().to_vec(),
+            &[],
             SourceInfo::default(),
         );
 
@@ -4180,6 +4238,7 @@ lda #Variable
     "#
             .as_bytes(),
             one_two().to_vec(),
+            &[],
             SourceInfo::default(),
         );
 
@@ -4206,6 +4265,7 @@ lda #Variable
         "#
             .as_bytes(),
             one_two().to_vec(),
+            &[],
             SourceInfo::default(),
         );
 
diff --git a/scripts/test-e2e.sh b/scripts/test-e2e.sh
index 9ac7d04..2185768 100755
--- a/scripts/test-e2e.sh
+++ b/scripts/test-e2e.sh
@@ -15,6 +15,17 @@ popd
 
 cargo build
 
+# Custom
+
+./target/debug/nasm -c empty -Werror -o tests/out/defines-undefined.nes tests/defines.s
+diff tests/out/defines-undefined.nes tests/expected/defines-undefined.nes
+
+./target/debug/nasm -D LALA -c empty -Werror -o tests/out/defines-one.nes tests/defines.s
+diff tests/out/defines-one.nes tests/expected/defines-one.nes
+
+./target/debug/nasm -D LALA=2 -c empty -Werror -o tests/out/defines-two.nes tests/defines.s
+diff tests/out/defines-two.nes tests/expected/defines-two.nes
+
 # code.nes
 
 ./target/debug/nasm -c nrom -Werror -o tests/out/sprite.nes tests/code.nes/basics/sprite.s
diff --git a/tests/defines.s b/tests/defines.s
new file mode 100644
index 0000000..6f92756
--- /dev/null
+++ b/tests/defines.s
@@ -0,0 +1,18 @@
+.segment "HEADER"
+    .byte 'N', 'E', 'S', $1A
+    .byte $02
+    .byte $01
+    .byte $00
+    .byte $00
+
+.segment "CODE"
+
+.ifdef LALA
+    .if LALA == 1
+        lda #0
+    .else
+        lda #1
+    .endif
+.else
+    lda #2
+.endif
diff --git a/tests/expected/defines-one.nes b/tests/expected/defines-one.nes
new file mode 100644
index 0000000..051a8a8
--- /dev/null
+++ b/tests/expected/defines-one.nes
diff --git a/tests/expected/defines-two.nes b/tests/expected/defines-two.nes
new file mode 100644
index 0000000..eb99eb1
--- /dev/null
+++ b/tests/expected/defines-two.nes
diff --git a/tests/expected/defines-undefined.nes b/tests/expected/defines-undefined.nes
new file mode 100644
index 0000000..40c3987
--- /dev/null
+++ b/tests/expected/defines-undefined.nes