Merge pull request #49 from dxrcy/stack-env-var

Require `LACE_STACK=1` to use non-standard 'stack' instruction set extension

src/env.rs

@@ -0,0 +1,50 @@
+use std::{cell::RefCell, ffi::OsStr};
+
+#[derive(Clone, Copy)]
+struct Env {
+    stack_enabled: bool,
+}
+
+thread_local! {
+    /// Must only be mutated within `set_env`
+    static ENV: RefCell<Option<Env>> = const { RefCell::new(None) };
+}
+
+pub fn init() {
+    let value = Env {
+        stack_enabled: var_is("LACE_STACK", "1"),
+    };
+    set_env(value);
+}
+
+pub fn is_stack_enabled() -> bool {
+    with_env(|env| env.stack_enabled)
+}
+
+fn set_env(value: Env) {
+    ENV.with(|env| {
+        let mut env = env.borrow_mut();
+        assert!(
+            env.is_none(),
+            "tried to initialize environment state multiple times"
+        );
+        *env = Some(value);
+    });
+}
+
+fn with_env<F, R>(callback: F) -> R
+where
+    F: Fn(&Env) -> R,
+{
+    ENV.with(|env| {
+        let env = env.borrow();
+        let env = env.unwrap_or_else(|| {
+            panic!("tried to access environment state before initialization");
+        });
+        callback(&env)
+    })
+}
+
+fn var_is(name: impl AsRef<OsStr>, value: impl AsRef<str>) -> bool {
+    std::env::var(name.as_ref()).is_ok_and(|v| &v == value.as_ref())
+}

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 {

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

src/lib.rs

@@ -14,3 +14,5 @@ pub use symbol::{reset_state, StaticSource};
 
 mod error;
 mod lexer;
+
+pub mod env;

src/main.rs

@@ -66,6 +66,7 @@ enum Command {
 fn main() -> miette::Result<()> {
     use MsgColor::*;
     let args = Args::parse();
+    lace::env::init();
 
     if let Some(command) = args.command {
         match command {

src/runtime.rs

@@ -5,7 +5,7 @@ use std::{
     u16, u32, u8, usize,
 };
 
-use crate::Air;
+use crate::{env, Air};
 use colored::Colorize;
 use console::Term;
 use miette::Result;
@@ -153,6 +153,17 @@ impl RunState {
     }
 
     fn stack(&mut self, instr: u16) {
+        if !env::is_stack_enabled() {
+            eprintln!(
+                "\
+                You called a reserved instruction.\n\
+                Note: Run with `LACE_STACK=1` to enable stack features.\n\
+                Halting...\
+                "
+            );
+            std::process::exit(1);
+        }
+
         // Bit to determine call/ret or push/pop
         if instr & 0x0800 != 0 {
             // Call
@@ -180,6 +191,10 @@ impl RunState {
     }
 
     fn push_val(&mut self, val: u16) {
+        debug_assert!(
+            env::is_stack_enabled(),
+            "caller should have ensured stack features are enabled",
+        );
         // Decrement stack
         *self.reg(7) -= 1;
         let sp = *self.reg(7);
@@ -188,6 +203,10 @@ impl RunState {
     }
 
     fn pop_val(&mut self) -> u16 {
+        debug_assert!(
+            env::is_stack_enabled(),
+            "caller should have ensured stack features are enabled",
+        );
         let sp = *self.reg(7);
         let val = *self.mem(sp);
         *self.reg(7) += 1;

tests/integration_tests.rs

@@ -23,6 +23,7 @@ fn runs_hello_world() {
 fn runs_stack_example() {
     let mut cmd = Command::cargo_bin("lace").unwrap();
     cmd.arg("run").arg("tests/files/stack.asm");
+    cmd.env("LACE_STACK", "1");
 
     cmd.assert()
         .success()
@@ -35,6 +36,7 @@ fn runs_stack_example() {
 fn runs_recursive_fibonacci_example() {
     let mut cmd = Command::cargo_bin("lace").unwrap();
     cmd.arg("run").arg("tests/files/fibonacci.asm");
+    cmd.env("LACE_STACK", "1");
 
     // Hardcoded 23rd fibonacci number
     cmd.assert()