Implement the asan:stack statement

This statement allows programmers to reserve a memory range for the
stack. This is useful to reserve a memory region which is not attached
to any variable, and:

1. It is a way to safely reserve space on the $0100 page.
2. The --stats/--write-info flags will be able to add up the memory
   space reserved for the stack.
3. Future tooling might be able to use this value as a way to detect
   stack overflows.

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

13 files changed, 348 insertions, 7 deletions

diff --git a/lib/xixanta/src/assembler.rs b/lib/xixanta/src/assembler.rs
index 1b963f2..2f0292e 100644
--- a/lib/xixanta/src/assembler.rs
+++ b/lib/xixanta/src/assembler.rs
@@ -146,6 +146,9 @@ struct Assembler<'a> {
     // The amount of bytes to reserve for the next variable assignment.
     asan_next_reserve: usize,
 
+    // The memory range as defined by an 'asan:stack' statement.
+    asan_stack_definition: Option<MemoryRange>,
+
     // Stack of macro expansions. That is, every time a macro is expanded
     // (i.e. via 'bundle_call'), the node that performed the expansion is
     // stacked here. This is then picked up via structs like 'PendingNode' and
@@ -305,6 +308,27 @@ pub fn assemble_with_mapping(
         };
     }
 
+    // We have enough information for asan:stack after evaluating the
+    // context. So, if asan is enabled, account for it now, otherwise push a
+    // warning.
+    if asm.asan_enabled {
+        match asm.asan_stack_definition {
+            Some(ref mr) => {
+                memory.memory_ranges.push(mr.clone());
+                memory.total_internal_ram += mr.range.len();
+            }
+            None => {
+                asm.warnings.push(Error {
+                    line: 0,
+                    message: "you have not defined a range for 'asan:stack'".to_string(),
+                    source: source.clone(),
+                    global: true,
+                    expanded_from: vec![],
+                });
+            }
+        }
+    }
+
     // Convert the relevant nodes into binary bundles which can be used by
     // the caller. This is done for most nodes, even if some of them will
     // have to be marked as pending, since they depend on knowing the exact
@@ -394,6 +418,7 @@ impl<'a> Assembler<'a> {
             asan_enabled: false,
             asan_next_ignore: false,
             asan_next_reserve: 1,
+            asan_stack_definition: None,
             macro_context: vec![],
             pending_defines: vec![],
         }
@@ -525,6 +550,28 @@ impl<'a> Assembler<'a> {
                         }
                     }
                 }
+                NodeType::Comment(CommentType::AsanStack(range)) => {
+                    match self.asan_stack_definition {
+                        Some(_) => {
+                            if self.asan_stack_definition.is_some() {
+                                errors.push(Error {
+                                    message: "you have already defined a value for 'asan:stack'"
+                                        .to_string(),
+                                    line: node.value.line,
+                                    source: self.source_for(node),
+                                    expanded_from: self.macro_context.clone(),
+                                    global: false,
+                                });
+                            }
+                        }
+                        None => {
+                            self.asan_stack_definition = Some(MemoryRange {
+                                range: range.clone(),
+                                name: "<stack>".to_string(),
+                            });
+                        }
+                    }
+                }
                 NodeType::Comment(CommentType::AsanIgnore) => {
                     self.asan_next_ignore = true;
                 }
@@ -3618,6 +3665,75 @@ mod tests {
     }
 
     #[test]
