Allow constants in asan:reserve statements

This way, if you define a constant like:

     MY_BUFFER_LEN_IN_BYTES = $10

You can then declare your buffer like so:

    zp_buffer = $00 ; asan:reserve MY_BUFFER_LEN_IN_BYTES

And then further in the code you can rely on just using the constant for
bound checking, and then the address sanitizer will check on bound
checks via static analysis as well.

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

6 files changed, 104 insertions, 29 deletions

diff --git a/lib/xixanta/src/assembler.rs b/lib/xixanta/src/assembler.rs
index 6b82645..8c40afc 100644
--- a/lib/xixanta/src/assembler.rs
+++ b/lib/xixanta/src/assembler.rs
@@ -480,6 +480,22 @@ impl<'a> Assembler<'a> {
                 NodeType::Comment(CommentType::AsanReserve(size)) => {
                     self.asan_next_reserve = *size;
                 }
+                NodeType::Comment(CommentType::AsanReserveIdentifier(id)) => {
+                    match self.context.get_variable(id, &self.mappings) {
+                        Ok(var) => self.asan_next_reserve = var.bundle.value() as usize,
+                        Err(_) => {
+                            errors.push(Error {
+                                line: node.value.line,
+                                message: format!(
+                                    "could not reserve for unknown value of '{}'",
+                                    id.value,
+                                ),
+                                source: self.source_for(node),
+                                global: false,
+                            });
+                        }
+                    }
+                }
                 NodeType::Comment(CommentType::AsanIgnore) => {
                     self.asan_next_ignore = true;
                 }
diff --git a/lib/xixanta/src/node.rs b/lib/xixanta/src/node.rs
index 65bc09e..3f1778a 100644
--- a/lib/xixanta/src/node.rs
+++ b/lib/xixanta/src/node.rs
@@ -259,6 +259,7 @@ pub enum OperationType {
 
 #[derive(Debug, Clone, PartialEq)]
 pub enum CommentType {
+    AsanReserveIdentifier(PString),
     AsanReserve(usize),
     AsanIgnore,
 }
@@ -356,6 +357,7 @@ impl fmt::Display for NodeType {
                 OperationType::Greater => write!(f, "greater"),
             },
             NodeType::Comment(ct) => match ct {
+                CommentType::AsanReserveIdentifier(_) => write!(f, ";; asan:reserve"),
                 CommentType::AsanReserve(_) => write!(f, ";; asan:reserve"),
                 CommentType::AsanIgnore => write!(f, ";; asan:ignore"),
             },
diff --git a/lib/xixanta/src/parser.rs b/lib/xixanta/src/parser.rs
index 68a80ad..cf01133 100644
--- a/lib/xixanta/src/parser.rs
+++ b/lib/xixanta/src/parser.rs
@@ -183,38 +183,12 @@ impl Parser {
                     arg.push(c);
                 }
 
-                // Argument validation.
-                if !arg.starts_with('$') || arg.len() > 5 {
-                    return Err(Error {
-                        line: self.line,
-                        global: false,
-                        source: self.sources[self.current_source].clone(),
-                        message: "expecting a number formatted with a leading '$' sign".to_string(),
-                    }
-                    .into());
-                }
-                let Ok(val) = usize::from_str_radix(arg.get(1..).unwrap_or("0000"), 16) else {
-                    return Err(Error {
-                        line: self.line,
-                        global: false,
-                        source: self.sources[self.current_source].clone(),
-                        message: "could not parse asan:reserve number".to_string(),
-                    }
-                    .into());
-                };
-                if val < 2 {
-                    return Err(Error {
-                        line: self.line,
-                        global: false,
-                        source: self.sources[self.current_source].clone(),
-                        message: "bad asan:reserve number, should be higher than $01".to_string(),
-                    }
-                    .into());
-                }
+                // Parse the given argument as build it as a NodeType.
+                let node_type = self.asan_reserve_from(arg)?;
 
                 // Push whatever was parsed.
                 self.nodes.last_mut().unwrap().push(PNode {
-                    node_type: NodeType::Comment(CommentType::AsanReserve(val)),
+                    node_type,
                     value: PString {
                         value: cmd,
                         line: self.line,
@@ -248,6 +222,63 @@ impl Parser {
         Ok(())
     }
 
+    // Returns the NodeType that can be filled with the given asan:reserve
+    // argument string in 'arg'.
+    fn asan_reserve_from(&mut self, arg: String) -> Result<NodeType, Error> {
+        if arg.starts_with('$') {
+            self.asan_reserve_from_numeric(arg)
+        } else {
+            let s = PString {
+                value: arg,
+                line: self.line,
+                start: 0,
+                end: 0,
+            };
+            if s.is_valid_identifier(true).is_ok() {
+                Ok(NodeType::Comment(CommentType::AsanReserveIdentifier(s)))
+            } else {
+                Err(Error {
+                    line: self.line,
+                    global: false,
+                    source: self.sources[self.current_source].clone(),
+                    message: "expecting a number formatted with a leading '$' sign".to_string(),
+                })
+            }
+        }
+    }
+
+    // Returns the NodeType that can be filled with the given asan:reserve
+    // argument which is assumed to be a numeric value.
+    fn asan_reserve_from_numeric(&mut self, arg: String) -> Result<NodeType, Error> {
+        // Argument validation.
+        if !arg.starts_with('$') || arg.len() > 5 {
+            return Err(Error {
+                line: self.line,
+                global: false,
+                source: self.sources[self.current_source].clone(),
+                message: "expecting a number formatted with a leading '$' sign".to_string(),
+            });
+        }
+        let Ok(val) = usize::from_str_radix(arg.get(1..).unwrap_or("0000"), 16) else {
+            return Err(Error {
+                line: self.line,
+                global: false,
+                source: self.sources[self.current_source].clone(),
+                message: "could not parse asan:reserve number".to_string(),
+            });
+        };
+        if val < 2 {
+            return Err(Error {
+                line: self.line,
+                global: false,
+                source: self.sources[self.current_source].clone(),
+                message: "bad asan:reserve number, should be higher than $01".to_string(),
+            });
+        }
+
+        Ok(NodeType::Comment(CommentType::AsanReserve(val)))
+    }
+
     // Parse a single `line` and push the parsed nodes into `self.nodes`.
     fn parse_line(&mut self, line: &str) -> Result<(), Vec<Error>> {
         self.column = 0;
diff --git a/scripts/test-e2e.sh b/scripts/test-e2e.sh
index 2d9370a..c9821c9 100755
--- a/scripts/test-e2e.sh
+++ b/scripts/test-e2e.sh
@@ -66,6 +66,11 @@ echo "test: custom => variable_names.nes"
 diff tests/out/variable_names.txt tests/expected/variable_names.txt
 exit_code=$((exit_code + $?))
 
+echo "test: custom => asan_reserve_constant.nes"
+./target/debug/nasm -c empty --asan tests/asan_reserve_constant.s 2>tests/out/asan_reserve_constant.txt
+diff tests/out/asan_reserve_constant.txt tests/expected/asan_reserve_constant.txt
+exit_code=$((exit_code + $?))
+
 ##
 # code.nes
 
diff --git a/tests/asan_reserve_constant.s b/tests/asan_reserve_constant.s
new file mode 100644
index 0000000..8af9cad
--- /dev/null
+++ b/tests/asan_reserve_constant.s
@@ -0,0 +1,20 @@
+.segment "HEADER"
+    .byte 'N', 'E', 'S', $1A
+    .byte $02
+    .byte $01
+    .byte $00
+    .byte $00
+
+.segment "CODE"
+
+CONST = $02
+
+zp_var         = $00    ; asan:reserve CONST
+zp_another     = $01
+zp_yet_another = $02
+zp_and_yet     = $03
+
+lda zp_var
+lda zp_another
+lda zp_yet_another
+lda zp_and_yet
diff --git a/tests/expected/asan_reserve_constant.txt b/tests/expected/asan_reserve_constant.txt
new file mode 100644
index 0000000..e08f84f
--- /dev/null
+++ b/tests/expected/asan_reserve_constant.txt
@@ -0,0 +1 @@
+error: The variable 'zp_another' ($01) conflicts with 'zp_var' ($00-$01) (asan_reserve_constant.s)