refactor: move stack feature check to lexer module

rozukke · Nov 29, 2024 · 6ac4c68 · 6ac4c68
1 parent fb7f162
commit 6ac4c68
Show file tree

Hide file tree

Showing 3 changed files with 46 additions and 47 deletions.
diff --git a/src/error.rs b/src/error.rs
@@ -54,6 +54,22 @@ pub fn lex_unknown(span: Span, src: &'static str) -> Report {
     .with_source_code(src)
 }
 
+pub fn lex_stack_extension_not_enabled(instr: &str, span: Span, src: &'static str) -> Report {
+    miette!(
+        severity = Severity::Error,
+        code = "lex::stack_extension_not_enabled",
+        help = "\
+        this instruction requires the non-standard 'stack' extension\n\
+        run with `LACE_STACK=1` to enable\n\
+        note: this identifier cannot be used as a label\
+        ",
+        labels = vec![LabeledSpan::at(span, "non-standard instruction")],
+        "Non-standard '{}' instruction used without 'stack' extension enabled",
+        instr
+    )
+    .with_source_code(src)
+}
+
 // Preprocessor errors
 
 pub fn preproc_bad_lit(span: Span, src: &'static str, is_present: bool) -> Report {
@@ -147,19 +163,3 @@ pub fn parse_lit_range(span: Span, src: &'static str, bits: Bits) -> Report {
     )
     .with_source_code(src)
 }
-
-pub fn parse_stack_extension_not_enabled(instr: &str, span: Span, src: &'static str) -> Report {
-    miette!(
-        severity = Severity::Error,
-        code = "parse::stack_extension_not_enabled",
-        help = "\
-        this instruction requires the non-standard 'stack' extension\n\
-        run with `LACE_STACK=1` to enable\n\
-        note: this identifier cannot be used as a label\
-        ",
-        labels = vec![LabeledSpan::at(span, "non-standard instruction")],
-        "Non-standard '{}' instruction used without 'stack' extension enabled",
-        instr
-    )
-    .with_source_code(src)
-}
diff --git a/src/lexer/mod.rs b/src/lexer/mod.rs
@@ -5,9 +5,9 @@ use std::{i16, u16};
 
 use miette::Result;
 
-use crate::error;
 use crate::lexer::cursor::Cursor;
 use crate::symbol::{DirKind, Flag, InstrKind, Register, Span, SrcOffset, TrapKind};
+use crate::{env, error};
 
 pub mod cursor;
 
@@ -146,7 +146,7 @@ impl Cursor<'_> {
                     self.bump();
                     self.hex()?
                 }
-                _ => self.ident(),
+                _ => self.ident()?,
             },
             // Register literals
             'r' | 'R' => match self.first() {
@@ -162,13 +162,13 @@ impl Cursor<'_> {
                         // SAFETY: c is always valid
                         TokenKind::Reg(Register::from_str(&c.to_string()).unwrap())
                     } else {
-                        self.ident()
+                        self.ident()?
                     }
                 }
-                _ => self.ident(),
+                _ => self.ident()?,
             },
             // Check only after other identifier-likes
-            c if is_id(c) => self.ident(),
+            c if is_id(c) => self.ident()?,
             // Decimal literal
             '#' => self.dec()?,
             // Directive
@@ -210,7 +210,7 @@ impl Cursor<'_> {
                             e,
                         ))
                     }
-                    _ => return Ok(self.ident()),
+                    _ => return Ok(self.ident()?),
                 },
             },
         };
@@ -282,18 +282,18 @@ impl Cursor<'_> {
         }
     }
 
-    fn ident(&mut self) -> TokenKind {
+    fn ident(&mut self) -> Result<TokenKind> {
         let ident_start = self.abs_pos() - 1;
         self.take_while(is_id);
         let ident = self
             .get_range(ident_start..self.abs_pos())
             .to_ascii_lowercase();
 
-        let mut token_kind = self.check_instruction(&ident);
+        let mut token_kind = self.check_instruction(&ident, ident_start)?;
         if token_kind == TokenKind::Label {
             token_kind = self.check_trap(&ident);
         }
-        token_kind
+        Ok(token_kind)
     }
 
     /// Expects lowercase
@@ -311,10 +311,21 @@ impl Cursor<'_> {
     }
 
     /// Expects lowercase