+    fn asan_no_stack() {
+        let line = "nop";
+        let real_line = minimal_header().to_string() + line;
+
+        let res = assemble_with_mapping(
+            real_line.as_bytes(),
+            empty(),
+            &[],
+            &SourceInfo::default(),
+            true,
+        );
+
+        assert_eq!(res.warnings.len(), 1);
+        assert_eq!(
+            res.warnings[0].to_string(),
+            "you have not defined a range for 'asan:stack'"
+        );
+    }
+
+    #[test]
+    fn asan_dup_stack() {
+        let line = r#";; asan:stack $00-$FF
+;; asan:stack $00-$FF
+"#;
+
+        let real_line = minimal_header().to_string() + line;
+
+        let res = assemble_with_mapping(
+            real_line.as_bytes(),
+            empty(),
+            &[],
+            &SourceInfo::default(),
+            true,
+        );
+
+        assert_eq!(res.errors.len(), 1);
+
+        // NOTE: 3 lines for the "minimal header" being prepended.
+        assert_eq!(
+            res.errors[0].to_string(),
+            "you have already defined a value for 'asan:stack' (line 5)"
+        );
+    }
+
+    #[test]
+    fn asan_stack() {
+        let line = "nop ;; asan:stack $00-$FF";
+        let real_line = minimal_header().to_string() + line;
+
+        let res = assemble_with_mapping(
+            real_line.as_bytes(),
+            empty(),
+            &[],
+            &SourceInfo::default(),
+            true,
+        );
+
+        assert_eq!(res.errors.len(), 0);
+        assert_eq!(res.warnings.len(), 0);
+        assert_eq!(res.memory.memory_ranges.len(), 1);
+
+        let mr = res.memory.memory_ranges.first().unwrap();
+        assert_eq!(mr.range, (0x100..0x200));
+        assert_eq!(mr.name, "<stack>".to_string());
+
+        assert_eq!(res.memory.total_internal_ram, 0x100);
+    }
+
+    #[test]
     fn asan_invalid_negative_variable() {
         let line = "m_var = -1";
         let real_line = minimal_header().to_string() + line;
diff --git a/lib/xixanta/src/node.rs b/lib/xixanta/src/node.rs
index 7a5cf3d..3803840 100644
--- a/lib/xixanta/src/node.rs
+++ b/lib/xixanta/src/node.rs
@@ -1,4 +1,5 @@
 use std::fmt;
+use std::ops::Range;
 
 /// Returns true if this String follows the naming
 /// conventions as expected by the address sanitizer.
@@ -264,6 +265,7 @@ pub enum OperationType {
 pub enum CommentType {
     AsanReserveIdentifier(PString),
     AsanReserve(usize),
+    AsanStack(Range<usize>),
     AsanIgnore,
 }
 
@@ -366,6 +368,7 @@ impl fmt::Display for NodeType {
             NodeType::Comment(ct) => match ct {
                 CommentType::AsanReserveIdentifier(_) => write!(f, ";; asan:reserve"),
                 CommentType::AsanReserve(_) => write!(f, ";; asan:reserve"),
+                CommentType::AsanStack(_) => write!(f, ";; asan:stack"),
                 CommentType::AsanIgnore => write!(f, ";; asan:ignore"),
             },
         }
diff --git a/lib/xixanta/src/parser.rs b/lib/xixanta/src/parser.rs
index aca5dba..9f8ffc7 100644
--- a/lib/xixanta/src/parser.rs
+++ b/lib/xixanta/src/parser.rs
@@ -201,6 +201,39 @@ impl Parser {
                     source: self.current_source,
                 });
             }
+            "asan:stack" => {
+                // Skip whitespaces.
+                for c in line.get(offset..).unwrap_or("").chars() {
+                    if c != ';' && !c.is_whitespace() {
+                        break;
+                    }
+                    offset += 1;
+                }
+
+                // Fetch the argument.
+                let mut arg = String::from("");
+                for c in line.get(offset..).unwrap_or("").chars() {
+                    if c.is_whitespace() {
+                        break;
+                    }
+                    arg.push(c);
+                }
+
+                let node_type = self.asan_stack_from(arg)?;
+                self.nodes.last_mut().unwrap().push(PNode {
+                    node_type,
+                    value: PString {
+                        value: cmd,
+                        line: self.line,
+                        start,
+                        end: offset - 1,
+                    },
+                    left: None,
+                    right: None,
+                    args: None,
+                    source: self.current_source,
+                });
+            }
             "asan:ignore" => {
                 self.nodes.last_mut().unwrap().push(PNode {
                     node_type: NodeType::Comment(CommentType::AsanIgnore),
@@ -241,13 +274,112 @@ impl Parser {
                     line: self.line,
                     global: false,
                     source: self.sources[self.current_source].clone(),
-                    message: "expecting a number formatted with a leading '$' sign".to_string(),
+                    message: "'asan:reserve' expects a number formatted with a leading '$' sign"
+                        .to_string(),
                     expanded_from: vec![],
                 })
             }
         }
     }
 
