From e1a423a35500256b09731b2cf467d31194f332a7 Mon Sep 17 00:00:00 2001
From: Erin <erin@erindesu.cz>
Date: Mon, 24 Jul 2023 20:41:10 +0200
Subject: [PATCH] Kekw

---
 hbasm/src/lib.rs         | 2 +-
 hbasm/src/macros/asm.rs  | 7 +++++++
 hbasm/src/macros/text.rs | 9 +++++++++
 hbvm/src/validate.rs     | 6 ++++--
 4 files changed, 21 insertions(+), 3 deletions(-)

diff --git a/hbasm/src/lib.rs b/hbasm/src/lib.rs
index 022f8d07..62d13717 100644
--- a/hbasm/src/lib.rs
+++ b/hbasm/src/lib.rs
@@ -24,7 +24,7 @@ macros::impl_both!(
         => [DIR, DIRF, FMAF],
     bbb(p0: R, p1: R, p2: R)
         => [ADD, SUB, MUL, AND, OR, XOR, SL, SR, SRS, CMP, CMPU, /*BRC,*/ ADDF, SUBF, MULF],
-    bbdh(p0: R, p1: R, p2: I, p3: u16)
+    bbdh(p0: R, p1: R, p2: I, p3: L)
         => [LD, ST],
     bbd(p0: R, p1: R, p2: I)
         => [ADDI, MULI, ANDI, ORI, XORI, CMPI, CMPUI, BMC, JAL, JEQ, JNE, JLT, JGT, JLTU,
diff --git a/hbasm/src/macros/asm.rs b/hbasm/src/macros/asm.rs
index 74e0e8f7..98db69f0 100644
--- a/hbasm/src/macros/asm.rs
+++ b/hbasm/src/macros/asm.rs
@@ -38,6 +38,12 @@ macro_rules! impl_asm_insert {
         Imm::insert(&$id, $self)
     };
 
+    // Length - cannot be more than 2048
+    ($self:expr, $id:ident, L) => {{
+        assert!($id <= 2048);
+        $self.buf.extend($id.to_le_bytes())
+    }};
+
     // Other numbers, just insert their bytes, little endian
     ($self:expr, $id:ident, $_:ident) => {
         $self.buf.extend($id.to_le_bytes())
@@ -76,6 +82,7 @@ macro_rules! impl_asm {
 macro_rules! ident_map_ty {
     (R)         => { u8 };       // Register is just u8
     (I)         => { impl Imm }; // Immediate is anything implementing the trait
+    (L)         => { u16 };      // Copy count
     ($id:ident) => { $id };      // Anything else → identity map
 }
 
diff --git a/hbasm/src/macros/text.rs b/hbasm/src/macros/text.rs
index d26e7e1f..22f82061 100644
--- a/hbasm/src/macros/text.rs
+++ b/hbasm/src/macros/text.rs
@@ -226,6 +226,15 @@ macro_rules! gen_extract {
                 extract_pat!($self, Token::Register($id));
             };
 
+            ($self:expr, L, $id:ident) => {
+                extract_pat!($self, Token::Integer($id));
+                if $id > 2048 {
+                    return Err(ErrorKind::InvalidToken);
+                }
+
+                let $id = u16::try_from($id).unwrap();
+            };
+
             // Immediate
             ($self:expr, I, $id:ident) => {
                 let $id = match $self.next()? {
diff --git a/hbvm/src/validate.rs b/hbvm/src/validate.rs
index 103a2d85..81f9a0a8 100644
--- a/hbvm/src/validate.rs
+++ b/hbvm/src/validate.rs
@@ -48,8 +48,10 @@ pub fn validate(mut program: &[u8]) -> Result<(), Error> {
         // Match on instruction types and perform necessary checks
         program = match program {
             [] => return Ok(()),
-            [LD..=ST, reg, _, _, _, _, _, _, _, _, _, count, ..]
-                if usize::from(*reg) * 8 + usize::from(*count) > 2048 =>
+            [LD..=ST, reg, _, _, _, _, _, _, _, _, count_0, count_1, ..]
+                if usize::from(*reg) * 8
+                    + usize::from(u16::from_le_bytes([*count_0, *count_1]))
+                    > 2048 =>
             {
                 return Err(Error {
                     kind:  ErrorKind::RegisterArrayOverflow,