From 961c9e998ac7f3111e91fcfb5a70906484a36076 Mon Sep 17 00:00:00 2001
From: darcy <acednes@gmail.com>
Date: Wed, 18 Dec 2024 11:12:02 +1100
Subject: [PATCH 1/3] feat: support literals for PC-offset operands

---
 src/parser.rs | 46 ++++++++++++++++++++++++++++++++--------------
 1 file changed, 32 insertions(+), 14 deletions(-)

diff --git a/src/parser.rs b/src/parser.rs
index 56124bb..e8f9149 100644
--- a/src/parser.rs
+++ b/src/parser.rs
@@ -234,8 +234,7 @@ impl AsmParser {
                 })
             }
             InstrKind::Br(flag) => {
-                let label_tok = self.expect(TokenKind::Label)?;
-                let dest_label = Label::try_fill(self.get_span(label_tok.span));
+                let dest_label = self.expect_lit_or_label(9)?;
                 Ok(AirStmt::Branch { flag, dest_label })
             }
             InstrKind::Jmp => {
@@ -243,8 +242,7 @@ impl AsmParser {
                 Ok(AirStmt::Jump { src_reg })
             }
             InstrKind::Jsr => {
-                let label_tok = self.expect(TokenKind::Label)?;
-                let dest_label = Label::try_fill(self.get_span(label_tok.span));
+                let dest_label = self.expect_lit_or_label(11)?;
                 Ok(AirStmt::JumbSub { dest_label })
             }
             InstrKind::Jsrr => {
@@ -253,14 +251,12 @@ impl AsmParser {
             }
             InstrKind::Ld => {
                 let dest = self.expect_reg()?;
-                let label_tok = self.expect(TokenKind::Label)?;
-                let src_label = Label::try_fill(self.get_span(label_tok.span));
+                let src_label = self.expect_lit_or_label(9)?;
                 Ok(AirStmt::Load { dest, src_label })
             }
             InstrKind::Ldi => {
                 let dest = self.expect_reg()?;
-                let label_tok = self.expect(TokenKind::Label)?;
-                let src_label = Label::try_fill(self.get_span(label_tok.span));
+                let src_label = self.expect_lit_or_label(9)?;
                 Ok(AirStmt::LoadInd { dest, src_label })
             }
             InstrKind::Ldr => {
@@ -275,8 +271,7 @@ impl AsmParser {
             }
             InstrKind::Lea => {
                 let dest = self.expect_reg()?;
-                let label_tok = self.expect(TokenKind::Label)?;
-                let src_label = Label::try_fill(self.get_span(label_tok.span));
+                let src_label = self.expect_lit_or_label(9)?;
                 Ok(AirStmt::LoadEAddr { dest, src_label })
             }
             InstrKind::Not => {
@@ -288,8 +283,7 @@ impl AsmParser {
             InstrKind::Rti => Ok(AirStmt::Interrupt),
             InstrKind::St => {
                 let src_reg = self.expect_reg()?;
-                let label_tok = self.expect(TokenKind::Label)?;
-                let dest_label = Label::try_fill(self.get_span(label_tok.span));
+                let dest_label = self.expect_lit_or_label(9)?;
                 Ok(AirStmt::Store {
                     src_reg,
                     dest_label,
@@ -297,8 +291,7 @@ impl AsmParser {
             }
             InstrKind::Sti => {
                 let src_reg = self.expect_reg()?;
-                let label_tok = self.expect(TokenKind::Label)?;
-                let dest_label = Label::try_fill(self.get_span(label_tok.span));
+                let dest_label = self.expect_lit_or_label(9)?;
                 Ok(AirStmt::StoreInd {
                     src_reg,
                     dest_label,
@@ -435,6 +428,31 @@ impl AsmParser {
             None => return Err(error::parse_eof(self.src)),
         }
     }
+
+    fn expect_lit_or_label(&mut self, bits: u8) -> Result<Label> {
+        match self.toks.peek() {
+            Some(tok) => match tok.kind {
+                TokenKind::Label => {
+                    let span = tok.span;
+                    let label = Label::try_fill(self.get_span(span));
+                    Ok(label)
+                }
+                TokenKind::Lit(_) => {
+                    let val = self.expect_lit(Bits::Signed(bits))?;
+                    let label = Label::Ref(self.line + 1 + val);
+                    Ok(label)
+                }
+                _ => {
+                    return Err(error::parse_generic_unexpected(
+                        self.src,
+                        "literal or label",
+                        *tok,
+                    ))
+                }
+            },
+            None => return Err(error::parse_eof(self.src)),
+        }
+    }
 }
 
 /// Convenient way to pass around bit limits

From 67d514ea6e441e4ca59985685bfc40b1c5b09c19 Mon Sep 17 00:00:00 2001
From: darcy <acednes@gmail.com>
Date: Wed, 18 Dec 2024 11:15:55 +1100
Subject: [PATCH 2/3] fix: increment cursor when parsing label operand

---
 src/parser.rs | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/parser.rs b/src/parser.rs
index e8f9149..0aacaf9 100644
--- a/src/parser.rs
+++ b/src/parser.rs
@@ -433,8 +433,8 @@ impl AsmParser {
         match self.toks.peek() {
             Some(tok) => match tok.kind {
                 TokenKind::Label => {
-                    let span = tok.span;
-                    let label = Label::try_fill(self.get_span(span));
+                    let label_tok = self.expect(TokenKind::Label)?;
+                    let label = Label::try_fill(self.get_span(label_tok.span));
                     Ok(label)
                 }
                 TokenKind::Lit(_) => {

From 967ec6791860d1b61677b7f19806189929dac788 Mon Sep 17 00:00:00 2001
From: darcy <acednes@gmail.com>
Date: Wed, 18 Dec 2024 11:23:12 +1100
Subject: [PATCH 3/3] chore: add tests for literal PC-offset operands

---
 src/parser.rs | 26 ++++++++++++++++++++++++++
 1 file changed, 26 insertions(+)

diff --git a/src/parser.rs b/src/parser.rs
index 0aacaf9..c0e4d2d 100644
--- a/src/parser.rs
+++ b/src/parser.rs
@@ -673,6 +673,21 @@ mod test {
         )
     }
 
+    #[test]
+    fn parse_branch_lit() {
+        let air = AsmParser::new("br x2").unwrap().parse().unwrap();
+        assert_eq!(
+            air.get(0),
+            &AsmLine {
+                line: 1,
+                stmt: AirStmt::Branch {
+                    flag: Flag::Nzp,
+                    dest_label: Label::Ref(0x2 + 0x2)
+                }
+            }
+        )
+    }
+
     #[test]
     fn parse_fill() {
         let air = AsmParser::new("label .fill x30").unwrap().parse().unwrap();
@@ -758,6 +773,7 @@ mod test {
         label add r0 r0 r0
               br label
               br not_existing
+              br x30
         "#,
         )
         .unwrap()
@@ -794,5 +810,15 @@ mod test {
                 }
             }
         );
+        assert_eq!(
+            air.get(3),
+            &AsmLine {
+                line: 4,
+                stmt: AirStmt::Branch {
+                    flag: Flag::Nzp,
+                    dest_label: Label::Ref(0x5 + 0x30),
+                }
+            }
+        );
     }
 }