+    // Returns the NodeType that can be filled with the given asan:stack
+    // argument which is assumed to be a memory range from the $0100 page.
+    fn asan_stack_from(&mut self, arg: String) -> Result<NodeType, Error> {
+        // Maybe it's the "full" shorthand. If that's the case return early.
+        if arg.to_lowercase() == "full" {
+            return Ok(NodeType::Comment(CommentType::AsanStack(0x100..0x200)));
+        }
+
+        // Validate that the first address is in hexadecimal format.
+        if !arg.starts_with('$') {
+            return Err(Error {
+                line: self.line,
+                global: false,
+                source: self.sources[self.current_source].clone(),
+                message: "'asan:stack' expects a number formatted with a leading '$' sign"
+                    .to_string(),
+                expanded_from: vec![],
+            });
+        }
+
+        // Ensure that it's a range of addresses and that the second address is
+        // in hexadecimal format.
+        let Some(sep) = arg.find('-') else {
+            return Err(Error {
+                line: self.line,
+                global: false,
+                source: self.sources[self.current_source].clone(),
+                message: "'asan:stack' expects a memory range".to_string(),
+                expanded_from: vec![],
+            });
+        };
+        let Some(next) = arg.get(sep..).unwrap_or("$0000").find('$') else {
+            return Err(Error {
+                line: self.line,
+                global: false,
+                source: self.sources[self.current_source].clone(),
+                message: "'asan:stack' expects a number formatted with a leading '$' sign"
+                    .to_string(),
+                expanded_from: vec![],
+            });
+        };
+
+        // Now we can extract the proper hexadecimal values from it.
+        let Ok(mut first) = usize::from_str_radix(arg.get(1..sep).unwrap_or("0000"), 16) else {
+            return Err(Error {
+                line: self.line,
+                global: false,
+                source: self.sources[self.current_source].clone(),
+                expanded_from: vec![],
+                message: "could not parse 'asan:stack' number".to_string(),
+            });
+        };
+        let Ok(mut second) = usize::from_str_radix(arg.get(sep + next + 1..).unwrap_or("0000"), 16)
+        else {
+            return Err(Error {
+                line: self.line,
+                global: false,
+                source: self.sources[self.current_source].clone(),
+                expanded_from: vec![],
+                message: "could not parse 'asan:stack' number".to_string(),
+            });
+        };
+
+        // If the given addresses are below what a byte can contain, assume that
+        // page $1xx was implicit.
+        if first <= u8::MAX.into() && second <= u8::MAX.into() {
+            first += 0x100;
+            second += 0x100;
+        }
+
+        // Validate that we are on the $0100 page.
+        if !(0x100..0x200).contains(&first) || !(0x100..0x200).contains(&second) {
+            return Err(Error {
+                line: self.line,
+                global: false,
+                source: self.sources[self.current_source].clone(),
+                expanded_from: vec![],
+                message: "the stack must be on the $0100 page".to_string(),
+            });
+        }
+
+        // Since the stack grows to lower addresses, we allow the programmer to
+        // put the memory range in any order they see fit.
+        if first > second {
+            Ok(NodeType::Comment(CommentType::AsanStack(second..first + 1)))
+        } else if first < second {
+            Ok(NodeType::Comment(CommentType::AsanStack(first..second + 1)))
+        } else {
+            Err(Error {
+                line: self.line,
+                global: false,
+                source: self.sources[self.current_source].clone(),
+                expanded_from: vec![],
+                message: "bad 'asan:stack' range".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> {
@@ -257,7 +389,8 @@ impl Parser {
                 line: self.line,
                 global: false,
                 source: self.sources[self.current_source].clone(),
-                message: "expecting a number formatted with a leading '$' sign".to_string(),
+                message: "'asan:reserve' expects a number formatted with a leading '$' sign"
+                    .to_string(),
                 expanded_from: vec![],
             });
         }
@@ -267,7 +400,7 @@ impl Parser {
                 global: false,
                 source: self.sources[self.current_source].clone(),
                 expanded_from: vec![],
-                message: "could not parse asan:reserve number".to_string(),
+                message: "could not parse 'asan:reserve' number".to_string(),
             });
         };
         if val < 2 {
@@ -276,7 +409,7 @@ impl Parser {
                 global: false,
                 source: self.sources[self.current_source].clone(),
                 expanded_from: vec![],
-                message: "bad asan:reserve number, should be higher than $01".to_string(),
+                message: "bad 'asan:reserve' number, should be higher than $01".to_string(),
             });
         }
 
@@ -3495,15 +3628,84 @@ VAR3 = $200 ;; asan:reserve $100
 
         assert_eq!(
             errors.first().unwrap().message,
-            "expecting a number formatted with a leading '$' sign"
+            "'asan:reserve' expects a number formatted with a leading '$' sign"
+        );
+        assert_eq!(
+            errors.get(1).unwrap().message,
+            "bad 'asan:reserve' number, should be higher than $01"
+        );
+        assert_eq!(
+            errors.get(2).unwrap().message,
+            "bad 'asan:reserve' number, should be higher than $01"
+        );
+    }
+
+    #[test]
+    fn parse_comment_asan_stack() {
+        let code = r#";; asan:stack full
+;; asan:stack $00-$FF
+;; asan:stack $100-$1FF
+"#;
+        let mut parser = Parser::default();
+        assert!(parser
+            .parse(code.as_bytes(), &SourceInfo::default())
+            .is_ok());
+
+        let nodes = parser.nodes();
+        assert_eq!(nodes.len(), 3);
+
+        assert_node(
+            nodes.first().unwrap(),
+            NodeType::Comment(CommentType::AsanStack(256..512)),
+            code,
+            "asan:stack",
+        );
+        assert_node(
+            nodes.get(1).unwrap(),
+            NodeType::Comment(CommentType::AsanStack(256..512)),
+            code,
+            "asan:stack",
+        );
+        assert_node(
+            nodes.get(2).unwrap(),
+            NodeType::Comment(CommentType::AsanStack(256..512)),
+            code,
+            "asan:stack",
+        );
+    }
+
+    #[test]
+    fn parse_bad_asan_stack() {
+        let code = r#";; asan:stack 02
+;; asan:stack $00
+;; asan:stack $01-02
+;; asan:stack $0200-$02FF
+;; asan:stack $00-$00
+"#;
+
+        let mut parser = Parser::default();
+        let res = parser.parse(code.as_bytes(), &SourceInfo::default());
+
+        assert!(res.is_err());
+        let errors = res.unwrap_err();
+        assert_eq!(errors.len(), 5);
+
+        assert_eq!(
+            errors.first().unwrap().message,
+            "'asan:stack' expects a number formatted with a leading '$' sign"
         );
         assert_eq!(
             errors.get(1).unwrap().message,
-            "bad asan:reserve number, should be higher than $01"
+            "'asan:stack' expects a memory range"
         );
         assert_eq!(
             errors.get(2).unwrap().message,
-            "bad asan:reserve number, should be higher than $01"
+            "'asan:stack' expects a number formatted with a leading '$' sign"
+        );
+        assert_eq!(
+            errors.get(3).unwrap().message,
+            "the stack must be on the $0100 page"
         );
+        assert_eq!(errors.get(4).unwrap().message, "bad 'asan:stack' range");
     }
 }
