I guess I improved local labels?

hbasm/src/lib.rs

@@ -67,6 +67,7 @@ pub enum ErrorKind {
     UnexpectedToken,
     InvalidToken,
     UnexpectedEnd,
+    InvalidSymbol,
 }
 
 #[derive(Clone, Debug, PartialEq, Eq)]
@@ -94,8 +95,8 @@ pub fn assembly(code: &str, buf: &mut Vec<u8>) -> Result<(), Error> {
     struct Assembler<'a> {
         lexer: Lexer<'a, Token>,
         buf: &'a mut Vec<u8>,
-        lblmap: HashMap<Spur, u64>,
-        subset: HashSet<usize>,
+        label_map: HashMap<Spur, u64>,
+        to_sub_label: HashMap<usize, Spur>,
     }
 
     impl<'a> Assembler<'a> {
@@ -129,7 +130,7 @@ pub fn assembly(code: &str, buf: &mut Vec<u8>) -> Result<(), Error> {
                         }
                     }
                     Some(Ok(Token::Label(lbl))) => {
-                        self.lblmap.insert(lbl, self.buf.len() as u64 + 1);
+                        self.label_map.insert(lbl, self.buf.len() as u64 + 1);
                     }
                     Some(Ok(Token::ISep)) => (),
                     Some(Ok(_)) => return Err(ErrorKind::UnexpectedToken),
@@ -139,6 +140,22 @@ pub fn assembly(code: &str, buf: &mut Vec<u8>) -> Result<(), Error> {
             }
         }
 
+        fn link_local_syms(&mut self) -> Result<(), ErrorKind> {
+            for (ix, sym) in &self.to_sub_label {
+                self.label_map
+                    .get(sym)
+                    .ok_or(ErrorKind::InvalidSymbol)?
+                    .to_le_bytes()
+                    .iter()
+                    .enumerate()
+                    .for_each(|(i, b)| {
+                        self.buf[ix + i] = *b;
+                    });
+            }
+
+            Ok(())
+        }
+
         fn rrr(&mut self) -> Result<(), ErrorKind> {
             expect_matches!(
                 self,
@@ -159,19 +176,9 @@ 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, imm @ (Token::Integer(_) | Token::Symbol(_)));
-
+            expect_matches!(self, Token::Register(r0), Token::PSep);
             self.buf.push(r0);
-            let imm = match imm {
-                Token::Integer(n) => n.to_le_bytes(),
-                Token::Symbol(s) => {
-                    self.subset.insert(self.buf.len());
-                    s.into_usize().to_le_bytes()
-                }
-                _ => unreachable!(),
-            };
-
-            self.buf.extend(imm);
+            self.insert_imm()?;
             Ok(())
         }
 
@@ -182,18 +189,21 @@ pub fn assembly(code: &str, buf: &mut Vec<u8>) -> Result<(), Error> {
                 Token::PSep,
                 Token::Register(r1),
                 Token::PSep,
-                imm @ (Token::Integer(_) | Token::Symbol(_)),
             );
             self.buf.extend([r0, r1]);
-            let imm = match imm {
-                Token::Integer(n) => n.to_le_bytes(),
-                Token::Symbol(s) => {
-                    self.subset.insert(self.buf.len());
-                    s.into_usize().to_le_bytes()
-                }
-                _ => unreachable!(),
-            };
+            self.insert_imm()?;
+            Ok(())
+        }
 
+        fn insert_imm(&mut self) -> Result<(), ErrorKind> {
+            let imm = match self.next()? {
+                Token::Integer(i) => i.to_le_bytes(),
+                Token::Symbol(s) => {
+                    self.to_sub_label.insert(self.buf.len(), s);
+                    [0; 8]
+                }
+                _ => return Err(ErrorKind::UnexpectedToken),
+            };
             self.buf.extend(imm);
             Ok(())
         }
@@ -201,8 +211,8 @@ pub fn assembly(code: &str, buf: &mut Vec<u8>) -> Result<(), Error> {
 
     let mut asm = Assembler {
         lexer: Token::lexer(code),
-        lblmap: Default::default(),
-        subset: Default::default(),
+        label_map: Default::default(),
+        to_sub_label: Default::default(),
         buf,
     };
 
@@ -211,20 +221,6 @@ pub fn assembly(code: &str, buf: &mut Vec<u8>) -> Result<(), Error> {
         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(())
+    asm.link_local_syms()
+        .map_err(|kind| Error { kind, span: 0..0 })
 }