Add support for the asan:fixed-segments comment

This magic comment allows for the definition of segments that are fixed
and for which references are always safe. This goes in tandem with the
asan:safe comment, which can then be used more sporadically.

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

8 files changed, 370 insertions, 24 deletions

diff --git a/crates/nasm/README.md b/crates/nasm/README.md
index 2d24011..2675bfb 100644
--- a/crates/nasm/README.md
+++ b/crates/nasm/README.md
@@ -312,6 +312,48 @@ now be written like so:
 jsr foo      ; check:safe
 ```
 
+Moreover, you can define whole segments as "fixed" ones with the
+`asan:fixed-segments` (or `check:fixed-segments`) comments. This way, you
+express to the assembler that references to addresses of these segments are
+guaranteed to always be valid. Consider the following example:
+
+```asm
+;;; asan:fixed-segments ONE, OTHER
+
+.segment "ONE"
+
+.proc foo
+    rts
+.endproc
+
+.segment "OTHER"
+
+.proc bar
+    rts
+.endproc
+
+.segment "FIXED"
+
+jsr foo
+jsr bar
+```
+
+In the code above, we state that both 'ONE' and 'OTHER' are guaranteed to have a
+stable address space (they are fixed, never to be re-mapped). Hence, the
+assembler won't spit any warning at the final two 'jsr' instructions. Also note
+that you can define this comment multiple times. So the code below achieves the
+same thing:
+
+```asm
+;;; asan:fixed-segments ONE
+.segment "ONE"
+;; bla bla
+
+;;; asan:fixed-segments OTHER
+.segment "OTHER"
+;; rest
+```
+
 ## Address sanitizer
 
 This assembler comes with a set of tools that builds up an "address
diff --git a/lib/xixanta/src/assembler.rs b/lib/xixanta/src/assembler.rs
index 98ae19d..26d12b2 100644
--- a/lib/xixanta/src/assembler.rs
+++ b/lib/xixanta/src/assembler.rs
@@ -184,6 +184,10 @@ struct Assembler<'a> {
     // is a tuple where the first member is the String identifier of the
     // variable/address/label/etc. that is being stored in the second member.
     objects_visited: Vec<(String, Object)>,
+
+    // List of mapping/segment indeces which the programmer asked to be
+    // considered safe.
+    safe_segments: Vec<(usize, usize)>,
 }
 
 /// The result to be given at the end of `assembler::assemble` and
@@ -453,6 +457,7 @@ impl<'a> Assembler<'a> {
             macro_context: vec![],
             pending_defines: vec![],
             objects_visited: vec![],
+            safe_segments: vec![],
         }
     }
 
@@ -950,6 +955,14 @@ impl<'a> Assembler<'a> {
                     next_safe = true;
                     self.asan_next_safe = true;
                 }
+                NodeType::Comment(CommentType::AsanFixedSegments(segments)) => {
+                    for name in segments {
+                        let (midx, sidx) = self.get_mapping_segment_by_name(node, name)?;
+                        if !self.is_in_safe_segment(midx, sidx) {
+                            self.safe_segments.push((midx, sidx));
+                        }
+                    }
+                }
                 NodeType::Instruction => {
                     self.literal_mode = None;
                     match self.evaluate_node(node) {
@@ -1107,7 +1120,7 @@ impl<'a> Assembler<'a> {
                     bundle.address = current.bundles[pn.bundle_index].address;
 
                     // The first check deals on whether there is any reference
-                    // that crosses the mapping boundary. There are two
+                    // that crosses the mapping boundary. There are some
                     // exceptions to this check:
                     //
                     //  1. We are currently on the Vector segment, which
@@ -1115,6 +1128,9 @@ impl<'a> Assembler<'a> {
                     //     that is to be expected and safe.
                     //  2. If the bundle was marked as "safe" by the programmer
                     //     via the asan:safe/check:safe comment.
+                    //  3. The object belongs to a segment that the programmer
+                    //     explicitely marked as safe via the
+                    //     asan:fixed-segments/check:fixed-segments comment.
                     //
                     // Otherwise iterate over the objects that were visited over
                     // the previous evaluate_node() call, and check for
@@ -1125,7 +1141,9 @@ impl<'a> Assembler<'a> {
                         && !matches!(self.mappings[pn.mapping].section_type, SectionType::Vector)
                     {
                         for (name, object) in &self.objects_visited {
-                            if pn.mapping != object.mapping {
+                            if pn.mapping != object.mapping
+                                && !self.is_in_safe_segment(object.mapping, object.segment)
+                            {
                                 let reference =
                                     &self.mappings[object.mapping].segments[object.segment].name;
 
@@ -2925,6 +2943,39 @@ impl<'a> Assembler<'a> {
         Ok(())
     }
 
+    // Returns true if the given mapping/segment pair were marked by the
+    // programmer to be safe.
+    fn is_in_safe_segment(&self, mapping: usize, segment: usize) -> bool {
+        self.safe_segments
+            .iter()
+            .any(|s| s.0 == mapping && s.1 == segment)
+    }
+
+    // Return the mapping and segment indeces which identify the given segment
+    // identified by 'name'. If it cannot be found, then it returns an error by
+    // pinning the location of the given 'node'.
+    fn get_mapping_segment_by_name(
+        &self,
+        node: &'a PNode,
+        name: &str,
+    ) -> Result<(usize, usize), Error> {
+        for (mapping_idx, mapping) in self.mappings.iter().enumerate() {
+            for (segment_idx, segment) in mapping.segments.iter().enumerate() {
+                if segment.name == name {
+                    return Ok((mapping_idx, segment_idx));
+                }
+            }
+        }
+
+        Err(Error {
+            line: node.value.line,
+            message: format!("unknown segment '{name}'"),
+            source: self.source_for(node),
+            expanded_from: self.macro_context.clone(),
+            global: false,
+        })
+    }
+
     // Change the current segment to the one referenced in `node`.
     fn switch_to_segment(&mut self, node: &'a PNode) -> Result<(), Error> {
         let name = self.fetch_quoted_first_argument(node)?;
@@ -2941,28 +2992,10 @@ impl<'a> Assembler<'a> {
             });
         }
 
-        // Find the segment being referenced and update the
-        // `self.current_segment` accordingly.
-        let mut found = false;
-        for (mapping_idx, mapping) in self.mappings.iter().enumerate() {
-            for (segment_idx, segment) in mapping.segments.iter().enumerate() {
-                if segment.name == name {
-                    self.current_mapping = mapping_idx;
-                    self.current_segment = segment_idx;
-                    found = true;
-                    break;
-                }
-            }
-        }
-        if !found {
-            return Err(Error {
-                line: node.value.line,
-                message: format!("unknown segment '{name}'"),
-                source: self.source_for(node),
-                expanded_from: self.macro_context.clone(),
-                global: false,
-            });
-        }
+        let (midx, sidx) = self.get_mapping_segment_by_name(node, name)?;
+        self.current_mapping = midx;
+        self.current_segment = sidx;
+
         Ok(())
     }
 
diff --git a/lib/xixanta/src/node.rs b/lib/xixanta/src/node.rs
index 4f25158..649f333 100644
--- a/lib/xixanta/src/node.rs
+++ b/lib/xixanta/src/node.rs
@@ -276,6 +276,7 @@ pub enum CommentType {
     AsanStack(Range<usize>),
     AsanIgnore,
     AsanSafe,
+    AsanFixedSegments(Vec<String>),
 }
 
 /// The PNode type.
@@ -380,6 +381,7 @@ impl fmt::Display for NodeType {
                 CommentType::AsanStack(_) => write!(f, ";; asan:stack"),
                 CommentType::AsanIgnore => write!(f, ";; asan:ignore"),
                 CommentType::AsanSafe => write!(f, ";; asan:safe"),
+                CommentType::AsanFixedSegments(_) => write!(f, ";; asan:fixed-segments"),
             },
         }
     }
diff --git a/lib/xixanta/src/parser.rs b/lib/xixanta/src/parser.rs
index 03039f1..abcf1e8 100644
--- a/lib/xixanta/src/parser.rs
+++ b/lib/xixanta/src/parser.rs
@@ -264,6 +264,77 @@ impl Parser {
                     source: self.current_source,
                 });
             }
+            "check:fixed-segments" | "asan:fixed-segments" => {
+                let mut args = vec![];
+                let mut eol = false;
+                let end = offset;
+
+                while !eol {
+                    // We will assume that this is the last argument unless the
+                    // last loop sees a comma.
+                    eol = true;
+
+                    // Skip initial whitespaces.
+                    for c in line.get(offset..).unwrap_or("").chars() {
+                        if !c.is_whitespace() {
+                            break;
+                        }
+                        offset += 1;
+                    }
+
+                    // Fetch the argument and push it unless it's empty.
+                    let mut arg = String::from("");
+                    for c in line.get(offset..).unwrap_or("").chars() {
+                        if c.is_whitespace() {
+                            offset += 1;
+                            break;
+                        } else if c == ',' {
+                            break;
+                        }
+                        offset += 1;
+
+                        arg.push(c);
+                    }
+                    if !arg.is_empty() {
+                        args.push(arg);
+                    }
+
+                    // Skip whitespaces until a comma if available.
+                    for c in line.get(offset..).unwrap_or("").chars() {
+                        offset += 1;
+
+                        if c == ',' {
+                            eol = false;
+                            break;
+                        }
+                    }
+                }
+
+                if args.is_empty() {
+                    return Err(Error {
+                        line: self.line,
+                        global: false,
+                        source: self.sources[self.current_source].clone(),
+                        message: "'asan:fixed-segments' expects at least one argument".to_string(),
+                        expanded_from: vec![],
+                    }
+                    .into());
+                }
+
+                self.nodes.last_mut().unwrap().push(PNode {
+                    node_type: NodeType::Comment(CommentType::AsanFixedSegments(args)),
+                    value: PString {
+                        value: cmd,
+                        line: self.line,
+                        start,
+                        end,
+                    },
+                    left: None,
+                    right: None,
+                    args: None,
+                    source: self.current_source,
+                });
+            }
             &_ => {}
         }
 
@@ -3843,4 +3914,63 @@ VAR3 = $200 ;; asan:reserve $100
         );
         assert_eq!(errors.get(4).unwrap().message, "bad 'asan:stack' range");
     }
+
+    #[test]
+    fn parse_fixed_segments() {
+        let code = r#";; asan:fixed-segments CODE
+;; asan:fixed-segments CODE , FIXED
+;; asan:fixed-segments TAIL, FIXED,
+"#;
+        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::AsanFixedSegments(vec!["CODE".to_string()])),
+            code,
+            "asan:fixed-segments",
+        );
+        assert_node(
+            nodes.get(1).unwrap(),
+            NodeType::Comment(CommentType::AsanFixedSegments(vec![
+                "CODE".to_string(),
+                "FIXED".to_string(),
+            ])),
+            code,
+            "asan:fixed-segments",
+        );
+        assert_node(
+            nodes.get(2).unwrap(),
+            NodeType::Comment(CommentType::AsanFixedSegments(vec![
+                "TAIL".to_string(),
+                "FIXED".to_string(),
+            ])),
+            code,
+            "asan:fixed-segments",
+        );
+    }
+
+    #[test]
+    fn parse_bad_asan_fixed_segments() {
+        let code = ";; asan:fixed-segments";
+
+        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(), 1);
+
+        assert_eq!(
+            errors.first().unwrap().message,
+            "'asan:fixed-segments' expects at least one argument"
+        );
+    }
 }
diff --git a/scripts/test-e2e.sh b/scripts/test-e2e.sh
index 43572b0..2a72d38 100755
--- a/scripts/test-e2e.sh
+++ b/scripts/test-e2e.sh
@@ -133,6 +133,13 @@ exit_code=$((exit_code + $?))
 diff tests/out/unrom.nes tests/expected/unrom.nes
 exit_code=$((exit_code + $?))
 
+echo "test: custom => fixed-segments.nes"
+./target/debug/nasm -c unrom tests/fixed-segments.s -o tests/out/fixed-segments.nes 2>tests/out/fixed-segments.txt
+diff tests/out/fixed-segments.txt tests/expected/fixed-segments.txt
+exit_code=$((exit_code + $?))
+diff tests/out/fixed-segments.nes tests/expected/fixed-segments.nes
+exit_code=$((exit_code + $?))
+
 ##
 # code.nes
 
diff --git a/tests/expected/fixed-segments.nes b/tests/expected/fixed-segments.nes
new file mode 100644
index 0000000..06959e0
--- /dev/null
+++ b/tests/expected/fixed-segments.nes
diff --git a/tests/expected/fixed-segments.txt b/tests/expected/fixed-segments.txt
new file mode 100644
index 0000000..9a93bf4
--- /dev/null
+++ b/tests/expected/fixed-segments.txt
@@ -0,0 +1 @@
+warning: referencing 'hello_bank1' from 'FIXED', which belongs to the 'BANK1' segment (fixed-segments.s: line 114)
diff --git a/tests/fixed-segments.s b/tests/fixed-segments.s
new file mode 100644
index 0000000..2d502a5
--- /dev/null
+++ b/tests/fixed-segments.s
@@ -0,0 +1,131 @@
+;;; asan:fixed-segments BANK0, FIXED
+
+.segment "HEADER"
+    .byte 'N', 'E', 'S', $1A
+    .byte $08                   ; 128KB  of PRG-ROM (8 x 16KB)
+    .byte $00                   ; No CHR-ROM. Games using this chip used RAM instead.
+
+    ;; On the 6th byte we have to set the lower nibble of the mapper (#%0010 for
+    ;; UNROM).
+    .byte $20
+
+    ;; Forcing iNES 2.0 format, which will help us for the next bytes.
+    .byte $08
+
+    ;; And now iNES 2.0-specific thingies.
+    .byte $00       ; No submapper
+    .byte $00       ; PRG ROM not 4 MiB or larger
+    .byte $00       ; No PRG RAM
+    .byte $07       ; 8192 (64 * 2^7) bytes CHR RAM, no battery
+    .byte $00       ; NTSC; use $01 for PAL
+    .byte $00       ; No special PPU
+
+.segment "VECTORS"
+    .addr nmi, reset, irq
+
+.segment "BANK0"
+
+hello_bank0:
+    lda #2
+    sta $10
+    rts
+
+.segment "BANK1"
+
+hello_bank1:
+    lda #3
+    sta $11
+    rts
+
+.segment "BANK2"
+.byte $00
+
+.segment "BANK3"
+.byte $00
+
+.segment "BANK4"
+.byte $00
+
+.segment "BANK5"
+.byte $00
+
+.segment "BANK6"
+.byte $00
+
+.segment "FIXED"
+
+banktable:
+  .byte $00, $01, $02, $03, $04, $05, $06
+
+zp_current_bank = $00
+
+bankswitch:
+    sty zp_current_bank
+    tya
+    sta banktable, y
+    rts
+
+.proc reset
+    sei
+    cld
+    ldx #$40
+    stx $4017
+
+    ldx #$FF
+    txs
+
+    inx
+    stx $2000
+    stx $2001
+    stx $4010
+
+    bit $2002
+@vblankwait1:
+    bit $2002
+    bpl @vblankwait1
+
+    ldx #0
+    lda #0
+@ram_reset_loop:
+    sta $000, x
+    sta $100, x
+    sta $200, x
+    sta $300, x
+    sta $400, x
+    sta $500, x
+    sta $600, x
+    sta $700, x
+    inx
+    bne @ram_reset_loop
+
+@vblankwait2:
+    bit $2002
+    bpl @vblankwait2
+
+    ;;;
+    ;; NOTE: configuration/reset is done, the code below is our actual program :D
+
+    ldy #0
+    jsr bankswitch
+    jsr hello_bank0
+
+    ldy #1
+    jsr bankswitch
+    jsr hello_bank1
+
+    lda $10
+    clc
+    adc $11
+    sta $12
+
+@loop:
+    jmp @loop
+.endproc
+
+.proc nmi
+    rti
+.endproc
+
+.proc irq
+    rti
+.endproc