diff --git a/tests/asan_reserve_constant.s b/tests/asan_reserve_constant.s
index 8af9cad..49f461b 100644
--- a/tests/asan_reserve_constant.s
+++ b/tests/asan_reserve_constant.s
@@ -18,3 +18,5 @@ lda zp_var
 lda zp_another
 lda zp_yet_another
 lda zp_and_yet
+
+;;; asan:stack full
diff --git a/tests/avoid_bad_macro.s b/tests/avoid_bad_macro.s
index f701c2b..cd310a8 100644
--- a/tests/avoid_bad_macro.s
+++ b/tests/avoid_bad_macro.s
@@ -14,3 +14,5 @@
 .endif
 
 lda #1
+
+;;; asan:stack full
diff --git a/tests/bad_fallthrough.s b/tests/bad_fallthrough.s
index 80c35c3..43dfa18 100644
--- a/tests/bad_fallthrough.s
+++ b/tests/bad_fallthrough.s
@@ -25,3 +25,5 @@ __fallthrough__ other
 .endproc
 
 __fallthrough__ other
+
+;;; asan:stack full
diff --git a/tests/bad_jal.s b/tests/bad_jal.s
index 693fa7d..5fa2316 100644
--- a/tests/bad_jal.s
+++ b/tests/bad_jal.s
@@ -16,3 +16,5 @@
 .proc bar
     rts
 .endproc
+
+;;; asan:stack full
diff --git a/tests/bare_accesses.s b/tests/bare_accesses.s
index cebdc10..8b05fb0 100644
--- a/tests/bare_accesses.s
+++ b/tests/bare_accesses.s
@@ -42,3 +42,5 @@ lda $200
 
 jmp $8000
 jmp ($8000)
+
+;;; asan:stack full
diff --git a/tests/delayed_macro_arguments.s b/tests/delayed_macro_arguments.s
index b91ec92..acd3910 100644
--- a/tests/delayed_macro_arguments.s
+++ b/tests/delayed_macro_arguments.s
@@ -47,3 +47,5 @@
         rts
     .endproc
 .endscope
+
+;;; asan:stack full
diff --git a/tests/fallthrough.s b/tests/fallthrough.s
index 6e95b5e..2243415 100644
--- a/tests/fallthrough.s
+++ b/tests/fallthrough.s
@@ -21,3 +21,5 @@ __fallthrough__ bar
 .proc other
     rts
 .endproc
+
+;;; asan:stack full
diff --git a/tests/indirect.s b/tests/indirect.s
index 9629d5f..2237b7e 100644
--- a/tests/indirect.s
+++ b/tests/indirect.s
@@ -16,3 +16,5 @@ jmp (zp_something)
 
 lala:
     rts
+
+;;; asan:stack full
diff --git a/tests/unused.s b/tests/unused.s
index c26ac42..938794d 100644
--- a/tests/unused.s
+++ b/tests/unused.s
@@ -16,3 +16,5 @@ zp_yet     = $04
 lda zp_used
 lda (zp_more), y
 lda #.lobyte(zp_yet)
+
+;;; asan:stack full
diff --git a/tests/variable_names.s b/tests/variable_names.s
index beda2b8..a1dc578 100644
--- a/tests/variable_names.s
+++ b/tests/variable_names.s
@@ -36,3 +36,5 @@ lda wr_out
 .endmacro
 
 INC_MOVEMENT_X zp_good + 1
+
+;;; asan:stack full