-    fn check_instruction(&self, ident: &str) -> TokenKind {
+    fn check_instruction(&self, ident: &str, start_pos: usize) -> Result<TokenKind> {
         use InstrKind::*;
         use TokenKind::Instr;
-        match ident {
+
+        if matches!(ident, "pop" | "push" | "call" | "rets") {
+            if !env::is_stack_enabled() {
+                return Err(error::lex_stack_extension_not_enabled(
+                    ident,
+                    Span::new(SrcOffset(start_pos), self.pos_in_token()),
+                    self.src(),
+                ));
+            }
+        }
+
+        Ok(match ident {
             "add" => Instr(Add),
             "and" => Instr(And),
             "br" => Instr(Br(Flag::Nzp)),
@@ -343,7 +354,7 @@ impl Cursor<'_> {
             "call" => Instr(Call),
             "rets" => Instr(Rets),
             _ => TokenKind::Label,
-        }
+        })
     }
 
     /// Expects lowercase

diff --git a/src/parser.rs b/src/parser.rs
@@ -3,7 +3,10 @@ use std::{borrow::Cow, fmt::Display, iter::Peekable, vec::IntoIter};
 use miette::Result;
 
 use crate::{
-    air::{Air, AirStmt, ImmediateOrReg, RawWord}, env, error, lexer::{cursor::Cursor, LiteralKind, Token, TokenKind}, symbol::{DirKind, InstrKind, Label, Register, Span, TrapKind}
+    air::{Air, AirStmt, ImmediateOrReg, RawWord},
+    error,
+    lexer::{cursor::Cursor, LiteralKind, Token, TokenKind},
+    symbol::{DirKind, InstrKind, Label, Register, Span, TrapKind},
 };
 
 /// Replaces raw value directives .fill, .blkw, .stringz with equivalent raw bytes
@@ -163,7 +166,7 @@ impl AsmParser {
                         self.air.set_orig(orig)?;
                         continue;
                     }
-                    TokenKind::Instr(instr_kind) => self.parse_instr(instr_kind, tok.span)?,
+                    TokenKind::Instr(instr_kind) => self.parse_instr(instr_kind)?,
                     TokenKind::Trap(trap_kind) => self.parse_trap(trap_kind)?,
                     TokenKind::Byte(val) => self.parse_byte(val),
                     // Does not exist in preprocessed token stream
@@ -193,29 +196,23 @@ impl AsmParser {
     }
 
     /// Process several tokens to form valid AIR statement
-    fn parse_instr(&mut self, kind: InstrKind, span: Span) -> Result<AirStmt> {
+    fn parse_instr(&mut self, kind: InstrKind) -> Result<AirStmt> {
         use crate::symbol::InstrKind;
         match kind {
             InstrKind::Push => {
-                self.expect_stack_enabled("push", span)?;
                 let src_reg = self.expect_reg()?;
                 Ok(AirStmt::Push { src_reg })
             }
             InstrKind::Pop => {
-                self.expect_stack_enabled("push", span)?;
                 let dest_reg = self.expect_reg()?;
                 Ok(AirStmt::Pop { dest_reg })
             }
             InstrKind::Call => {
-                self.expect_stack_enabled("push", span)?;
                 let label_tok = self.expect(TokenKind::Label)?;
                 let dest_label = Label::try_fill(self.get_span(label_tok.span));
                 Ok(AirStmt::Call { dest_label })
             }
-            InstrKind::Rets => {
-                self.expect_stack_enabled("push", span)?;
-                Ok(AirStmt::Rets)
-            }
+            InstrKind::Rets => Ok(AirStmt::Rets),
             InstrKind::Add => {
                 let dest = self.expect_reg()?;
                 let src_reg = self.expect_reg()?;
@@ -438,15 +435,6 @@ impl AsmParser {
             None => return Err(error::parse_eof(self.src)),
         }
     }
-
-    fn expect_stack_enabled(&self, instr_name: &str, span: Span) -> Result<()> {
-        if !env::is_stack_enabled() {
-            return Err(error::parse_stack_extension_not_enabled(
-                instr_name, span, self.src,
-            ));
-        }
-        Ok(())
-    }
 }
 
 /// Convenient way to pass around bit limits