From 2a08362b52ec9c2f0a3721e13dbd9f8e9e9d6ead Mon Sep 17 00:00:00 2001
From: Erin <erin@erindesu.cz>
Date: Fri, 9 Jun 2023 02:42:00 +0200
Subject: [PATCH] Very quick and dirty labels.
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

I wanna sleep now, so... will fix later™
---
 Cargo.lock       | 10 ++++++
 hbasm/Cargo.toml |  1 +
 hbasm/src/lib.rs | 83 +++++++++++++++++++++++++++++++++++++-----------
 3 files changed, 76 insertions(+), 18 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index abe9a78e..9ba3ea20 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -49,6 +49,7 @@ name = "hbasm"
 version = "0.1.0"
 dependencies = [
  "hbbytecode",
+ "lasso",
  "logos",
  "paste",
 ]
@@ -67,6 +68,15 @@ dependencies = [
  "static_assertions",
 ]
 
+[[package]]
+name = "lasso"
+version = "0.7.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "4644821e1c3d7a560fe13d842d13f587c07348a1a05d3a797152d41c90c56df2"
+dependencies = [
+ "hashbrown",
+]
+
 [[package]]
 name = "log"
 version = "0.4.17"
diff --git a/hbasm/Cargo.toml b/hbasm/Cargo.toml
index c9add3ca..1ab3f41f 100644
--- a/hbasm/Cargo.toml
+++ b/hbasm/Cargo.toml
@@ -5,6 +5,7 @@ edition = "2021"
 
 [dependencies]
 hbbytecode = { path = "../hbbytecode" }
+lasso      = "0.7"
 paste      = "1.0"
 
 [dependencies.logos]
diff --git a/hbasm/src/lib.rs b/hbasm/src/lib.rs
index 7bd12af4..67f6d2af 100644
--- a/hbasm/src/lib.rs
+++ b/hbasm/src/lib.rs
@@ -1,16 +1,18 @@
+use std::collections::{HashMap, HashSet};
+
+use lasso::Key;
+
 use {
+    lasso::{Rodeo, Spur},
     logos::{Lexer, Logos, Span},
-    std::{
-        fmt::{Display, Formatter},
-        ops::Range,
-        str::FromStr,
-    },
+    std::fmt::{Display, Formatter},
 };
 
 macro_rules! tokendef {
     ($($opcode:literal),* $(,)?) => {
         paste::paste! {
             #[derive(Clone, Copy, Debug, PartialEq, Eq, Logos)]
+            #[logos(extras = Rodeo)]
             #[logos(skip r"[ \t\f]+")]
             pub enum Token {
                 $(#[token($opcode, |_| hbbytecode::opcode::[<$opcode:upper>])])*
@@ -32,9 +34,18 @@ macro_rules! tokendef {
                     },
                 )] Register(u8),
 
+                #[regex(
+                    r"\p{XID_Start}\p{XID_Continue}*:",
+                    |lexer| lexer.extras.get_or_intern(&lexer.slice()[..lexer.slice().len() - 1]),
+                )] Label(Spur),
+
+                #[regex(
+                    r"\p{XID_Start}\p{XID_Continue}*",
+                    |lexer| lexer.extras.get_or_intern(lexer.slice()),
+                )] Symbol(Spur),
+
                 #[token("\n")]
                 #[token(";")] ISep,
-
                 #[token(",")] PSep,
             }
         }
@@ -80,8 +91,10 @@ macro_rules! expect_matches {
 
 pub fn assembly(code: &str, buf: &mut Vec<u8>) -> Result<(), Error> {
     struct Assembler<'a> {
-        lexer: Lexer<'a, Token>,
-        buf:   &'a mut Vec<u8>,
+        lexer:  Lexer<'a, Token>,
+        buf:    &'a mut Vec<u8>,
+        lblmap: HashMap<Spur, u64>,
+        subset: HashSet<usize>,
     }
 
     impl<'a> Assembler<'a> {
@@ -114,6 +127,9 @@ pub fn assembly(code: &str, buf: &mut Vec<u8>) -> Result<(), Error> {
                             Err(e) => return Err(e),
                         }
                     }
+                    Some(Ok(Token::Label(lbl))) => {
+                        self.lblmap.insert(lbl, self.buf.len() as u64 + 1);
+                    }
                     Some(Ok(_)) => return Err(ErrorKind::UnexpectedToken),
                     Some(Err(())) => return Err(ErrorKind::InvalidToken),
                     None => return Ok(()),
@@ -141,10 +157,19 @@ pub fn assembly(code: &str, buf: &mut Vec<u8>) -> Result<(), Error> {
         }
 
         fn ri(&mut self) -> Result<(), ErrorKind> {
-            expect_matches!(self, Token::Register(r0), Token::PSep, Token::Integer(r1),);
+            expect_matches!(self, Token::Register(r0), Token::PSep, imm @ (Token::Integer(_) | Token::Symbol(_)));
 
             self.buf.push(r0);
-            self.buf.extend(r1.to_le_bytes());
+            let imm = match imm {
+                Token::Integer(n) => n.to_le_bytes(),
+                Token::Symbol(s) => {
+                    self.subset.insert(self.buf.len() - 1);
+                    s.into_usize().to_le_bytes()
+                }
+                _ => unreachable!(),
+            };
+
+            self.buf.extend(imm);
             Ok(())
         }
 
@@ -155,27 +180,49 @@ pub fn assembly(code: &str, buf: &mut Vec<u8>) -> Result<(), Error> {
                 Token::PSep,
                 Token::Register(r1),
                 Token::PSep,
-                Token::Integer(imm),
+                imm @ (Token::Integer(_) | Token::Symbol(_)),
             );
             self.buf.extend([r0, r1]);
-            self.buf.extend(imm.to_le_bytes());
-            Ok(())
-        }
+            let imm = match imm {
+                Token::Integer(n) => n.to_le_bytes(),
+                Token::Symbol(s) => {
+                    self.subset.insert(self.buf.len() - 1);
+                    (s.into_inner().get() as usize).to_le_bytes()
+                }
+                _ => unreachable!(),
+            };
 
-        fn i(&mut self) -> Result<(), ErrorKind> {
-            expect_matches!(self, Token::Integer(imm),);
-            self.buf.extend(imm.to_le_bytes());
+            self.buf.extend(imm);
             Ok(())
         }
     }
 
     let mut asm = Assembler {
         lexer: Token::lexer(code),
+        lblmap: Default::default(),
+        subset: Default::default(),
         buf,
     };
 
     asm.assemble().map_err(|kind| Error {
         kind,
         span: asm.lexer.span(),
-    })
+    })?;
+
+    for sub in asm.subset {
+        asm.lblmap[&Spur::try_from_usize(usize::from_le_bytes(
+            asm.buf[sub..sub + core::mem::size_of::<usize>()]
+                .try_into()
+                .expect("Expected location"),
+        ))
+        .expect("Expected valid intern key")]
+            .to_le_bytes()
+            .iter()
+            .enumerate()
+            .for_each(|(i, b)| {
+                asm.buf[sub + i] = *b;
+            });
+    }
+
+    Ok(())
 }