diff --git a/crates/polkavm-assembler/src/amd64.rs b/crates/polkavm-assembler/src/amd64.rs
index ba4831db..01a52bd0 100644
--- a/crates/polkavm-assembler/src/amd64.rs
+++ b/crates/polkavm-assembler/src/amd64.rs
@@ -573,16 +573,16 @@ impl Inst {
 
                 self.modrm |= base.modrm_rm_bits();
 
-                if offset != 0 || matches!(base, Reg::rbp | Reg::r13) {
-                    if offset <= i8::MAX as i32 && offset >= i8::MIN as i32 {
-                        self.modrm |= 0b01000000;
-                        self.displacement = offset as u8 as u32;
-                        self.displacement_length = 8;
-                    } else {
-                        self.modrm |= 0b10000000;
-                        self.displacement = offset as u32;
-                        self.displacement_length = 32;
-                    }
+                let set_displacement = (offset != 0) | matches!(base, Reg::rbp | Reg::r13);
+                let set_displacement_mask = (-(set_displacement as i32)) as u32;
+                if offset <= i8::MAX as i32 && offset >= i8::MIN as i32 {
+                    self.modrm |= 0b01000000 & set_displacement_mask as u8;
+                    self.displacement = (offset as u8 as u32) & set_displacement_mask;
+                    self.displacement_length = 8 & set_displacement_mask;
+                } else {
+                    self.modrm |= 0b10000000 & set_displacement_mask as u8;
+                    self.displacement = (offset as u32) & set_displacement_mask;
+                    self.displacement_length = 32 & set_displacement_mask;
                 }
 
                 self.override_segment = segment;
@@ -602,16 +602,16 @@ impl Inst {
                 self.sib |= base.modrm_rm_bits();
                 self.sib |= ((scale as usize) << 6) as u8;
 
-                if offset != 0 || matches!(base, Reg::rbp | Reg::r13) {
-                    if offset <= i8::MAX as i32 && offset >= i8::MIN as i32 {
-                        self.modrm |= 0b01000000;
-                        self.displacement = offset as u8 as u32;
-                        self.displacement_length = 8;
-                    } else {
-                        self.modrm |= 0b10000000;
-                        self.displacement = offset as u32;
-                        self.displacement_length = 32;
-                    }
+                let set_displacement = (offset != 0) | matches!(base, Reg::rbp | Reg::r13);
+                let set_displacement_mask = (-(set_displacement as i32)) as u32;
+                if offset <= i8::MAX as i32 && offset >= i8::MIN as i32 {
+                    self.modrm |= 0b01000000 & set_displacement_mask as u8;
+                    self.displacement = (offset as u8 as u32) & set_displacement_mask;
+                    self.displacement_length = 8 & set_displacement_mask;
+                } else {
+                    self.modrm |= 0b10000000 & set_displacement_mask as u8;
+                    self.displacement = (offset as u32) & set_displacement_mask;
+                    self.displacement_length = 32 & set_displacement_mask;
                 }
 
                 self.override_segment = segment;
diff --git a/crates/polkavm-assembler/src/assembler.rs b/crates/polkavm-assembler/src/assembler.rs
index 4919eabf..209e5fd7 100644
--- a/crates/polkavm-assembler/src/assembler.rs
+++ b/crates/polkavm-assembler/src/assembler.rs
@@ -48,6 +48,107 @@ impl<'a> Drop for AssembledCode<'a> {
     }
 }
 
+/// # Safety
+///
+/// `VALUE` must be non-zero, and `Self::Next::VALUE` must be `VALUE - 1` if `Self::Next` is `NonZero`.
+pub unsafe trait NonZero {
+    const VALUE: usize;
+    type Next;
+}
+
+pub struct U0;
+
+macro_rules! impl_non_zero {
+    ($(($name:ident = $value:expr, $next:ident))*) => {
+        $(
+            pub struct $name;
+
+            const _: () = {
+                assert!($value != 0);
+            };
+
+            /// SAFETY: `VALUE` is non-zero.
+            unsafe impl NonZero for $name {
+                const VALUE: usize = $value;
+                type Next = $next;
+            }
+        )*
+    }
+}
+
+impl_non_zero! {
+    (U1 = 1, U0)
+    (U2 = 2, U1)
+    (U3 = 3, U2)
+    (U4 = 4, U3)
+    (U5 = 5, U4)
+    (U6 = 6, U5)
+}
+
+#[repr(transparent)]
+pub struct ReservedAssembler<'a, R>(&'a mut Assembler, core::marker::PhantomData<R>);
+
+impl<'a> ReservedAssembler<'a, U0> {
+    #[allow(clippy::unused_self)]
+    #[cfg_attr(not(debug_assertions), inline(always))]
+    pub fn assert_reserved_exactly_as_needed(self) {}
+}
+
+impl<'a, R> ReservedAssembler<'a, R> {
+    #[cfg_attr(not(debug_assertions), inline(always))]
+    pub fn push<T>(self, instruction: Instruction<T>) -> ReservedAssembler<'a, R::Next>
+    where
+        R: NonZero,
+        T: core::fmt::Display,
+    {
+        // SAFETY: `R: NonZero`, so we still have space in the buffer.
+        unsafe {
+            self.0.push_unchecked(instruction);
+        }
+
+        ReservedAssembler(self.0, core::marker::PhantomData)
+    }
+
+    #[cfg_attr(not(debug_assertions), inline(always))]
+    pub fn push_if<T>(self, condition: bool, instruction: Instruction<T>) -> ReservedAssembler<'a, R::Next>
+    where
+        R: NonZero,
+        T: core::fmt::Display,
+    {
+        if condition {
+            // SAFETY: `R: NonZero`, so we still have space in the buffer.
+            unsafe {
+                self.0.push_unchecked(instruction);
+            }
+        }
+
+        ReservedAssembler(self.0, core::marker::PhantomData)
+    }
+
+    #[cfg_attr(not(debug_assertions), inline(always))]
+    pub fn push_none(self) -> ReservedAssembler<'a, R::Next>
+    where
+        R: NonZero,
+    {
+        ReservedAssembler(self.0, core::marker::PhantomData)
+    }
+
+    #[cfg_attr(not(debug_assertions), inline(always))]
+    pub fn get_label_origin_offset(&self, label: Label) -> Option<isize> {
+        self.0.get_label_origin_offset(label)
+    }
+
+    #[cfg_attr(not(debug_assertions), inline(always))]
+    pub fn len(&self) -> usize {
+        self.0.len()
+    }
+
+    #[cfg_attr(not(debug_assertions), inline(always))]
+    pub fn is_empty(&self) -> bool {
+        self.0.is_empty()
+    }
+}
+
 impl Assembler {
     pub const fn new() -> Self {
         Assembler {
@@ -140,9 +241,14 @@ impl Assembler {
     }
 
     #[inline(always)]
-    pub fn reserve<const INSTRUCTIONS: usize>(&mut self) {
-        InstBuf::reserve::<INSTRUCTIONS>(&mut self.code);
-        self.guaranteed_capacity = INSTRUCTIONS;
+    pub fn reserve<T>(&mut self) -> ReservedAssembler<T>
+    where
+        T: NonZero,
+    {
+        InstBuf::reserve(&mut self.code, T::VALUE);
+
+        self.guaranteed_capacity = T::VALUE;
+        ReservedAssembler(self, core::marker::PhantomData)
     }
 
     #[cfg_attr(not(debug_assertions), inline(always))]
@@ -150,17 +256,28 @@ impl Assembler {
     where
         T: core::fmt::Display,
     {
-        #[cfg(debug_assertions)]
-        log::trace!("{:08x}: {}", self.origin + self.code.len() as u64, instruction);
-
         if self.guaranteed_capacity == 0 {
-            InstBuf::reserve::<1>(&mut self.code);
+            InstBuf::reserve_const::<1>(&mut self.code);
             self.guaranteed_capacity = 1;
         }
 
+        // SAFETY: We've reserved space for at least one instruction.
+        unsafe { self.push_unchecked(instruction) }
+    }
+
+    // SAFETY: The buffer *must* have space for at least one instruction.
+    #[cfg_attr(not(debug_assertions), inline(always))]
+    unsafe fn push_unchecked<T>(&mut self, instruction: Instruction<T>) -> &mut Self
+    where
+        T: core::fmt::Display,
+    {
+        #[cfg(debug_assertions)]
+        log::trace!("{:08x}: {}", self.origin + self.code.len() as u64, instruction);
+
+        debug_assert!(self.guaranteed_capacity > 0);
         let instruction_offset = self.code.len();
 
-        // SAFETY: We've reserved space for at least one instruction.
+        // SAFETY: The caller reserved space for at least one instruction.
         unsafe {
             instruction.bytes.encode_into_vec_unsafe(&mut self.code);
         }
@@ -404,8 +521,13 @@ impl InstBuf {
     }
 
     #[inline(always)]
-    fn reserve<const INSTRUCTIONS: usize>(output: &mut Vec<u8>) {
-        let count = INSTRUCTIONS.checked_mul(MAXIMUM_INSTRUCTION_SIZE).unwrap();
+    fn reserve_const<const INSTRUCTIONS: usize>(output: &mut Vec<u8>) {
+        Self::reserve(output, INSTRUCTIONS);
+    }
+
+    #[inline(always)]
+    fn reserve(output: &mut Vec<u8>, count: usize) {
+        let count = count.checked_mul(MAXIMUM_INSTRUCTION_SIZE).unwrap();
         if output.spare_capacity_mut().len() < count {
             Self::reserve_impl(output, count);
             if output.spare_capacity_mut().len() < count {
diff --git a/crates/polkavm-assembler/src/lib.rs b/crates/polkavm-assembler/src/lib.rs
index ff87dbc8..72034215 100644
--- a/crates/polkavm-assembler/src/lib.rs
+++ b/crates/polkavm-assembler/src/lib.rs
@@ -8,4 +8,4 @@ mod assembler;
 
 extern crate alloc;
 
-pub use crate::assembler::{Assembler, Instruction, Label};
+pub use crate::assembler::{Assembler, Instruction, Label, NonZero, ReservedAssembler, U0, U1, U2, U3, U4, U5, U6};
diff --git a/crates/polkavm-common/src/program.rs b/crates/polkavm-common/src/program.rs
index c23647e9..84835e11 100644
--- a/crates/polkavm-common/src/program.rs
+++ b/crates/polkavm-common/src/program.rs
@@ -1,8 +1,69 @@
 use crate::abi::{VM_CODE_ADDRESS_ALIGNMENT, VM_MAXIMUM_CODE_SIZE, VM_MAXIMUM_IMPORT_COUNT, VM_MAXIMUM_JUMP_TABLE_ENTRIES};
 use crate::utils::CowBytes;
-use crate::varint::{read_simple_varint_fast, read_varint, read_varint_fast, write_simple_varint, write_varint, MAX_VARINT_LENGTH};
+use crate::varint::{read_simple_varint, read_simple_varint_length_minus_1, read_varint, write_simple_varint, MAX_VARINT_LENGTH};
 use core::ops::Range;
 
+#[derive(Copy, Clone)]
+#[repr(transparent)]
+pub struct RawReg(u32);
+
+impl Eq for RawReg {}
+impl PartialEq for RawReg {
+    fn eq(&self, rhs: &Self) -> bool {
+        self.get() == rhs.get()
+    }
+}
+
+impl RawReg {
+    #[inline]
+    pub const fn get(self) -> Reg {
+        let mut value = self.0 & 0b1111;
+        if value > 12 {
+            value = 12;
+        }
+
+        match value {
+            0 => Reg::RA,
+            1 => Reg::SP,
+            2 => Reg::T0,
+            3 => Reg::T1,
+            4 => Reg::T2,
+            5 => Reg::S0,
+            6 => Reg::S1,
+            7 => Reg::A0,
+            8 => Reg::A1,
+            9 => Reg::A2,
+            10 => Reg::A3,
+            11 => Reg::A4,
+            12 => Reg::A5,
+            _ => unreachable!(),
+        }
+    }
+
+    #[inline]
+    pub const fn raw_unparsed(self) -> u32 {
+        self.0
+    }
+}
+
+impl From<Reg> for RawReg {
+    fn from(reg: Reg) -> Self {
+        Self(reg as u32)
+    }
+}
+
+impl core::fmt::Debug for RawReg {
+    fn fmt(&self, fmt: &mut core::fmt::Formatter) -> core::fmt::Result {
+        write!(fmt, "{} (0x{:x})", self.get(), self.0)
+    }
+}
+
+impl core::fmt::Display for RawReg {
+    fn fmt(&self, fmt: &mut core::fmt::Formatter) -> core::fmt::Result {
+        self.get().fmt(fmt)
+    }
+}
+
 #[derive(Copy, Clone, PartialEq, Eq, Debug, Hash)]
 #[repr(u32)]
 pub enum Reg {
@@ -23,23 +84,8 @@ pub enum Reg {
 
 impl Reg {
     #[inline]
-    pub const fn from_u8(value: u8) -> Option<Reg> {
-        match value {
-            0 => Some(Reg::RA),
-            1 => Some(Reg::SP),
-            2 => Some(Reg::T0),
-            3 => Some(Reg::T1),
-            4 => Some(Reg::T2),
-            5 => Some(Reg::S0),
-            6 => Some(Reg::S1),
-            7 => Some(Reg::A0),
-            8 => Some(Reg::A1),
-            9 => Some(Reg::A2),
-            10 => Some(Reg::A3),
-            11 => Some(Reg::A4),
-            12 => Some(Reg::A5),
-            _ => None,
-        }
+    pub const fn raw(self) -> RawReg {
+        RawReg(self as u32)
     }
 
     pub const fn name(self) -> &'static str {
@@ -83,48 +129,7 @@ impl core::fmt::Display for Reg {
 #[allow(clippy::partial_pub_fields)]
 #[doc(hidden)]
 pub struct VisitorHelper<T> {
-    chunk: u128,
     pub visitor: T,
-    args_length: usize,
-    instruction_offset: u32,
-    opcode: u8,
-}
-
-macro_rules! skip {
-    ($chunk:expr, $count:expr) => {
-        $chunk >> ($count << 3)
-    };
-}
-
-macro_rules! read_reg {
-    ($chunk:ident) => {{
-        let reg = Reg::from_u8($chunk as u8)?;
-        (reg, skip!($chunk, 1))
-    }};
-}
-
-macro_rules! read_reg2 {
-    ($chunk:ident) => {{
-        let value = $chunk as u8;
-        let reg1 = Reg::from_u8(value & 0b1111)?;
-        let reg2 = Reg::from_u8(value >> 4)?;
-        (reg1, reg2, skip!($chunk, 1))
-    }};
-}
-
-macro_rules! read_simple_varint {
-    ($chunk:expr, $length:expr) => {{
-        let imm = read_simple_varint_fast($chunk as u32, $length as u32);
-        (imm, skip!($chunk, $length))
-    }};
-}
-
-macro_rules! read_varint {
-    ($chunk:ident) => {{
-        let (imm_length, imm) = read_varint_fast($chunk as u64)?;
-        let imm_length = imm_length as usize;
-        (imm, skip!($chunk, imm_length), imm_length)
-    }};
 }
 
 impl<T> VisitorHelper<T> {
@@ -134,28 +139,28 @@ impl<T> VisitorHelper<T> {
     fn step_slow(
         &mut self,
         code: &[u8],
+        bitmask: &[u8],
         instruction_offset: usize,
-        instruction_length: usize,
-        decode_table: &[fn(state: &mut Self) -> <T as InstructionVisitor>::ReturnTy],
-    ) -> Option<<T as InstructionVisitor>::ReturnTy>
+        decode_table: &[fn(state: &mut Self, chunk: u128, instruction_offset: u32, args_length: u32) -> <T as ParsingVisitor>::ReturnTy],
+    ) -> Option<(usize, <T as ParsingVisitor>::ReturnTy)>
     where
         T: ParsingVisitor,
     {
-        let chunk = code.get(instruction_offset..instruction_offset + instruction_length)?;
+        let (next_offset, args_length) = parse_bitmask_slow(bitmask, instruction_offset)?;
+        let chunk_length = core::cmp::min(16, args_length + 1);
+        let chunk = code.get(instruction_offset..instruction_offset + chunk_length)?;
         let opcode = chunk[0];
-        let args = &chunk[1..instruction_length];
+
         let mut t: [u8; 16] = [0; 16];
-        t[..instruction_length].copy_from_slice(chunk);
-        self.chunk = u128::from_le_bytes([
+        t[..chunk_length].copy_from_slice(chunk);
+        let chunk = u128::from_le_bytes([
             t[0], t[1], t[2], t[3], t[4], t[5], t[6], t[7], t[8], t[9], t[10], t[11], t[12], t[13], t[14], t[15],
         ]) >> 8;
 
-        self.opcode = opcode;
-        self.args_length = args.len();
-        self.instruction_offset = instruction_offset as u32;
-
-        self.visitor.on_pre_visit(instruction_offset, instruction_length, opcode);
-        Some(decode_table[opcode as usize](self))
+        Some((
+            next_offset,
+            decode_table[opcode as usize](self, chunk, instruction_offset as u32, args_length as u32),
+        ))
     }
 
     #[allow(clippy::type_complexity)]
@@ -164,51 +169,46 @@ impl<T> VisitorHelper<T> {
         &mut self,
         code: &[u8],
         bitmask: &[u8],
-        offset: &mut usize,
-        decode_table: &[fn(state: &mut Self) -> <T as InstructionVisitor>::ReturnTy],
-    ) -> Option<<T as InstructionVisitor>::ReturnTy>
+        instruction_offset: usize,
+        decode_table: &[fn(state: &mut Self, chunk: u128, instruction_offset: u32, args_length: u32) -> <T as ParsingVisitor>::ReturnTy],
+    ) -> Option<(usize, <T as ParsingVisitor>::ReturnTy)>
     where
         T: ParsingVisitor,
     {
-        let instruction_offset = *offset;
-        let args_length = parse_bitmask(bitmask, offset)?;
-        let instruction_length = args_length + 1;
-
-        let padded_length = core::cmp::max(instruction_length, 16);
-        let Some(chunk) = code.get(instruction_offset..instruction_offset + padded_length) else {
-            return self.step_slow(code, instruction_offset, instruction_length, decode_table);
-        };
-
-        assert!(chunk.len() >= 16);
-        let opcode = chunk[0];
-
-        // NOTE: This should produce the same assembly as the unsafe `read_unaligned`.
-        self.chunk = u128::from_le_bytes([
-            chunk[0], chunk[1], chunk[2], chunk[3], chunk[4], chunk[5], chunk[6], chunk[7], chunk[8], chunk[9], chunk[10], chunk[11],
-            chunk[12], chunk[13], chunk[14], chunk[15],
-        ]) >> 8;
-
-        self.opcode = opcode;
-        self.args_length = instruction_length - 1;
-        self.instruction_offset = instruction_offset as u32;
+        if let Some((next_offset, args_length)) = parse_bitmask_fast(bitmask, instruction_offset) {
+            debug_assert!(args_length <= 31);
+            if let Some(chunk) = code.get(instruction_offset..instruction_offset + 32) {
+                assert!(chunk.len() >= 32);
+                let opcode = chunk[0];
+
+                // NOTE: This should produce the same assembly as the unsafe `read_unaligned`.
+                let chunk = u128::from_le_bytes([
+                    chunk[1], chunk[2], chunk[3], chunk[4], chunk[5], chunk[6], chunk[7], chunk[8], chunk[9], chunk[10], chunk[11],
+                    chunk[12], chunk[13], chunk[14], chunk[15], chunk[16],
+                ]);
+
+                return Some((
+                    next_offset,
+                    decode_table[opcode as usize](self, chunk, instruction_offset as u32, args_length as u32),
+                ));
+            }
+        }
 
-        self.visitor.on_pre_visit(instruction_offset, instruction_length, opcode);
-        Some(decode_table[opcode as usize](self))
+        self.step_slow(code, bitmask, instruction_offset, decode_table)
     }
 
     #[inline]
     pub fn new(visitor: T) -> Self {
-        VisitorHelper {
-            chunk: 0,
-            visitor,
-            args_length: 0,
-            instruction_offset: 0,
-            opcode: 0,
-        }
+        VisitorHelper { visitor }
     }
 
+    #[allow(clippy::type_complexity)]
     #[inline]
-    pub fn run(mut self, blob: &ProgramBlob, decode_table: &[fn(&mut Self) -> <T as InstructionVisitor>::ReturnTy; 256]) -> T
+    pub fn run(
+        mut self,
+        blob: &ProgramBlob,
+        decode_table: &[fn(&mut Self, chunk: u128, instruction_offset: u32, args_length: u32) -> <T as ParsingVisitor>::ReturnTy; 256],
+    ) -> T
     where
         T: ParsingVisitor<ReturnTy = ()>,
     {
@@ -217,100 +217,127 @@ impl<T> VisitorHelper<T> {
         debug_assert_eq!(bitmask[0] & 0b1, 1);
 
         let mut offset = 0;
-        while self.step(code, bitmask, &mut offset, decode_table).is_some() {}
-        self.visitor
-    }
+        loop {
+            let Some((next_offset, ())) = self.step(code, bitmask, offset, decode_table) else {
+                break;
+            };
+            offset = next_offset;
+        }
 
-    #[inline]
-    pub fn opcode(&self) -> u8 {
-        self.opcode
+        self.visitor
     }
 
     #[inline(always)]
-    pub fn read_args_offset(&mut self) -> Option<u32> {
-        let (imm, _) = read_simple_varint!(self.chunk, self.args_length);
-        Some(self.instruction_offset.wrapping_add(imm))
+    pub fn read_args_offset(&mut self, chunk: u128, instruction_offset: u32, args_length: u32) -> Option<u32> {
+        let imm = read_simple_varint(chunk as u32, args_length);
+        Some(instruction_offset.wrapping_add(imm))
     }
 
     #[inline(always)]
-    pub fn read_args_imm(&mut self) -> Option<u32> {
-        let (imm, _) = read_simple_varint!(self.chunk, self.args_length);
+    pub fn read_args_imm(&mut self, chunk: u128, args_length: u32) -> Option<u32> {
+        let imm = read_simple_varint(chunk as u32, args_length);
         Some(imm)
     }
 
     #[inline(always)]
-    pub fn read_args_imm2(&mut self) -> Option<(u32, u32)> {
-        let chunk = self.chunk;
-        let (imm1, chunk, imm1_length) = read_varint!(chunk);
-        let (imm2, _) = read_simple_varint!(chunk, self.args_length - imm1_length);
+    pub fn read_args_imm2(&mut self, chunk: u128, args_length: u32) -> Option<(u32, u32)> {
+        let imm1_length = u32::from(chunk as u8);
+        let chunk = chunk >> 8;
+        let chunk = chunk as u64;
+        let imm1 = read_simple_varint(chunk as u32, imm1_length);
+        let chunk = chunk >> (imm1_length << 3);
+        let imm2 = read_simple_varint_length_minus_1(chunk as u32, args_length.wrapping_sub(imm1_length));
         Some((imm1, imm2))
     }
 
     #[inline(always)]
-    pub fn read_args_reg_imm(&mut self) -> Option<(Reg, u32)> {
-        let chunk = self.chunk;
-        let (reg, chunk) = read_reg!(chunk);
-        let (imm, _) = read_simple_varint!(chunk, self.args_length - 1);
+    pub fn read_args_reg_imm(&mut self, chunk: u128, args_length: u32) -> Option<(RawReg, u32)> {
+        let chunk = chunk as u64;
+        let reg = RawReg(chunk as u32);
+        let chunk = chunk >> 8;
+        let imm = read_simple_varint_length_minus_1(chunk as u32, args_length);
         Some((reg, imm))
     }
 
     #[inline(always)]
-    pub fn read_args_reg_imm_offset(&mut self) -> Option<(Reg, u32, u32)> {
-        let chunk = self.chunk;
-        let (reg, chunk) = read_reg!(chunk);
-        let (imm1, chunk, imm1_length) = read_varint!(chunk);
-        let (imm2, _) = read_simple_varint!(chunk, self.args_length - 1 - imm1_length);
-        let imm2 = self.instruction_offset.wrapping_add(imm2);
+    pub fn read_args_reg_imm_offset(&mut self, chunk: u128, instruction_offset: u32, args_length: u32) -> Option<(RawReg, u32, u32)> {
+        let (reg, imm1_length) = {
+            let value = chunk as u32;
+            (RawReg(value), (value >> 4) & 0b1111)
+        };
+        let chunk = chunk >> 8;
+        let chunk = chunk as u64;
+        let imm1 = read_simple_varint(chunk as u32, imm1_length);
+        let chunk = chunk >> (imm1_length << 3);
+        let imm2 = read_simple_varint_length_minus_1(chunk as u32, args_length.wrapping_sub(imm1_length));
+        let imm2 = instruction_offset.wrapping_add(imm2);
         Some((reg, imm1, imm2))
     }
 
     #[inline(always)]
-    pub fn read_args_reg_imm2(&mut self) -> Option<(Reg, u32, u32)> {
-        let chunk = self.chunk;
-        let (reg, chunk) = read_reg!(chunk);
-        let (imm1, chunk, imm1_length) = read_varint!(chunk);
-        let (imm2, _) = read_simple_varint!(chunk, self.args_length - 1 - imm1_length);
+    pub fn read_args_reg_imm2(&mut self, chunk: u128, args_length: u32) -> Option<(RawReg, u32, u32)> {
+        let (reg, imm1_length) = {
+            let value = chunk as u32;
+            (RawReg(value), (value >> 4) & 0b1111)
+        };
+        let chunk = chunk >> 8;
+        let chunk = chunk as u64;
+        let imm1 = read_simple_varint(chunk as u32, imm1_length);
+        let chunk = chunk >> (imm1_length << 3);
+        let imm2 = read_simple_varint_length_minus_1(chunk as u32, args_length.wrapping_sub(imm1_length));
         Some((reg, imm1, imm2))
     }
 
     #[inline(always)]
-    pub fn read_args_regs2_imm2(&mut self) -> Option<(Reg, Reg, u32, u32)> {
-        let chunk = self.chunk;
-        let (reg1, reg2, chunk) = read_reg2!(chunk);
-        let (imm1, chunk, imm1_length) = read_varint!(chunk);
-        let (imm2, _) = read_simple_varint!(chunk, self.args_length - 1 - imm1_length);
+    pub fn read_args_regs2_imm2(&mut self, chunk: u128, args_length: u32) -> Option<(RawReg, RawReg, u32, u32)> {
+        let (reg1, reg2, imm1_length) = {
+            let value = chunk as u32;
+            (RawReg(value), RawReg(value >> 4), u32::from((value >> 8) as u8))
+        };
+        let chunk = chunk >> 16;
+        let chunk = chunk as u64;
+        let imm1 = read_simple_varint(chunk as u32, imm1_length);
+        let chunk = chunk >> (imm1_length << 3);
+        let imm2 = read_simple_varint_length_minus_1(chunk as u32, args_length.wrapping_sub(imm1_length + 1));
         Some((reg1, reg2, imm1, imm2))
     }
 
     #[inline(always)]
-    pub fn read_args_regs2_imm(&mut self) -> Option<(Reg, Reg, u32)> {
-        let chunk = self.chunk;
-        let (reg1, reg2, chunk) = read_reg2!(chunk);
-        let (imm, _) = read_simple_varint!(chunk, self.args_length - 1);
+    pub fn read_args_regs2_imm(&mut self, chunk: u128, args_length: u32) -> Option<(RawReg, RawReg, u32)> {
+        let chunk = chunk as u64;
+        let (reg1, reg2) = {
+            let value = chunk as u32;
+            (RawReg(value), RawReg(value >> 4))
+        };
+        let chunk = chunk >> 8;
+        let imm = read_simple_varint_length_minus_1(chunk as u32, args_length);
         Some((reg1, reg2, imm))
     }
 
     #[inline(always)]
-    pub fn read_args_regs2_offset(&mut self) -> Option<(Reg, Reg, u32)> {
-        let chunk = self.chunk;
-        let (reg1, reg2, chunk) = read_reg2!(chunk);
-        let (imm, _) = read_simple_varint!(chunk, self.args_length - 1);
-        let imm = self.instruction_offset.wrapping_add(imm);
+    pub fn read_args_regs2_offset(&mut self, chunk: u128, instruction_offset: u32, args_length: u32) -> Option<(RawReg, RawReg, u32)> {
+        let chunk = chunk as u64;
+        let (reg1, reg2) = {
+            let value = chunk as u32;
+            (RawReg(value), RawReg(value >> 4))
+        };
+        let chunk = chunk >> 8;
+        let imm = read_simple_varint_length_minus_1(chunk as u32, args_length);
+        let imm = instruction_offset.wrapping_add(imm);
         Some((reg1, reg2, imm))
     }
 
     #[inline(always)]
-    pub fn read_args_regs3(&mut self) -> Option<(Reg, Reg, Reg)> {
-        let chunk = self.chunk;
-        let (reg1, reg2, chunk) = read_reg2!(chunk);
-        let (reg3, _) = read_reg!(chunk);
+    pub fn read_args_regs3(&mut self, chunk: u128) -> Option<(RawReg, RawReg, RawReg)> {
+        let chunk = chunk as u32;
+        let (reg1, reg2, reg3) = (RawReg(chunk), RawReg(chunk >> 4), RawReg(chunk >> 8));
         Some((reg1, reg2, reg3))
     }
 
     #[inline(always)]
-    pub fn read_args_regs2(&mut self) -> Option<(Reg, Reg)> {
-        let chunk = self.chunk;
-        let (reg1, reg2, _) = read_reg2!(chunk);
+    pub fn read_args_regs2(&mut self, chunk: u128) -> Option<(RawReg, RawReg)> {
+        let chunk = chunk as u32;
+        let (reg1, reg2) = (RawReg(chunk), RawReg(chunk >> 4));
         Some((reg1, reg2))
     }
 }
@@ -386,98 +413,69 @@ macro_rules! define_opcodes {
         [$($name_reg_reg:ident = $value_reg_reg:expr,)+]
         [$($name_reg_reg_imm_imm:ident = $value_reg_reg_imm_imm:expr,)+]
     ) => {
-        pub trait ParsingVisitor: InstructionVisitor {
-            fn on_pre_visit(&mut self, _offset: usize, _length: usize, _opcode: u8);
+        pub trait ParsingVisitor {
+            type ReturnTy;
+
+            $(fn $name_argless(&mut self, offset: u32, args_length: u32) -> Self::ReturnTy;)+
+            $(fn $name_reg_imm(&mut self, offset: u32, args_length: u32, reg: RawReg, imm: u32) -> Self::ReturnTy;)+
+            $(fn $name_reg_imm_offset(&mut self, offset: u32, args_length: u32, reg: RawReg, imm1: u32, imm2: u32) -> Self::ReturnTy;)+
+            $(fn $name_reg_imm_imm(&mut self, offset: u32, args_length: u32, reg: RawReg, imm1: u32, imm2: u32) -> Self::ReturnTy;)+
+            $(fn $name_reg_reg_imm(&mut self, offset: u32, args_length: u32, reg1: RawReg, reg2: RawReg, imm: u32) -> Self::ReturnTy;)+
+            $(fn $name_reg_reg_offset(&mut self, offset: u32, args_length: u32, reg1: RawReg, reg2: RawReg, imm: u32) -> Self::ReturnTy;)+
+            $(fn $name_reg_reg_reg(&mut self, offset: u32, args_length: u32, reg1: RawReg, reg2: RawReg, reg3: RawReg) -> Self::ReturnTy;)+
+            $(fn $name_offset(&mut self, offset: u32, args_length: u32, imm: u32) -> Self::ReturnTy;)+
+            $(fn $name_imm(&mut self, offset: u32, args_length: u32, imm: u32) -> Self::ReturnTy;)+
+            $(fn $name_imm_imm(&mut self, offset: u32, args_length: u32, imm1: u32, imm2: u32) -> Self::ReturnTy;)+
+            $(fn $name_reg_reg(&mut self, offset: u32, args_length: u32, reg1: RawReg, reg2: RawReg) -> Self::ReturnTy;)+
+            $(fn $name_reg_reg_imm_imm(&mut self, offset: u32, args_length: u32, reg1: RawReg, reg2: RawReg, imm1: u32, imm2: u32) -> Self::ReturnTy;)+
+
+            #[inline(never)]
+            #[cold]
+            fn invalid(&mut self, offset: u32, args_length: u32) -> Self::ReturnTy {
+                self.trap(offset, args_length)
+            }
         }
 
         pub trait InstructionVisitor {
             type ReturnTy;
 
             $(fn $name_argless(&mut self) -> Self::ReturnTy;)+
-            $(fn $name_reg_imm(&mut self, reg: Reg, imm: u32) -> Self::ReturnTy;)+
-            $(fn $name_reg_imm_offset(&mut self, reg: Reg, imm1: u32, imm2: u32) -> Self::ReturnTy;)+
-            $(fn $name_reg_imm_imm(&mut self, reg: Reg, imm1: u32, imm2: u32) -> Self::ReturnTy;)+
-            $(fn $name_reg_reg_imm(&mut self, reg1: Reg, reg2: Reg, imm: u32) -> Self::ReturnTy;)+
-            $(fn $name_reg_reg_offset(&mut self, reg1: Reg, reg2: Reg, imm: u32) -> Self::ReturnTy;)+
-            $(fn $name_reg_reg_reg(&mut self, reg1: Reg, reg2: Reg, reg3: Reg) -> Self::ReturnTy;)+
+            $(fn $name_reg_imm(&mut self, reg: RawReg, imm: u32) -> Self::ReturnTy;)+
+            $(fn $name_reg_imm_offset(&mut self, reg: RawReg, imm1: u32, imm2: u32) -> Self::ReturnTy;)+
+            $(fn $name_reg_imm_imm(&mut self, reg: RawReg, imm1: u32, imm2: u32) -> Self::ReturnTy;)+
+            $(fn $name_reg_reg_imm(&mut self, reg1: RawReg, reg2: RawReg, imm: u32) -> Self::ReturnTy;)+
+            $(fn $name_reg_reg_offset(&mut self, reg1: RawReg, reg2: RawReg, imm: u32) -> Self::ReturnTy;)+
+            $(fn $name_reg_reg_reg(&mut self, reg1: RawReg, reg2: RawReg, reg3: RawReg) -> Self::ReturnTy;)+
             $(fn $name_offset(&mut self, imm: u32) -> Self::ReturnTy;)+
             $(fn $name_imm(&mut self, imm: u32) -> Self::ReturnTy;)+
             $(fn $name_imm_imm(&mut self, imm1: u32, imm2: u32) -> Self::ReturnTy;)+
-            $(fn $name_reg_reg(&mut self, reg1: Reg, reg2: Reg) -> Self::ReturnTy;)+
-            $(fn $name_reg_reg_imm_imm(&mut self, reg1: Reg, reg2: Reg, imm1: u32, imm2: u32) -> Self::ReturnTy;)+
+            $(fn $name_reg_reg(&mut self, reg1: RawReg, reg2: RawReg) -> Self::ReturnTy;)+
+            $(fn $name_reg_reg_imm_imm(&mut self, reg1: RawReg, reg2: RawReg, imm1: u32, imm2: u32) -> Self::ReturnTy;)+
 
             #[inline(never)]
             #[cold]
-            fn invalid(&mut self, _opcode: u8) -> Self::ReturnTy {
+            fn invalid(&mut self) -> Self::ReturnTy {
                 self.trap()
             }
         }
 
-        #[macro_export]
-        macro_rules! implement_instruction_visitor {
-            (impl<$d($visitor_ty_params:tt),*> $visitor_ty:ty, $method:ident) => {
-                impl<$d($visitor_ty_params),*> polkavm_common::program::InstructionVisitor for $visitor_ty {
-                    type ReturnTy = ();
-
-                    $(fn $name_argless(&mut self) -> Self::ReturnTy {
-                        self.$method(polkavm_common::program::Instruction::$name_argless);
-                    })+
-                    $(fn $name_reg_imm(&mut self, reg: Reg, imm: u32) -> Self::ReturnTy {
-                        self.$method(polkavm_common::program::Instruction::$name_reg_imm(reg, imm));
-                    })+
-                    $(fn $name_reg_imm_offset(&mut self, reg: Reg, imm1: u32, imm2: u32) -> Self::ReturnTy {
-                        self.$method(polkavm_common::program::Instruction::$name_reg_imm_offset(reg, imm1, imm2));
-                    })+
-                    $(fn $name_reg_imm_imm(&mut self, reg: Reg, imm1: u32, imm2: u32) -> Self::ReturnTy {
-                        self.$method(polkavm_common::program::Instruction::$name_reg_imm_imm(reg, imm1, imm2));
-                    })+
-                    $(fn $name_reg_reg_imm(&mut self, reg1: Reg, reg2: Reg, imm: u32) -> Self::ReturnTy {
-                        self.$method(polkavm_common::program::Instruction::$name_reg_reg_imm(reg1, reg2, imm));
-                    })+
-                    $(fn $name_reg_reg_offset(&mut self, reg1: Reg, reg2: Reg, imm: u32) -> Self::ReturnTy {
-                        self.$method(polkavm_common::program::Instruction::$name_reg_reg_offset(reg1, reg2, imm));
-                    })+
-                    $(fn $name_reg_reg_reg(&mut self, reg1: Reg, reg2: Reg, reg3: Reg) -> Self::ReturnTy {
-                        self.$method(polkavm_common::program::Instruction::$name_reg_reg_reg(reg1, reg2, reg3));
-                    })+
-                    $(fn $name_offset(&mut self, imm: u32) -> Self::ReturnTy {
-                        self.$method(polkavm_common::program::Instruction::$name_offset(imm));
-                    })+
-                    $(fn $name_imm(&mut self, imm: u32) -> Self::ReturnTy {
-                        self.$method(polkavm_common::program::Instruction::$name_imm(imm));
-                    })+
-                    $(fn $name_imm_imm(&mut self, imm1: u32, imm2: u32) -> Self::ReturnTy {
-                        self.$method(polkavm_common::program::Instruction::$name_imm_imm(imm1, imm2));
-                    })+
-                    $(fn $name_reg_reg(&mut self, reg1: Reg, reg2: Reg) -> Self::ReturnTy {
-                        self.$method(polkavm_common::program::Instruction::$name_reg_reg(reg1, reg2));
-                    })+
-                    $(fn $name_reg_reg_imm_imm(&mut self, reg1: Reg, reg2: Reg, imm1: u32, imm2: u32) -> Self::ReturnTy {
-                        self.$method(polkavm_common::program::Instruction::$name_reg_reg_imm_imm(reg1, reg2, imm1, imm2));
-                    })+
-                }
-            }
-        }
-
-        pub use implement_instruction_visitor;
-
         #[derive(Copy, Clone, PartialEq, Eq, Debug)]
         #[allow(non_camel_case_types)]
         #[repr(u32)]
         pub enum Instruction {
             $($name_argless = $value_argless,)+
-            $($name_reg_imm(Reg, u32) = $value_reg_imm,)+
-            $($name_reg_imm_offset(Reg, u32, u32) = $value_reg_imm_offset,)+
-            $($name_reg_imm_imm(Reg, u32, u32) = $value_reg_imm_imm,)+
-            $($name_reg_reg_imm(Reg, Reg, u32) = $value_reg_reg_imm,)+
-            $($name_reg_reg_offset(Reg, Reg, u32) = $value_reg_reg_offset,)+
-            $($name_reg_reg_reg(Reg, Reg, Reg) = $value_reg_reg_reg,)+
+            $($name_reg_imm(RawReg, u32) = $value_reg_imm,)+
+            $($name_reg_imm_offset(RawReg, u32, u32) = $value_reg_imm_offset,)+
+            $($name_reg_imm_imm(RawReg, u32, u32) = $value_reg_imm_imm,)+
+            $($name_reg_reg_imm(RawReg, RawReg, u32) = $value_reg_reg_imm,)+
+            $($name_reg_reg_offset(RawReg, RawReg, u32) = $value_reg_reg_offset,)+
+            $($name_reg_reg_reg(RawReg, RawReg, RawReg) = $value_reg_reg_reg,)+
             $($name_offset(u32) = $value_offset,)+
             $($name_imm(u32) = $value_imm,)+
             $($name_imm_imm(u32, u32) = $value_imm_imm,)+
-            $($name_reg_reg(Reg, Reg) = $value_reg_reg,)+
-            $($name_reg_reg_imm_imm(Reg, Reg, u32, u32) = $value_reg_reg_imm_imm,)+
-            invalid(u8) = 88,
+            $($name_reg_reg(RawReg, RawReg) = $value_reg_reg,)+
+            $($name_reg_reg_imm_imm(RawReg, RawReg, u32, u32) = $value_reg_reg_imm_imm,)+
+            invalid = 88,
         }
 
         impl Instruction {
@@ -495,7 +493,7 @@ macro_rules! define_opcodes {
                     $(Self::$name_imm_imm(imm1, imm2) => visitor.$name_imm_imm(imm1, imm2),)+
                     $(Self::$name_reg_reg(reg1, reg2) => visitor.$name_reg_reg(reg1, reg2),)+
                     $(Self::$name_reg_reg_imm_imm(reg1, reg2, imm1, imm2) => visitor.$name_reg_reg_imm_imm(reg1, reg2, imm1, imm2),)+
-                    Self::invalid(opcode) => visitor.invalid(opcode),
+                    Self::invalid => visitor.invalid(),
                 }
             }
 
@@ -513,7 +511,7 @@ macro_rules! define_opcodes {
                     $(Self::$name_imm_imm(imm1, imm2) => Self::serialize_imm_imm(buffer, Opcode::$name_imm_imm, imm1, imm2),)+
                     $(Self::$name_reg_reg(reg1, reg2) => Self::serialize_reg_reg(buffer, Opcode::$name_reg_reg, reg1, reg2),)+
                     $(Self::$name_reg_reg_imm_imm(reg1, reg2, imm1, imm2) => Self::serialize_reg_reg_imm_imm(buffer, Opcode::$name_reg_reg_imm_imm, reg1, reg2, imm1, imm2),)+
-                    Self::invalid(..) => Self::serialize_argless(buffer, Opcode::trap),
+                    Self::invalid => Self::serialize_argless(buffer, Opcode::trap),
 
                 }
             }
@@ -532,17 +530,11 @@ macro_rules! define_opcodes {
                     $(Self::$name_imm_imm(..) => Opcode::$name_imm_imm,)+
                     $(Self::$name_reg_reg(..) => Opcode::$name_reg_reg,)+
                     $(Self::$name_reg_reg_imm_imm(..) => Opcode::$name_reg_reg_imm_imm,)+
-                    Self::invalid(..) => Opcode::trap,
+                    Self::invalid => Opcode::trap,
                 }
             }
         }
 
-        impl core::fmt::Display for Instruction {
-            fn fmt(&self, fmt: &mut core::fmt::Formatter) -> core::fmt::Result {
-                self.visit(fmt)
-            }
-        }
-
         pub mod asm {
             use super::{Instruction, Reg};
 
@@ -554,37 +546,37 @@ macro_rules! define_opcodes {
 
             $(
                 pub fn $name_reg_imm(reg: Reg, imm: u32) -> Instruction {
-                    Instruction::$name_reg_imm(reg, imm)
+                    Instruction::$name_reg_imm(reg.into(), imm)
                 }
             )+
 
             $(
                 pub fn $name_reg_imm_offset(reg: Reg, imm1: u32, imm2: u32) -> Instruction {
-                    Instruction::$name_reg_imm_offset(reg, imm1, imm2)
+                    Instruction::$name_reg_imm_offset(reg.into(), imm1, imm2)
                 }
             )+
 
             $(
                 pub fn $name_reg_imm_imm(reg: Reg, imm1: u32, imm2: u32) -> Instruction {
-                    Instruction::$name_reg_imm_imm(reg, imm1, imm2)
+                    Instruction::$name_reg_imm_imm(reg.into(), imm1, imm2)
                 }
             )+
 
             $(
                 pub fn $name_reg_reg_imm(reg1: Reg, reg2: Reg, imm: u32) -> Instruction {
-                    Instruction::$name_reg_reg_imm(reg1, reg2, imm)
+                    Instruction::$name_reg_reg_imm(reg1.into(), reg2.into(), imm)
                 }
             )+
 
             $(
                 pub fn $name_reg_reg_offset(reg1: Reg, reg2: Reg, imm: u32) -> Instruction {
-                    Instruction::$name_reg_reg_offset(reg1, reg2, imm)
+                    Instruction::$name_reg_reg_offset(reg1.into(), reg2.into(), imm)
                 }
             )+
 
             $(
                 pub fn $name_reg_reg_reg(reg1: Reg, reg2: Reg, reg3: Reg) -> Instruction {
-                    Instruction::$name_reg_reg_reg(reg1, reg2, reg3)
+                    Instruction::$name_reg_reg_reg(reg1.into(), reg2.into(), reg3.into())
                 }
             )+
 
@@ -608,13 +600,13 @@ macro_rules! define_opcodes {
 
             $(
                 pub fn $name_reg_reg(reg1: Reg, reg2: Reg) -> Instruction {
-                    Instruction::$name_reg_reg(reg1, reg2)
+                    Instruction::$name_reg_reg(reg1.into(), reg2.into())
                 }
             )+
 
             $(
                 pub fn $name_reg_reg_imm_imm(reg1: Reg, reg2: Reg, imm1: u32, imm2: u32) -> Instruction {
-                    Instruction::$name_reg_reg_imm_imm(reg1, reg2, imm1, imm2)
+                    Instruction::$name_reg_reg_imm_imm(reg1.into(), reg2.into(), imm1, imm2)
                 }
             )+
 
@@ -627,12 +619,12 @@ macro_rules! define_opcodes {
         macro_rules! prepare_visitor {
             (@define_table $table_name:ident, $visitor_ty:ident<$d($visitor_ty_params:tt),*>) => {
                 use $crate::program::{
-                    InstructionVisitor,
+                    ParsingVisitor,
                     VisitorHelper,
                 };
 
-                type ReturnTy<$d($visitor_ty_params),*> = <$visitor_ty<$d($visitor_ty_params),*> as InstructionVisitor>::ReturnTy;
-                type VisitFn<$d($visitor_ty_params),*> = fn(state: &mut VisitorHelper<$visitor_ty<$d($visitor_ty_params),*>>) -> ReturnTy<$d($visitor_ty_params),*>;
+                type ReturnTy<$d($visitor_ty_params),*> = <$visitor_ty<$d($visitor_ty_params),*> as ParsingVisitor>::ReturnTy;
+                type VisitFn<$d($visitor_ty_params),*> = fn(state: &mut VisitorHelper<$visitor_ty<$d($visitor_ty_params),*>>, chunk: u128, instruction_offset: u32, args_length: u32) -> ReturnTy<$d($visitor_ty_params),*>;
 
                 #[allow(unsafe_code)]
                 static $table_name: [VisitFn; 256] = {
@@ -643,8 +635,8 @@ macro_rules! define_opcodes {
                         // compiler and the linker to put all of this code near each other, minimizing
                         // instruction cache misses.
                         #[cfg_attr(target_os = "linux", link_section = concat!(".text.", stringify!($table_name)))]
-                        fn $name_argless<$d($visitor_ty_params),*>(state: &mut VisitorHelper<$visitor_ty<$d($visitor_ty_params),*>>) -> ReturnTy<$d($visitor_ty_params),*>{
-                            state.visitor.$name_argless()
+                        fn $name_argless<$d($visitor_ty_params),*>(state: &mut VisitorHelper<$visitor_ty<$d($visitor_ty_params),*>>, _chunk: u128, instruction_offset: u32, args_length: u32) -> ReturnTy<$d($visitor_ty_params),*>{
+                            state.visitor.$name_argless(instruction_offset, args_length)
                         }
 
                         table[$value_argless] = $name_argless;
@@ -652,12 +644,9 @@ macro_rules! define_opcodes {
 
                     $({
                         #[cfg_attr(target_os = "linux", link_section = concat!(".text.", stringify!($table_name)))]
-                        fn $name_reg_imm<$d($visitor_ty_params),*>(state: &mut VisitorHelper<$visitor_ty<$d($visitor_ty_params),*>>) -> ReturnTy<$d($visitor_ty_params),*>{
-                            if let Some((reg, imm)) = state.read_args_reg_imm() {
-                                return state.visitor.$name_reg_imm(reg, imm);
-                            };
-
-                            state.visitor.invalid($value_reg_imm)
+                        fn $name_reg_imm<$d($visitor_ty_params),*>(state: &mut VisitorHelper<$visitor_ty<$d($visitor_ty_params),*>>, chunk: u128, instruction_offset: u32, args_length: u32) -> ReturnTy<$d($visitor_ty_params),*>{
+                            let Some((reg, imm)) = state.read_args_reg_imm(chunk, args_length) else { return state.visitor.invalid(instruction_offset, args_length) };
+                            state.visitor.$name_reg_imm(instruction_offset, args_length, reg, imm)
                         }
 
                         table[$value_reg_imm] = $name_reg_imm;
@@ -665,12 +654,9 @@ macro_rules! define_opcodes {
 
                     $({
                         #[cfg_attr(target_os = "linux", link_section = concat!(".text.", stringify!($table_name)))]
-                        fn $name_reg_imm_offset<$d($visitor_ty_params),*>(state: &mut VisitorHelper<$visitor_ty<$d($visitor_ty_params),*>>) -> ReturnTy<$d($visitor_ty_params),*>{
-                            if let Some((reg, imm1, imm2)) = state.read_args_reg_imm_offset() {
-                                return state.visitor.$name_reg_imm_offset(reg, imm1, imm2);
-                            }
-
-                            state.visitor.invalid($value_reg_imm_offset)
+                        fn $name_reg_imm_offset<$d($visitor_ty_params),*>(state: &mut VisitorHelper<$visitor_ty<$d($visitor_ty_params),*>>, chunk: u128, instruction_offset: u32, args_length: u32) -> ReturnTy<$d($visitor_ty_params),*>{
+                            let Some((reg, imm1, imm2)) = state.read_args_reg_imm_offset(chunk, instruction_offset, args_length) else { return state.visitor.invalid(instruction_offset, args_length) };
+                            state.visitor.$name_reg_imm_offset(instruction_offset, args_length, reg, imm1, imm2)
                         }
 
                         table[$value_reg_imm_offset] = $name_reg_imm_offset;
@@ -678,12 +664,9 @@ macro_rules! define_opcodes {
 
                     $({
                         #[cfg_attr(target_os = "linux", link_section = concat!(".text.", stringify!($table_name)))]
-                        fn $name_reg_imm_imm<$d($visitor_ty_params),*>(state: &mut VisitorHelper<$visitor_ty<$d($visitor_ty_params),*>>) -> ReturnTy<$d($visitor_ty_params),*>{
-                            if let Some((reg, imm1, imm2)) = state.read_args_reg_imm2() {
-                                return state.visitor.$name_reg_imm_imm(reg, imm1, imm2);
-                            }
-
-                            state.visitor.invalid($value_reg_imm_imm)
+                        fn $name_reg_imm_imm<$d($visitor_ty_params),*>(state: &mut VisitorHelper<$visitor_ty<$d($visitor_ty_params),*>>, chunk: u128, instruction_offset: u32, args_length: u32) -> ReturnTy<$d($visitor_ty_params),*>{
+                            let Some((reg, imm1, imm2)) = state.read_args_reg_imm2(chunk, args_length) else { return state.visitor.invalid(instruction_offset, args_length) };
+                            state.visitor.$name_reg_imm_imm(instruction_offset, args_length, reg, imm1, imm2)
                         }
 
                         table[$value_reg_imm_imm] = $name_reg_imm_imm;
@@ -691,12 +674,9 @@ macro_rules! define_opcodes {
 
                     $({
                         #[cfg_attr(target_os = "linux", link_section = concat!(".text.", stringify!($table_name)))]
-                        fn $name_reg_reg_imm<$d($visitor_ty_params),*>(state: &mut VisitorHelper<$visitor_ty<$d($visitor_ty_params),*>>) -> ReturnTy<$d($visitor_ty_params),*>{
-                            if let Some((reg1, reg2, imm)) = state.read_args_regs2_imm() {
-                                return state.visitor.$name_reg_reg_imm(reg1, reg2, imm);
-                            }
-
-                            state.visitor.invalid($value_reg_reg_imm)
+                        fn $name_reg_reg_imm<$d($visitor_ty_params),*>(state: &mut VisitorHelper<$visitor_ty<$d($visitor_ty_params),*>>, chunk: u128, instruction_offset: u32, args_length: u32) -> ReturnTy<$d($visitor_ty_params),*>{
+                            let Some((reg1, reg2, imm)) = state.read_args_regs2_imm(chunk, args_length) else { return state.visitor.invalid(instruction_offset, args_length) };
+                            state.visitor.$name_reg_reg_imm(instruction_offset, args_length, reg1, reg2, imm)
                         }
 
                         table[$value_reg_reg_imm] = $name_reg_reg_imm;
@@ -704,12 +684,9 @@ macro_rules! define_opcodes {
 
                     $({
                         #[cfg_attr(target_os = "linux", link_section = concat!(".text.", stringify!($table_name)))]
-                        fn $name_reg_reg_offset<$d($visitor_ty_params),*>(state: &mut VisitorHelper<$visitor_ty<$d($visitor_ty_params),*>>) -> ReturnTy<$d($visitor_ty_params),*>{
-                            if let Some((reg1, reg2, imm)) = state.read_args_regs2_offset() {
-                                return state.visitor.$name_reg_reg_offset(reg1, reg2, imm);
-                            }
-
-                            state.visitor.invalid($value_reg_reg_offset)
+                        fn $name_reg_reg_offset<$d($visitor_ty_params),*>(state: &mut VisitorHelper<$visitor_ty<$d($visitor_ty_params),*>>, chunk: u128, instruction_offset: u32, args_length: u32) -> ReturnTy<$d($visitor_ty_params),*>{
+                            let Some((reg1, reg2, imm)) = state.read_args_regs2_offset(chunk, instruction_offset, args_length) else { return state.visitor.invalid(instruction_offset, args_length) };
+                            state.visitor.$name_reg_reg_offset(instruction_offset, args_length, reg1, reg2, imm)
                         }
 
                         table[$value_reg_reg_offset] = $name_reg_reg_offset;
@@ -717,12 +694,9 @@ macro_rules! define_opcodes {
 
                     $({
                         #[cfg_attr(target_os = "linux", link_section = concat!(".text.", stringify!($table_name)))]
-                        fn $name_reg_reg_reg<$d($visitor_ty_params),*>(state: &mut VisitorHelper<$visitor_ty<$d($visitor_ty_params),*>>) -> ReturnTy<$d($visitor_ty_params),*>{
-                            if let Some((reg1, reg2, reg3)) = state.read_args_regs3() {
-                                return state.visitor.$name_reg_reg_reg(reg1, reg2, reg3);
-                            }
-
-                            state.visitor.invalid($value_reg_reg_reg)
+                        fn $name_reg_reg_reg<$d($visitor_ty_params),*>(state: &mut VisitorHelper<$visitor_ty<$d($visitor_ty_params),*>>, chunk: u128, instruction_offset: u32, args_length: u32) -> ReturnTy<$d($visitor_ty_params),*>{
+                            let Some((reg1, reg2, reg3)) = state.read_args_regs3(chunk) else { return state.visitor.invalid(instruction_offset, args_length) };
+                            state.visitor.$name_reg_reg_reg(instruction_offset, args_length, reg1, reg2, reg3)
                         }
 
                         table[$value_reg_reg_reg] = $name_reg_reg_reg;
@@ -730,12 +704,9 @@ macro_rules! define_opcodes {
 
                     $({
                         #[cfg_attr(target_os = "linux", link_section = concat!(".text.", stringify!($table_name)))]
-                        fn $name_offset<$d($visitor_ty_params),*>(state: &mut VisitorHelper<$visitor_ty<$d($visitor_ty_params),*>>) -> ReturnTy<$d($visitor_ty_params),*>{
-                            if let Some(imm) = state.read_args_offset() {
-                                return state.visitor.$name_offset(imm);
-                            }
-
-                            state.visitor.invalid($value_offset)
+                        fn $name_offset<$d($visitor_ty_params),*>(state: &mut VisitorHelper<$visitor_ty<$d($visitor_ty_params),*>>, chunk: u128, instruction_offset: u32, args_length: u32) -> ReturnTy<$d($visitor_ty_params),*>{
+                            let Some(imm) = state.read_args_offset(chunk, instruction_offset, args_length) else { return state.visitor.invalid(instruction_offset, args_length) };
+                            state.visitor.$name_offset(instruction_offset, args_length, imm)
                         }
 
                         table[$value_offset] = $name_offset;
@@ -743,12 +714,9 @@ macro_rules! define_opcodes {
 
                     $({
                         #[cfg_attr(target_os = "linux", link_section = concat!(".text.", stringify!($table_name)))]
-                        fn $name_imm<$d($visitor_ty_params),*>(state: &mut VisitorHelper<$visitor_ty<$d($visitor_ty_params),*>>) -> ReturnTy<$d($visitor_ty_params),*>{
-                            if let Some(imm) = state.read_args_imm() {
-                                return state.visitor.$name_imm(imm);
-                            }
-
-                            state.visitor.invalid($value_imm)
+                        fn $name_imm<$d($visitor_ty_params),*>(state: &mut VisitorHelper<$visitor_ty<$d($visitor_ty_params),*>>, chunk: u128, instruction_offset: u32, args_length: u32) -> ReturnTy<$d($visitor_ty_params),*>{
+                            let Some(imm) = state.read_args_imm(chunk, args_length) else { return state.visitor.invalid(instruction_offset, args_length) };
+                            state.visitor.$name_imm(instruction_offset, args_length, imm)
                         }
 
                         table[$value_imm] = $name_imm;
@@ -756,12 +724,9 @@ macro_rules! define_opcodes {
 
                     $({
                         #[cfg_attr(target_os = "linux", link_section = concat!(".text.", stringify!($table_name)))]
-                        fn $name_imm_imm<$d($visitor_ty_params),*>(state: &mut VisitorHelper<$visitor_ty<$d($visitor_ty_params),*>>) -> ReturnTy<$d($visitor_ty_params),*>{
-                            if let Some((imm1, imm2)) = state.read_args_imm2() {
-                                return state.visitor.$name_imm_imm(imm1, imm2);
-                            }
-
-                            state.visitor.invalid($value_imm_imm)
+                        fn $name_imm_imm<$d($visitor_ty_params),*>(state: &mut VisitorHelper<$visitor_ty<$d($visitor_ty_params),*>>, chunk: u128, instruction_offset: u32, args_length: u32) -> ReturnTy<$d($visitor_ty_params),*>{
+                            let Some((imm1, imm2)) = state.read_args_imm2(chunk, args_length) else { return state.visitor.invalid(instruction_offset, args_length) };
+                            state.visitor.$name_imm_imm(instruction_offset, args_length, imm1, imm2)
                         }
 
                         table[$value_imm_imm] = $name_imm_imm;
@@ -769,12 +734,9 @@ macro_rules! define_opcodes {
 
                     $({
                         #[cfg_attr(target_os = "linux", link_section = concat!(".text.", stringify!($table_name)))]
-                        fn $name_reg_reg<$d($visitor_ty_params),*>(state: &mut VisitorHelper<$visitor_ty<$d($visitor_ty_params),*>>) -> ReturnTy<$d($visitor_ty_params),*>{
-                            if let Some((reg1, reg2)) = state.read_args_regs2() {
-                                return state.visitor.$name_reg_reg(reg1, reg2);
-                            }
-
-                            state.visitor.invalid($value_reg_reg)
+                        fn $name_reg_reg<$d($visitor_ty_params),*>(state: &mut VisitorHelper<$visitor_ty<$d($visitor_ty_params),*>>, chunk: u128, instruction_offset: u32, args_length: u32) -> ReturnTy<$d($visitor_ty_params),*>{
+                            let Some((reg1, reg2)) = state.read_args_regs2(chunk) else { return state.visitor.invalid(instruction_offset, args_length) };
+                            state.visitor.$name_reg_reg(instruction_offset, args_length, reg1, reg2)
                         }
 
                         table[$value_reg_reg] = $name_reg_reg;
@@ -782,12 +744,9 @@ macro_rules! define_opcodes {
 
                     $({
                         #[cfg_attr(target_os = "linux", link_section = concat!(".text.", stringify!($table_name)))]
-                        fn $name_reg_reg_imm_imm<$d($visitor_ty_params),*>(state: &mut VisitorHelper<$visitor_ty<$d($visitor_ty_params),*>>) -> ReturnTy<$d($visitor_ty_params),*>{
-                            if let Some((reg1, reg2, imm1, imm2)) = state.read_args_regs2_imm2() {
-                                return state.visitor.$name_reg_reg_imm_imm(reg1, reg2, imm1, imm2);
-                            }
-
-                            state.visitor.invalid($value_reg_reg_imm_imm)
+                        fn $name_reg_reg_imm_imm<$d($visitor_ty_params),*>(state: &mut VisitorHelper<$visitor_ty<$d($visitor_ty_params),*>>, chunk: u128, instruction_offset: u32, args_length: u32) -> ReturnTy<$d($visitor_ty_params),*>{
+                            let Some((reg1, reg2, imm1, imm2)) = state.read_args_regs2_imm2(chunk, args_length) else { return state.visitor.invalid(instruction_offset, args_length) };
+                            state.visitor.$name_reg_reg_imm_imm(instruction_offset, args_length, reg1, reg2, imm1, imm2)
                         }
 
                         table[$value_reg_reg_imm_imm] = $name_reg_reg_imm_imm;
@@ -795,8 +754,8 @@ macro_rules! define_opcodes {
 
                     #[cfg_attr(target_os = "linux", link_section = concat!(".text.", stringify!($table_name)))]
                     #[cold]
-                    fn invalid_instruction<$d($visitor_ty_params),*>(state: &mut VisitorHelper<$visitor_ty<$d($visitor_ty_params),*>>) -> ReturnTy<$d($visitor_ty_params),*>{
-                        state.visitor.invalid(state.opcode())
+                    fn invalid_instruction<$d($visitor_ty_params),*>(state: &mut VisitorHelper<$visitor_ty<$d($visitor_ty_params),*>>, _chunk: u128, instruction_offset: u32, args_length: u32) -> ReturnTy<$d($visitor_ty_params),*>{
+                        state.visitor.invalid(instruction_offset, args_length)
                     }
 
                     table
@@ -828,75 +787,59 @@ macro_rules! define_opcodes {
 
         pub use prepare_visitor;
 
-        struct ToEnumVisitor<'a>(core::marker::PhantomData<&'a ()>);
-
         impl<'a> ParsingVisitor for ToEnumVisitor<'a> {
-            fn on_pre_visit(&mut self, _offset: usize, _length: usize, _opcode: u8) {}
-        }
-
-        impl<'a> InstructionVisitor for ToEnumVisitor<'a> {
             type ReturnTy = Instruction;
 
-            $(fn $name_argless(&mut self) -> Self::ReturnTy {
+            $(fn $name_argless(&mut self, _offset: u32, args_length: u32) -> Self::ReturnTy {
+                self.args_length = args_length;
                 Instruction::$name_argless
             })+
-            $(fn $name_reg_imm(&mut self, reg: Reg, imm: u32) -> Self::ReturnTy {
+            $(fn $name_reg_imm(&mut self, _offset: u32, args_length: u32, reg: RawReg, imm: u32) -> Self::ReturnTy {
+                self.args_length = args_length;
                 Instruction::$name_reg_imm(reg, imm)
             })+
-            $(fn $name_reg_imm_offset(&mut self, reg: Reg, imm1: u32, imm2: u32) -> Self::ReturnTy {
+            $(fn $name_reg_imm_offset(&mut self, _offset: u32, args_length: u32, reg: RawReg, imm1: u32, imm2: u32) -> Self::ReturnTy {
+                self.args_length = args_length;
                 Instruction::$name_reg_imm_offset(reg, imm1, imm2)
             })+
-            $(fn $name_reg_imm_imm(&mut self, reg: Reg, imm1: u32, imm2: u32) -> Self::ReturnTy {
+            $(fn $name_reg_imm_imm(&mut self, _offset: u32, args_length: u32, reg: RawReg, imm1: u32, imm2: u32) -> Self::ReturnTy {
+                self.args_length = args_length;
                 Instruction::$name_reg_imm_imm(reg, imm1, imm2)
             })+
-            $(fn $name_reg_reg_imm(&mut self, reg1: Reg, reg2: Reg, imm: u32) -> Self::ReturnTy {
+            $(fn $name_reg_reg_imm(&mut self, _offset: u32, args_length: u32, reg1: RawReg, reg2: RawReg, imm: u32) -> Self::ReturnTy {
+                self.args_length = args_length;
                 Instruction::$name_reg_reg_imm(reg1, reg2, imm)
             })+
-            $(fn $name_reg_reg_offset(&mut self, reg1: Reg, reg2: Reg, imm: u32) -> Self::ReturnTy {
+            $(fn $name_reg_reg_offset(&mut self, _offset: u32, args_length: u32, reg1: RawReg, reg2: RawReg, imm: u32) -> Self::ReturnTy {
+                self.args_length = args_length;
                 Instruction::$name_reg_reg_offset(reg1, reg2, imm)
             })+
-            $(fn $name_reg_reg_reg(&mut self, reg1: Reg, reg2: Reg, reg3: Reg) -> Self::ReturnTy {
+            $(fn $name_reg_reg_reg(&mut self, _offset: u32, args_length: u32, reg1: RawReg, reg2: RawReg, reg3: RawReg) -> Self::ReturnTy {
+                self.args_length = args_length;
                 Instruction::$name_reg_reg_reg(reg1, reg2, reg3)
             })+
-            $(fn $name_offset(&mut self, imm: u32) -> Self::ReturnTy {
+            $(fn $name_offset(&mut self, _offset: u32, args_length: u32, imm: u32) -> Self::ReturnTy {
+                self.args_length = args_length;
                 Instruction::$name_offset(imm)
             })+
-            $(fn $name_imm(&mut self, imm: u32) -> Self::ReturnTy {
+            $(fn $name_imm(&mut self, _offset: u32, args_length: u32, imm: u32) -> Self::ReturnTy {
+                self.args_length = args_length;
                 Instruction::$name_imm(imm)
             })+
-            $(fn $name_imm_imm(&mut self, imm1: u32, imm2: u32) -> Self::ReturnTy {
+            $(fn $name_imm_imm(&mut self, _offset: u32, args_length: u32, imm1: u32, imm2: u32) -> Self::ReturnTy {
+                self.args_length = args_length;
                 Instruction::$name_imm_imm(imm1, imm2)
             })+
-            $(fn $name_reg_reg(&mut self, reg1: Reg, reg2: Reg) -> Self::ReturnTy {
+            $(fn $name_reg_reg(&mut self, _offset: u32, args_length: u32, reg1: RawReg, reg2: RawReg) -> Self::ReturnTy {
+                self.args_length = args_length;
                 Instruction::$name_reg_reg(reg1, reg2)
             })+
-            $(fn $name_reg_reg_imm_imm(&mut self, reg1: Reg, reg2: Reg, imm1: u32, imm2: u32) -> Self::ReturnTy {
+            $(fn $name_reg_reg_imm_imm(&mut self, _offset: u32, args_length: u32, reg1: RawReg, reg2: RawReg, imm1: u32, imm2: u32) -> Self::ReturnTy {
+                self.args_length = args_length;
                 Instruction::$name_reg_reg_imm_imm(reg1, reg2, imm1, imm2)
             })+
         }
 
-        #[inline]
-        fn parse_instruction(code: &[u8], bitmask: &[u8], offset: &mut usize) -> Option<ParsedInstruction> {
-            prepare_visitor!(@define_table TO_ENUM_VISITOR, ToEnumVisitor<'a>);
-
-            let decode_table: &[VisitFn; 256] = &TO_ENUM_VISITOR;
-
-            #[allow(unsafe_code)]
-            // SAFETY: Here we transmute the lifetimes which were unnecessarily extended to be 'static due to the table here being a `static`.
-            let decode_table: &[VisitFn; 256] = unsafe { core::mem::transmute(decode_table) };
-
-            let mut helper = VisitorHelper::new(ToEnumVisitor(core::marker::PhantomData));
-            let origin = *offset;
-            let instruction = helper.step(code, bitmask, offset, decode_table)?;
-            let length = helper.args_length + 1;
-
-            Some(ParsedInstruction {
-                kind: instruction,
-                offset: origin as u32,
-                length: length as u32,
-            })
-        }
-
         define_opcodes!(
             @impl_shared
             $($name_argless = $value_argless,)+
@@ -915,6 +858,147 @@ macro_rules! define_opcodes {
     }
 }
 
+struct ToEnumVisitor<'a> {
+    args_length: u32,
+    phantom: core::marker::PhantomData<&'a ()>,
+}
+
+#[inline]
+fn parse_instruction(code: &[u8], bitmask: &[u8], offset: &mut usize) -> Option<ParsedInstruction> {
+    prepare_visitor!(@define_table TO_ENUM_VISITOR, ToEnumVisitor<'a>);
+
+    let decode_table: &[VisitFn; 256] = &TO_ENUM_VISITOR;
+
+    #[allow(unsafe_code)]
+    // SAFETY: Here we transmute the lifetimes which were unnecessarily extended to be 'static due to the table here being a `static`.
+    let decode_table: &[VisitFn; 256] = unsafe { core::mem::transmute(decode_table) };
+
+    let mut helper = VisitorHelper::new(ToEnumVisitor {
+        args_length: 0,
+        phantom: core::marker::PhantomData,
+    });
+
+    let origin = *offset;
+    let (next_offset, instruction) = helper.step(code, bitmask, origin, decode_table)?;
+    *offset = next_offset;
+
+    let length = helper.visitor.args_length + 1;
+    Some(ParsedInstruction {
+        kind: instruction,
+        offset: origin as u32,
+        length,
+    })
+}
+
+#[test]
+fn test_parse_instruction() {
+    // Instruction with no arguments.
+    assert_eq!(
+        parse_instruction(&[Opcode::fallthrough as u8], &[0b11111111], &mut 0),
+        Some(ParsedInstruction {
+            kind: Instruction::fallthrough,
+            offset: 0,
+            length: 1
+        })
+    );
+
+    // Instruction with no arguments, overparametrized.
+    assert_eq!(
+        parse_instruction(&[Opcode::fallthrough as u8, 0xff], &[0b00000101], &mut 0),
+        Some(ParsedInstruction {
+            kind: Instruction::fallthrough,
+            offset: 0,
+            length: 2
+        })
+    );
+
+    // Instruction with no arguments, overparametrized, truncated code.
+    assert_eq!(parse_instruction(&[Opcode::fallthrough as u8], &[0b00000101], &mut 0), None);
+
+    // Instruction with no arguments, overparametrized until end of code.
+    assert_eq!(
+        parse_instruction(
+            &[Opcode::fallthrough as u8, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff],
+            &[0b00000001],
+            &mut 0
+        ),
+        Some(ParsedInstruction {
+            kind: Instruction::fallthrough,
+            offset: 0,
+            length: 8
+        })
+    );
+
+    // Instruction with one immediate argument.
+    assert_eq!(
+        parse_instruction(&[Opcode::ecalli as u8], &[0b00000011], &mut 0),
+        Some(ParsedInstruction {
+            kind: Instruction::ecalli(0),
+            offset: 0,
+            length: 1
+        })
+    );
+
+    assert_eq!(
+        parse_instruction(&[Opcode::ecalli as u8, 0xff, 0xff, 0xff, 0xff], &[0b00100001], &mut 0),
+        Some(ParsedInstruction {
+            kind: Instruction::ecalli(0x80000000),
+            offset: 0,
+            length: 5
+        })
+    );
+
+    // Instruction with one immediate argument, overparametrized.
+    assert_eq!(
+        parse_instruction(&[Opcode::ecalli as u8, 0xff, 0xff, 0xff, 0xff, 0x66], &[0b01000001], &mut 0),
+        Some(ParsedInstruction {
+            kind: Instruction::ecalli(0x80000000),
+            offset: 0,
+            length: 6
+        })
+    );
+
+    // Instruction with two registers and one immediate argument.
+    assert_eq!(
+        parse_instruction(&[Opcode::add_imm as u8, 0x00, 0xff, 0xff, 0xff, 0xff], &[0b01000001], &mut 0),
+        Some(ParsedInstruction {
+            kind: Instruction::add_imm(Reg::RA.into(), Reg::RA.into(), 0x80000000),
+            offset: 0,
+            length: 6
+        })
+    );
+
+    // Instruction with two registers and one immediate argument, overparametrized.
+    assert_eq!(
+        parse_instruction(&[Opcode::add_imm as u8, 0x00, 0xff, 0xff, 0xff, 0xff, 0x66], &[0b10000001], &mut 0),
+        Some(ParsedInstruction {
+            kind: Instruction::add_imm(Reg::RA.into(), Reg::RA.into(), 0x80000000),
+            offset: 0,
+            length: 7
+        })
+    );
+
+    extern crate alloc;
+    use alloc::vec;
+    use alloc::vec::Vec;
+
+    let length = 512;
+    let mut bitmask = vec![0; length / 8];
+    bitmask[0] = 0b00000001;
+    let mut code = Vec::new();
+    code.resize(length, 0xff);
+    code[0] = Opcode::add_imm as u8;
+    code[1] = 0x00;
+    assert_eq!(
+        parse_instruction(&code, &bitmask, &mut 0),
+        Some(ParsedInstruction {
+            kind: Instruction::add_imm(Reg::RA.into(), Reg::RA.into(), 0x80000000),
+            offset: 0,
+            length: 512
+        })
+    );
+}
+
 // NOTE: The opcodes here are assigned roughly in the order of how common a given instruction is,
 // except the `trap` which is deliberately hardcoded as zero.
 define_opcodes! {
@@ -1089,6 +1173,12 @@ impl Opcode {
     }
 }
 
+impl core::fmt::Display for Instruction {
+    fn fmt(&self, fmt: &mut core::fmt::Formatter) -> core::fmt::Result {
+        self.visit(fmt)
+    }
+}
+
 impl Instruction {
     pub fn starts_new_basic_block(self) -> bool {
         self.opcode().starts_new_basic_block()
@@ -1099,56 +1189,60 @@ impl Instruction {
         1
     }
 
-    fn serialize_reg_imm_offset(buffer: &mut [u8], position: u32, opcode: Opcode, reg: Reg, imm1: u32, imm2: u32) -> usize {
+    fn serialize_reg_imm_offset(buffer: &mut [u8], position: u32, opcode: Opcode, reg: RawReg, imm1: u32, imm2: u32) -> usize {
         let imm2 = imm2.wrapping_sub(position);
         buffer[0] = opcode as u8;
-        buffer[1] = reg as u8;
         let mut position = 2;
-        position += write_varint(imm1, &mut buffer[position..]);
+        let imm1_length = write_simple_varint(imm1, &mut buffer[position..]);
+        position += imm1_length;
+        buffer[1] = reg.0 as u8 | (imm1_length << 4) as u8;
         position += write_simple_varint(imm2, &mut buffer[position..]);
         position
     }
 
-    fn serialize_reg_imm_imm(buffer: &mut [u8], opcode: Opcode, reg: Reg, imm1: u32, imm2: u32) -> usize {
+    fn serialize_reg_imm_imm(buffer: &mut [u8], opcode: Opcode, reg: RawReg, imm1: u32, imm2: u32) -> usize {
         buffer[0] = opcode as u8;
-        buffer[1] = reg as u8;
         let mut position = 2;
-        position += write_varint(imm1, &mut buffer[position..]);
+        let imm1_length = write_simple_varint(imm1, &mut buffer[position..]);
+        position += imm1_length;
+        buffer[1] = reg.0 as u8 | (imm1_length << 4) as u8;
         position += write_simple_varint(imm2, &mut buffer[position..]);
         position
     }
-    fn serialize_reg_reg_imm_imm(buffer: &mut [u8], opcode: Opcode, reg1: Reg, reg2: Reg, imm1: u32, imm2: u32) -> usize {
+    fn serialize_reg_reg_imm_imm(buffer: &mut [u8], opcode: Opcode, reg1: RawReg, reg2: RawReg, imm1: u32, imm2: u32) -> usize {
         buffer[0] = opcode as u8;
-        buffer[1] = reg1 as u8 | (reg2 as u8) << 4;
-        let mut position = 2;
-        position += write_varint(imm1, &mut buffer[position..]);
+        buffer[1] = reg1.0 as u8 | (reg2.0 as u8) << 4;
+        let mut position = 3;
+        let imm1_length = write_simple_varint(imm1, &mut buffer[position..]);
+        buffer[2] = imm1_length as u8;
+        position += imm1_length;
         position += write_simple_varint(imm2, &mut buffer[position..]);
         position
     }
 
-    fn serialize_reg_reg_reg(buffer: &mut [u8], opcode: Opcode, reg1: Reg, reg2: Reg, reg3: Reg) -> usize {
+    fn serialize_reg_reg_reg(buffer: &mut [u8], opcode: Opcode, reg1: RawReg, reg2: RawReg, reg3: RawReg) -> usize {
         buffer[0] = opcode as u8;
-        buffer[1] = reg1 as u8 | (reg2 as u8) << 4;
-        buffer[2] = reg3 as u8;
+        buffer[1] = reg1.0 as u8 | (reg2.0 as u8) << 4;
+        buffer[2] = reg3.0 as u8;
         3
     }
 
-    fn serialize_reg_reg_imm(buffer: &mut [u8], opcode: Opcode, reg1: Reg, reg2: Reg, imm: u32) -> usize {
+    fn serialize_reg_reg_imm(buffer: &mut [u8], opcode: Opcode, reg1: RawReg, reg2: RawReg, imm: u32) -> usize {
         buffer[0] = opcode as u8;
-        buffer[1] = reg1 as u8 | (reg2 as u8) << 4;
+        buffer[1] = reg1.0 as u8 | (reg2.0 as u8) << 4;
         write_simple_varint(imm, &mut buffer[2..]) + 2
     }
 
-    fn serialize_reg_reg_offset(buffer: &mut [u8], position: u32, opcode: Opcode, reg1: Reg, reg2: Reg, imm: u32) -> usize {
+    fn serialize_reg_reg_offset(buffer: &mut [u8], position: u32, opcode: Opcode, reg1: RawReg, reg2: RawReg, imm: u32) -> usize {
         let imm = imm.wrapping_sub(position);
         buffer[0] = opcode as u8;
-        buffer[1] = reg1 as u8 | (reg2 as u8) << 4;
+        buffer[1] = reg1.0 as u8 | (reg2.0 as u8) << 4;
         write_simple_varint(imm, &mut buffer[2..]) + 2
     }
 
-    fn serialize_reg_imm(buffer: &mut [u8], opcode: Opcode, reg: Reg, imm: u32) -> usize {
+    fn serialize_reg_imm(buffer: &mut [u8], opcode: Opcode, reg: RawReg, imm: u32) -> usize {
         buffer[0] = opcode as u8;
-        buffer[1] = reg as u8;
+        buffer[1] = reg.0 as u8;
         write_simple_varint(imm, &mut buffer[2..]) + 2
     }
 
@@ -1165,15 +1259,17 @@ impl Instruction {
 
     fn serialize_imm_imm(buffer: &mut [u8], opcode: Opcode, imm1: u32, imm2: u32) -> usize {
         buffer[0] = opcode as u8;
-        let mut position = 1;
-        position += write_varint(imm1, &mut buffer[position..]);
+        let mut position = 2;
+        let imm1_length = write_simple_varint(imm1, &mut buffer[position..]);
+        buffer[1] = imm1_length as u8;
+        position += imm1_length;
         position += write_simple_varint(imm2, &mut buffer[position..]);
         position
     }
 
-    fn serialize_reg_reg(buffer: &mut [u8], opcode: Opcode, reg1: Reg, reg2: Reg) -> usize {
+    fn serialize_reg_reg(buffer: &mut [u8], opcode: Opcode, reg1: RawReg, reg2: RawReg) -> usize {
         buffer[0] = opcode as u8;
-        buffer[1] = reg1 as u8 | (reg2 as u8) << 4;
+        buffer[1] = reg1.0 as u8 | (reg2.0 as u8) << 4;
         2
     }
 }
@@ -1191,7 +1287,7 @@ impl<'a> InstructionVisitor for core::fmt::Formatter<'a> {
         write!(self, "fallthrough")
     }
 
-    fn sbrk(&mut self, d: Reg, s: Reg) -> Self::ReturnTy {
+    fn sbrk(&mut self, d: RawReg, s: RawReg) -> Self::ReturnTy {
         write!(self, "{d} = sbrk {s}")
     }
 
@@ -1199,167 +1295,167 @@ impl<'a> InstructionVisitor for core::fmt::Formatter<'a> {
         write!(self, "ecalli {nth_import}")
     }
 
-    fn set_less_than_unsigned(&mut self, d: Reg, s1: Reg, s2: Reg) -> Self::ReturnTy {
+    fn set_less_than_unsigned(&mut self, d: RawReg, s1: RawReg, s2: RawReg) -> Self::ReturnTy {
         write!(self, "{d} = {s1} <u {s2}")
     }
 
-    fn set_less_than_signed(&mut self, d: Reg, s1: Reg, s2: Reg) -> Self::ReturnTy {
+    fn set_less_than_signed(&mut self, d: RawReg, s1: RawReg, s2: RawReg) -> Self::ReturnTy {
         write!(self, "{d} = {s1} <s {s2}")
     }
 
-    fn shift_logical_right(&mut self, d: Reg, s1: Reg, s2: Reg) -> Self::ReturnTy {
+    fn shift_logical_right(&mut self, d: RawReg, s1: RawReg, s2: RawReg) -> Self::ReturnTy {
         write!(self, "{d} = {s1} >> {s2}")
     }
 
-    fn shift_arithmetic_right(&mut self, d: Reg, s1: Reg, s2: Reg) -> Self::ReturnTy {
+    fn shift_arithmetic_right(&mut self, d: RawReg, s1: RawReg, s2: RawReg) -> Self::ReturnTy {
         write!(self, "{d} = {s1} >>a {s2}")
     }
 
-    fn shift_logical_left(&mut self, d: Reg, s1: Reg, s2: Reg) -> Self::ReturnTy {
+    fn shift_logical_left(&mut self, d: RawReg, s1: RawReg, s2: RawReg) -> Self::ReturnTy {
         write!(self, "{d} = {s1} << {s2}")
     }
 
-    fn xor(&mut self, d: Reg, s1: Reg, s2: Reg) -> Self::ReturnTy {
+    fn xor(&mut self, d: RawReg, s1: RawReg, s2: RawReg) -> Self::ReturnTy {
         write!(self, "{d} = {s1} ^ {s2}")
     }
 
-    fn and(&mut self, d: Reg, s1: Reg, s2: Reg) -> Self::ReturnTy {
+    fn and(&mut self, d: RawReg, s1: RawReg, s2: RawReg) -> Self::ReturnTy {
         write!(self, "{d} = {s1} & {s2}")
     }
 
-    fn or(&mut self, d: Reg, s1: Reg, s2: Reg) -> Self::ReturnTy {
+    fn or(&mut self, d: RawReg, s1: RawReg, s2: RawReg) -> Self::ReturnTy {
         write!(self, "{d} = {s1} | {s2}")
     }
 
-    fn add(&mut self, d: Reg, s1: Reg, s2: Reg) -> Self::ReturnTy {
+    fn add(&mut self, d: RawReg, s1: RawReg, s2: RawReg) -> Self::ReturnTy {
         write!(self, "{d} = {s1} + {s2}")
     }
 
-    fn sub(&mut self, d: Reg, s1: Reg, s2: Reg) -> Self::ReturnTy {
+    fn sub(&mut self, d: RawReg, s1: RawReg, s2: RawReg) -> Self::ReturnTy {
         write!(self, "{d} = {s1} - {s2}")
     }
 
-    fn mul(&mut self, d: Reg, s1: Reg, s2: Reg) -> Self::ReturnTy {
+    fn mul(&mut self, d: RawReg, s1: RawReg, s2: RawReg) -> Self::ReturnTy {
         write!(self, "{d} = {s1} * {s2}")
     }
 
-    fn mul_imm(&mut self, d: Reg, s1: Reg, s2: u32) -> Self::ReturnTy {
+    fn mul_imm(&mut self, d: RawReg, s1: RawReg, s2: u32) -> Self::ReturnTy {
         write!(self, "{d} = {s1} * {s2}")
     }
 
-    fn mul_upper_signed_signed(&mut self, d: Reg, s1: Reg, s2: Reg) -> Self::ReturnTy {
+    fn mul_upper_signed_signed(&mut self, d: RawReg, s1: RawReg, s2: RawReg) -> Self::ReturnTy {
         write!(self, "{d} = ({s1} as i64 * {s2} as i64) >> 32")
     }
 
-    fn mul_upper_signed_signed_imm(&mut self, d: Reg, s1: Reg, s2: u32) -> Self::ReturnTy {
+    fn mul_upper_signed_signed_imm(&mut self, d: RawReg, s1: RawReg, s2: u32) -> Self::ReturnTy {
         write!(self, "{d} = ({s1} as i64 * {s2} as i64) >> 32", s2 = s2 as i32)
     }
 
-    fn mul_upper_unsigned_unsigned(&mut self, d: Reg, s1: Reg, s2: Reg) -> Self::ReturnTy {
+    fn mul_upper_unsigned_unsigned(&mut self, d: RawReg, s1: RawReg, s2: RawReg) -> Self::ReturnTy {
         write!(self, "{d} = ({s1} as u64 * {s2} as u64) >> 32")
     }
 
-    fn mul_upper_unsigned_unsigned_imm(&mut self, d: Reg, s1: Reg, s2: u32) -> Self::ReturnTy {
+    fn mul_upper_unsigned_unsigned_imm(&mut self, d: RawReg, s1: RawReg, s2: u32) -> Self::ReturnTy {
         write!(self, "{d} = ({s1} as u64 * {s2} as u64) >> 32")
     }
 
-    fn mul_upper_signed_unsigned(&mut self, d: Reg, s1: Reg, s2: Reg) -> Self::ReturnTy {
+    fn mul_upper_signed_unsigned(&mut self, d: RawReg, s1: RawReg, s2: RawReg) -> Self::ReturnTy {
         write!(self, "{d} = ({s1} as i64 * {s2} as u64) >> 32")
     }
 
-    fn div_unsigned(&mut self, d: Reg, s1: Reg, s2: Reg) -> Self::ReturnTy {
+    fn div_unsigned(&mut self, d: RawReg, s1: RawReg, s2: RawReg) -> Self::ReturnTy {
         write!(self, "{d} = {s1} /u {s2}")
     }
 
-    fn div_signed(&mut self, d: Reg, s1: Reg, s2: Reg) -> Self::ReturnTy {
+    fn div_signed(&mut self, d: RawReg, s1: RawReg, s2: RawReg) -> Self::ReturnTy {
         write!(self, "{d} = {s1} /s {s2}")
     }
 
-    fn rem_unsigned(&mut self, d: Reg, s1: Reg, s2: Reg) -> Self::ReturnTy {
+    fn rem_unsigned(&mut self, d: RawReg, s1: RawReg, s2: RawReg) -> Self::ReturnTy {
         write!(self, "{d} = {s1} %u {s2}")
     }
 
-    fn rem_signed(&mut self, d: Reg, s1: Reg, s2: Reg) -> Self::ReturnTy {
+    fn rem_signed(&mut self, d: RawReg, s1: RawReg, s2: RawReg) -> Self::ReturnTy {
         write!(self, "{d} = {s1} %s {s2}")
     }
 
-    fn set_less_than_unsigned_imm(&mut self, d: Reg, s1: Reg, s2: u32) -> Self::ReturnTy {
+    fn set_less_than_unsigned_imm(&mut self, d: RawReg, s1: RawReg, s2: u32) -> Self::ReturnTy {
         write!(self, "{d} = {s1} <u 0x{s2:x}")
     }
 
-    fn set_greater_than_unsigned_imm(&mut self, d: Reg, s1: Reg, s2: u32) -> Self::ReturnTy {
+    fn set_greater_than_unsigned_imm(&mut self, d: RawReg, s1: RawReg, s2: u32) -> Self::ReturnTy {
         write!(self, "{d} = {s1} >u 0x{s2:x}")
     }
 
-    fn set_less_than_signed_imm(&mut self, d: Reg, s1: Reg, s2: u32) -> Self::ReturnTy {
+    fn set_less_than_signed_imm(&mut self, d: RawReg, s1: RawReg, s2: u32) -> Self::ReturnTy {
         write!(self, "{d} = {s1} <s {s2}", s2 = s2 as i32)
     }
 
-    fn set_greater_than_signed_imm(&mut self, d: Reg, s1: Reg, s2: u32) -> Self::ReturnTy {
+    fn set_greater_than_signed_imm(&mut self, d: RawReg, s1: RawReg, s2: u32) -> Self::ReturnTy {
         write!(self, "{d} = {s1} >s {s2}", s2 = s2 as i32)
     }
 
-    fn shift_logical_right_imm(&mut self, d: Reg, s1: Reg, s2: u32) -> Self::ReturnTy {
+    fn shift_logical_right_imm(&mut self, d: RawReg, s1: RawReg, s2: u32) -> Self::ReturnTy {
         write!(self, "{d} = {s1} >> {s2}")
     }
 
-    fn shift_logical_right_imm_alt(&mut self, d: Reg, s2: Reg, s1: u32) -> Self::ReturnTy {
+    fn shift_logical_right_imm_alt(&mut self, d: RawReg, s2: RawReg, s1: u32) -> Self::ReturnTy {
         write!(self, "{d} = {s1} >> {s2}")
     }
 
-    fn shift_arithmetic_right_imm(&mut self, d: Reg, s1: Reg, s2: u32) -> Self::ReturnTy {
+    fn shift_arithmetic_right_imm(&mut self, d: RawReg, s1: RawReg, s2: u32) -> Self::ReturnTy {
         write!(self, "{d} = {s1} >>a {s2}")
     }
 
-    fn shift_arithmetic_right_imm_alt(&mut self, d: Reg, s2: Reg, s1: u32) -> Self::ReturnTy {
+    fn shift_arithmetic_right_imm_alt(&mut self, d: RawReg, s2: RawReg, s1: u32) -> Self::ReturnTy {
         write!(self, "{d} = {s1} >>a {s2}")
     }
 
-    fn shift_logical_left_imm(&mut self, d: Reg, s1: Reg, s2: u32) -> Self::ReturnTy {
+    fn shift_logical_left_imm(&mut self, d: RawReg, s1: RawReg, s2: u32) -> Self::ReturnTy {
         write!(self, "{d} = {s1} << {s2}")
     }
 
-    fn shift_logical_left_imm_alt(&mut self, d: Reg, s2: Reg, s1: u32) -> Self::ReturnTy {
+    fn shift_logical_left_imm_alt(&mut self, d: RawReg, s2: RawReg, s1: u32) -> Self::ReturnTy {
         write!(self, "{d} = {s1} << {s2}")
     }
 
-    fn or_imm(&mut self, d: Reg, s1: Reg, s2: u32) -> Self::ReturnTy {
+    fn or_imm(&mut self, d: RawReg, s1: RawReg, s2: u32) -> Self::ReturnTy {
         write!(self, "{d} = {s1} | 0x{s2:x}")
     }
 
-    fn and_imm(&mut self, d: Reg, s1: Reg, s2: u32) -> Self::ReturnTy {
+    fn and_imm(&mut self, d: RawReg, s1: RawReg, s2: u32) -> Self::ReturnTy {
         write!(self, "{d} = {s1} & 0x{s2:x}")
     }
 
-    fn xor_imm(&mut self, d: Reg, s1: Reg, s2: u32) -> Self::ReturnTy {
+    fn xor_imm(&mut self, d: RawReg, s1: RawReg, s2: u32) -> Self::ReturnTy {
         write!(self, "{d} = {s1} ^ 0x{s2:x}")
     }
 
-    fn load_imm(&mut self, d: Reg, a: u32) -> Self::ReturnTy {
+    fn load_imm(&mut self, d: RawReg, a: u32) -> Self::ReturnTy {
         write!(self, "{d} = 0x{a:x}")
     }
 
-    fn move_reg(&mut self, d: Reg, s: Reg) -> Self::ReturnTy {
+    fn move_reg(&mut self, d: RawReg, s: RawReg) -> Self::ReturnTy {
         write!(self, "{d} = {s}")
     }
 
-    fn cmov_if_zero(&mut self, d: Reg, s: Reg, c: Reg) -> Self::ReturnTy {
+    fn cmov_if_zero(&mut self, d: RawReg, s: RawReg, c: RawReg) -> Self::ReturnTy {
         write!(self, "{d} = {s} if {c} == 0")
     }
 
-    fn cmov_if_not_zero(&mut self, d: Reg, s: Reg, c: Reg) -> Self::ReturnTy {
+    fn cmov_if_not_zero(&mut self, d: RawReg, s: RawReg, c: RawReg) -> Self::ReturnTy {
         write!(self, "{d} = {s} if {c} != 0")
     }
 
-    fn cmov_if_zero_imm(&mut self, d: Reg, c: Reg, s: u32) -> Self::ReturnTy {
+    fn cmov_if_zero_imm(&mut self, d: RawReg, c: RawReg, s: u32) -> Self::ReturnTy {
         write!(self, "{d} = {s} if {c} == 0")
     }
 
-    fn cmov_if_not_zero_imm(&mut self, d: Reg, c: Reg, s: u32) -> Self::ReturnTy {
+    fn cmov_if_not_zero_imm(&mut self, d: RawReg, c: RawReg, s: u32) -> Self::ReturnTy {
         write!(self, "{d} = {s} if {c} != 0")
     }
 
-    fn add_imm(&mut self, d: Reg, s1: Reg, s2: u32) -> Self::ReturnTy {
+    fn add_imm(&mut self, d: RawReg, s1: RawReg, s2: u32) -> Self::ReturnTy {
         if (s2 as i32) < 0 && (s2 as i32) > -4096 {
             write!(self, "{d} = {s1} - {s2}", s2 = -(s2 as i32))
         } else {
@@ -1367,7 +1463,7 @@ impl<'a> InstructionVisitor for core::fmt::Formatter<'a> {
         }
     }
 
-    fn negate_and_add_imm(&mut self, d: Reg, s1: Reg, s2: u32) -> Self::ReturnTy {
+    fn negate_and_add_imm(&mut self, d: RawReg, s1: RawReg, s2: u32) -> Self::ReturnTy {
         if s2 == 0 {
             write!(self, "{d} = -{s1}")
         } else {
@@ -1375,19 +1471,19 @@ impl<'a> InstructionVisitor for core::fmt::Formatter<'a> {
         }
     }
 
-    fn store_imm_indirect_u8(&mut self, base: Reg, offset: u32, value: u32) -> Self::ReturnTy {
+    fn store_imm_indirect_u8(&mut self, base: RawReg, offset: u32, value: u32) -> Self::ReturnTy {
         write!(self, "u8 [{base} + {offset}] = {value}")
     }
 
-    fn store_imm_indirect_u16(&mut self, base: Reg, offset: u32, value: u32) -> Self::ReturnTy {
+    fn store_imm_indirect_u16(&mut self, base: RawReg, offset: u32, value: u32) -> Self::ReturnTy {
         write!(self, "u16 [{base} + {offset}] = {value}")
     }
 
-    fn store_imm_indirect_u32(&mut self, base: Reg, offset: u32, value: u32) -> Self::ReturnTy {
+    fn store_imm_indirect_u32(&mut self, base: RawReg, offset: u32, value: u32) -> Self::ReturnTy {
         write!(self, "u32 [{base} + {offset}] = {value}")
     }
 
-    fn store_indirect_u8(&mut self, src: Reg, base: Reg, offset: u32) -> Self::ReturnTy {
+    fn store_indirect_u8(&mut self, src: RawReg, base: RawReg, offset: u32) -> Self::ReturnTy {
         if offset != 0 {
             write!(self, "u8 [{base} + {offset}] = {src}")
         } else {
@@ -1395,7 +1491,7 @@ impl<'a> InstructionVisitor for core::fmt::Formatter<'a> {
         }
     }
 
-    fn store_indirect_u16(&mut self, src: Reg, base: Reg, offset: u32) -> Self::ReturnTy {
+    fn store_indirect_u16(&mut self, src: RawReg, base: RawReg, offset: u32) -> Self::ReturnTy {
         if offset != 0 {
             write!(self, "u16 [{base} + {offset}] = {src}")
         } else {
@@ -1403,7 +1499,7 @@ impl<'a> InstructionVisitor for core::fmt::Formatter<'a> {
         }
     }
 
-    fn store_indirect_u32(&mut self, src: Reg, base: Reg, offset: u32) -> Self::ReturnTy {
+    fn store_indirect_u32(&mut self, src: RawReg, base: RawReg, offset: u32) -> Self::ReturnTy {
         if offset != 0 {
             write!(self, "u32 [{base} + {offset}] = {src}")
         } else {
@@ -1411,31 +1507,31 @@ impl<'a> InstructionVisitor for core::fmt::Formatter<'a> {
         }
     }
 
-    fn store_imm_u8(&mut self, value: u32, offset: u32) -> Self::ReturnTy {
+    fn store_imm_u8(&mut self, offset: u32, value: u32) -> Self::ReturnTy {
         write!(self, "u8 [0x{offset:x}] = {value}")
     }
 
-    fn store_imm_u16(&mut self, value: u32, offset: u32) -> Self::ReturnTy {
+    fn store_imm_u16(&mut self, offset: u32, value: u32) -> Self::ReturnTy {
         write!(self, "u16 [0x{offset:x}] = {value}")
     }
 
-    fn store_imm_u32(&mut self, value: u32, offset: u32) -> Self::ReturnTy {
+    fn store_imm_u32(&mut self, offset: u32, value: u32) -> Self::ReturnTy {
         write!(self, "u32 [0x{offset:x}] = {value}")
     }
 
-    fn store_u8(&mut self, src: Reg, offset: u32) -> Self::ReturnTy {
+    fn store_u8(&mut self, src: RawReg, offset: u32) -> Self::ReturnTy {
         write!(self, "u8 [0x{offset:x}] = {src}")
     }
 
-    fn store_u16(&mut self, src: Reg, offset: u32) -> Self::ReturnTy {
+    fn store_u16(&mut self, src: RawReg, offset: u32) -> Self::ReturnTy {
         write!(self, "u16 [0x{offset:x}] = {src}")
     }
 
-    fn store_u32(&mut self, src: Reg, offset: u32) -> Self::ReturnTy {
+    fn store_u32(&mut self, src: RawReg, offset: u32) -> Self::ReturnTy {
         write!(self, "u32 [0x{offset:x}] = {src}")
     }
 
-    fn load_indirect_u8(&mut self, dst: Reg, base: Reg, offset: u32) -> Self::ReturnTy {
+    fn load_indirect_u8(&mut self, dst: RawReg, base: RawReg, offset: u32) -> Self::ReturnTy {
         if offset != 0 {
             write!(self, "{} = u8 [{} + {}]", dst, base, offset)
         } else {
@@ -1443,7 +1539,7 @@ impl<'a> InstructionVisitor for core::fmt::Formatter<'a> {
         }
     }
 
-    fn load_indirect_i8(&mut self, dst: Reg, base: Reg, offset: u32) -> Self::ReturnTy {
+    fn load_indirect_i8(&mut self, dst: RawReg, base: RawReg, offset: u32) -> Self::ReturnTy {
         if offset != 0 {
             write!(self, "{} = i8 [{} + {}]", dst, base, offset)
         } else {
@@ -1451,7 +1547,7 @@ impl<'a> InstructionVisitor for core::fmt::Formatter<'a> {
         }
     }
 
-    fn load_indirect_u16(&mut self, dst: Reg, base: Reg, offset: u32) -> Self::ReturnTy {
+    fn load_indirect_u16(&mut self, dst: RawReg, base: RawReg, offset: u32) -> Self::ReturnTy {
         if offset != 0 {
             write!(self, "{} = u16 [{} + {}]", dst, base, offset)
         } else {
@@ -1459,7 +1555,7 @@ impl<'a> InstructionVisitor for core::fmt::Formatter<'a> {
         }
     }
 
-    fn load_indirect_i16(&mut self, dst: Reg, base: Reg, offset: u32) -> Self::ReturnTy {
+    fn load_indirect_i16(&mut self, dst: RawReg, base: RawReg, offset: u32) -> Self::ReturnTy {
         if offset != 0 {
             write!(self, "{} = i16 [{} + {}]", dst, base, offset)
         } else {
@@ -1467,7 +1563,7 @@ impl<'a> InstructionVisitor for core::fmt::Formatter<'a> {
         }
     }
 
-    fn load_indirect_u32(&mut self, dst: Reg, base: Reg, offset: u32) -> Self::ReturnTy {
+    fn load_indirect_u32(&mut self, dst: RawReg, base: RawReg, offset: u32) -> Self::ReturnTy {
         if offset != 0 {
             write!(self, "{} = u32 [{} + {}]", dst, base, offset)
         } else {
@@ -1475,87 +1571,87 @@ impl<'a> InstructionVisitor for core::fmt::Formatter<'a> {
         }
     }
 
-    fn load_u8(&mut self, dst: Reg, offset: u32) -> Self::ReturnTy {
+    fn load_u8(&mut self, dst: RawReg, offset: u32) -> Self::ReturnTy {
         write!(self, "{} = u8 [0x{:x}]", dst, offset)
     }
 
-    fn load_i8(&mut self, dst: Reg, offset: u32) -> Self::ReturnTy {
+    fn load_i8(&mut self, dst: RawReg, offset: u32) -> Self::ReturnTy {
         write!(self, "{} = i8 [0x{:x}]", dst, offset)
     }
 
-    fn load_u16(&mut self, dst: Reg, offset: u32) -> Self::ReturnTy {
+    fn load_u16(&mut self, dst: RawReg, offset: u32) -> Self::ReturnTy {
         write!(self, "{} = u16 [0x{:x}]", dst, offset)
     }
 
-    fn load_i16(&mut self, dst: Reg, offset: u32) -> Self::ReturnTy {
+    fn load_i16(&mut self, dst: RawReg, offset: u32) -> Self::ReturnTy {
         write!(self, "{} = i16 [0x{:x}]", dst, offset)
     }
 
-    fn load_u32(&mut self, dst: Reg, offset: u32) -> Self::ReturnTy {
+    fn load_u32(&mut self, dst: RawReg, offset: u32) -> Self::ReturnTy {
         write!(self, "{} = u32 [0x{:x}]", dst, offset)
     }
 
-    fn branch_less_unsigned(&mut self, s1: Reg, s2: Reg, imm: u32) -> Self::ReturnTy {
+    fn branch_less_unsigned(&mut self, s1: RawReg, s2: RawReg, imm: u32) -> Self::ReturnTy {
         write!(self, "jump {} if {} <u {}", imm, s1, s2)
     }
 
-    fn branch_less_signed(&mut self, s1: Reg, s2: Reg, imm: u32) -> Self::ReturnTy {
+    fn branch_less_signed(&mut self, s1: RawReg, s2: RawReg, imm: u32) -> Self::ReturnTy {
         write!(self, "jump {} if {} <s {}", imm, s1, s2)
     }
 
-    fn branch_less_unsigned_imm(&mut self, s1: Reg, s2: u32, imm: u32) -> Self::ReturnTy {
+    fn branch_less_unsigned_imm(&mut self, s1: RawReg, s2: u32, imm: u32) -> Self::ReturnTy {
         write!(self, "jump {} if {} <u {}", imm, s1, s2)
     }
 
-    fn branch_less_signed_imm(&mut self, s1: Reg, s2: u32, imm: u32) -> Self::ReturnTy {
+    fn branch_less_signed_imm(&mut self, s1: RawReg, s2: u32, imm: u32) -> Self::ReturnTy {
         write!(self, "jump {} if {} <s {}", imm, s1, s2)
     }
 
-    fn branch_greater_or_equal_unsigned(&mut self, s1: Reg, s2: Reg, imm: u32) -> Self::ReturnTy {
+    fn branch_greater_or_equal_unsigned(&mut self, s1: RawReg, s2: RawReg, imm: u32) -> Self::ReturnTy {
         write!(self, "jump {} if {} >=u {}", imm, s1, s2)
     }
 
-    fn branch_greater_or_equal_signed(&mut self, s1: Reg, s2: Reg, imm: u32) -> Self::ReturnTy {
+    fn branch_greater_or_equal_signed(&mut self, s1: RawReg, s2: RawReg, imm: u32) -> Self::ReturnTy {
         write!(self, "jump {} if {} >=s {}", imm, s1, s2)
     }
 
-    fn branch_greater_or_equal_unsigned_imm(&mut self, s1: Reg, s2: u32, imm: u32) -> Self::ReturnTy {
+    fn branch_greater_or_equal_unsigned_imm(&mut self, s1: RawReg, s2: u32, imm: u32) -> Self::ReturnTy {
         write!(self, "jump {} if {} >=u {}", imm, s1, s2)
     }
 
-    fn branch_greater_or_equal_signed_imm(&mut self, s1: Reg, s2: u32, imm: u32) -> Self::ReturnTy {
+    fn branch_greater_or_equal_signed_imm(&mut self, s1: RawReg, s2: u32, imm: u32) -> Self::ReturnTy {
         write!(self, "jump {} if {} >=s {}", imm, s1, s2)
     }
 
-    fn branch_eq(&mut self, s1: Reg, s2: Reg, imm: u32) -> Self::ReturnTy {
+    fn branch_eq(&mut self, s1: RawReg, s2: RawReg, imm: u32) -> Self::ReturnTy {
         write!(self, "jump {} if {} == {}", imm, s1, s2)
     }
 
-    fn branch_not_eq(&mut self, s1: Reg, s2: Reg, imm: u32) -> Self::ReturnTy {
+    fn branch_not_eq(&mut self, s1: RawReg, s2: RawReg, imm: u32) -> Self::ReturnTy {
         write!(self, "jump {} if {} != {}", imm, s1, s2)
     }
 
-    fn branch_eq_imm(&mut self, s1: Reg, s2: u32, imm: u32) -> Self::ReturnTy {
+    fn branch_eq_imm(&mut self, s1: RawReg, s2: u32, imm: u32) -> Self::ReturnTy {
         write!(self, "jump {} if {} == {}", imm, s1, s2)
     }
 
-    fn branch_not_eq_imm(&mut self, s1: Reg, s2: u32, imm: u32) -> Self::ReturnTy {
+    fn branch_not_eq_imm(&mut self, s1: RawReg, s2: u32, imm: u32) -> Self::ReturnTy {
         write!(self, "jump {} if {} != {}", imm, s1, s2)
     }
 
-    fn branch_less_or_equal_unsigned_imm(&mut self, s1: Reg, s2: u32, imm: u32) -> Self::ReturnTy {
+    fn branch_less_or_equal_unsigned_imm(&mut self, s1: RawReg, s2: u32, imm: u32) -> Self::ReturnTy {
         write!(self, "jump {} if {} <=u {}", imm, s1, s2)
     }
 
-    fn branch_less_or_equal_signed_imm(&mut self, s1: Reg, s2: u32, imm: u32) -> Self::ReturnTy {
+    fn branch_less_or_equal_signed_imm(&mut self, s1: RawReg, s2: u32, imm: u32) -> Self::ReturnTy {
         write!(self, "jump {} if {} <=s {}", imm, s1, s2)
     }
 
-    fn branch_greater_unsigned_imm(&mut self, s1: Reg, s2: u32, imm: u32) -> Self::ReturnTy {
+    fn branch_greater_unsigned_imm(&mut self, s1: RawReg, s2: u32, imm: u32) -> Self::ReturnTy {
         write!(self, "jump {} if {} >u {}", imm, s1, s2)
     }
 
-    fn branch_greater_signed_imm(&mut self, s1: Reg, s2: u32, imm: u32) -> Self::ReturnTy {
+    fn branch_greater_signed_imm(&mut self, s1: RawReg, s2: u32, imm: u32) -> Self::ReturnTy {
         write!(self, "jump {} if {} >s {}", imm, s1, s2)
     }
 
@@ -1563,20 +1659,19 @@ impl<'a> InstructionVisitor for core::fmt::Formatter<'a> {
         write!(self, "jump {}", target)
     }
 
-    fn load_imm_and_jump(&mut self, ra: Reg, value: u32, target: u32) -> Self::ReturnTy {
+    fn load_imm_and_jump(&mut self, ra: RawReg, value: u32, target: u32) -> Self::ReturnTy {
         write!(self, "{ra} = {value}, jump {target}")
     }
 
-    fn jump_indirect(&mut self, base: Reg, offset: u32) -> Self::ReturnTy {
-        use Reg::*;
+    fn jump_indirect(&mut self, base: RawReg, offset: u32) -> Self::ReturnTy {
         match (base, offset) {
-            (RA, 0) => write!(self, "ret"),
+            (_, 0) if base == Reg::RA.into() => write!(self, "ret"),
             (_, 0) => write!(self, "jump [{}]", base),
             (_, _) => write!(self, "jump [{} + {}]", base, offset),
         }
     }
 
-    fn load_imm_and_jump_indirect(&mut self, ra: Reg, base: Reg, value: u32, offset: u32) -> Self::ReturnTy {
+    fn load_imm_and_jump_indirect(&mut self, ra: RawReg, base: RawReg, value: u32, offset: u32) -> Self::ReturnTy {
         if offset == 0 {
             write!(self, "jump [{base}], {ra} = {value}")
         } else {
@@ -1584,8 +1679,8 @@ impl<'a> InstructionVisitor for core::fmt::Formatter<'a> {
         }
     }
 
-    fn invalid(&mut self, opcode: u8) -> Self::ReturnTy {
-        write!(self, "invalid 0x{opcode:02x}")
+    fn invalid(&mut self) -> Self::ReturnTy {
+        write!(self, "invalid")
     }
 }
 
@@ -2033,12 +2128,12 @@ impl<'a> Iterator for JumpTableIter<'a> {
 
 #[inline(never)]
 #[cold]
-fn parse_bitmask_slow(bitmask: &[u8], code_offset: &mut usize) -> Option<usize> {
+fn parse_bitmask_slow(bitmask: &[u8], mut offset: usize) -> Option<(usize, usize)> {
     if bitmask.is_empty() {
         return None;
     }
 
-    let mut offset = *code_offset + 1;
+    offset += 1;
     let mut args_length = 0;
     while let Some(&byte) = bitmask.get(offset >> 3) {
         let shift = offset & 7;
@@ -2059,47 +2154,54 @@ fn parse_bitmask_slow(bitmask: &[u8], code_offset: &mut usize) -> Option<usize>
         offset += length;
     }
 
-    *code_offset = offset;
-    Some(args_length)
+    Some((offset, args_length))
 }
 
 #[cfg_attr(not(debug_assertions), inline(always))]
-pub(crate) fn parse_bitmask(bitmask: &[u8], code_offset: &mut usize) -> Option<usize> {
-    let mut offset = *code_offset + 1;
+pub(crate) fn parse_bitmask_fast(bitmask: &[u8], mut offset: usize) -> Option<(usize, usize)> {
+    offset += 1;
     let Some(bitmask) = bitmask.get(offset >> 3..(offset >> 3) + 4) else {
-        return parse_bitmask_slow(bitmask, code_offset);
+        return None;
     };
 
     let shift = offset & 7;
     let mask = u32::from_le_bytes([bitmask[0], bitmask[1], bitmask[2], bitmask[3]]) >> shift;
-
     if mask == 0 {
-        return parse_bitmask_slow(bitmask, code_offset);
+        return None;
     }
 
     let args_length = mask.trailing_zeros() as usize;
     offset += args_length;
-    *code_offset = offset;
 
-    Some(args_length)
+    Some((offset, args_length))
 }
 
 #[test]
 fn test_parse_bitmask() {
-    fn p(bitmask: &[u8]) -> (Option<usize>, usize) {
-        let mut offset = 0;
-        (parse_bitmask(bitmask, &mut offset), offset)
+    fn p(bitmask: &[u8]) -> Option<(usize, usize)> {
+        let result_fast = parse_bitmask_fast(bitmask, 0);
+        let result_slow = parse_bitmask_slow(bitmask, 0);
+        if result_fast.is_some() {
+            assert_eq!(result_fast, result_slow);
+        }
+
+        result_slow
     }
 
-    assert_eq!(p(&[0b00000001, 0, 0, 0]), (Some(31), 32));
-    assert_eq!(p(&[0b00000001, 0, 0]), (Some(23), 24));
-    assert_eq!(p(&[0b00000001, 0]), (Some(15), 16));
-    assert_eq!(p(&[0b00000001]), (Some(7), 8));
-    assert_eq!(p(&[0b00000011]), (Some(0), 1));
-    assert_eq!(p(&[0b00000011, 0]), (Some(0), 1));
-    assert_eq!(p(&[0b00000101]), (Some(1), 2));
-    assert_eq!(p(&[0b10000001]), (Some(6), 7));
-    assert_eq!(p(&[0b00000001, 1]), (Some(7), 8));
+    assert_eq!(p(&[0b00000001, 0, 0, 0]), Some((32, 31)));
+    assert_eq!(p(&[0b00000001, 0, 0]), Some((24, 23)));
+    assert_eq!(p(&[0b00000001, 0]), Some((16, 15)));
+    assert_eq!(p(&[0b00000001]), Some((8, 7)));
+    assert_eq!(p(&[0b00000011]), Some((1, 0)));
+    assert_eq!(p(&[0b00000011, 0]), Some((1, 0)));
+    assert_eq!(p(&[0b00000101]), Some((2, 1)));
+    assert_eq!(p(&[0b10000001]), Some((7, 6)));
+    assert_eq!(p(&[0b00000001, 1]), Some((8, 7)));
+
+    assert_eq!(parse_bitmask_fast(&[1, 0, 1, 0, 0, 0, 0, 0], 0), Some((16, 15)));
+    assert_eq!(parse_bitmask_fast(&[1, 0, 0, 1, 0, 0, 0, 0], 0), Some((24, 23)));
+    assert_eq!(parse_bitmask_fast(&[1, 0, 0, 0b10000000, 0, 0, 0, 0], 0), Some((31, 30)));
+    assert_eq!(parse_bitmask_fast(&[1, 0, 0, 0, 1, 0, 0, 0], 0), None);
 }
 
 #[derive(Clone)]
@@ -2109,7 +2211,7 @@ pub struct Instructions<'a> {
     offset: usize,
 }
 
-#[derive(Copy, Clone, Debug)]
+#[derive(Copy, Clone, PartialEq, Eq, Debug)]
 pub struct ParsedInstruction {
     pub kind: Instruction,
     pub offset: u32,
diff --git a/crates/polkavm-common/src/varint.rs b/crates/polkavm-common/src/varint.rs
index e3f24002..ba12127f 100644
--- a/crates/polkavm-common/src/varint.rs
+++ b/crates/polkavm-common/src/varint.rs
@@ -7,20 +7,6 @@ fn get_varint_length(leading_zeros: u32) -> u32 {
 
 pub const MAX_VARINT_LENGTH: usize = 5;
 
-#[inline(always)]
-pub(crate) fn read_varint_fast(chunk: u64) -> Option<(u32, u32)> {
-    let first_byte = chunk as u8;
-    let length = (!first_byte).leading_zeros();
-    if length > 4 {
-        return None;
-    }
-
-    let upper_mask = 0b11111111_u32 >> length;
-    let lower_mask = ((1_u64 << (length << 3)) - 1) as u32;
-    let value = (upper_mask & u32::from(first_byte)).wrapping_shl(length * 8) | (((chunk >> 8) as u32) & lower_mask);
-    Some((length + 1, value))
-}
-
 #[inline]
 pub(crate) fn read_varint(input: &[u8], first_byte: u8) -> Option<(usize, u32)> {
     let length = (!first_byte).leading_zeros();
@@ -90,24 +76,71 @@ proptest::proptest! {
 }
 
 #[inline(always)]
-pub(crate) fn read_simple_varint_fast(chunk: u32, length: u32) -> u32 {
-    let mask = ((1_u64 << (length << 3)) - 1) as u32;
-    let value = chunk & mask;
+fn apply_zigzag(value: u32) -> u32 {
+    let value = value as i32;
+    let value = (value << 1) ^ (value >> 31);
+    value as u32
+}
+
+#[inline(always)]
+fn undo_zigzag(value: u32) -> u32 {
     (value >> 1) ^ (-((value & 1) as i32)) as u32
 }
 
+const fn bit_length_to_mask_u32_slow(length: u32) -> u32 {
+    if length >= 32 {
+        0xffffffff
+    } else {
+        (1_u32 << length) - 1
+    }
+}
+
+static LENGTH_MINUS_1_TO_MASK: [u32; 256] = {
+    let mut output = [0; 256];
+    let mut length = 0_u32;
+    while length < 256 {
+        let length_effective = length.wrapping_sub(1);
+        let length_effective = if length_effective < 4 { length_effective } else { 4 };
+        output[length as usize] = bit_length_to_mask_u32_slow(length_effective * 8);
+        length += 1;
+    }
+    output
+};
+
+#[inline(always)]
+pub(crate) fn read_simple_varint_length_minus_1(chunk: u32, length: u32) -> u32 {
+    let mask = LENGTH_MINUS_1_TO_MASK[length as u8 as usize];
+    undo_zigzag(chunk & mask)
+}
+
+static LENGTH_TO_MASK: [u32; 256] = {
+    let mut output = [0; 256];
+    let mut length = 0_u32;
+    while length < 256 {
+        let length_effective = if length < 4 { length } else { 4 };
+        output[length as usize] = bit_length_to_mask_u32_slow(length_effective * 8);
+        length += 1;
+    }
+    output
+};
+
+#[inline(always)]
+pub(crate) fn read_simple_varint(chunk: u32, length: u32) -> u32 {
+    let mask = LENGTH_TO_MASK[length as u8 as usize];
+    undo_zigzag(chunk & mask)
+}
+
 #[inline]
-fn get_simple_varint_length(leading_zeros: u32) -> u32 {
+fn get_bytes_required(value: u32) -> u32 {
+    let leading_zeros = value.leading_zeros();
     let bits_required = 32 - leading_zeros;
     bits_required / 8 + (if leading_zeros % 8 > 0 { 1 } else { 0 })
 }
 
 #[inline]
 pub(crate) fn write_simple_varint(value: u32, buffer: &mut [u8]) -> usize {
-    let value = value as i32;
-    let value = (value << 1) ^ (value >> 31);
-    let value = value as u32;
-    let varint_length = get_simple_varint_length(value.leading_zeros());
+    let value = apply_zigzag(value);
+    let varint_length = get_bytes_required(value);
     match varint_length {
         0 => {}
         1 => {
@@ -142,7 +175,7 @@ proptest::proptest! {
     #[allow(clippy::ignored_unit_patterns)]
     #[test]
     fn simple_varint_serialization(value in 0u32..=0xffffffff) {
-        fn read_simple_varint(input: &[u8]) -> Option<u32> {
+        fn read_simple_varint_slow(input: &[u8]) -> Option<u32> {
             let value = match input.len() {
                 0 => 0,
                 1 => u32::from(input[0]),
@@ -156,9 +189,16 @@ proptest::proptest! {
             Some(value)
         }
 
-        let mut buffer = [0; MAX_VARINT_LENGTH];
-        let length = write_simple_varint(value, &mut buffer);
-        let parsed_value = read_simple_varint(&buffer[..length]).unwrap();
-        assert_eq!(parsed_value, value, "value mismatch");
+        let mut t = [0; 4];
+        let length = write_simple_varint(value, &mut t);
+        assert_eq!(read_simple_varint_slow(&t[..length]).unwrap(), value, "value mismatch");
+
+        let chunk = u32::from_le_bytes([t[0], t[1], t[2], t[3]]);
+        assert_eq!(read_simple_varint(chunk, length as u32), value);
+        assert_eq!(read_simple_varint(chunk, length as u32 + 1), value);
+        assert_eq!(read_simple_varint(chunk, length as u32 + 2), value);
+        assert_eq!(read_simple_varint(chunk, length as u32 + 3), value);
+        assert_eq!(read_simple_varint(chunk, length as u32 + 4), value);
+        assert_eq!(read_simple_varint(chunk, 0xffffffff), value);
     }
 }
diff --git a/crates/polkavm-common/src/writer.rs b/crates/polkavm-common/src/writer.rs
index 09701a46..e4a6e2f5 100644
--- a/crates/polkavm-common/src/writer.rs
+++ b/crates/polkavm-common/src/writer.rs
@@ -271,7 +271,7 @@ impl ProgramBlobBuilder {
 
             assert_eq!(parsed.len(), instructions.len());
             for ((offset, mut instruction), entry) in parsed.into_iter().zip(instructions.into_iter()) {
-                assert_eq!(instruction, entry.instruction);
+                assert_eq!(instruction, entry.instruction, "broken serialization: {:?}", entry.bytes.bytes);
                 assert_eq!(entry.position, offset);
                 if let Some(target) = instruction.target_mut() {
                     assert!(offsets.contains(target));
diff --git a/crates/polkavm-disassembler/src/lib.rs b/crates/polkavm-disassembler/src/lib.rs
index b4f6841c..c7f36d16 100644
--- a/crates/polkavm-disassembler/src/lib.rs
+++ b/crates/polkavm-disassembler/src/lib.rs
@@ -426,7 +426,7 @@ mod tests {
         builder.add_import("hostcall".into());
         builder.set_code(
             &[
-                asm::store_imm_u32(0x12345678, memory_map.rw_data_address()),
+                asm::store_imm_u32(memory_map.rw_data_address(), 0x12345678),
                 asm::add(S0, A0, A1),
                 asm::ecalli(0),
                 asm::add(A0, A0, S0),
diff --git a/crates/polkavm-linker/src/program_from_elf.rs b/crates/polkavm-linker/src/program_from_elf.rs
index ea75fb50..077af93c 100644
--- a/crates/polkavm-linker/src/program_from_elf.rs
+++ b/crates/polkavm-linker/src/program_from_elf.rs
@@ -5104,7 +5104,7 @@ fn emit_code(
     is_optimized: bool,
 ) -> Result<Vec<(SourceStack, Instruction)>, ProgramFromElfError> {
     use polkavm_common::program::Reg as PReg;
-    fn conv_reg(reg: Reg) -> PReg {
+    fn conv_reg(reg: Reg) -> polkavm_common::program::RawReg {
         match reg {
             Reg::RA => PReg::RA,
             Reg::SP => PReg::SP,
@@ -5123,6 +5123,7 @@ fn emit_code(
                 unreachable!("internal error: temporary register was not spilled into memory");
             }
         }
+        .into()
     }
 
     let can_fallthrough_to_next_block = calculate_whether_can_fallthrough(all_blocks, used_blocks);
@@ -5216,7 +5217,7 @@ fn emit_code(
                         }
                         RegImm::Imm(value) => {
                             codegen! {
-                                args = (value, target),
+                                args = (target, value),
                                 kind = kind,
                                 {
                                     StoreKind::U32 => store_imm_u32,
diff --git a/crates/polkavm/src/compiler.rs b/crates/polkavm/src/compiler.rs
index a1299eda..34f0fcb8 100644
--- a/crates/polkavm/src/compiler.rs
+++ b/crates/polkavm/src/compiler.rs
@@ -4,7 +4,7 @@ use std::sync::Mutex;
 use core::marker::PhantomData;
 
 use polkavm_assembler::{Assembler, Label};
-use polkavm_common::program::{ParsedInstruction, ProgramExport, Instructions, JumpTable, Reg};
+use polkavm_common::program::{ParsedInstruction, ProgramExport, Instructions, JumpTable, RawReg, InstructionVisitor};
 use polkavm_common::zygote::{
     VM_COMPILER_MAXIMUM_EPILOGUE_LENGTH, VM_COMPILER_MAXIMUM_INSTRUCTION_LENGTH,
 };
@@ -33,10 +33,6 @@ struct Cache {
 pub(crate) struct CompilerCache(Mutex<Cache>);
 
 pub(crate) struct CompilerVisitor<'a, S> where S: Sandbox {
-    current_code_offset: u32,
-    next_code_offset: u32,
-    is_last_instruction: bool,
-
     init: GuestInit<'a>,
     jump_table: JumpTable<'a>,
     code: &'a [u8],
@@ -133,9 +129,6 @@ impl<'a, S> CompilerVisitor<'a, S> where S: Sandbox {
         asm.set_origin(native_code_address);
 
         let mut visitor = CompilerVisitor {
-            current_code_offset: 0,
-            next_code_offset: 0,
-            is_last_instruction: false,
             gas_visitor: GasVisitor::default(),
             asm,
             exports,
@@ -167,7 +160,7 @@ impl<'a, S> CompilerVisitor<'a, S> where S: Sandbox {
             ArchVisitor(&mut visitor).emit_trace_trampoline();
         }
 
-        visitor.start_new_basic_block::<true>();
+        visitor.force_start_new_basic_block(0);
         Ok((visitor, address_space))
     }
 
@@ -272,58 +265,61 @@ impl<'a, S> CompilerVisitor<'a, S> where S: Sandbox {
     }
 
     #[inline(always)]
-    fn start_new_basic_block<const IS_INITIAL: bool>(&mut self) {
-        if !IS_INITIAL && self.gas_metering.is_some() {
+    fn start_new_basic_block(&mut self, code_offset: u32, args_length: u32) {
+        if self.gas_metering.is_some() {
             let cost = self.gas_visitor.take_block_cost().unwrap();
             self.gas_cost_for_basic_block.push(cost);
         }
 
-        if !self.is_last_instruction {
-            let code_offset = self.next_code_offset;
+        let next_code_offset = code_offset + args_length + 1;
+        let is_last_instruction = next_code_offset as usize >= self.code.len();
+        if !is_last_instruction {
+            self.force_start_new_basic_block(next_code_offset);
+        }
+    }
 
-            log::trace!("Starting new basic block at: {code_offset}");
-            if let Some(label) = self.code_offset_to_label.get(code_offset) {
-                log::trace!("Label: {label} -> {code_offset} -> {:08x}", self.asm.current_address());
-                self.asm.define_label(label);
-            } else {
-                let label = self.asm.create_label();
-                log::trace!("Label: {label} -> {code_offset} -> {:08x}", self.asm.current_address());
-                self.code_offset_to_label.insert(code_offset, label);
-            }
+    #[inline(always)]
+    fn force_start_new_basic_block(&mut self, next_code_offset: u32) {
+        log::trace!("Starting new basic block at: {next_code_offset}");
+        if let Some(label) = self.code_offset_to_label.get(next_code_offset) {
+            log::trace!("Label: {label} -> {next_code_offset} -> {:08x}", self.asm.current_address());
+            self.asm.define_label(label);
+        } else {
+            let label = self.asm.create_label();
+            log::trace!("Label: {label} -> {next_code_offset} -> {:08x}", self.asm.current_address());
+            self.code_offset_to_label.insert(next_code_offset, label);
+        }
 
-            if let Some(gas_metering) = self.gas_metering {
-                self.gas_metering_stub_offsets.push(self.asm.len());
-                ArchVisitor(self).emit_gas_metering_stub(gas_metering);
-            }
+        if let Some(gas_metering) = self.gas_metering {
+            self.gas_metering_stub_offsets.push(self.asm.len());
+            ArchVisitor(self).emit_gas_metering_stub(gas_metering);
         }
     }
 
     #[inline(always)]
-    fn before_instruction(&mut self) {
-        self.code_offset_to_native_code_offset.push((self.current_code_offset, self.asm.len() as u32));
+    fn before_instruction(&mut self, code_offset: u32) {
+        self.code_offset_to_native_code_offset.push((code_offset, self.asm.len() as u32));
 
         if log::log_enabled!(log::Level::Trace) {
-            self.trace_compiled_instruction();
+            self.trace_compiled_instruction(code_offset);
         }
 
         if self.debug_trace_execution {
-            ArchVisitor(self).trace_execution();
+            ArchVisitor(self).trace_execution(code_offset);
         }
-
-        self.asm.reserve::<8>();
     }
 
-    fn after_instruction(&mut self) {
+    fn after_instruction(&mut self, code_offset: u32) {
         if cfg!(debug_assertions) && !self.debug_trace_execution {
             let offset = self.code_offset_to_native_code_offset.last().unwrap().1 as usize;
             let instruction_length = self.asm.len() - offset;
             if instruction_length > VM_COMPILER_MAXIMUM_INSTRUCTION_LENGTH as usize {
-                self.panic_on_too_long_instruction(instruction_length)
+                self.panic_on_too_long_instruction(code_offset, instruction_length)
             }
         }
     }
 
-    fn current_instruction(&self) -> impl core::fmt::Display {
+    fn current_instruction(&self, code_offset: u32) -> impl core::fmt::Display {
         struct MaybeInstruction(Option<ParsedInstruction>);
         impl core::fmt::Display for MaybeInstruction {
             fn fmt(&self, fmt: &mut core::fmt::Formatter) -> core::fmt::Result {
@@ -334,34 +330,38 @@ impl<'a, S> CompilerVisitor<'a, S> where S: Sandbox {
                 }
             }
         }
-        MaybeInstruction(Instructions::new(self.code, self.bitmask, self.current_code_offset).next())
+        MaybeInstruction(Instructions::new(self.code, self.bitmask, code_offset).next())
     }
 
     #[cold]
-    fn panic_on_too_long_instruction(&self, instruction_length: usize) -> ! {
+    fn panic_on_too_long_instruction(&self, code_offset: u32, instruction_length: usize) -> ! {
         panic!(
             "maximum instruction length of {} exceeded with {} bytes for instruction: {}",
             VM_COMPILER_MAXIMUM_INSTRUCTION_LENGTH,
             instruction_length,
-            self.current_instruction(),
+            self.current_instruction(code_offset),
         );
     }
 
     #[inline(never)]
     #[cold]
-    fn trace_compiled_instruction(&self) {
-        log::trace!("Compiling {}", self.current_instruction());
+    fn trace_compiled_instruction(&self, code_offset: u32) {
+        log::trace!("Compiling {}", self.current_instruction(code_offset));
     }
 
-    fn get_or_forward_declare_label(&mut self, code_offset: u32) -> Label {
+    fn get_or_forward_declare_label(&mut self, code_offset: u32) -> Option<Label> {
         match self.code_offset_to_label.get(code_offset) {
-            Some(label) => label,
+            Some(label) => Some(label),
             None => {
+                if code_offset > self.code_offset_to_label.len() {
+                    return None;
+                }
+
                 let label = self.asm.forward_declare_label();
                 log::trace!("Label: {label} -> {code_offset} (forward declare)");
 
                 self.code_offset_to_label.insert(code_offset, label);
-                label
+                Some(label)
             }
         }
     }
@@ -373,742 +373,732 @@ impl<'a, S> CompilerVisitor<'a, S> where S: Sandbox {
 }
 
 impl<'a, S> polkavm_common::program::ParsingVisitor for CompilerVisitor<'a, S> where S: Sandbox {
-    #[cfg_attr(not(debug_assertions), inline)]
-    fn on_pre_visit(&mut self, offset: usize, length: usize, _opcode: u8) -> Self::ReturnTy {
-        self.current_code_offset = offset as u32;
-        let next_code_offset = offset + length;
-        self.next_code_offset = next_code_offset as u32;
-        self.is_last_instruction = next_code_offset == self.code.len();
-    }
-}
-
-impl<'a, S> polkavm_common::program::InstructionVisitor for CompilerVisitor<'a, S> where S: Sandbox {
     type ReturnTy = ();
 
     #[inline(always)]
-    fn trap(&mut self) -> Self::ReturnTy {
-        self.before_instruction();
+    fn trap(&mut self, code_offset: u32, args_length: u32) -> Self::ReturnTy {
+        self.before_instruction(code_offset);
         self.gas_visitor.trap();
         ArchVisitor(self).trap();
-        self.after_instruction();
-        self.start_new_basic_block::<false>();
+        self.after_instruction(code_offset);
+        self.start_new_basic_block(code_offset, args_length);
     }
 
     #[inline(always)]
-    fn fallthrough(&mut self) -> Self::ReturnTy {
-        self.before_instruction();
+    fn fallthrough(&mut self, code_offset: u32, args_length: u32) -> Self::ReturnTy {
+        self.before_instruction(code_offset);
         self.gas_visitor.fallthrough();
         ArchVisitor(self).fallthrough();
-        self.after_instruction();
-        self.start_new_basic_block::<false>();
+        self.after_instruction(code_offset);
+        self.start_new_basic_block(code_offset, args_length);
     }
 
     #[inline(always)]
-    fn sbrk(&mut self, d: Reg, s: Reg) -> Self::ReturnTy {
-        self.before_instruction();
+    fn sbrk(&mut self, code_offset: u32, _args_length: u32, d: RawReg, s: RawReg) -> Self::ReturnTy {
+        self.before_instruction(code_offset);
         self.gas_visitor.sbrk(d, s);
         ArchVisitor(self).sbrk(d, s);
-        self.after_instruction();
+        self.after_instruction(code_offset);
     }
 
     #[inline(always)]
-    fn ecalli(&mut self, imm: u32) -> Self::ReturnTy {
-        self.before_instruction();
+    fn ecalli(&mut self, code_offset: u32, _args_length: u32, imm: u32) -> Self::ReturnTy {
+        self.before_instruction(code_offset);
         self.gas_visitor.ecalli(imm);
         ArchVisitor(self).ecalli(imm);
-        self.after_instruction();
+        self.after_instruction(code_offset);
     }
 
     #[inline(always)]
-    fn set_less_than_unsigned(&mut self, d: Reg, s1: Reg, s2: Reg) -> Self::ReturnTy {
-        self.before_instruction();
+    fn set_less_than_unsigned(&mut self, code_offset: u32, _args_length: u32, d: RawReg, s1: RawReg, s2: RawReg) -> Self::ReturnTy {
+        self.before_instruction(code_offset);
         self.gas_visitor.set_less_than_unsigned(d, s1, s2);
         ArchVisitor(self).set_less_than_unsigned(d, s1, s2);
-        self.after_instruction();
+        self.after_instruction(code_offset);
     }
 
     #[inline(always)]
-    fn set_less_than_signed(&mut self, d: Reg, s1: Reg, s2: Reg) -> Self::ReturnTy {
-        self.before_instruction();
+    fn set_less_than_signed(&mut self, code_offset: u32, _args_length: u32, d: RawReg, s1: RawReg, s2: RawReg) -> Self::ReturnTy {
+        self.before_instruction(code_offset);
         self.gas_visitor.set_less_than_signed(d, s1, s2);
         ArchVisitor(self).set_less_than_signed(d, s1, s2);
-        self.after_instruction();
+        self.after_instruction(code_offset);
     }
 
     #[inline(always)]
-    fn shift_logical_right(&mut self, d: Reg, s1: Reg, s2: Reg) -> Self::ReturnTy {
-        self.before_instruction();
+    fn shift_logical_right(&mut self, code_offset: u32, _args_length: u32, d: RawReg, s1: RawReg, s2: RawReg) -> Self::ReturnTy {
+        self.before_instruction(code_offset);
         self.gas_visitor.shift_logical_right(d, s1, s2);
         ArchVisitor(self).shift_logical_right(d, s1, s2);
-        self.after_instruction();
+        self.after_instruction(code_offset);
     }
 
     #[inline(always)]
-    fn shift_arithmetic_right(&mut self, d: Reg, s1: Reg, s2: Reg) -> Self::ReturnTy {
-        self.before_instruction();
+    fn shift_arithmetic_right(&mut self, code_offset: u32, _args_length: u32, d: RawReg, s1: RawReg, s2: RawReg) -> Self::ReturnTy {
+        self.before_instruction(code_offset);
         self.gas_visitor.shift_arithmetic_right(d, s1, s2);
         ArchVisitor(self).shift_arithmetic_right(d, s1, s2);
-        self.after_instruction();
+        self.after_instruction(code_offset);
     }
 
     #[inline(always)]
-    fn shift_logical_left(&mut self, d: Reg, s1: Reg, s2: Reg) -> Self::ReturnTy {
-        self.before_instruction();
+    fn shift_logical_left(&mut self, code_offset: u32, _args_length: u32, d: RawReg, s1: RawReg, s2: RawReg) -> Self::ReturnTy {
+        self.before_instruction(code_offset);
         self.gas_visitor.shift_logical_left(d, s1, s2);
         ArchVisitor(self).shift_logical_left(d, s1, s2);
-        self.after_instruction();
+        self.after_instruction(code_offset);
     }
 
     #[inline(always)]
-    fn xor(&mut self, d: Reg, s1: Reg, s2: Reg) -> Self::ReturnTy {
-        self.before_instruction();
+    fn xor(&mut self, code_offset: u32, _args_length: u32, d: RawReg, s1: RawReg, s2: RawReg) -> Self::ReturnTy {
+        self.before_instruction(code_offset);
         self.gas_visitor.xor(d, s1, s2);
         ArchVisitor(self).xor(d, s1, s2);
-        self.after_instruction();
+        self.after_instruction(code_offset);
     }
 
     #[inline(always)]
-    fn and(&mut self, d: Reg, s1: Reg, s2: Reg) -> Self::ReturnTy {
-        self.before_instruction();
+    fn and(&mut self, code_offset: u32, _args_length: u32, d: RawReg, s1: RawReg, s2: RawReg) -> Self::ReturnTy {
+        self.before_instruction(code_offset);
         self.gas_visitor.and(d, s1, s2);
         ArchVisitor(self).and(d, s1, s2);
-        self.after_instruction();
+        self.after_instruction(code_offset);
     }
 
     #[inline(always)]
-    fn or(&mut self, d: Reg, s1: Reg, s2: Reg) -> Self::ReturnTy {
-        self.before_instruction();
+    fn or(&mut self, code_offset: u32, _args_length: u32, d: RawReg, s1: RawReg, s2: RawReg) -> Self::ReturnTy {
+        self.before_instruction(code_offset);
         self.gas_visitor.or(d, s1, s2);
         ArchVisitor(self).or(d, s1, s2);
-        self.after_instruction();
+        self.after_instruction(code_offset);
     }
 
     #[inline(always)]
-    fn add(&mut self, d: Reg, s1: Reg, s2: Reg) -> Self::ReturnTy {
-        self.before_instruction();
+    fn add(&mut self, code_offset: u32, _args_length: u32, d: RawReg, s1: RawReg, s2: RawReg) -> Self::ReturnTy {
+        self.before_instruction(code_offset);
         self.gas_visitor.add(d, s1, s2);
         ArchVisitor(self).add(d, s1, s2);
-        self.after_instruction();
+        self.after_instruction(code_offset);
     }
 
     #[inline(always)]
-    fn sub(&mut self, d: Reg, s1: Reg, s2: Reg) -> Self::ReturnTy {
-        self.before_instruction();
+    fn sub(&mut self, code_offset: u32, _args_length: u32, d: RawReg, s1: RawReg, s2: RawReg) -> Self::ReturnTy {
+        self.before_instruction(code_offset);
         self.gas_visitor.sub(d, s1, s2);
         ArchVisitor(self).sub(d, s1, s2);
-        self.after_instruction();
+        self.after_instruction(code_offset);
     }
 
     #[inline(always)]
-    fn mul(&mut self, d: Reg, s1: Reg, s2: Reg) -> Self::ReturnTy {
-        self.before_instruction();
+    fn mul(&mut self, code_offset: u32, _args_length: u32, d: RawReg, s1: RawReg, s2: RawReg) -> Self::ReturnTy {
+        self.before_instruction(code_offset);
         self.gas_visitor.mul(d, s1, s2);
         ArchVisitor(self).mul(d, s1, s2);
-        self.after_instruction();
+        self.after_instruction(code_offset);
     }
 
     #[inline(always)]
-    fn mul_upper_signed_signed(&mut self, d: Reg, s1: Reg, s2: Reg) -> Self::ReturnTy {
-        self.before_instruction();
+    fn mul_upper_signed_signed(&mut self, code_offset: u32, _args_length: u32, d: RawReg, s1: RawReg, s2: RawReg) -> Self::ReturnTy {
+        self.before_instruction(code_offset);
         self.gas_visitor.mul_upper_signed_signed(d, s1, s2);
         ArchVisitor(self).mul_upper_signed_signed(d, s1, s2);
-        self.after_instruction();
+        self.after_instruction(code_offset);
     }
 
     #[inline(always)]
-    fn mul_upper_unsigned_unsigned(&mut self, d: Reg, s1: Reg, s2: Reg) -> Self::ReturnTy {
-        self.before_instruction();
+    fn mul_upper_unsigned_unsigned(&mut self, code_offset: u32, _args_length: u32, d: RawReg, s1: RawReg, s2: RawReg) -> Self::ReturnTy {
+        self.before_instruction(code_offset);
         self.gas_visitor.mul_upper_unsigned_unsigned(d, s1, s2);
         ArchVisitor(self).mul_upper_unsigned_unsigned(d, s1, s2);
-        self.after_instruction();
+        self.after_instruction(code_offset);
     }
 
     #[inline(always)]
-    fn mul_upper_signed_unsigned(&mut self, d: Reg, s1: Reg, s2: Reg) -> Self::ReturnTy {
-        self.before_instruction();
+    fn mul_upper_signed_unsigned(&mut self, code_offset: u32, _args_length: u32, d: RawReg, s1: RawReg, s2: RawReg) -> Self::ReturnTy {
+        self.before_instruction(code_offset);
         self.gas_visitor.mul_upper_signed_unsigned(d, s1, s2);
         ArchVisitor(self).mul_upper_signed_unsigned(d, s1, s2);
-        self.after_instruction();
+        self.after_instruction(code_offset);
     }
 
     #[inline(always)]
-    fn div_unsigned(&mut self, d: Reg, s1: Reg, s2: Reg) -> Self::ReturnTy {
-        self.before_instruction();
+    fn div_unsigned(&mut self, code_offset: u32, _args_length: u32, d: RawReg, s1: RawReg, s2: RawReg) -> Self::ReturnTy {
+        self.before_instruction(code_offset);
         self.gas_visitor.div_unsigned(d, s1, s2);
         ArchVisitor(self).div_unsigned(d, s1, s2);
-        self.after_instruction();
+        self.after_instruction(code_offset);
     }
 
     #[inline(always)]
-    fn div_signed(&mut self, d: Reg, s1: Reg, s2: Reg) -> Self::ReturnTy {
-        self.before_instruction();
+    fn div_signed(&mut self, code_offset: u32, _args_length: u32, d: RawReg, s1: RawReg, s2: RawReg) -> Self::ReturnTy {
+        self.before_instruction(code_offset);
         self.gas_visitor.div_signed(d, s1, s2);
         ArchVisitor(self).div_signed(d, s1, s2);
-        self.after_instruction();
+        self.after_instruction(code_offset);
     }
 
     #[inline(always)]
-    fn rem_unsigned(&mut self, d: Reg, s1: Reg, s2: Reg) -> Self::ReturnTy {
-        self.before_instruction();
+    fn rem_unsigned(&mut self, code_offset: u32, _args_length: u32, d: RawReg, s1: RawReg, s2: RawReg) -> Self::ReturnTy {
+        self.before_instruction(code_offset);
         self.gas_visitor.rem_unsigned(d, s1, s2);
         ArchVisitor(self).rem_unsigned(d, s1, s2);
-        self.after_instruction();
+        self.after_instruction(code_offset);
     }
 
     #[inline(always)]
-    fn rem_signed(&mut self, d: Reg, s1: Reg, s2: Reg) -> Self::ReturnTy {
-        self.before_instruction();
+    fn rem_signed(&mut self, code_offset: u32, _args_length: u32, d: RawReg, s1: RawReg, s2: RawReg) -> Self::ReturnTy {
+        self.before_instruction(code_offset);
         self.gas_visitor.rem_signed(d, s1, s2);
         ArchVisitor(self).rem_signed(d, s1, s2);
-        self.after_instruction();
+        self.after_instruction(code_offset);
     }
 
     #[inline(always)]
-    fn mul_imm(&mut self, d: Reg, s1: Reg, s2: u32) -> Self::ReturnTy {
-        self.before_instruction();
+    fn mul_imm(&mut self, code_offset: u32, _args_length: u32, d: RawReg, s1: RawReg, s2: u32) -> Self::ReturnTy {
+        self.before_instruction(code_offset);
         self.gas_visitor.mul_imm(d, s1, s2);
         ArchVisitor(self).mul_imm(d, s1, s2);
-        self.after_instruction();
+        self.after_instruction(code_offset);
     }
 
     #[inline(always)]
-    fn mul_upper_signed_signed_imm(&mut self, d: Reg, s1: Reg, s2: u32) -> Self::ReturnTy {
-        self.before_instruction();
+    fn mul_upper_signed_signed_imm(&mut self, code_offset: u32, _args_length: u32, d: RawReg, s1: RawReg, s2: u32) -> Self::ReturnTy {
+        self.before_instruction(code_offset);
         self.gas_visitor.mul_upper_signed_signed_imm(d, s1, s2);
         ArchVisitor(self).mul_upper_signed_signed_imm(d, s1, s2);
-        self.after_instruction();
+        self.after_instruction(code_offset);
     }
 
     #[inline(always)]
-    fn mul_upper_unsigned_unsigned_imm(&mut self, d: Reg, s1: Reg, s2: u32) -> Self::ReturnTy {
-        self.before_instruction();
+    fn mul_upper_unsigned_unsigned_imm(&mut self, code_offset: u32, _args_length: u32, d: RawReg, s1: RawReg, s2: u32) -> Self::ReturnTy {
+        self.before_instruction(code_offset);
         self.gas_visitor.mul_upper_unsigned_unsigned_imm(d, s1, s2);
         ArchVisitor(self).mul_upper_unsigned_unsigned_imm(d, s1, s2);
-        self.after_instruction();
+        self.after_instruction(code_offset);
     }
 
     #[inline(always)]
-    fn set_less_than_unsigned_imm(&mut self, d: Reg, s1: Reg, s2: u32) -> Self::ReturnTy {
-        self.before_instruction();
+    fn set_less_than_unsigned_imm(&mut self, code_offset: u32, _args_length: u32, d: RawReg, s1: RawReg, s2: u32) -> Self::ReturnTy {
+        self.before_instruction(code_offset);
         self.gas_visitor.set_less_than_unsigned_imm(d, s1, s2);
         ArchVisitor(self).set_less_than_unsigned_imm(d, s1, s2);
-        self.after_instruction();
+        self.after_instruction(code_offset);
     }
 
     #[inline(always)]
-    fn set_less_than_signed_imm(&mut self, d: Reg, s1: Reg, s2: u32) -> Self::ReturnTy {
-        self.before_instruction();
+    fn set_less_than_signed_imm(&mut self, code_offset: u32, _args_length: u32, d: RawReg, s1: RawReg, s2: u32) -> Self::ReturnTy {
+        self.before_instruction(code_offset);
         self.gas_visitor.set_less_than_signed_imm(d, s1, s2);
         ArchVisitor(self).set_less_than_signed_imm(d, s1, s2);
-        self.after_instruction();
+        self.after_instruction(code_offset);
     }
 
     #[inline(always)]
-    fn set_greater_than_unsigned_imm(&mut self, d: Reg, s1: Reg, s2: u32) -> Self::ReturnTy {
-        self.before_instruction();
+    fn set_greater_than_unsigned_imm(&mut self, code_offset: u32, _args_length: u32, d: RawReg, s1: RawReg, s2: u32) -> Self::ReturnTy {
+        self.before_instruction(code_offset);
         self.gas_visitor.set_greater_than_unsigned_imm(d, s1, s2);
         ArchVisitor(self).set_greater_than_unsigned_imm(d, s1, s2);
-        self.after_instruction();
+        self.after_instruction(code_offset);
     }
 
     #[inline(always)]
-    fn set_greater_than_signed_imm(&mut self, d: Reg, s1: Reg, s2: u32) -> Self::ReturnTy {
-        self.before_instruction();
+    fn set_greater_than_signed_imm(&mut self, code_offset: u32, _args_length: u32, d: RawReg, s1: RawReg, s2: u32) -> Self::ReturnTy {
+        self.before_instruction(code_offset);
         self.gas_visitor.set_greater_than_signed_imm(d, s1, s2);
         ArchVisitor(self).set_greater_than_signed_imm(d, s1, s2);
-        self.after_instruction();
+        self.after_instruction(code_offset);
     }
 
     #[inline(always)]
-    fn shift_logical_right_imm(&mut self, d: Reg, s1: Reg, s2: u32) -> Self::ReturnTy {
-        self.before_instruction();
+    fn shift_logical_right_imm(&mut self, code_offset: u32, _args_length: u32, d: RawReg, s1: RawReg, s2: u32) -> Self::ReturnTy {
+        self.before_instruction(code_offset);
         self.gas_visitor.shift_logical_right_imm(d, s1, s2);
         ArchVisitor(self).shift_logical_right_imm(d, s1, s2);
-        self.after_instruction();
+        self.after_instruction(code_offset);
     }
 
     #[inline(always)]
-    fn shift_arithmetic_right_imm(&mut self, d: Reg, s1: Reg, s2: u32) -> Self::ReturnTy {
-        self.before_instruction();
+    fn shift_arithmetic_right_imm(&mut self, code_offset: u32, _args_length: u32, d: RawReg, s1: RawReg, s2: u32) -> Self::ReturnTy {
+        self.before_instruction(code_offset);
         self.gas_visitor.shift_arithmetic_right_imm(d, s1, s2);
         ArchVisitor(self).shift_arithmetic_right_imm(d, s1, s2);
-        self.after_instruction();
+        self.after_instruction(code_offset);
     }
 
     #[inline(always)]
-    fn shift_logical_left_imm(&mut self, d: Reg, s1: Reg, s2: u32) -> Self::ReturnTy {
-        self.before_instruction();
+    fn shift_logical_left_imm(&mut self, code_offset: u32, _args_length: u32, d: RawReg, s1: RawReg, s2: u32) -> Self::ReturnTy {
+        self.before_instruction(code_offset);
         self.gas_visitor.shift_logical_left_imm(d, s1, s2);
         ArchVisitor(self).shift_logical_left_imm(d, s1, s2);
-        self.after_instruction();
+        self.after_instruction(code_offset);
     }
 
     #[inline(always)]
-    fn shift_logical_right_imm_alt(&mut self, d: Reg, s2: Reg, s1: u32) -> Self::ReturnTy {
-        self.before_instruction();
+    fn shift_logical_right_imm_alt(&mut self, code_offset: u32, _args_length: u32, d: RawReg, s2: RawReg, s1: u32) -> Self::ReturnTy {
+        self.before_instruction(code_offset);
         self.gas_visitor.shift_logical_right_imm_alt(d, s2, s1);
         ArchVisitor(self).shift_logical_right_imm_alt(d, s2, s1);
-        self.after_instruction();
+        self.after_instruction(code_offset);
     }
 
     #[inline(always)]
-    fn shift_arithmetic_right_imm_alt(&mut self, d: Reg, s2: Reg, s1: u32) -> Self::ReturnTy {
-        self.before_instruction();
+    fn shift_arithmetic_right_imm_alt(&mut self, code_offset: u32, _args_length: u32, d: RawReg, s2: RawReg, s1: u32) -> Self::ReturnTy {
+        self.before_instruction(code_offset);
         self.gas_visitor.shift_arithmetic_right_imm_alt(d, s2, s1);
         ArchVisitor(self).shift_arithmetic_right_imm_alt(d, s2, s1);
-        self.after_instruction();
+        self.after_instruction(code_offset);
     }
 
     #[inline(always)]
-    fn shift_logical_left_imm_alt(&mut self, d: Reg, s2: Reg, s1: u32) -> Self::ReturnTy {
-        self.before_instruction();
+    fn shift_logical_left_imm_alt(&mut self, code_offset: u32, _args_length: u32, d: RawReg, s2: RawReg, s1: u32) -> Self::ReturnTy {
+        self.before_instruction(code_offset);
         self.gas_visitor.shift_logical_left_imm_alt(d, s2, s1);
         ArchVisitor(self).shift_logical_left_imm_alt(d, s2, s1);
-        self.after_instruction();
+        self.after_instruction(code_offset);
     }
 
     #[inline(always)]
-    fn or_imm(&mut self, d: Reg, s: Reg, imm: u32) -> Self::ReturnTy {
-        self.before_instruction();
+    fn or_imm(&mut self, code_offset: u32, _args_length: u32, d: RawReg, s: RawReg, imm: u32) -> Self::ReturnTy {
+        self.before_instruction(code_offset);
         self.gas_visitor.or_imm(d, s, imm);
         ArchVisitor(self).or_imm(d, s, imm);
-        self.after_instruction();
+        self.after_instruction(code_offset);
     }
 
     #[inline(always)]
-    fn and_imm(&mut self, d: Reg, s: Reg, imm: u32) -> Self::ReturnTy {
-        self.before_instruction();
+    fn and_imm(&mut self, code_offset: u32, _args_length: u32, d: RawReg, s: RawReg, imm: u32) -> Self::ReturnTy {
+        self.before_instruction(code_offset);
         self.gas_visitor.and_imm(d, s, imm);
         ArchVisitor(self).and_imm(d, s, imm);
-        self.after_instruction();
+        self.after_instruction(code_offset);
     }
 
     #[inline(always)]
-    fn xor_imm(&mut self, d: Reg, s: Reg, imm: u32) -> Self::ReturnTy {
-        self.before_instruction();
+    fn xor_imm(&mut self, code_offset: u32, _args_length: u32, d: RawReg, s: RawReg, imm: u32) -> Self::ReturnTy {
+        self.before_instruction(code_offset);
         self.gas_visitor.xor_imm(d, s, imm);
         ArchVisitor(self).xor_imm(d, s, imm);
-        self.after_instruction();
+        self.after_instruction(code_offset);
     }
 
     #[inline(always)]
-    fn move_reg(&mut self, d: Reg, s: Reg) -> Self::ReturnTy {
-        self.before_instruction();
+    fn move_reg(&mut self, code_offset: u32, _args_length: u32, d: RawReg, s: RawReg) -> Self::ReturnTy {
+        self.before_instruction(code_offset);
         self.gas_visitor.move_reg(d, s);
         ArchVisitor(self).move_reg(d, s);
-        self.after_instruction();
+        self.after_instruction(code_offset);
     }
 
     #[inline(always)]
-    fn cmov_if_zero(&mut self, d: Reg, s: Reg, c: Reg) -> Self::ReturnTy {
-        self.before_instruction();
+    fn cmov_if_zero(&mut self, code_offset: u32, _args_length: u32, d: RawReg, s: RawReg, c: RawReg) -> Self::ReturnTy {
+        self.before_instruction(code_offset);
         self.gas_visitor.cmov_if_zero(d, s, c);
         ArchVisitor(self).cmov_if_zero(d, s, c);
-        self.after_instruction();
+        self.after_instruction(code_offset);
     }
 
     #[inline(always)]
-    fn cmov_if_not_zero(&mut self, d: Reg, s: Reg, c: Reg) -> Self::ReturnTy {
-        self.before_instruction();
+    fn cmov_if_not_zero(&mut self, code_offset: u32, _args_length: u32, d: RawReg, s: RawReg, c: RawReg) -> Self::ReturnTy {
+        self.before_instruction(code_offset);
         self.gas_visitor.cmov_if_not_zero(d, s, c);
         ArchVisitor(self).cmov_if_not_zero(d, s, c);
-        self.after_instruction();
+        self.after_instruction(code_offset);
     }
 
     #[inline(always)]
-    fn cmov_if_zero_imm(&mut self, d: Reg, c: Reg, s: u32) -> Self::ReturnTy {
-        self.before_instruction();
+    fn cmov_if_zero_imm(&mut self, code_offset: u32, _args_length: u32, d: RawReg, c: RawReg, s: u32) -> Self::ReturnTy {
+        self.before_instruction(code_offset);
         self.gas_visitor.cmov_if_zero_imm(d, c, s);
         ArchVisitor(self).cmov_if_zero_imm(d, c, s);
-        self.after_instruction();
+        self.after_instruction(code_offset);
     }
 
     #[inline(always)]
-    fn cmov_if_not_zero_imm(&mut self, d: Reg, c: Reg, s: u32) -> Self::ReturnTy {
-        self.before_instruction();
+    fn cmov_if_not_zero_imm(&mut self, code_offset: u32, _args_length: u32, d: RawReg, c: RawReg, s: u32) -> Self::ReturnTy {
+        self.before_instruction(code_offset);
         self.gas_visitor.cmov_if_not_zero_imm(d, c, s);
         ArchVisitor(self).cmov_if_not_zero_imm(d, c, s);
-        self.after_instruction();
+        self.after_instruction(code_offset);
     }
 
     #[inline(always)]
-    fn add_imm(&mut self, d: Reg, s: Reg, imm: u32) -> Self::ReturnTy {
-        self.before_instruction();
+    fn add_imm(&mut self, code_offset: u32, _args_length: u32, d: RawReg, s: RawReg, imm: u32) -> Self::ReturnTy {
+        self.before_instruction(code_offset);
         self.gas_visitor.add_imm(d, s, imm);
         ArchVisitor(self).add_imm(d, s, imm);
-        self.after_instruction();
+        self.after_instruction(code_offset);
     }
 
     #[inline(always)]
-    fn negate_and_add_imm(&mut self, d: Reg, s1: Reg, s2: u32) -> Self::ReturnTy {
-        self.before_instruction();
+    fn negate_and_add_imm(&mut self, code_offset: u32, _args_length: u32, d: RawReg, s1: RawReg, s2: u32) -> Self::ReturnTy {
+        self.before_instruction(code_offset);
         self.gas_visitor.negate_and_add_imm(d, s1, s2);
         ArchVisitor(self).negate_and_add_imm(d, s1, s2);
-        self.after_instruction();
+        self.after_instruction(code_offset);
     }
 
     #[inline(always)]
-    fn store_imm_indirect_u8(&mut self, base: Reg, offset: u32, value: u32) -> Self::ReturnTy {
-        self.before_instruction();
+    fn store_imm_indirect_u8(&mut self, code_offset: u32, _args_length: u32, base: RawReg, offset: u32, value: u32) -> Self::ReturnTy {
+        self.before_instruction(code_offset);
         self.gas_visitor.store_imm_indirect_u8(base, offset, value);
         ArchVisitor(self).store_imm_indirect_u8(base, offset, value);
-        self.after_instruction();
+        self.after_instruction(code_offset);
     }
 
     #[inline(always)]
-    fn store_imm_indirect_u16(&mut self, base: Reg, offset: u32, value: u32) -> Self::ReturnTy {
-        self.before_instruction();
+    fn store_imm_indirect_u16(&mut self, code_offset: u32, _args_length: u32, base: RawReg, offset: u32, value: u32) -> Self::ReturnTy {
+        self.before_instruction(code_offset);
         self.gas_visitor.store_imm_indirect_u16(base, offset, value);
         ArchVisitor(self).store_imm_indirect_u16(base, offset, value);
-        self.after_instruction();
+        self.after_instruction(code_offset);
     }
 
     #[inline(always)]
-    fn store_imm_indirect_u32(&mut self, base: Reg, offset: u32, value: u32) -> Self::ReturnTy {
-        self.before_instruction();
+    fn store_imm_indirect_u32(&mut self, code_offset: u32, _args_length: u32, base: RawReg, offset: u32, value: u32) -> Self::ReturnTy {
+        self.before_instruction(code_offset);
         self.gas_visitor.store_imm_indirect_u32(base, offset, value);
         ArchVisitor(self).store_imm_indirect_u32(base, offset, value);
-        self.after_instruction();
+        self.after_instruction(code_offset);
     }
 
     #[inline(always)]
-    fn store_indirect_u8(&mut self, src: Reg, base: Reg, offset: u32) -> Self::ReturnTy {
-        self.before_instruction();
+    fn store_indirect_u8(&mut self, code_offset: u32, _args_length: u32, src: RawReg, base: RawReg, offset: u32) -> Self::ReturnTy {
+        self.before_instruction(code_offset);
         self.gas_visitor.store_indirect_u8(src, base, offset);
         ArchVisitor(self).store_indirect_u8(src, base, offset);
-        self.after_instruction();
+        self.after_instruction(code_offset);
     }
 
     #[inline(always)]
-    fn store_indirect_u16(&mut self, src: Reg, base: Reg, offset: u32) -> Self::ReturnTy {
-        self.before_instruction();
+    fn store_indirect_u16(&mut self, code_offset: u32, _args_length: u32, src: RawReg, base: RawReg, offset: u32) -> Self::ReturnTy {
+        self.before_instruction(code_offset);
         self.gas_visitor.store_indirect_u16(src, base, offset);
         ArchVisitor(self).store_indirect_u16(src, base, offset);
-        self.after_instruction();
+        self.after_instruction(code_offset);
     }
 
     #[inline(always)]
-    fn store_indirect_u32(&mut self, src: Reg, base: Reg, offset: u32) -> Self::ReturnTy {
-        self.before_instruction();
+    fn store_indirect_u32(&mut self, code_offset: u32, _args_length: u32, src: RawReg, base: RawReg, offset: u32) -> Self::ReturnTy {
+        self.before_instruction(code_offset);
         self.gas_visitor.store_indirect_u32(src, base, offset);
         ArchVisitor(self).store_indirect_u32(src, base, offset);
-        self.after_instruction();
+        self.after_instruction(code_offset);
     }
 
     #[inline(always)]
-    fn store_imm_u8(&mut self, value: u32, offset: u32) -> Self::ReturnTy {
-        self.before_instruction();
+    fn store_imm_u8(&mut self, code_offset: u32, _args_length: u32, value: u32, offset: u32) -> Self::ReturnTy {
+        self.before_instruction(code_offset);
         self.gas_visitor.store_imm_u8(value, offset);
         ArchVisitor(self).store_imm_u8(value, offset);
-        self.after_instruction();
+        self.after_instruction(code_offset);
     }
 
     #[inline(always)]
-    fn store_imm_u16(&mut self, value: u32, offset: u32) -> Self::ReturnTy {
-        self.before_instruction();
+    fn store_imm_u16(&mut self, code_offset: u32, _args_length: u32, value: u32, offset: u32) -> Self::ReturnTy {
+        self.before_instruction(code_offset);
         self.gas_visitor.store_imm_u16(value, offset);
         ArchVisitor(self).store_imm_u16(value, offset);
-        self.after_instruction();
+        self.after_instruction(code_offset);
     }
 
     #[inline(always)]
-    fn store_imm_u32(&mut self, value: u32, offset: u32) -> Self::ReturnTy {
-        self.before_instruction();
+    fn store_imm_u32(&mut self, code_offset: u32, _args_length: u32, value: u32, offset: u32) -> Self::ReturnTy {
+        self.before_instruction(code_offset);
         self.gas_visitor.store_imm_u32(value, offset);
         ArchVisitor(self).store_imm_u32(value, offset);
-        self.after_instruction();
+        self.after_instruction(code_offset);
     }
 
     #[inline(always)]
-    fn store_u8(&mut self, src: Reg, offset: u32) -> Self::ReturnTy {
-        self.before_instruction();
+    fn store_u8(&mut self, code_offset: u32, _args_length: u32, src: RawReg, offset: u32) -> Self::ReturnTy {
+        self.before_instruction(code_offset);
         self.gas_visitor.store_u8(src, offset);
         ArchVisitor(self).store_u8(src, offset);
-        self.after_instruction();
+        self.after_instruction(code_offset);
     }
 
     #[inline(always)]
-    fn store_u16(&mut self, src: Reg, offset: u32) -> Self::ReturnTy {
-        self.before_instruction();
+    fn store_u16(&mut self, code_offset: u32, _args_length: u32, src: RawReg, offset: u32) -> Self::ReturnTy {
+        self.before_instruction(code_offset);
         self.gas_visitor.store_u16(src, offset);
         ArchVisitor(self).store_u16(src, offset);
-        self.after_instruction();
+        self.after_instruction(code_offset);
     }
 
     #[inline(always)]
-    fn store_u32(&mut self, src: Reg, offset: u32) -> Self::ReturnTy {
-        self.before_instruction();
+    fn store_u32(&mut self, code_offset: u32, _args_length: u32, src: RawReg, offset: u32) -> Self::ReturnTy {
+        self.before_instruction(code_offset);
         self.gas_visitor.store_u32(src, offset);
         ArchVisitor(self).store_u32(src, offset);
-        self.after_instruction();
+        self.after_instruction(code_offset);
     }
 
     #[inline(always)]
-    fn load_indirect_u8(&mut self, dst: Reg, base: Reg, offset: u32) -> Self::ReturnTy {
-        self.before_instruction();
+    fn load_indirect_u8(&mut self, code_offset: u32, _args_length: u32, dst: RawReg, base: RawReg, offset: u32) -> Self::ReturnTy {
+        self.before_instruction(code_offset);
         self.gas_visitor.load_indirect_u8(dst, base, offset);
         ArchVisitor(self).load_indirect_u8(dst, base, offset);
-        self.after_instruction();
+        self.after_instruction(code_offset);
     }
 
     #[inline(always)]
-    fn load_indirect_i8(&mut self, dst: Reg, base: Reg, offset: u32) -> Self::ReturnTy {
-        self.before_instruction();
+    fn load_indirect_i8(&mut self, code_offset: u32, _args_length: u32, dst: RawReg, base: RawReg, offset: u32) -> Self::ReturnTy {
+        self.before_instruction(code_offset);
         self.gas_visitor.load_indirect_i8(dst, base, offset);
         ArchVisitor(self).load_indirect_i8(dst, base, offset);
-        self.after_instruction();
+        self.after_instruction(code_offset);
     }
 
     #[inline(always)]
-    fn load_indirect_u16(&mut self, dst: Reg, base: Reg, offset: u32) -> Self::ReturnTy {
-        self.before_instruction();
+    fn load_indirect_u16(&mut self, code_offset: u32, _args_length: u32, dst: RawReg, base: RawReg, offset: u32) -> Self::ReturnTy {
+        self.before_instruction(code_offset);
         self.gas_visitor.load_indirect_u16(dst, base, offset);
         ArchVisitor(self).load_indirect_u16(dst, base, offset);
-        self.after_instruction();
+        self.after_instruction(code_offset);
     }
 
     #[inline(always)]
-    fn load_indirect_i16(&mut self, dst: Reg, base: Reg, offset: u32) -> Self::ReturnTy {
-        self.before_instruction();
+    fn load_indirect_i16(&mut self, code_offset: u32, _args_length: u32, dst: RawReg, base: RawReg, offset: u32) -> Self::ReturnTy {
+        self.before_instruction(code_offset);
         self.gas_visitor.load_indirect_i16(dst, base, offset);
         ArchVisitor(self).load_indirect_i16(dst, base, offset);
-        self.after_instruction();
+        self.after_instruction(code_offset);
     }
 
     #[inline(always)]
-    fn load_indirect_u32(&mut self, dst: Reg, base: Reg, offset: u32) -> Self::ReturnTy {
-        self.before_instruction();
+    fn load_indirect_u32(&mut self, code_offset: u32, _args_length: u32, dst: RawReg, base: RawReg, offset: u32) -> Self::ReturnTy {
+        self.before_instruction(code_offset);
         self.gas_visitor.load_indirect_u32(dst, base, offset);
         ArchVisitor(self).load_indirect_u32(dst, base, offset);
-        self.after_instruction();
+        self.after_instruction(code_offset);
     }
 
     #[inline(always)]
-    fn load_u8(&mut self, dst: Reg, offset: u32) -> Self::ReturnTy {
-        self.before_instruction();
+    fn load_u8(&mut self, code_offset: u32, _args_length: u32, dst: RawReg, offset: u32) -> Self::ReturnTy {
+        self.before_instruction(code_offset);
         self.gas_visitor.load_u8(dst, offset);
         ArchVisitor(self).load_u8(dst, offset);
-        self.after_instruction();
+        self.after_instruction(code_offset);
     }
 
     #[inline(always)]
-    fn load_i8(&mut self, dst: Reg, offset: u32) -> Self::ReturnTy {
-        self.before_instruction();
+    fn load_i8(&mut self, code_offset: u32, _args_length: u32, dst: RawReg, offset: u32) -> Self::ReturnTy {
+        self.before_instruction(code_offset);
         self.gas_visitor.load_i8(dst, offset);
         ArchVisitor(self).load_i8(dst, offset);
-        self.after_instruction();
+        self.after_instruction(code_offset);
     }
 
     #[inline(always)]
-    fn load_u16(&mut self, dst: Reg, offset: u32) -> Self::ReturnTy {
-        self.before_instruction();
+    fn load_u16(&mut self, code_offset: u32, _args_length: u32, dst: RawReg, offset: u32) -> Self::ReturnTy {
+        self.before_instruction(code_offset);
         self.gas_visitor.load_u16(dst, offset);
         ArchVisitor(self).load_u16(dst, offset);
-        self.after_instruction();
+        self.after_instruction(code_offset);
     }
 
     #[inline(always)]
-    fn load_i16(&mut self, dst: Reg, offset: u32) -> Self::ReturnTy {
-        self.before_instruction();
+    fn load_i16(&mut self, code_offset: u32, _args_length: u32, dst: RawReg, offset: u32) -> Self::ReturnTy {
+        self.before_instruction(code_offset);
         self.gas_visitor.load_i16(dst, offset);
         ArchVisitor(self).load_i16(dst, offset);
-        self.after_instruction();
+        self.after_instruction(code_offset);
     }
 
     #[inline(always)]
-    fn load_u32(&mut self, dst: Reg, offset: u32) -> Self::ReturnTy {
-        self.before_instruction();
+    fn load_u32(&mut self, code_offset: u32, _args_length: u32, dst: RawReg, offset: u32) -> Self::ReturnTy {
+        self.before_instruction(code_offset);
         self.gas_visitor.load_u32(dst, offset);
         ArchVisitor(self).load_u32(dst, offset);
-        self.after_instruction();
+        self.after_instruction(code_offset);
     }
 
     #[inline(always)]
-    fn branch_less_unsigned(&mut self, s1: Reg, s2: Reg, imm: u32) -> Self::ReturnTy {
-        self.before_instruction();
+    fn branch_less_unsigned(&mut self, code_offset: u32, args_length: u32, s1: RawReg, s2: RawReg, imm: u32) -> Self::ReturnTy {
+        self.before_instruction(code_offset);
         self.gas_visitor.branch_less_unsigned(s1, s2, imm);
         ArchVisitor(self).branch_less_unsigned(s1, s2, imm);
-        self.after_instruction();
-        self.start_new_basic_block::<false>();
+        self.after_instruction(code_offset);
+        self.start_new_basic_block(code_offset, args_length);
     }
 
     #[inline(always)]
-    fn branch_less_signed(&mut self, s1: Reg, s2: Reg, imm: u32) -> Self::ReturnTy {
-        self.before_instruction();
+    fn branch_less_signed(&mut self, code_offset: u32, args_length: u32, s1: RawReg, s2: RawReg, imm: u32) -> Self::ReturnTy {
+        self.before_instruction(code_offset);
         self.gas_visitor.branch_less_signed(s1, s2, imm);
         ArchVisitor(self).branch_less_signed(s1, s2, imm);
-        self.after_instruction();
-        self.start_new_basic_block::<false>();
+        self.after_instruction(code_offset);
+        self.start_new_basic_block(code_offset, args_length);
     }
 
     #[inline(always)]
-    fn branch_greater_or_equal_unsigned(&mut self, s1: Reg, s2: Reg, imm: u32) -> Self::ReturnTy {
-        self.before_instruction();
+    fn branch_greater_or_equal_unsigned(&mut self, code_offset: u32, args_length: u32, s1: RawReg, s2: RawReg, imm: u32) -> Self::ReturnTy {
+        self.before_instruction(code_offset);
         self.gas_visitor.branch_greater_or_equal_unsigned(s1, s2, imm);
         ArchVisitor(self).branch_greater_or_equal_unsigned(s1, s2, imm);
-        self.after_instruction();
-        self.start_new_basic_block::<false>();
+        self.after_instruction(code_offset);
+        self.start_new_basic_block(code_offset, args_length);
     }
 
     #[inline(always)]
-    fn branch_greater_or_equal_signed(&mut self, s1: Reg, s2: Reg, imm: u32) -> Self::ReturnTy {
-        self.before_instruction();
+    fn branch_greater_or_equal_signed(&mut self, code_offset: u32, args_length: u32, s1: RawReg, s2: RawReg, imm: u32) -> Self::ReturnTy {
+        self.before_instruction(code_offset);
         self.gas_visitor.branch_greater_or_equal_signed(s1, s2, imm);
         ArchVisitor(self).branch_greater_or_equal_signed(s1, s2, imm);
-        self.after_instruction();
-        self.start_new_basic_block::<false>();
+        self.after_instruction(code_offset);
+        self.start_new_basic_block(code_offset, args_length);
     }
 
     #[inline(always)]
-    fn branch_eq(&mut self, s1: Reg, s2: Reg, imm: u32) -> Self::ReturnTy {
-        self.before_instruction();
+    fn branch_eq(&mut self, code_offset: u32, args_length: u32, s1: RawReg, s2: RawReg, imm: u32) -> Self::ReturnTy {
+        self.before_instruction(code_offset);
         self.gas_visitor.branch_eq(s1, s2, imm);
         ArchVisitor(self).branch_eq(s1, s2, imm);
-        self.after_instruction();
-        self.start_new_basic_block::<false>();
+        self.after_instruction(code_offset);
+        self.start_new_basic_block(code_offset, args_length);
     }
 
     #[inline(always)]
-    fn branch_not_eq(&mut self, s1: Reg, s2: Reg, imm: u32) -> Self::ReturnTy {
-        self.before_instruction();
+    fn branch_not_eq(&mut self, code_offset: u32, args_length: u32, s1: RawReg, s2: RawReg, imm: u32) -> Self::ReturnTy {
+        self.before_instruction(code_offset);
         self.gas_visitor.branch_not_eq(s1, s2, imm);
         ArchVisitor(self).branch_not_eq(s1, s2, imm);
-        self.after_instruction();
-        self.start_new_basic_block::<false>();
+        self.after_instruction(code_offset);
+        self.start_new_basic_block(code_offset, args_length);
     }
 
     #[inline(always)]
-    fn branch_eq_imm(&mut self, s1: Reg, s2: u32, imm: u32) -> Self::ReturnTy {
-        self.before_instruction();
+    fn branch_eq_imm(&mut self, code_offset: u32, args_length: u32, s1: RawReg, s2: u32, imm: u32) -> Self::ReturnTy {
+        self.before_instruction(code_offset);
         self.gas_visitor.branch_eq_imm(s1, s2, imm);
         ArchVisitor(self).branch_eq_imm(s1, s2, imm);
-        self.after_instruction();
-        self.start_new_basic_block::<false>();
+        self.after_instruction(code_offset);
+        self.start_new_basic_block(code_offset, args_length);
     }
 
     #[inline(always)]
-    fn branch_not_eq_imm(&mut self, s1: Reg, s2: u32, imm: u32) -> Self::ReturnTy {
-        self.before_instruction();
+    fn branch_not_eq_imm(&mut self, code_offset: u32, args_length: u32, s1: RawReg, s2: u32, imm: u32) -> Self::ReturnTy {
+        self.before_instruction(code_offset);
         self.gas_visitor.branch_not_eq_imm(s1, s2, imm);
         ArchVisitor(self).branch_not_eq_imm(s1, s2, imm);
-        self.after_instruction();
-        self.start_new_basic_block::<false>();
+        self.after_instruction(code_offset);
+        self.start_new_basic_block(code_offset, args_length);
     }
 
     #[inline(always)]
-    fn branch_less_unsigned_imm(&mut self, s1: Reg, s2: u32, imm: u32) -> Self::ReturnTy {
-        self.before_instruction();
+    fn branch_less_unsigned_imm(&mut self, code_offset: u32, args_length: u32, s1: RawReg, s2: u32, imm: u32) -> Self::ReturnTy {
+        self.before_instruction(code_offset);
         self.gas_visitor.branch_less_unsigned_imm(s1, s2, imm);
         ArchVisitor(self).branch_less_unsigned_imm(s1, s2, imm);
-        self.after_instruction();
-        self.start_new_basic_block::<false>();
+        self.after_instruction(code_offset);
+        self.start_new_basic_block(code_offset, args_length);
     }
 
     #[inline(always)]
-    fn branch_less_signed_imm(&mut self, s1: Reg, s2: u32, imm: u32) -> Self::ReturnTy {
-        self.before_instruction();
+    fn branch_less_signed_imm(&mut self, code_offset: u32, args_length: u32, s1: RawReg, s2: u32, imm: u32) -> Self::ReturnTy {
+        self.before_instruction(code_offset);
         self.gas_visitor.branch_less_signed_imm(s1, s2, imm);
         ArchVisitor(self).branch_less_signed_imm(s1, s2, imm);
-        self.after_instruction();
-        self.start_new_basic_block::<false>();
+        self.after_instruction(code_offset);
+        self.start_new_basic_block(code_offset, args_length);
     }
 
     #[inline(always)]
-    fn branch_greater_or_equal_unsigned_imm(&mut self, s1: Reg, s2: u32, imm: u32) -> Self::ReturnTy {
-        self.before_instruction();
+    fn branch_greater_or_equal_unsigned_imm(&mut self, code_offset: u32, args_length: u32, s1: RawReg, s2: u32, imm: u32) -> Self::ReturnTy {
+        self.before_instruction(code_offset);
         self.gas_visitor.branch_greater_or_equal_unsigned_imm(s1, s2, imm);
         ArchVisitor(self).branch_greater_or_equal_unsigned_imm(s1, s2, imm);
-        self.after_instruction();
-        self.start_new_basic_block::<false>();
+        self.after_instruction(code_offset);
+        self.start_new_basic_block(code_offset, args_length);
     }
 
     #[inline(always)]
-    fn branch_greater_or_equal_signed_imm(&mut self, s1: Reg, s2: u32, imm: u32) -> Self::ReturnTy {
-        self.before_instruction();
+    fn branch_greater_or_equal_signed_imm(&mut self, code_offset: u32, args_length: u32, s1: RawReg, s2: u32, imm: u32) -> Self::ReturnTy {
+        self.before_instruction(code_offset);
         self.gas_visitor.branch_greater_or_equal_signed_imm(s1, s2, imm);
         ArchVisitor(self).branch_greater_or_equal_signed_imm(s1, s2, imm);
-        self.after_instruction();
-        self.start_new_basic_block::<false>();
+        self.after_instruction(code_offset);
+        self.start_new_basic_block(code_offset, args_length);
     }
 
     #[inline(always)]
-    fn branch_less_or_equal_unsigned_imm(&mut self, s1: Reg, s2: u32, imm: u32) -> Self::ReturnTy {
-        self.before_instruction();
+    fn branch_less_or_equal_unsigned_imm(&mut self, code_offset: u32, args_length: u32, s1: RawReg, s2: u32, imm: u32) -> Self::ReturnTy {
+        self.before_instruction(code_offset);
         self.gas_visitor.branch_less_or_equal_unsigned_imm(s1, s2, imm);
         ArchVisitor(self).branch_less_or_equal_unsigned_imm(s1, s2, imm);
-        self.after_instruction();
-        self.start_new_basic_block::<false>();
+        self.after_instruction(code_offset);
+        self.start_new_basic_block(code_offset, args_length);
     }
 
     #[inline(always)]
-    fn branch_less_or_equal_signed_imm(&mut self, s1: Reg, s2: u32, imm: u32) -> Self::ReturnTy {
-        self.before_instruction();
+    fn branch_less_or_equal_signed_imm(&mut self, code_offset: u32, args_length: u32, s1: RawReg, s2: u32, imm: u32) -> Self::ReturnTy {
+        self.before_instruction(code_offset);
         self.gas_visitor.branch_less_or_equal_signed_imm(s1, s2, imm);
         ArchVisitor(self).branch_less_or_equal_signed_imm(s1, s2, imm);
-        self.after_instruction();
-        self.start_new_basic_block::<false>();
+        self.after_instruction(code_offset);
+        self.start_new_basic_block(code_offset, args_length);
     }
 
     #[inline(always)]
-    fn branch_greater_unsigned_imm(&mut self, s1: Reg, s2: u32, imm: u32) -> Self::ReturnTy {
-        self.before_instruction();
+    fn branch_greater_unsigned_imm(&mut self, code_offset: u32, args_length: u32, s1: RawReg, s2: u32, imm: u32) -> Self::ReturnTy {
+        self.before_instruction(code_offset);
         self.gas_visitor.branch_greater_unsigned_imm(s1, s2, imm);
         ArchVisitor(self).branch_greater_unsigned_imm(s1, s2, imm);
-        self.after_instruction();
-        self.start_new_basic_block::<false>();
+        self.after_instruction(code_offset);
+        self.start_new_basic_block(code_offset, args_length);
     }
 
     #[inline(always)]
-    fn branch_greater_signed_imm(&mut self, s1: Reg, s2: u32, imm: u32) -> Self::ReturnTy {
-        self.before_instruction();
+    fn branch_greater_signed_imm(&mut self, code_offset: u32, args_length: u32, s1: RawReg, s2: u32, imm: u32) -> Self::ReturnTy {
+        self.before_instruction(code_offset);
         self.gas_visitor.branch_greater_signed_imm(s1, s2, imm);
         ArchVisitor(self).branch_greater_signed_imm(s1, s2, imm);
-        self.after_instruction();
-        self.start_new_basic_block::<false>();
+        self.after_instruction(code_offset);
+        self.start_new_basic_block(code_offset, args_length);
     }
 
     #[inline(always)]
-    fn load_imm(&mut self, dst: Reg, value: u32) -> Self::ReturnTy {
-        self.before_instruction();
+    fn load_imm(&mut self, code_offset: u32, _args_length: u32, dst: RawReg, value: u32) -> Self::ReturnTy {
+        self.before_instruction(code_offset);
         self.gas_visitor.load_imm(dst, value);
         ArchVisitor(self).load_imm(dst, value);
-        self.after_instruction();
+        self.after_instruction(code_offset);
     }
 
     #[inline(always)]
-    fn load_imm_and_jump(&mut self, ra: Reg, value: u32, target: u32) -> Self::ReturnTy {
-        self.before_instruction();
+    fn load_imm_and_jump(&mut self, code_offset: u32, args_length: u32, ra: RawReg, value: u32, target: u32) -> Self::ReturnTy {
+        self.before_instruction(code_offset);
         self.gas_visitor.load_imm_and_jump(ra, value, target);
         ArchVisitor(self).load_imm_and_jump(ra, value, target);
-        self.after_instruction();
-        self.start_new_basic_block::<false>();
+        self.after_instruction(code_offset);
+        self.start_new_basic_block(code_offset, args_length);
     }
 
     #[inline(always)]
-    fn load_imm_and_jump_indirect(&mut self, ra: Reg, base: Reg, value: u32, offset: u32) -> Self::ReturnTy {
-        self.before_instruction();
+    fn load_imm_and_jump_indirect(&mut self, code_offset: u32, args_length: u32, ra: RawReg, base: RawReg, value: u32, offset: u32) -> Self::ReturnTy {
+        self.before_instruction(code_offset);
         self.gas_visitor.load_imm_and_jump_indirect(ra, base, value, offset);
         ArchVisitor(self).load_imm_and_jump_indirect(ra, base, value, offset);
-        self.after_instruction();
-        self.start_new_basic_block::<false>();
+        self.after_instruction(code_offset);
+        self.start_new_basic_block(code_offset, args_length);
     }
 
     #[inline(always)]
-    fn jump(&mut self, target: u32) -> Self::ReturnTy {
-        self.before_instruction();
+    fn jump(&mut self, code_offset: u32, args_length: u32, target: u32) -> Self::ReturnTy {
+        self.before_instruction(code_offset);
         self.gas_visitor.jump(target);
         ArchVisitor(self).jump(target);
-        self.after_instruction();
-        self.start_new_basic_block::<false>();
+        self.after_instruction(code_offset);
+        self.start_new_basic_block(code_offset, args_length);
     }
 
     #[inline(always)]
-    fn jump_indirect(&mut self, base: Reg, offset: u32) -> Self::ReturnTy {
-        self.before_instruction();
+    fn jump_indirect(&mut self, code_offset: u32, args_length: u32, base: RawReg, offset: u32) -> Self::ReturnTy {
+        self.before_instruction(code_offset);
         self.gas_visitor.jump_indirect(base, offset);
         ArchVisitor(self).jump_indirect(base, offset);
-        self.after_instruction();
-        self.start_new_basic_block::<false>();
+        self.after_instruction(code_offset);
+        self.start_new_basic_block(code_offset, args_length);
     }
 }
 
diff --git a/crates/polkavm/src/compiler/amd64.rs b/crates/polkavm/src/compiler/amd64.rs
index 4bc17fd7..d6785b37 100644
--- a/crates/polkavm/src/compiler/amd64.rs
+++ b/crates/polkavm/src/compiler/amd64.rs
@@ -4,9 +4,9 @@ use polkavm_assembler::amd64::RegIndex as NativeReg;
 use polkavm_assembler::amd64::RegIndex::*;
 use polkavm_assembler::amd64::Reg::rsp;
 use polkavm_assembler::amd64::{Condition, LoadKind, RegSize, Size, MemOp};
-use polkavm_assembler::Label;
+use polkavm_assembler::{ReservedAssembler, Label, U1, U2, U3, U4, NonZero};
 
-use polkavm_common::program::{InstructionVisitor, Reg};
+use polkavm_common::program::{InstructionVisitor, Reg, RawReg};
 use polkavm_common::zygote::VM_ADDR_VMCTX;
 
 use crate::config::GasMeteringKind;
@@ -42,8 +42,8 @@ const fn conv_reg_const(reg: Reg) -> NativeReg {
     }
 }
 
-static REG_MAP: [NativeReg; Reg::ALL.len()] = {
-    let mut output = [NativeReg::rcx; Reg::ALL.len()];
+static REG_MAP: [NativeReg; 16] = {
+    let mut output = [conv_reg_const(Reg::T2); 16];
     let mut index = 0;
     while index < Reg::ALL.len() {
         assert!(Reg::ALL[index] as usize == index);
@@ -54,14 +54,17 @@ static REG_MAP: [NativeReg; Reg::ALL.len()] = {
 };
 
 #[inline]
-fn conv_reg(reg: Reg) -> NativeReg {
-    REG_MAP[reg as usize]
+fn conv_reg(reg: RawReg) -> NativeReg {
+    let native_reg = REG_MAP[reg.raw_unparsed() as usize & 0b1111];
+    debug_assert_eq!(native_reg, conv_reg_const(reg.get()));
+
+    native_reg
 }
 
 #[test]
 fn test_conv_reg() {
     for reg in Reg::ALL {
-        assert_eq!(conv_reg(reg), conv_reg_const(reg));
+        assert_eq!(conv_reg(reg.into()), conv_reg_const(reg));
     }
 }
 
@@ -149,6 +152,55 @@ enum ShiftKind {
     ArithmeticRight,
 }
 
+#[cfg_attr(not(debug_assertions), inline(always))]
+fn calculate_label_offset(asm_len: usize, rel8_len: usize, rel32_len: usize, offset: isize) -> Result<i8, i32> {
+    let offset_near = offset - (asm_len as isize + rel8_len as isize);
+    if offset_near <= i8::MAX as isize && offset_near >= i8::MIN as isize {
+        Ok(offset_near as i8)
+    } else {
+        let offset = offset - (asm_len as isize + rel32_len as isize);
+        Err(offset as i32)
+    }
+}
+
+#[cfg_attr(not(debug_assertions), inline(always))]
+fn branch_to_label<R>(asm: ReservedAssembler<R>, condition: Condition, label: Label) -> ReservedAssembler<R::Next> where R: NonZero {
+    if let Some(offset) = asm.get_label_origin_offset(label) {
+        let offset = calculate_label_offset(
+            asm.len(),
+            jcc_rel8(condition, i8::MAX).len(),
+            jcc_rel32(condition, i32::MAX).len(),
+            offset
+        );
+
+        match offset {
+            Ok(offset) => asm.push(jcc_rel8(condition, offset)),
+            Err(offset) => asm.push(jcc_rel32(condition, offset))
+        }
+    } else {
+        asm.push(jcc_label32(condition, label))
+    }
+}
+
+#[cfg_attr(not(debug_assertions), inline(always))]
+fn jump_to_label<R>(asm: ReservedAssembler<R>, label: Label) -> ReservedAssembler<R::Next> where R: NonZero {
+    if let Some(offset) = asm.get_label_origin_offset(label) {
+        let offset = calculate_label_offset(
+            asm.len(),
+            jmp_rel8(i8::MAX).len(),
+            jmp_rel32(i32::MAX).len(),
+            offset
+        );
+
+        match offset {
+            Ok(offset) => asm.push(jmp_rel8(offset)),
+            Err(offset) => asm.push(jmp_rel32(offset))
+        }
+    } else {
+        asm.push(jmp_label32(label))
+    }
+}
+
 impl<'r, 'a, S> ArchVisitor<'r, 'a, S> where S: Sandbox {
     pub const PADDING_BYTE: u8 = 0x90; // NOP
 
@@ -164,12 +216,12 @@ impl<'r, 'a, S> ArchVisitor<'r, 'a, S> where S: Sandbox {
     }
 
     #[cfg_attr(not(debug_assertions), inline(always))]
-    fn load_immediate(&mut self, dst: Reg, value: u32) {
+    fn load_immediate(&mut self, dst: RawReg, value: u32) {
         self.push(mov_imm(conv_reg(dst), imm32(value)));
     }
 
     #[cfg_attr(not(debug_assertions), inline(always))]
-    fn store(&mut self, src: impl Into<RegImm>, base: Option<Reg>, offset: u32, kind: Size) {
+    fn store(&mut self, src: impl Into<RegImm>, base: Option<RawReg>, offset: u32, kind: Size) {
         let src = src.into();
         load_store_operand!(self, S::KIND, base, offset, |dst| {
             match src {
@@ -187,20 +239,20 @@ impl<'r, 'a, S> ArchVisitor<'r, 'a, S> where S: Sandbox {
     }
 
     #[cfg_attr(not(debug_assertions), inline(always))]
-    fn load(&mut self, dst: Reg, base: Option<Reg>, offset: u32, kind: LoadKind) {
+    fn load(&mut self, dst: RawReg, base: Option<RawReg>, offset: u32, kind: LoadKind) {
         load_store_operand!(self, S::KIND, base, offset, |src| {
             self.push(load(kind, conv_reg(dst), src));
         });
     }
 
     #[cfg_attr(not(debug_assertions), inline(always))]
-    fn clear_reg(&mut self, reg: Reg) {
+    fn clear_reg(&mut self, reg: RawReg) {
         let reg = conv_reg(reg);
         self.push(xor((RegSize::R32, reg, reg)));
     }
 
     #[cfg_attr(not(debug_assertions), inline(always))]
-    fn fill_with_ones(&mut self, reg: Reg) {
+    fn fill_with_ones(&mut self, reg: RawReg) {
         match self.reg_size() {
             RegSize::R32 => {
                 self.push(mov_imm(conv_reg(reg), imm32(0xffffffff)));
@@ -213,125 +265,121 @@ impl<'r, 'a, S> ArchVisitor<'r, 'a, S> where S: Sandbox {
     }
 
     #[cfg_attr(not(debug_assertions), inline(always))]
-    fn compare_reg_reg(&mut self, d: Reg, s1: Reg, s2: Reg, condition: Condition) {
+    fn compare_reg_reg(&mut self, d: RawReg, s1: RawReg, s2: RawReg, condition: Condition) {
+        let reg_size = self.reg_size();
+        let s1 = conv_reg(s1);
+        let s2 = conv_reg(s2);
+        let d = conv_reg(d);
+        let asm = self.asm.reserve::<U3>();
         if d == s1 || d == s2 {
-            self.push(cmp((self.reg_size(), conv_reg(s1), conv_reg(s2))));
-            self.push(setcc(condition, conv_reg(d)));
-            self.push(and((conv_reg(d), imm32(1))));
+            let asm = asm.push(cmp((reg_size, s1, s2)));
+            let asm = asm.push(setcc(condition, d));
+            asm.push(and((d, imm32(1))));
         } else {
-            self.clear_reg(d);
-            self.push(cmp((self.reg_size(), conv_reg(s1), conv_reg(s2))));
-            self.push(setcc(condition, conv_reg(d)));
+            let asm = asm.push(xor((RegSize::R32, d, d)));
+            let asm = asm.push(cmp((reg_size, s1, s2)));
+            asm.push(setcc(condition, d));
         }
     }
 
     #[cfg_attr(not(debug_assertions), inline(always))]
-    fn compare_reg_imm(&mut self, d: Reg, s1: Reg, s2: u32, condition: Condition) {
-        if d != s1 {
-            self.clear_reg(d);
-        }
+    fn compare_reg_imm(&mut self, d: RawReg, s1: RawReg, s2: u32, condition: Condition) {
+        let reg_size = self.reg_size();
+        let d = conv_reg(d);
+        let s1 = conv_reg(s1);
+
+        let asm = self.asm.reserve::<U4>();
+        let asm = if d != s1 {
+            asm.push(xor((RegSize::R32, d, d)))
+        } else {
+            asm.push_none()
+        };
 
-        if condition == Condition::Below && s2 == 1 {
+        let asm = if condition == Condition::Below && s2 == 1 {
             // d = s1 <u 1  =>  d = s1 == 0
-            self.push(test((self.reg_size(), conv_reg(s1), conv_reg(s1))));
-            self.push(setcc(Condition::Equal, conv_reg(d)));
+            let asm = asm.push(test((reg_size, s1, s1)));
+            asm.push(setcc(Condition::Equal, d))
         } else if condition == Condition::Above && s2 == 0 {
             // d = s1 >u 0  =>  d = s1 != 0
-            self.push(test((self.reg_size(), conv_reg(s1), conv_reg(s1))));
-            self.push(setcc(Condition::NotEqual, conv_reg(d)));
+            let asm = asm.push(test((reg_size, s1, s1)));
+            asm.push(setcc(Condition::NotEqual, d))
         } else {
-            match self.reg_size() {
+            let asm = match reg_size {
                 RegSize::R32 => {
-                    self.push(cmp((conv_reg(s1), imm32(s2))));
+                    asm.push(cmp((s1, imm32(s2))))
                 },
                 RegSize::R64 => {
-                    self.push(cmp((conv_reg(s1), imm64(s2 as i32))));
+                    asm.push(cmp((s1, imm64(s2 as i32))))
                 }
-            }
-            self.push(setcc(condition, conv_reg(d)));
-        }
+            };
+            asm.push(setcc(condition, d))
+        };
 
-        if d == s1 {
-            self.push(and((conv_reg(d), imm32(1))));
-        }
+        let asm = asm.push_if(d == s1, and((d, imm32(1))));
+        asm.assert_reserved_exactly_as_needed();
     }
 
     #[cfg_attr(not(debug_assertions), inline(always))]
-    fn shift_imm(&mut self, d: Reg, s1: Reg, s2: u32, kind: ShiftKind) {
+    fn shift_imm(&mut self, d: RawReg, s1: RawReg, s2: u32, kind: ShiftKind) {
+        let reg_size = self.reg_size();
+        let d = conv_reg(d);
+        let s1 = conv_reg(s1);
+        let asm = self.asm.reserve::<polkavm_assembler::U2>();
+
         if s2 >= 32 {
             // d = s << 32+
-            self.clear_reg(d);
+            asm.push(xor((RegSize::R32, d, d)));
             return;
         }
 
-        if d != s1 {
-            self.mov(d, s1);
-        }
+        let asm = asm.push_if(d != s1, mov(reg_size, d, s1));
 
         // d = d << s2
-        match kind {
-            ShiftKind::LogicalLeft => self.push(shl_imm(self.reg_size(), conv_reg(d), s2 as u8)),
-            ShiftKind::LogicalRight => self.push(shr_imm(self.reg_size(), conv_reg(d), s2 as u8)),
-            ShiftKind::ArithmeticRight => self.push(sar_imm(self.reg_size(), conv_reg(d), s2 as u8)),
-        }
+        let asm = match kind {
+            ShiftKind::LogicalLeft => asm.push(shl_imm(reg_size, d, s2 as u8)),
+            ShiftKind::LogicalRight => asm.push(shr_imm(reg_size, d, s2 as u8)),
+            ShiftKind::ArithmeticRight => asm.push(sar_imm(reg_size, d, s2 as u8)),
+        };
+
+        asm.assert_reserved_exactly_as_needed();
     }
 
     #[cfg_attr(not(debug_assertions), inline(always))]
-    fn shift(&mut self, d: Reg, s1: impl Into<RegImm>, s2: Reg, kind: ShiftKind) {
-        // TODO: Consider using shlx/shrx/sarx when BMI2 is available.
-        self.push(mov(self.reg_size(), rcx, conv_reg(s2)));
+    fn shift(&mut self, d: RawReg, s1: impl Into<RegImm>, s2: RawReg, kind: ShiftKind) {
+        let reg_size = self.reg_size();
+        let d = conv_reg(d);
+        let s2 = conv_reg(s2);
+        let asm = self.asm.reserve::<polkavm_assembler::U3>();
 
-        match s1.into() {
+        // TODO: Consider using shlx/shrx/sarx when BMI2 is available.
+        let asm = asm.push(mov(reg_size, rcx, s2));
+        let asm = match s1.into() {
             RegImm::Reg(s1) => {
-                if s1 != d {
-                    self.mov(d, s1);
-                }
+                let s1 = conv_reg(s1);
+                asm.push_if(d != s1, mov(reg_size, d, s1))
             },
             RegImm::Imm(s1) => {
-                self.load_immediate(d, s1);
+                asm.push(mov_imm(d, imm32(s1)))
             }
-        }
+        };
 
         // d = d << s2
         match kind {
-            ShiftKind::LogicalLeft => self.push(shl_cl(self.reg_size(), conv_reg(d))),
-            ShiftKind::LogicalRight => self.push(shr_cl(self.reg_size(), conv_reg(d))),
-            ShiftKind::ArithmeticRight => self.push(sar_cl(self.reg_size(), conv_reg(d))),
-        }
+            ShiftKind::LogicalLeft => asm.push(shl_cl(reg_size, d)),
+            ShiftKind::LogicalRight => asm.push(shr_cl(reg_size, d)),
+            ShiftKind::ArithmeticRight => asm.push(sar_cl(reg_size, d)),
+        }.assert_reserved_exactly_as_needed()
     }
 
     #[cfg_attr(not(debug_assertions), inline(always))]
-    fn mov(&mut self, dst: Reg, src: Reg) {
+    fn mov(&mut self, dst: RawReg, src: RawReg) {
         self.push(mov(self.reg_size(), conv_reg(dst), conv_reg(src)))
     }
 
-    #[cfg_attr(not(debug_assertions), inline(always))]
-    fn calculate_label_offset(&self, rel8_len: usize, rel32_len: usize, offset: isize) -> Result<i8, i32> {
-        let offset_near = offset - (self.asm.len() as isize + rel8_len as isize);
-        if offset_near <= i8::MAX as isize && offset_near >= i8::MIN as isize {
-            Ok(offset_near as i8)
-        } else {
-            let offset = offset - (self.asm.len() as isize + rel32_len as isize);
-            Err(offset as i32)
-        }
-    }
-
     #[cfg_attr(not(debug_assertions), inline(always))]
     fn jump_to_label(&mut self, label: Label) {
-        if let Some(offset) = self.asm.get_label_origin_offset(label) {
-            let offset = self.calculate_label_offset(
-                jmp_rel8(i8::MAX).len(),
-                jmp_rel32(i32::MAX).len(),
-                offset
-            );
-
-            match offset {
-                Ok(offset) => self.push(jmp_rel8(offset)),
-                Err(offset) => self.push(jmp_rel32(offset))
-            }
-        } else {
-            self.push(jmp_label32(label));
-        }
+        let asm = self.asm.reserve::<U1>();
+        jump_to_label(asm, label).assert_reserved_exactly_as_needed();
     }
 
     #[cfg_attr(not(debug_assertions), inline(always))]
@@ -341,58 +389,56 @@ impl<'r, 'a, S> ArchVisitor<'r, 'a, S> where S: Sandbox {
 
     #[cfg_attr(not(debug_assertions), inline(always))]
     fn branch_to_label(&mut self, condition: Condition, label: Label) {
-        if let Some(offset) = self.asm.get_label_origin_offset(label) {
-            let offset = self.calculate_label_offset(
-                jcc_rel8(condition, i8::MAX).len(),
-                jcc_rel32(condition, i32::MAX).len(),
-                offset
-            );
-
-            match offset {
-                Ok(offset) => self.push(jcc_rel8(condition, offset)),
-                Err(offset) => self.push(jcc_rel32(condition, offset))
-            }
-        } else {
-            self.push(jcc_label32(condition, label));
-        }
+        let asm = self.asm.reserve::<U1>();
+        branch_to_label(asm, condition, label).assert_reserved_exactly_as_needed();
     }
 
     #[cfg_attr(not(debug_assertions), inline(always))]
-    fn branch(&mut self, s1: Reg, s2: impl Into<RegImm>, target: u32, condition: Condition) {
-        match s2.into() {
-            RegImm::Reg(s2) => self.push(cmp((self.reg_size(), conv_reg(s1), conv_reg(s2)))),
-            RegImm::Imm(s2) => self.push(cmp((conv_reg(s1), imm32(s2)))),
-        }
+    fn branch(&mut self, s1: RawReg, s2: impl Into<RegImm>, target: u32, condition: Condition) {
+        let reg_size = self.reg_size();
+        let s1 = conv_reg(s1);
+        let label = self.get_or_forward_declare_label(target).unwrap_or(self.trap_label);
+
+        let asm = self.asm.reserve::<U2>();
+        let asm = match s2.into() {
+            RegImm::Reg(s2) => asm.push(cmp((reg_size, s1, conv_reg(s2)))),
+            RegImm::Imm(s2) => asm.push(cmp((s1, imm32(s2)))),
+        };
 
-        let label = self.get_or_forward_declare_label(target);
-        self.branch_to_label(condition, label);
+        branch_to_label(asm, condition, label).assert_reserved_exactly_as_needed();
     }
 
     #[cfg_attr(not(debug_assertions), inline(always))]
-    fn cmov(&mut self, d: Reg, s: Reg, c: Reg, condition: Condition) {
+    fn cmov(&mut self, d: RawReg, s: RawReg, c: RawReg, condition: Condition) {
         if d == s {
             return;
         }
 
+        let reg_size = self.reg_size();
         let d = conv_reg(d);
         let s = conv_reg(s);
         let c = conv_reg(c);
 
-        self.push(test((self.reg_size(), c, c)));
-        self.push(cmov(condition, self.reg_size(), d, s));
+        let asm = self.asm.reserve::<U2>();
+        let asm = asm.push(test((reg_size, c, c)));
+        let asm = asm.push(cmov(condition, reg_size, d, s));
+        asm.assert_reserved_exactly_as_needed();
     }
 
     #[cfg_attr(not(debug_assertions), inline(always))]
-    fn cmov_imm(&mut self, d: Reg, s: u32, c: Reg, condition: Condition) {
+    fn cmov_imm(&mut self, d: RawReg, s: u32, c: RawReg, condition: Condition) {
+        let reg_size = self.reg_size();
         let d = conv_reg(d);
         let c = conv_reg(c);
 
-        self.push(test((self.reg_size(), c, c)));
-        self.push(mov_imm(TMP_REG, imm32(s)));
-        self.push(cmov(condition, self.reg_size(), d, TMP_REG));
+        let asm = self.asm.reserve::<U3>();
+        let asm = asm.push(test((reg_size, c, c)));
+        let asm = asm.push(mov_imm(TMP_REG, imm32(s)));
+        let asm = asm.push(cmov(condition, reg_size, d, TMP_REG));
+        asm.assert_reserved_exactly_as_needed();
     }
 
-    fn div_rem(&mut self, d: Reg, s1: Reg, s2: Reg, div_rem: DivRem, kind: Signedness) {
+    fn div_rem(&mut self, d: RawReg, s1: RawReg, s2: RawReg, div_rem: DivRem, kind: Signedness) {
         // Unlike most other architectures RISC-V doesn't trap on division by zero
         // nor on division with overflow, and has well defined results in such cases.
 
@@ -505,14 +551,14 @@ impl<'r, 'a, S> ArchVisitor<'r, 'a, S> where S: Sandbox {
     fn save_registers_to_vmctx(&mut self) {
         let regs_base = self.load_vmctx_field_address(S::vmctx_regs_offset());
         for (nth, reg) in Reg::ALL.iter().copied().enumerate() {
-            self.push(store(Size::U32, reg_indirect(RegSize::R64, regs_base + nth as i32 * 4), conv_reg(reg)));
+            self.push(store(Size::U32, reg_indirect(RegSize::R64, regs_base + nth as i32 * 4), conv_reg(reg.into())));
         }
     }
 
     fn restore_registers_from_vmctx(&mut self) {
         let regs_base = self.load_vmctx_field_address(S::vmctx_regs_offset());
         for (nth, reg) in Reg::ALL.iter().copied().enumerate() {
-            self.push(load(LoadKind::U32, conv_reg(reg), reg_indirect(RegSize::R64, regs_base + nth as i32 * 4)));
+            self.push(load(LoadKind::U32, conv_reg(reg.into()), reg_indirect(RegSize::R64, regs_base + nth as i32 * 4)));
         }
     }
 
@@ -534,7 +580,7 @@ impl<'r, 'a, S> ArchVisitor<'r, 'a, S> where S: Sandbox {
                 self.branch_to_label(Condition::Sign, self.trap_label);
             }
 
-            let target_label = self.get_or_forward_declare_label(export.target_code_offset());
+            let target_label = self.get_or_forward_declare_label(export.target_code_offset()).unwrap_or(self.trap_label);
             self.jump_to_label(target_label);
         }
     }
@@ -608,8 +654,7 @@ impl<'r, 'a, S> ArchVisitor<'r, 'a, S> where S: Sandbox {
 
     #[inline(never)]
     #[cold]
-    pub(crate) fn trace_execution(&mut self) {
-        let code_offset = self.current_code_offset;
+    pub(crate) fn trace_execution(&mut self, code_offset: u32) {
         self.push(mov_imm(TMP_REG, imm32(code_offset)));
         self.call_to_label(self.trace_label);
     }
@@ -629,23 +674,27 @@ impl<'r, 'a, S> ArchVisitor<'r, 'a, S> where S: Sandbox {
     }
 
     #[cfg_attr(not(debug_assertions), inline(always))]
-    fn jump_indirect_impl(&mut self, load_imm: Option<(Reg, u32)>, base: Reg, offset: u32) {
+    fn jump_indirect_impl(&mut self, load_imm: Option<(RawReg, u32)>, base: RawReg, offset: u32) {
         match S::KIND {
             SandboxKind::Linux => {
                 use polkavm_assembler::amd64::{SegReg, Scale};
 
-                let target = if offset != 0 || load_imm.map_or(false, |(t, _)| t == base) {
-                    self.push(lea(RegSize::R32, TMP_REG, reg_indirect(RegSize::R32, conv_reg(base) + offset as i32)));
-                    TMP_REG
+                let asm = self.asm.reserve::<U3>();
+                let (asm, target) = if offset != 0 || load_imm.map_or(false, |(t, _)| t == base) {
+                    let asm = asm.push(lea(RegSize::R32, TMP_REG, reg_indirect(RegSize::R32, conv_reg(base) + offset as i32)));
+                    (asm, TMP_REG)
                 } else {
-                    conv_reg(base)
+                    (asm.push_none(), conv_reg(base))
                 };
 
-                if let Some((return_register, return_address)) = load_imm {
-                    self.load_immediate(return_register, return_address);
-                }
+                let asm = if let Some((return_register, return_address)) = load_imm {
+                    asm.push(mov_imm(conv_reg(return_register), imm32(return_address)))
+                } else {
+                    asm.push_none()
+                };
 
-                self.asm.push(jmp(MemOp::IndexScaleOffset(Some(SegReg::gs), RegSize::R64, target, Scale::x8, 0)));
+                let asm = asm.push(jmp(MemOp::IndexScaleOffset(Some(SegReg::gs), RegSize::R64, target, Scale::x8, 0)));
+                asm.assert_reserved_exactly_as_needed();
             },
             SandboxKind::Generic => {
                 // TODO: This also could be more efficient.
@@ -680,8 +729,8 @@ impl<'r, 'a, S> InstructionVisitor for ArchVisitor<'r, 'a, S> where S: Sandbox {
 
     #[inline(always)]
     #[cold]
-    fn invalid(&mut self, opcode: u8) -> Self::ReturnTy {
-        log::debug!("Encountered invalid instruction: opcode = {opcode}");
+    fn invalid(&mut self) -> Self::ReturnTy {
+        log::debug!("Encountered invalid instruction");
         self.trap()
     }
 
@@ -690,7 +739,7 @@ impl<'r, 'a, S> InstructionVisitor for ArchVisitor<'r, 'a, S> where S: Sandbox {
     }
 
     #[inline(always)]
-    fn sbrk(&mut self, dst: Reg, size: Reg) -> Self::ReturnTy {
+    fn sbrk(&mut self, dst: RawReg, size: RawReg) -> Self::ReturnTy {
         let label_bump_only = self.asm.forward_declare_label();
         let label_continue = self.asm.forward_declare_label();
         let sbrk_label = self.sbrk_label;
@@ -728,242 +777,285 @@ impl<'r, 'a, S> InstructionVisitor for ArchVisitor<'r, 'a, S> where S: Sandbox {
 
     #[inline(always)]
     fn ecalli(&mut self, imm: u32) -> Self::ReturnTy {
-        self.push(mov_imm(TMP_REG, imm32(imm)));
-        self.call_to_label(self.ecall_label);
+        let ecall_label = self.ecall_label;
+        let asm = self.asm.reserve::<U2>();
+        let asm = asm.push(mov_imm(TMP_REG, imm32(imm)));
+        let asm = asm.push(call_label32(ecall_label));
+        asm.assert_reserved_exactly_as_needed();
     }
 
     #[inline(always)]
-    fn set_less_than_unsigned(&mut self, d: Reg, s1: Reg, s2: Reg) -> Self::ReturnTy {
+    fn set_less_than_unsigned(&mut self, d: RawReg, s1: RawReg, s2: RawReg) -> Self::ReturnTy {
         self.compare_reg_reg(d, s1, s2, Condition::Below);
     }
 
     #[inline(always)]
-    fn set_less_than_unsigned_imm(&mut self, d: Reg, s1: Reg, s2: u32) -> Self::ReturnTy {
+    fn set_less_than_unsigned_imm(&mut self, d: RawReg, s1: RawReg, s2: u32) -> Self::ReturnTy {
         self.compare_reg_imm(d, s1, s2, Condition::Below);
     }
 
     #[inline(always)]
-    fn set_greater_than_unsigned_imm(&mut self, d: Reg, s1: Reg, s2: u32) -> Self::ReturnTy {
+    fn set_greater_than_unsigned_imm(&mut self, d: RawReg, s1: RawReg, s2: u32) -> Self::ReturnTy {
         self.compare_reg_imm(d, s1, s2, Condition::Above);
     }
 
     #[inline(always)]
-    fn set_less_than_signed(&mut self, d: Reg, s1: Reg, s2: Reg) -> Self::ReturnTy {
+    fn set_less_than_signed(&mut self, d: RawReg, s1: RawReg, s2: RawReg) -> Self::ReturnTy {
         self.compare_reg_reg(d, s1, s2, Condition::Less);
     }
 
     #[inline(always)]
-    fn set_less_than_signed_imm(&mut self, d: Reg, s1: Reg, s2: u32) -> Self::ReturnTy {
+    fn set_less_than_signed_imm(&mut self, d: RawReg, s1: RawReg, s2: u32) -> Self::ReturnTy {
         self.compare_reg_imm(d, s1, s2, Condition::Less);
     }
 
     #[inline(always)]
-    fn set_greater_than_signed_imm(&mut self, d: Reg, s1: Reg, s2: u32) -> Self::ReturnTy {
+    fn set_greater_than_signed_imm(&mut self, d: RawReg, s1: RawReg, s2: u32) -> Self::ReturnTy {
         self.compare_reg_imm(d, s1, s2, Condition::Greater);
     }
 
     #[inline(always)]
-    fn shift_logical_right(&mut self, d: Reg, s1: Reg, s2: Reg) -> Self::ReturnTy {
+    fn shift_logical_right(&mut self, d: RawReg, s1: RawReg, s2: RawReg) -> Self::ReturnTy {
         self.shift(d, s1, s2, ShiftKind::LogicalRight);
     }
 
     #[inline(always)]
-    fn shift_arithmetic_right(&mut self, d: Reg, s1: Reg, s2: Reg) -> Self::ReturnTy {
+    fn shift_arithmetic_right(&mut self, d: RawReg, s1: RawReg, s2: RawReg) -> Self::ReturnTy {
         self.shift(d, s1, s2, ShiftKind::ArithmeticRight);
     }
 
     #[inline(always)]
-    fn shift_logical_left(&mut self, d: Reg, s1: Reg, s2: Reg) -> Self::ReturnTy {
+    fn shift_logical_left(&mut self, d: RawReg, s1: RawReg, s2: RawReg) -> Self::ReturnTy {
         self.shift(d, s1, s2, ShiftKind::LogicalLeft);
     }
 
     #[inline(always)]
-    fn shift_logical_right_imm_alt(&mut self, d: Reg, s2: Reg, s1: u32) -> Self::ReturnTy {
+    fn shift_logical_right_imm_alt(&mut self, d: RawReg, s2: RawReg, s1: u32) -> Self::ReturnTy {
         self.shift(d, s1, s2, ShiftKind::LogicalRight);
     }
 
     #[inline(always)]
-    fn shift_arithmetic_right_imm_alt(&mut self, d: Reg, s2: Reg, s1: u32) -> Self::ReturnTy  {
+    fn shift_arithmetic_right_imm_alt(&mut self, d: RawReg, s2: RawReg, s1: u32) -> Self::ReturnTy  {
         self.shift(d, s1, s2, ShiftKind::ArithmeticRight);
     }
 
     #[inline(always)]
-    fn shift_logical_left_imm_alt(&mut self, d: Reg, s2: Reg, s1: u32) -> Self::ReturnTy  {
+    fn shift_logical_left_imm_alt(&mut self, d: RawReg, s2: RawReg, s1: u32) -> Self::ReturnTy  {
         self.shift(d, s1, s2, ShiftKind::LogicalLeft);
     }
 
     #[inline(always)]
-    fn xor(&mut self, d: Reg, s1: Reg, s2: Reg) -> Self::ReturnTy {
+    fn xor(&mut self, d: RawReg, s1: RawReg, s2: RawReg) -> Self::ReturnTy {
         let reg_size = self.reg_size();
+        let d = conv_reg(d);
+        let s1 = conv_reg(s1);
+        let s2 = conv_reg(s2);
+
+        let asm = self.asm.reserve::<U2>();
         match (d, s1, s2) {
             // d = d ^ s2
-            (_, _, _) if d == s1 => self.push(xor((reg_size, conv_reg(d), conv_reg(s2)))),
+            (_, _, _) if d == s1 => asm.push(xor((reg_size, d, s2))).push_none(),
             // d = s1 ^ d
-            (_, _, _) if d == s2 => self.push(xor((reg_size, conv_reg(d), conv_reg(s1)))),
+            (_, _, _) if d == s2 => asm.push(xor((reg_size, d, s1))).push_none(),
             // d = s1 ^ s2
             _ => {
-                self.mov(d, s1);
-                self.push(xor((reg_size, conv_reg(d), conv_reg(s2))));
+                let asm = asm.push(mov(reg_size, d, s1));
+                asm.push(xor((reg_size, d, s2)))
             }
-        }
+        }.assert_reserved_exactly_as_needed();
     }
 
     #[inline(always)]
-    fn and(&mut self, d: Reg, s1: Reg, s2: Reg) -> Self::ReturnTy {
+    fn and(&mut self, d: RawReg, s1: RawReg, s2: RawReg) -> Self::ReturnTy {
         let reg_size = self.reg_size();
+        let d = conv_reg(d);
+        let s1 = conv_reg(s1);
+        let s2 = conv_reg(s2);
+
+        let asm = self.asm.reserve::<U2>();
         match (d, s1, s2) {
             // d = d & s2
-            (_, _, _) if d == s1 => self.push(and((reg_size, conv_reg(d), conv_reg(s2)))),
+            (_, _, _) if d == s1 => asm.push(and((reg_size, d, s2))).push_none(),
             // d = s1 & d
-            (_, _, _) if d == s2 => self.push(and((reg_size, conv_reg(d), conv_reg(s1)))),
+            (_, _, _) if d == s2 => asm.push(and((reg_size, d, s1))).push_none(),
             // d = s1 & s2
             _ => {
-                self.mov(d, s1);
-                self.push(and((reg_size, conv_reg(d), conv_reg(s2))));
+                let asm = asm.push(mov(reg_size, d, s1));
+                asm.push(and((reg_size, d, s2)))
             }
-        }
+        }.assert_reserved_exactly_as_needed();
     }
 
     #[inline(always)]
-    fn or(&mut self, d: Reg, s1: Reg, s2: Reg) -> Self::ReturnTy {
+    fn or(&mut self, d: RawReg, s1: RawReg, s2: RawReg) -> Self::ReturnTy {
         let reg_size = self.reg_size();
+        let d = conv_reg(d);
+        let s1 = conv_reg(s1);
+        let s2 = conv_reg(s2);
+
+        let asm = self.asm.reserve::<U2>();
         match (d, s1, s2) {
             // d = d | s2
-            (_, _, _) if d == s1 => self.push(or((reg_size, conv_reg(d), conv_reg(s2)))),
+            (_, _, _) if d == s1 => asm.push(or((reg_size, d, s2))).push_none(),
             // d = s1 | d
-            (_, _, _) if d == s2 => self.push(or((reg_size, conv_reg(d), conv_reg(s1)))),
+            (_, _, _) if d == s2 => asm.push(or((reg_size, d, s1))).push_none(),
             // d = s1 | s2
             _ => {
-                self.mov(d, s1);
-                self.push(or((reg_size, conv_reg(d), conv_reg(s2))));
+                let asm = asm.push(mov(reg_size, d, s1));
+                asm.push(or((reg_size, d, s2)))
             }
-        }
+        }.assert_reserved_exactly_as_needed();
     }
 
     #[inline(always)]
-    fn add(&mut self, d: Reg, s1: Reg, s2: Reg) -> Self::ReturnTy {
+    fn add(&mut self, d: RawReg, s1: RawReg, s2: RawReg) -> Self::ReturnTy {
         let reg_size = self.reg_size();
+        let d = conv_reg(d);
+        let s1 = conv_reg(s1);
+        let s2 = conv_reg(s2);
+
+        let asm = self.asm.reserve::<U2>();
         match (d, s1, s2) {
             // d = d + s2
-            (_, _, _) if d == s1 => self.push(add((reg_size, conv_reg(d), conv_reg(s2)))),
+            (_, _, _) if d == s1 => asm.push(add((reg_size, d, s2))).push_none(),
             // d = s1 + d
-            (_, _, _) if d == s2 => self.push(add((reg_size, conv_reg(d), conv_reg(s1)))),
+            (_, _, _) if d == s2 => asm.push(add((reg_size, d, s1))).push_none(),
             // d = s1 + s2
             _ => {
-                if d != s1 {
-                    self.mov(d, s1);
-                }
-                self.push(add((reg_size, conv_reg(d), conv_reg(s2))));
+                let asm = asm.push_if(d != s1, mov(reg_size, d, s1));
+                asm.push(add((reg_size, d, s2)))
             }
-        }
+        }.assert_reserved_exactly_as_needed();
     }
 
     #[inline(always)]
-    fn sub(&mut self, d: Reg, s1: Reg, s2: Reg) -> Self::ReturnTy {
+    fn sub(&mut self, d: RawReg, s1: RawReg, s2: RawReg) -> Self::ReturnTy {
         let reg_size = self.reg_size();
+        let d = conv_reg(d);
+        let s1 = conv_reg(s1);
+        let s2 = conv_reg(s2);
+
+        let asm = self.asm.reserve::<U2>();
         match (d, s1, s2) {
             // d = d - s2
-            (_, _, _) if d == s1 => self.push(sub((reg_size, conv_reg(d), conv_reg(s2)))),
+            (_, _, _) if d == s1 => asm.push(sub((reg_size, d, s2))).push_none(),
             // d = s1 - d
             (_, _, _) if d == s2 => {
-                self.push(neg(reg_size, conv_reg(d)));
-                self.push(add((reg_size, conv_reg(d), conv_reg(s1))));
+                let asm = asm.push(neg(reg_size, d));
+                asm.push(add((reg_size, d, s1)))
             }
             // d = s1 - s2
             _ => {
-                self.mov(d, s1);
-                self.push(sub((reg_size, conv_reg(d), conv_reg(s2))));
+                let asm = asm.push(mov(reg_size, d, s1));
+                asm.push(sub((reg_size, d, s2)))
             }
-        }
+        }.assert_reserved_exactly_as_needed();
     }
 
     #[inline(always)]
-    fn negate_and_add_imm(&mut self, d: Reg, s1: Reg, s2: u32) -> Self::ReturnTy {
+    fn negate_and_add_imm(&mut self, d: RawReg, s1: RawReg, s2: u32) -> Self::ReturnTy {
+        let reg_size = self.reg_size();
+        let d = conv_reg(d);
+        let s1 = conv_reg(s1);
+
+        let asm = self.asm.reserve::<U2>();
         if d == s1 {
             // d = -d + s2
-            self.push(neg(RegSize::R32, conv_reg(d)));
+            let asm = asm.push(neg(RegSize::R32, d));
             if s2 != 0 {
-                self.push(add((conv_reg(d), imm32(s2))));
+                asm.push(add((d, imm32(s2))))
+            } else {
+                asm.push_none()
             }
         } else {
             // d = -s1 + s2  =>  d = s2 - s1
             if s2 == 0 {
-                self.mov(d, s1);
-                self.push(neg(RegSize::R32, conv_reg(d)));
+                let asm = asm.push(mov(reg_size, d, s1));
+                asm.push(neg(RegSize::R32, d))
             } else {
-                self.push(mov_imm(conv_reg(d), imm32(s2)));
-                self.push(sub((RegSize::R32, conv_reg(d), conv_reg(s1))));
+                let asm = asm.push(mov_imm(d, imm32(s2)));
+                asm.push(sub((RegSize::R32, d, s1)))
             }
-        }
+        }.assert_reserved_exactly_as_needed();
     }
 
     #[inline(always)]
-    fn mul(&mut self, d: Reg, s1: Reg, s2: Reg) -> Self::ReturnTy {
+    fn mul(&mut self, d: RawReg, s1: RawReg, s2: RawReg) -> Self::ReturnTy {
         let reg_size = self.reg_size();
+        let d = conv_reg(d);
+        let s1 = conv_reg(s1);
+        let s2 = conv_reg(s2);
+
+        let asm = self.asm.reserve::<U2>();
         if d == s1 {
             // d = d * s2
-            self.push(imul(reg_size, conv_reg(d), conv_reg(s2)))
+            asm.push(imul(reg_size, d, s2)).push_none()
         } else if d == s2 {
             // d = s1 * d
-            self.push(imul(reg_size, conv_reg(d), conv_reg(s1)))
+            asm.push(imul(reg_size, d, s1)).push_none()
         } else {
             // d = s1 * s2
-            self.mov(d, s1);
-            self.push(imul(reg_size, conv_reg(d), conv_reg(s2)));
-        }
+            let asm = asm.push(mov(reg_size, d, s1));
+            asm.push(imul(reg_size, d, s2))
+        }.assert_reserved_exactly_as_needed();
     }
 
     #[inline(always)]
-    fn mul_imm(&mut self, d: Reg, s1: Reg, s2: u32) -> Self::ReturnTy {
+    fn mul_imm(&mut self, d: RawReg, s1: RawReg, s2: u32) -> Self::ReturnTy {
         self.push(imul_imm(RegSize::R32, conv_reg(d), conv_reg(s1), s2 as i32));
     }
 
     #[inline(always)]
-    fn mul_upper_signed_signed(&mut self, d: Reg, s1: Reg, s2: Reg) -> Self::ReturnTy {
-        self.push(movsxd_32_to_64(TMP_REG, conv_reg(s2)));
-        self.push(movsxd_32_to_64(conv_reg(d), conv_reg(s1)));
-        self.push(imul(RegSize::R64, conv_reg(d), TMP_REG));
-        self.push(shr_imm(RegSize::R64, conv_reg(d), 32));
+    fn mul_upper_signed_signed(&mut self, d: RawReg, s1: RawReg, s2: RawReg) -> Self::ReturnTy {
+        let asm = self.asm.reserve::<U4>();
+        let asm = asm.push(movsxd_32_to_64(TMP_REG, conv_reg(s2)));
+        let asm = asm.push(movsxd_32_to_64(conv_reg(d), conv_reg(s1)));
+        let asm = asm.push(imul(RegSize::R64, conv_reg(d), TMP_REG));
+        let asm = asm.push(shr_imm(RegSize::R64, conv_reg(d), 32));
+        asm.assert_reserved_exactly_as_needed();
     }
 
     #[inline(always)]
-    fn mul_upper_signed_signed_imm(&mut self, d: Reg, s1: Reg, s2: u32) -> Self::ReturnTy {
-        self.push(mov_imm(TMP_REG, imm64(s2 as i32)));
-        self.push(movsxd_32_to_64(conv_reg(d), conv_reg(s1)));
-        self.push(imul(RegSize::R64, conv_reg(d), TMP_REG));
-        self.push(shr_imm(RegSize::R64, conv_reg(d), 32));
+    fn mul_upper_signed_signed_imm(&mut self, d: RawReg, s1: RawReg, s2: u32) -> Self::ReturnTy {
+        let asm = self.asm.reserve::<U4>();
+        let asm = asm.push(mov_imm(TMP_REG, imm64(s2 as i32)));
+        let asm = asm.push(movsxd_32_to_64(conv_reg(d), conv_reg(s1)));
+        let asm = asm.push(imul(RegSize::R64, conv_reg(d), TMP_REG));
+        let asm = asm.push(shr_imm(RegSize::R64, conv_reg(d), 32));
+        asm.assert_reserved_exactly_as_needed();
     }
 
     #[inline(always)]
-    fn mul_upper_unsigned_unsigned(&mut self, d: Reg, s1: Reg, s2: Reg) -> Self::ReturnTy {
-        if d == s1 {
+    fn mul_upper_unsigned_unsigned(&mut self, d: RawReg, s1: RawReg, s2: RawReg) -> Self::ReturnTy {
+        let asm = self.asm.reserve::<U3>();
+        let asm = if d == s1 {
             // d = d * s2
-            self.push(imul(RegSize::R64, conv_reg(d), conv_reg(s2)));
+            asm.push(imul(RegSize::R64, conv_reg(d), conv_reg(s2))).push_none()
         } else if d == s2 {
             // d = s1 * d
-            self.push(imul(RegSize::R64, conv_reg(d), conv_reg(s1)));
+            asm.push(imul(RegSize::R64, conv_reg(d), conv_reg(s1))).push_none()
         } else {
             // d = s1 * s2
-            self.push(mov(RegSize::R32, conv_reg(d), conv_reg(s1)));
-            self.push(imul(RegSize::R64, conv_reg(d), conv_reg(s2)));
-        }
+            let asm = asm.push(mov(RegSize::R32, conv_reg(d), conv_reg(s1)));
+            asm.push(imul(RegSize::R64, conv_reg(d), conv_reg(s2)))
+        };
 
-        self.push(shr_imm(RegSize::R64, conv_reg(d), 32));
+        let asm = asm.push(shr_imm(RegSize::R64, conv_reg(d), 32));
+        asm.assert_reserved_exactly_as_needed();
     }
 
     #[inline(always)]
-    fn mul_upper_unsigned_unsigned_imm(&mut self, d: Reg, s1: Reg, s2: u32) -> Self::ReturnTy {
-        self.push(mov_imm(TMP_REG, imm32(s2)));
-        if d != s1 {
-            self.push(mov(RegSize::R32, conv_reg(d), conv_reg(s1)));
-        }
-
-        self.push(imul(RegSize::R64, conv_reg(d), TMP_REG));
-        self.push(shr_imm(RegSize::R64, conv_reg(d), 32));
+    fn mul_upper_unsigned_unsigned_imm(&mut self, d: RawReg, s1: RawReg, s2: u32) -> Self::ReturnTy {
+        let asm = self.asm.reserve::<U4>();
+        let asm = asm.push(mov_imm(TMP_REG, imm32(s2)));
+        let asm = asm.push_if(d != s1, mov(RegSize::R32, conv_reg(d), conv_reg(s1)));
+        let asm = asm.push(imul(RegSize::R64, conv_reg(d), TMP_REG));
+        let asm = asm.push(shr_imm(RegSize::R64, conv_reg(d), 32));
+        asm.assert_reserved_exactly_as_needed();
     }
 
     #[inline(always)]
-    fn mul_upper_signed_unsigned(&mut self, d: Reg, s1: Reg, s2: Reg) -> Self::ReturnTy {
+    fn mul_upper_signed_unsigned(&mut self, d: RawReg, s1: RawReg, s2: RawReg) -> Self::ReturnTy {
         // This instruction is equivalent to:
         //   let s1: i32;
         //   let s2: u32;
@@ -977,348 +1069,363 @@ impl<'r, 'a, S> InstructionVisitor for ArchVisitor<'r, 'a, S> where S: Sandbox {
         //   3) multiply,
         //   4) return the upper 32-bits.
 
-        if d == s2 {
+        let asm = self.asm.reserve::<U4>();
+        let asm = if d == s2 {
             // d = s1 * d
-            self.push(mov(RegSize::R32, TMP_REG, conv_reg(s2)));
-            self.push(movsxd_32_to_64(conv_reg(d), conv_reg(s1)));
-            self.push(imul(RegSize::R64, conv_reg(d), TMP_REG));
+            let asm = asm.push(mov(RegSize::R32, TMP_REG, conv_reg(s2)));
+            let asm = asm.push(movsxd_32_to_64(conv_reg(d), conv_reg(s1)));
+            asm.push(imul(RegSize::R64, conv_reg(d), TMP_REG))
         } else {
             // d = s1 * s2
-            self.push(movsxd_32_to_64(conv_reg(d), conv_reg(s1)));
-            self.push(imul(RegSize::R64, conv_reg(d), conv_reg(s2)));
-        }
+            let asm = asm.push(movsxd_32_to_64(conv_reg(d), conv_reg(s1)));
+            let asm = asm.push(imul(RegSize::R64, conv_reg(d), conv_reg(s2)));
+            asm.push_none()
+        };
 
-        self.push(shr_imm(RegSize::R64, conv_reg(d), 32));
+        let asm = asm.push(shr_imm(RegSize::R64, conv_reg(d), 32));
+        asm.assert_reserved_exactly_as_needed();
     }
 
     #[inline(always)]
-    fn div_unsigned(&mut self, d: Reg, s1: Reg, s2: Reg) -> Self::ReturnTy {
+    fn div_unsigned(&mut self, d: RawReg, s1: RawReg, s2: RawReg) -> Self::ReturnTy {
         self.div_rem(d, s1, s2, DivRem::Div, Signedness::Unsigned);
     }
 
     #[inline(always)]
-    fn div_signed(&mut self, d: Reg, s1: Reg, s2: Reg) -> Self::ReturnTy {
+    fn div_signed(&mut self, d: RawReg, s1: RawReg, s2: RawReg) -> Self::ReturnTy {
         self.div_rem(d, s1, s2, DivRem::Div, Signedness::Signed);
     }
 
     #[inline(always)]
-    fn rem_unsigned(&mut self, d: Reg, s1: Reg, s2: Reg) -> Self::ReturnTy {
+    fn rem_unsigned(&mut self, d: RawReg, s1: RawReg, s2: RawReg) -> Self::ReturnTy {
         self.div_rem(d, s1, s2, DivRem::Rem, Signedness::Unsigned);
     }
 
     #[inline(always)]
-    fn rem_signed(&mut self, d: Reg, s1: Reg, s2: Reg) -> Self::ReturnTy {
+    fn rem_signed(&mut self, d: RawReg, s1: RawReg, s2: RawReg) -> Self::ReturnTy {
         self.div_rem(d, s1, s2, DivRem::Rem, Signedness::Signed);
     }
 
     #[inline(always)]
-    fn shift_logical_right_imm(&mut self, d: Reg, s1: Reg, s2: u32) -> Self::ReturnTy {
+    fn shift_logical_right_imm(&mut self, d: RawReg, s1: RawReg, s2: u32) -> Self::ReturnTy {
         self.shift_imm(d, s1, s2, ShiftKind::LogicalRight);
     }
 
     #[inline(always)]
-    fn shift_arithmetic_right_imm(&mut self, d: Reg, s1: Reg, s2: u32) -> Self::ReturnTy {
+    fn shift_arithmetic_right_imm(&mut self, d: RawReg, s1: RawReg, s2: u32) -> Self::ReturnTy {
         self.shift_imm(d, s1, s2, ShiftKind::ArithmeticRight);
     }
 
     #[inline(always)]
-    fn shift_logical_left_imm(&mut self, d: Reg, s1: Reg, s2: u32) -> Self::ReturnTy {
+    fn shift_logical_left_imm(&mut self, d: RawReg, s1: RawReg, s2: u32) -> Self::ReturnTy {
         self.shift_imm(d, s1, s2, ShiftKind::LogicalLeft);
     }
 
     #[inline(always)]
-    fn or_imm(&mut self, d: Reg, s1: Reg, s2: u32) -> Self::ReturnTy {
-        if d != s1 {
-            self.mov(d, s1);
-        }
+    fn or_imm(&mut self, d: RawReg, s1: RawReg, s2: u32) -> Self::ReturnTy {
+        let reg_size = self.reg_size();
+        let d = conv_reg(d);
+        let s1 = conv_reg(s1);
+
+        let asm = self.asm.reserve::<U2>();
+        let asm = asm.push_if(d != s1, mov(reg_size, d, s1));
 
         // d = s1 | s2
-        self.push(or((conv_reg(d), imm32(s2))));
+        let asm = asm.push(or((d, imm32(s2))));
+        asm.assert_reserved_exactly_as_needed();
     }
 
     #[inline(always)]
-    fn and_imm(&mut self, d: Reg, s1: Reg, s2: u32) -> Self::ReturnTy {
-        if d != s1 {
-            self.mov(d, s1);
-        }
+    fn and_imm(&mut self, d: RawReg, s1: RawReg, s2: u32) -> Self::ReturnTy {
+        let reg_size = self.reg_size();
+        let d = conv_reg(d);
+        let s1 = conv_reg(s1);
+
+        let asm = self.asm.reserve::<U2>();
+        let asm = asm.push_if(d != s1, mov(reg_size, d, s1));
 
         // d = s1 & s2
-        self.push(and((conv_reg(d), imm32(s2))));
+        let asm = asm.push(and((d, imm32(s2))));
+        asm.assert_reserved_exactly_as_needed();
     }
 
     #[inline(always)]
-    fn xor_imm(&mut self, d: Reg, s1: Reg, s2: u32) -> Self::ReturnTy {
+    fn xor_imm(&mut self, d: RawReg, s1: RawReg, s2: u32) -> Self::ReturnTy {
         let reg_size = self.reg_size();
-        if d != s1 {
-            self.mov(d, s1);
-        }
+        let d = conv_reg(d);
+        let s1 = conv_reg(s1);
+        let asm = self.asm.reserve::<U2>();
+        let asm = asm.push_if(d != s1, mov(reg_size, d, s1));
 
         if s2 != !0 {
             // d = s1 ^ s2
-            self.push(xor((conv_reg(d), imm32(s2))));
+            asm.push(xor((d, imm32(s2))))
         } else {
             // d = s1 ^ 0xfffffff
-            self.push(not(reg_size, conv_reg(d)));
-        }
+            asm.push(not(reg_size, d))
+        }.assert_reserved_exactly_as_needed();
     }
 
     #[inline(always)]
-    fn load_imm(&mut self, dst: Reg, s2: u32) -> Self::ReturnTy {
+    fn load_imm(&mut self, dst: RawReg, s2: u32) -> Self::ReturnTy {
         self.load_immediate(dst, s2);
     }
 
     #[inline(always)]
-    fn move_reg(&mut self, d: Reg, s: Reg) -> Self::ReturnTy {
+    fn move_reg(&mut self, d: RawReg, s: RawReg) -> Self::ReturnTy {
         self.mov(d, s);
     }
 
     #[inline(always)]
-    fn cmov_if_zero(&mut self, d: Reg, s: Reg, c: Reg) -> Self::ReturnTy {
+    fn cmov_if_zero(&mut self, d: RawReg, s: RawReg, c: RawReg) -> Self::ReturnTy {
         self.cmov(d, s, c, Condition::Equal);
     }
 
     #[inline(always)]
-    fn cmov_if_not_zero(&mut self, d: Reg, s: Reg, c: Reg) -> Self::ReturnTy {
+    fn cmov_if_not_zero(&mut self, d: RawReg, s: RawReg, c: RawReg) -> Self::ReturnTy {
         self.cmov(d, s, c, Condition::NotEqual);
     }
 
     #[inline(always)]
-    fn cmov_if_zero_imm(&mut self, d: Reg, c: Reg, s: u32) -> Self::ReturnTy {
+    fn cmov_if_zero_imm(&mut self, d: RawReg, c: RawReg, s: u32) -> Self::ReturnTy {
         self.cmov_imm(d, s, c, Condition::Equal);
     }
 
     #[inline(always)]
-    fn cmov_if_not_zero_imm(&mut self, d: Reg, c: Reg, s: u32) -> Self::ReturnTy {
+    fn cmov_if_not_zero_imm(&mut self, d: RawReg, c: RawReg, s: u32) -> Self::ReturnTy {
         self.cmov_imm(d, s, c, Condition::NotEqual);
     }
 
     #[inline(always)]
-    fn add_imm(&mut self, d: Reg, s1: Reg, s2: u32) -> Self::ReturnTy {
+    fn add_imm(&mut self, d: RawReg, s1: RawReg, s2: u32) -> Self::ReturnTy {
         let reg_size = self.reg_size();
         let d = conv_reg(d);
         let s1 = conv_reg(s1);
-        match (d, s1, s2) {
-            // d = d + 1
-            (_, _, 1) if d == s1 => self.push(inc(reg_size, d)),
-            // d = d + s2
-            (_, _, _) if d == s1 => self.push(add((d, imm32(s2)))),
-            // d = s1 + s2
-            (_, _, _) => {
-                self.push(lea(reg_size, d, reg_indirect(reg_size, s1 + s2 as i32)));
+
+        let asm = self.asm.reserve::<U1>();
+        if d == s1 {
+            if s2 == 1 {
+                asm.push(inc(reg_size, d))
+            } else {
+                asm.push(add((d, imm32(s2))))
             }
-        }
+        } else {
+            asm.push(lea(reg_size, d, reg_indirect(reg_size, s1 + s2 as i32)))
+        }.assert_reserved_exactly_as_needed();
     }
 
     #[inline(always)]
-    fn store_u8(&mut self, src: Reg, offset: u32) -> Self::ReturnTy {
+    fn store_u8(&mut self, src: RawReg, offset: u32) -> Self::ReturnTy {
         self.store(src, None, offset, Size::U8);
     }
 
     #[inline(always)]
-    fn store_u16(&mut self, src: Reg, offset: u32) -> Self::ReturnTy {
+    fn store_u16(&mut self, src: RawReg, offset: u32) -> Self::ReturnTy {
         self.store(src, None, offset, Size::U16);
     }
 
     #[inline(always)]
-    fn store_u32(&mut self, src: Reg, offset: u32) -> Self::ReturnTy {
+    fn store_u32(&mut self, src: RawReg, offset: u32) -> Self::ReturnTy {
         self.store(src, None, offset, Size::U32);
     }
 
     #[inline(always)]
-    fn store_indirect_u8(&mut self, src: Reg, base: Reg, offset: u32) -> Self::ReturnTy {
+    fn store_indirect_u8(&mut self, src: RawReg, base: RawReg, offset: u32) -> Self::ReturnTy {
         self.store(src, Some(base), offset, Size::U8);
     }
 
     #[inline(always)]
-    fn store_indirect_u16(&mut self, src: Reg, base: Reg, offset: u32) -> Self::ReturnTy {
+    fn store_indirect_u16(&mut self, src: RawReg, base: RawReg, offset: u32) -> Self::ReturnTy {
         self.store(src, Some(base), offset, Size::U16);
     }
 
     #[inline(always)]
-    fn store_indirect_u32(&mut self, src: Reg, base: Reg, offset: u32) -> Self::ReturnTy {
+    fn store_indirect_u32(&mut self, src: RawReg, base: RawReg, offset: u32) -> Self::ReturnTy {
         self.store(src, Some(base), offset, Size::U32);
     }
 
     #[inline(always)]
-    fn store_imm_indirect_u8(&mut self, base: Reg, offset: u32, value: u32) -> Self::ReturnTy {
+    fn store_imm_indirect_u8(&mut self, base: RawReg, offset: u32, value: u32) -> Self::ReturnTy {
         self.store(value, Some(base), offset, Size::U8);
     }
 
     #[inline(always)]
-    fn store_imm_indirect_u16(&mut self, base: Reg, offset: u32, value: u32) -> Self::ReturnTy {
+    fn store_imm_indirect_u16(&mut self, base: RawReg, offset: u32, value: u32) -> Self::ReturnTy {
         self.store(value, Some(base), offset, Size::U16);
     }
 
     #[inline(always)]
-    fn store_imm_indirect_u32(&mut self, base: Reg, offset: u32, value: u32) -> Self::ReturnTy {
+    fn store_imm_indirect_u32(&mut self, base: RawReg, offset: u32, value: u32) -> Self::ReturnTy {
         self.store(value, Some(base), offset, Size::U32);
     }
 
     #[inline(always)]
-    fn store_imm_u8(&mut self, value: u32, offset: u32) -> Self::ReturnTy {
+    fn store_imm_u8(&mut self, offset: u32, value: u32) -> Self::ReturnTy {
         self.store(value, None, offset, Size::U8);
     }
 
     #[inline(always)]
-    fn store_imm_u16(&mut self, value: u32, offset: u32) -> Self::ReturnTy {
+    fn store_imm_u16(&mut self, offset: u32, value: u32) -> Self::ReturnTy {
         self.store(value, None, offset, Size::U16);
     }
 
     #[inline(always)]
-    fn store_imm_u32(&mut self, value: u32, offset: u32) -> Self::ReturnTy {
+    fn store_imm_u32(&mut self, offset: u32, value: u32) -> Self::ReturnTy {
         self.store(value, None, offset, Size::U32);
     }
 
     #[inline(always)]
-    fn load_indirect_u8(&mut self, dst: Reg, base: Reg, offset: u32) -> Self::ReturnTy {
+    fn load_indirect_u8(&mut self, dst: RawReg, base: RawReg, offset: u32) -> Self::ReturnTy {
         self.load(dst, Some(base), offset, LoadKind::U8);
     }
 
     #[inline(always)]
-    fn load_indirect_i8(&mut self, dst: Reg, base: Reg, offset: u32) -> Self::ReturnTy {
+    fn load_indirect_i8(&mut self, dst: RawReg, base: RawReg, offset: u32) -> Self::ReturnTy {
         self.load(dst, Some(base), offset, LoadKind::I8);
     }
 
     #[inline(always)]
-    fn load_indirect_u16(&mut self, dst: Reg, base: Reg, offset: u32) -> Self::ReturnTy {
+    fn load_indirect_u16(&mut self, dst: RawReg, base: RawReg, offset: u32) -> Self::ReturnTy {
         self.load(dst, Some(base), offset, LoadKind::U16);
     }
 
     #[inline(always)]
-    fn load_indirect_i16(&mut self, dst: Reg, base: Reg, offset: u32) -> Self::ReturnTy {
+    fn load_indirect_i16(&mut self, dst: RawReg, base: RawReg, offset: u32) -> Self::ReturnTy {
         self.load(dst, Some(base), offset, LoadKind::I16);
     }
 
     #[inline(always)]
-    fn load_indirect_u32(&mut self, dst: Reg, base: Reg, offset: u32) -> Self::ReturnTy {
+    fn load_indirect_u32(&mut self, dst: RawReg, base: RawReg, offset: u32) -> Self::ReturnTy {
         self.load(dst, Some(base), offset, LoadKind::U32);
     }
 
     #[inline(always)]
-    fn load_u8(&mut self, dst: Reg, offset: u32) -> Self::ReturnTy {
+    fn load_u8(&mut self, dst: RawReg, offset: u32) -> Self::ReturnTy {
         self.load(dst, None, offset, LoadKind::U8);
     }
 
     #[inline(always)]
-    fn load_i8(&mut self, dst: Reg, offset: u32) -> Self::ReturnTy {
+    fn load_i8(&mut self, dst: RawReg, offset: u32) -> Self::ReturnTy {
         self.load(dst, None, offset, LoadKind::I8);
     }
 
     #[inline(always)]
-    fn load_u16(&mut self, dst: Reg, offset: u32) -> Self::ReturnTy {
+    fn load_u16(&mut self, dst: RawReg, offset: u32) -> Self::ReturnTy {
         self.load(dst, None, offset, LoadKind::U16);
     }
 
     #[inline(always)]
-    fn load_i16(&mut self, dst: Reg, offset: u32) -> Self::ReturnTy {
+    fn load_i16(&mut self, dst: RawReg, offset: u32) -> Self::ReturnTy {
         self.load(dst, None, offset, LoadKind::I16);
     }
 
     #[inline(always)]
-    fn load_u32(&mut self, dst: Reg, offset: u32) -> Self::ReturnTy {
+    fn load_u32(&mut self, dst: RawReg, offset: u32) -> Self::ReturnTy {
         self.load(dst, None, offset, LoadKind::U32);
     }
 
     #[inline(always)]
-    fn branch_less_unsigned(&mut self, s1: Reg, s2: Reg, target: u32) -> Self::ReturnTy {
+    fn branch_less_unsigned(&mut self, s1: RawReg, s2: RawReg, target: u32) -> Self::ReturnTy {
         self.branch(s1, s2, target, Condition::Below);
     }
 
     #[inline(always)]
-    fn branch_less_signed(&mut self, s1: Reg, s2: Reg, target: u32) -> Self::ReturnTy {
+    fn branch_less_signed(&mut self, s1: RawReg, s2: RawReg, target: u32) -> Self::ReturnTy {
         self.branch(s1, s2, target, Condition::Less);
     }
 
     #[inline(always)]
-    fn branch_greater_or_equal_unsigned(&mut self, s1: Reg, s2: Reg, target: u32) -> Self::ReturnTy {
+    fn branch_greater_or_equal_unsigned(&mut self, s1: RawReg, s2: RawReg, target: u32) -> Self::ReturnTy {
         self.branch(s1, s2, target, Condition::AboveOrEqual);
     }
 
     #[inline(always)]
-    fn branch_greater_or_equal_signed(&mut self, s1: Reg, s2: Reg, target: u32) -> Self::ReturnTy {
+    fn branch_greater_or_equal_signed(&mut self, s1: RawReg, s2: RawReg, target: u32) -> Self::ReturnTy {
         self.branch(s1, s2, target, Condition::GreaterOrEqual);
     }
 
     #[inline(always)]
-    fn branch_eq(&mut self, s1: Reg, s2: Reg, target: u32) -> Self::ReturnTy {
+    fn branch_eq(&mut self, s1: RawReg, s2: RawReg, target: u32) -> Self::ReturnTy {
         self.branch(s1, s2, target, Condition::Equal);
     }
 
     #[inline(always)]
-    fn branch_not_eq(&mut self, s1: Reg, s2: Reg, target: u32) -> Self::ReturnTy {
+    fn branch_not_eq(&mut self, s1: RawReg, s2: RawReg, target: u32) -> Self::ReturnTy {
         self.branch(s1, s2, target, Condition::NotEqual);
     }
 
     #[inline(always)]
-    fn branch_eq_imm(&mut self, s1: Reg, s2: u32, target: u32) -> Self::ReturnTy {
+    fn branch_eq_imm(&mut self, s1: RawReg, s2: u32, target: u32) -> Self::ReturnTy {
         self.branch(s1, s2, target, Condition::Equal);
     }
 
     #[inline(always)]
-    fn branch_not_eq_imm(&mut self, s1: Reg, s2: u32, target: u32) -> Self::ReturnTy {
+    fn branch_not_eq_imm(&mut self, s1: RawReg, s2: u32, target: u32) -> Self::ReturnTy {
         self.branch(s1, s2, target, Condition::NotEqual);
     }
 
     #[inline(always)]
-    fn branch_less_unsigned_imm(&mut self, s1: Reg, s2: u32, target: u32) -> Self::ReturnTy {
+    fn branch_less_unsigned_imm(&mut self, s1: RawReg, s2: u32, target: u32) -> Self::ReturnTy {
         self.branch(s1, s2, target, Condition::Below);
     }
 
     #[inline(always)]
-    fn branch_less_signed_imm(&mut self, s1: Reg, s2: u32, target: u32) -> Self::ReturnTy {
+    fn branch_less_signed_imm(&mut self, s1: RawReg, s2: u32, target: u32) -> Self::ReturnTy {
         self.branch(s1, s2, target, Condition::Less);
     }
 
     #[inline(always)]
-    fn branch_greater_or_equal_unsigned_imm(&mut self, s1: Reg, s2: u32, target: u32) -> Self::ReturnTy {
+    fn branch_greater_or_equal_unsigned_imm(&mut self, s1: RawReg, s2: u32, target: u32) -> Self::ReturnTy {
         self.branch(s1, s2, target, Condition::AboveOrEqual);
     }
 
     #[inline(always)]
-    fn branch_greater_or_equal_signed_imm(&mut self, s1: Reg, s2: u32, target: u32) -> Self::ReturnTy {
+    fn branch_greater_or_equal_signed_imm(&mut self, s1: RawReg, s2: u32, target: u32) -> Self::ReturnTy {
         self.branch(s1, s2, target, Condition::GreaterOrEqual);
     }
 
     #[inline(always)]
-    fn branch_less_or_equal_unsigned_imm(&mut self, s1: Reg, s2: u32, target: u32) -> Self::ReturnTy {
+    fn branch_less_or_equal_unsigned_imm(&mut self, s1: RawReg, s2: u32, target: u32) -> Self::ReturnTy {
         self.branch(s1, s2, target, Condition::BelowOrEqual);
     }
 
     #[inline(always)]
-    fn branch_less_or_equal_signed_imm(&mut self, s1: Reg, s2: u32, target: u32) -> Self::ReturnTy {
+    fn branch_less_or_equal_signed_imm(&mut self, s1: RawReg, s2: u32, target: u32) -> Self::ReturnTy {
         self.branch(s1, s2, target, Condition::LessOrEqual);
     }
 
     #[inline(always)]
-    fn branch_greater_unsigned_imm(&mut self, s1: Reg, s2: u32, target: u32) -> Self::ReturnTy {
+    fn branch_greater_unsigned_imm(&mut self, s1: RawReg, s2: u32, target: u32) -> Self::ReturnTy {
         self.branch(s1, s2, target, Condition::Above);
     }
 
     #[inline(always)]
-    fn branch_greater_signed_imm(&mut self, s1: Reg, s2: u32, target: u32) -> Self::ReturnTy {
+    fn branch_greater_signed_imm(&mut self, s1: RawReg, s2: u32, target: u32) -> Self::ReturnTy {
         self.branch(s1, s2, target, Condition::Greater);
     }
 
     #[inline(always)]
     fn jump(&mut self, target: u32) -> Self::ReturnTy {
-        let label = self.get_or_forward_declare_label(target);
+        let label = self.get_or_forward_declare_label(target).unwrap_or(self.trap_label);
         self.jump_to_label(label);
     }
 
     #[inline(always)]
-    fn load_imm_and_jump(&mut self, ra: Reg, value: u32, target: u32) -> Self::ReturnTy {
-        let label = self.get_or_forward_declare_label(target);
-        self.load_immediate(ra, value);
-        self.jump_to_label(label);
+    fn load_imm_and_jump(&mut self, ra: RawReg, value: u32, target: u32) -> Self::ReturnTy {
+        let label = self.get_or_forward_declare_label(target).unwrap_or(self.trap_label);
+        let asm = self.asm.reserve::<U2>();
+        let asm = asm.push(mov_imm(conv_reg(ra), imm32(value)));
+        let asm = jump_to_label(asm, label);
+        asm.assert_reserved_exactly_as_needed();
     }
 
     #[inline(always)]
-    fn jump_indirect(&mut self, base: Reg, offset: u32) -> Self::ReturnTy {
+    fn jump_indirect(&mut self, base: RawReg, offset: u32) -> Self::ReturnTy {
         self.jump_indirect_impl(None, base, offset)
     }
 
     #[inline(always)]
-    fn load_imm_and_jump_indirect(&mut self, ra: Reg, base: Reg, value: u32, offset: u32) -> Self::ReturnTy {
+    fn load_imm_and_jump_indirect(&mut self, ra: RawReg, base: RawReg, value: u32, offset: u32) -> Self::ReturnTy {
         self.jump_indirect_impl(Some((ra, value)), base, offset)
     }
 }
diff --git a/crates/polkavm/src/gas.rs b/crates/polkavm/src/gas.rs
index 47b3213a..8343940b 100644
--- a/crates/polkavm/src/gas.rs
+++ b/crates/polkavm/src/gas.rs
@@ -1,4 +1,4 @@
-use polkavm_common::program::{Instructions, Reg};
+use polkavm_common::program::{Instructions, RawReg};
 
 // TODO: Come up with a better cost model.
 #[derive(Default)]
@@ -36,7 +36,7 @@ impl polkavm_common::program::InstructionVisitor for GasVisitor {
     }
 
     #[inline(always)]
-    fn sbrk(&mut self, _d: Reg, _s: Reg) -> Self::ReturnTy {
+    fn sbrk(&mut self, _d: RawReg, _s: RawReg) -> Self::ReturnTy {
         self.cost += 1;
     }
 
@@ -46,429 +46,429 @@ impl polkavm_common::program::InstructionVisitor for GasVisitor {
     }
 
     #[inline(always)]
-    fn set_less_than_unsigned(&mut self, _d: Reg, _s1: Reg, _s2: Reg) -> Self::ReturnTy {
+    fn set_less_than_unsigned(&mut self, _d: RawReg, _s1: RawReg, _s2: RawReg) -> Self::ReturnTy {
         self.cost += 1;
     }
 
     #[inline(always)]
-    fn set_less_than_signed(&mut self, _d: Reg, _s1: Reg, _s2: Reg) -> Self::ReturnTy {
+    fn set_less_than_signed(&mut self, _d: RawReg, _s1: RawReg, _s2: RawReg) -> Self::ReturnTy {
         self.cost += 1;
     }
 
     #[inline(always)]
-    fn shift_logical_right(&mut self, _d: Reg, _s1: Reg, _s2: Reg) -> Self::ReturnTy {
+    fn shift_logical_right(&mut self, _d: RawReg, _s1: RawReg, _s2: RawReg) -> Self::ReturnTy {
         self.cost += 1;
     }
 
     #[inline(always)]
-    fn shift_arithmetic_right(&mut self, _d: Reg, _s1: Reg, _s2: Reg) -> Self::ReturnTy {
+    fn shift_arithmetic_right(&mut self, _d: RawReg, _s1: RawReg, _s2: RawReg) -> Self::ReturnTy {
         self.cost += 1;
     }
 
     #[inline(always)]
-    fn shift_logical_left(&mut self, _d: Reg, _s1: Reg, _s2: Reg) -> Self::ReturnTy {
+    fn shift_logical_left(&mut self, _d: RawReg, _s1: RawReg, _s2: RawReg) -> Self::ReturnTy {
         self.cost += 1;
     }
 
     #[inline(always)]
-    fn xor(&mut self, _d: Reg, _s1: Reg, _s2: Reg) -> Self::ReturnTy {
+    fn xor(&mut self, _d: RawReg, _s1: RawReg, _s2: RawReg) -> Self::ReturnTy {
         self.cost += 1;
     }
 
     #[inline(always)]
-    fn and(&mut self, _d: Reg, _s1: Reg, _s2: Reg) -> Self::ReturnTy {
+    fn and(&mut self, _d: RawReg, _s1: RawReg, _s2: RawReg) -> Self::ReturnTy {
         self.cost += 1;
     }
 
     #[inline(always)]
-    fn or(&mut self, _d: Reg, _s1: Reg, _s2: Reg) -> Self::ReturnTy {
+    fn or(&mut self, _d: RawReg, _s1: RawReg, _s2: RawReg) -> Self::ReturnTy {
         self.cost += 1;
     }
 
     #[inline(always)]
-    fn add(&mut self, _d: Reg, _s1: Reg, _s2: Reg) -> Self::ReturnTy {
+    fn add(&mut self, _d: RawReg, _s1: RawReg, _s2: RawReg) -> Self::ReturnTy {
         self.cost += 1;
     }
 
     #[inline(always)]
-    fn sub(&mut self, _d: Reg, _s1: Reg, _s2: Reg) -> Self::ReturnTy {
+    fn sub(&mut self, _d: RawReg, _s1: RawReg, _s2: RawReg) -> Self::ReturnTy {
         self.cost += 1;
     }
 
     #[inline(always)]
-    fn mul(&mut self, _d: Reg, _s1: Reg, _s2: Reg) -> Self::ReturnTy {
+    fn mul(&mut self, _d: RawReg, _s1: RawReg, _s2: RawReg) -> Self::ReturnTy {
         self.cost += 1;
     }
 
     #[inline(always)]
-    fn mul_upper_signed_signed(&mut self, _d: Reg, _s1: Reg, _s2: Reg) -> Self::ReturnTy {
+    fn mul_upper_signed_signed(&mut self, _d: RawReg, _s1: RawReg, _s2: RawReg) -> Self::ReturnTy {
         self.cost += 1;
     }
 
     #[inline(always)]
-    fn mul_upper_unsigned_unsigned(&mut self, _d: Reg, _s1: Reg, _s2: Reg) -> Self::ReturnTy {
+    fn mul_upper_unsigned_unsigned(&mut self, _d: RawReg, _s1: RawReg, _s2: RawReg) -> Self::ReturnTy {
         self.cost += 1;
     }
 
     #[inline(always)]
-    fn mul_upper_signed_unsigned(&mut self, _d: Reg, _s1: Reg, _s2: Reg) -> Self::ReturnTy {
+    fn mul_upper_signed_unsigned(&mut self, _d: RawReg, _s1: RawReg, _s2: RawReg) -> Self::ReturnTy {
         self.cost += 1;
     }
 
     #[inline(always)]
-    fn div_unsigned(&mut self, _d: Reg, _s1: Reg, _s2: Reg) -> Self::ReturnTy {
+    fn div_unsigned(&mut self, _d: RawReg, _s1: RawReg, _s2: RawReg) -> Self::ReturnTy {
         self.cost += 1;
     }
 
     #[inline(always)]
-    fn div_signed(&mut self, _d: Reg, _s1: Reg, _s2: Reg) -> Self::ReturnTy {
+    fn div_signed(&mut self, _d: RawReg, _s1: RawReg, _s2: RawReg) -> Self::ReturnTy {
         self.cost += 1;
     }
 
     #[inline(always)]
-    fn rem_unsigned(&mut self, _d: Reg, _s1: Reg, _s2: Reg) -> Self::ReturnTy {
+    fn rem_unsigned(&mut self, _d: RawReg, _s1: RawReg, _s2: RawReg) -> Self::ReturnTy {
         self.cost += 1;
     }
 
     #[inline(always)]
-    fn rem_signed(&mut self, _d: Reg, _s1: Reg, _s2: Reg) -> Self::ReturnTy {
+    fn rem_signed(&mut self, _d: RawReg, _s1: RawReg, _s2: RawReg) -> Self::ReturnTy {
         self.cost += 1;
     }
 
     #[inline(always)]
-    fn mul_imm(&mut self, _d: Reg, _s1: Reg, _s2: u32) -> Self::ReturnTy {
+    fn mul_imm(&mut self, _d: RawReg, _s1: RawReg, _s2: u32) -> Self::ReturnTy {
         self.cost += 1;
     }
 
     #[inline(always)]
-    fn mul_upper_signed_signed_imm(&mut self, _d: Reg, _s1: Reg, _s2: u32) -> Self::ReturnTy {
+    fn mul_upper_signed_signed_imm(&mut self, _d: RawReg, _s1: RawReg, _s2: u32) -> Self::ReturnTy {
         self.cost += 1;
     }
 
     #[inline(always)]
-    fn mul_upper_unsigned_unsigned_imm(&mut self, _d: Reg, _s1: Reg, _s2: u32) -> Self::ReturnTy {
+    fn mul_upper_unsigned_unsigned_imm(&mut self, _d: RawReg, _s1: RawReg, _s2: u32) -> Self::ReturnTy {
         self.cost += 1;
     }
 
     #[inline(always)]
-    fn set_less_than_unsigned_imm(&mut self, _d: Reg, _s1: Reg, _s2: u32) -> Self::ReturnTy {
+    fn set_less_than_unsigned_imm(&mut self, _d: RawReg, _s1: RawReg, _s2: u32) -> Self::ReturnTy {
         self.cost += 1;
     }
 
     #[inline(always)]
-    fn set_less_than_signed_imm(&mut self, _d: Reg, _s1: Reg, _s2: u32) -> Self::ReturnTy {
+    fn set_less_than_signed_imm(&mut self, _d: RawReg, _s1: RawReg, _s2: u32) -> Self::ReturnTy {
         self.cost += 1;
     }
 
     #[inline(always)]
-    fn set_greater_than_unsigned_imm(&mut self, _d: Reg, _s1: Reg, _s2: u32) -> Self::ReturnTy {
+    fn set_greater_than_unsigned_imm(&mut self, _d: RawReg, _s1: RawReg, _s2: u32) -> Self::ReturnTy {
         self.cost += 1;
     }
 
     #[inline(always)]
-    fn set_greater_than_signed_imm(&mut self, _d: Reg, _s1: Reg, _s2: u32) -> Self::ReturnTy {
+    fn set_greater_than_signed_imm(&mut self, _d: RawReg, _s1: RawReg, _s2: u32) -> Self::ReturnTy {
         self.cost += 1;
     }
 
     #[inline(always)]
-    fn shift_logical_right_imm(&mut self, _d: Reg, _s1: Reg, _s2: u32) -> Self::ReturnTy {
+    fn shift_logical_right_imm(&mut self, _d: RawReg, _s1: RawReg, _s2: u32) -> Self::ReturnTy {
         self.cost += 1;
     }
 
     #[inline(always)]
-    fn shift_arithmetic_right_imm(&mut self, _d: Reg, _s1: Reg, _s2: u32) -> Self::ReturnTy {
+    fn shift_arithmetic_right_imm(&mut self, _d: RawReg, _s1: RawReg, _s2: u32) -> Self::ReturnTy {
         self.cost += 1;
     }
 
     #[inline(always)]
-    fn shift_logical_left_imm(&mut self, _d: Reg, _s1: Reg, _s2: u32) -> Self::ReturnTy {
+    fn shift_logical_left_imm(&mut self, _d: RawReg, _s1: RawReg, _s2: u32) -> Self::ReturnTy {
         self.cost += 1;
     }
 
     #[inline(always)]
-    fn shift_logical_right_imm_alt(&mut self, _d: Reg, _s2: Reg, _s1: u32) -> Self::ReturnTy {
+    fn shift_logical_right_imm_alt(&mut self, _d: RawReg, _s2: RawReg, _s1: u32) -> Self::ReturnTy {
         self.cost += 1;
     }
 
     #[inline(always)]
-    fn shift_arithmetic_right_imm_alt(&mut self, _d: Reg, _s2: Reg, _s1: u32) -> Self::ReturnTy {
+    fn shift_arithmetic_right_imm_alt(&mut self, _d: RawReg, _s2: RawReg, _s1: u32) -> Self::ReturnTy {
         self.cost += 1;
     }
 
     #[inline(always)]
-    fn shift_logical_left_imm_alt(&mut self, _d: Reg, _s2: Reg, _s1: u32) -> Self::ReturnTy {
+    fn shift_logical_left_imm_alt(&mut self, _d: RawReg, _s2: RawReg, _s1: u32) -> Self::ReturnTy {
         self.cost += 1;
     }
 
     #[inline(always)]
-    fn or_imm(&mut self, _d: Reg, _s: Reg, _imm: u32) -> Self::ReturnTy {
+    fn or_imm(&mut self, _d: RawReg, _s: RawReg, _imm: u32) -> Self::ReturnTy {
         self.cost += 1;
     }
 
     #[inline(always)]
-    fn and_imm(&mut self, _d: Reg, _s: Reg, _imm: u32) -> Self::ReturnTy {
+    fn and_imm(&mut self, _d: RawReg, _s: RawReg, _imm: u32) -> Self::ReturnTy {
         self.cost += 1;
     }
 
     #[inline(always)]
-    fn xor_imm(&mut self, _d: Reg, _s: Reg, _imm: u32) -> Self::ReturnTy {
+    fn xor_imm(&mut self, _d: RawReg, _s: RawReg, _imm: u32) -> Self::ReturnTy {
         self.cost += 1;
     }
 
     #[inline(always)]
-    fn move_reg(&mut self, _d: Reg, _s: Reg) -> Self::ReturnTy {
+    fn move_reg(&mut self, _d: RawReg, _s: RawReg) -> Self::ReturnTy {
         self.cost += 1;
     }
 
     #[inline(always)]
-    fn cmov_if_zero(&mut self, _d: Reg, _s: Reg, _c: Reg) -> Self::ReturnTy {
+    fn cmov_if_zero(&mut self, _d: RawReg, _s: RawReg, _c: RawReg) -> Self::ReturnTy {
         self.cost += 1;
     }
 
     #[inline(always)]
-    fn cmov_if_not_zero(&mut self, _d: Reg, _s: Reg, _c: Reg) -> Self::ReturnTy {
+    fn cmov_if_not_zero(&mut self, _d: RawReg, _s: RawReg, _c: RawReg) -> Self::ReturnTy {
         self.cost += 1;
     }
 
     #[inline(always)]
-    fn cmov_if_zero_imm(&mut self, _d: Reg, _c: Reg, _s: u32) -> Self::ReturnTy {
+    fn cmov_if_zero_imm(&mut self, _d: RawReg, _c: RawReg, _s: u32) -> Self::ReturnTy {
         self.cost += 1;
     }
 
     #[inline(always)]
-    fn cmov_if_not_zero_imm(&mut self, _d: Reg, _c: Reg, _s: u32) -> Self::ReturnTy {
+    fn cmov_if_not_zero_imm(&mut self, _d: RawReg, _c: RawReg, _s: u32) -> Self::ReturnTy {
         self.cost += 1;
     }
 
     #[inline(always)]
-    fn add_imm(&mut self, _d: Reg, _s: Reg, _imm: u32) -> Self::ReturnTy {
+    fn add_imm(&mut self, _d: RawReg, _s: RawReg, _imm: u32) -> Self::ReturnTy {
         self.cost += 1;
     }
 
     #[inline(always)]
-    fn negate_and_add_imm(&mut self, _d: Reg, _s1: Reg, _s2: u32) -> Self::ReturnTy {
+    fn negate_and_add_imm(&mut self, _d: RawReg, _s1: RawReg, _s2: u32) -> Self::ReturnTy {
         self.cost += 1;
     }
 
     #[inline(always)]
-    fn store_imm_indirect_u8(&mut self, _base: Reg, _offset: u32, _value: u32) -> Self::ReturnTy {
+    fn store_imm_indirect_u8(&mut self, _base: RawReg, _offset: u32, _value: u32) -> Self::ReturnTy {
         self.cost += 1;
     }
 
     #[inline(always)]
-    fn store_imm_indirect_u16(&mut self, _base: Reg, _offset: u32, _value: u32) -> Self::ReturnTy {
+    fn store_imm_indirect_u16(&mut self, _base: RawReg, _offset: u32, _value: u32) -> Self::ReturnTy {
         self.cost += 1;
     }
 
     #[inline(always)]
-    fn store_imm_indirect_u32(&mut self, _base: Reg, _offset: u32, _value: u32) -> Self::ReturnTy {
+    fn store_imm_indirect_u32(&mut self, _base: RawReg, _offset: u32, _value: u32) -> Self::ReturnTy {
         self.cost += 1;
     }
 
     #[inline(always)]
-    fn store_indirect_u8(&mut self, _src: Reg, _base: Reg, _offset: u32) -> Self::ReturnTy {
+    fn store_indirect_u8(&mut self, _src: RawReg, _base: RawReg, _offset: u32) -> Self::ReturnTy {
         self.cost += 1;
     }
 
     #[inline(always)]
-    fn store_indirect_u16(&mut self, _src: Reg, _base: Reg, _offset: u32) -> Self::ReturnTy {
+    fn store_indirect_u16(&mut self, _src: RawReg, _base: RawReg, _offset: u32) -> Self::ReturnTy {
         self.cost += 1;
     }
 
     #[inline(always)]
-    fn store_indirect_u32(&mut self, _src: Reg, _base: Reg, _offset: u32) -> Self::ReturnTy {
+    fn store_indirect_u32(&mut self, _src: RawReg, _base: RawReg, _offset: u32) -> Self::ReturnTy {
         self.cost += 1;
     }
 
     #[inline(always)]
-    fn store_imm_u8(&mut self, _value: u32, _offset: u32) -> Self::ReturnTy {
+    fn store_imm_u8(&mut self, _offset: u32, _value: u32) -> Self::ReturnTy {
         self.cost += 1;
     }
 
     #[inline(always)]
-    fn store_imm_u16(&mut self, _value: u32, _offset: u32) -> Self::ReturnTy {
+    fn store_imm_u16(&mut self, _offset: u32, _value: u32) -> Self::ReturnTy {
         self.cost += 1;
     }
 
     #[inline(always)]
-    fn store_imm_u32(&mut self, _value: u32, _offset: u32) -> Self::ReturnTy {
+    fn store_imm_u32(&mut self, _offset: u32, _value: u32) -> Self::ReturnTy {
         self.cost += 1;
     }
 
     #[inline(always)]
-    fn store_u8(&mut self, _src: Reg, _offset: u32) -> Self::ReturnTy {
+    fn store_u8(&mut self, _src: RawReg, _offset: u32) -> Self::ReturnTy {
         self.cost += 1;
     }
 
     #[inline(always)]
-    fn store_u16(&mut self, _src: Reg, _offset: u32) -> Self::ReturnTy {
+    fn store_u16(&mut self, _src: RawReg, _offset: u32) -> Self::ReturnTy {
         self.cost += 1;
     }
 
     #[inline(always)]
-    fn store_u32(&mut self, _src: Reg, _offset: u32) -> Self::ReturnTy {
+    fn store_u32(&mut self, _src: RawReg, _offset: u32) -> Self::ReturnTy {
         self.cost += 1;
     }
 
     #[inline(always)]
-    fn load_indirect_u8(&mut self, _dst: Reg, _base: Reg, _offset: u32) -> Self::ReturnTy {
+    fn load_indirect_u8(&mut self, _dst: RawReg, _base: RawReg, _offset: u32) -> Self::ReturnTy {
         self.cost += 1;
     }
 
     #[inline(always)]
-    fn load_indirect_i8(&mut self, _dst: Reg, _base: Reg, _offset: u32) -> Self::ReturnTy {
+    fn load_indirect_i8(&mut self, _dst: RawReg, _base: RawReg, _offset: u32) -> Self::ReturnTy {
         self.cost += 1;
     }
 
     #[inline(always)]
-    fn load_indirect_u16(&mut self, _dst: Reg, _base: Reg, _offset: u32) -> Self::ReturnTy {
+    fn load_indirect_u16(&mut self, _dst: RawReg, _base: RawReg, _offset: u32) -> Self::ReturnTy {
         self.cost += 1;
     }
 
     #[inline(always)]
-    fn load_indirect_i16(&mut self, _dst: Reg, _base: Reg, _offset: u32) -> Self::ReturnTy {
+    fn load_indirect_i16(&mut self, _dst: RawReg, _base: RawReg, _offset: u32) -> Self::ReturnTy {
         self.cost += 1;
     }
 
     #[inline(always)]
-    fn load_indirect_u32(&mut self, _dst: Reg, _base: Reg, _offset: u32) -> Self::ReturnTy {
+    fn load_indirect_u32(&mut self, _dst: RawReg, _base: RawReg, _offset: u32) -> Self::ReturnTy {
         self.cost += 1;
     }
 
     #[inline(always)]
-    fn load_u8(&mut self, _dst: Reg, _offset: u32) -> Self::ReturnTy {
+    fn load_u8(&mut self, _dst: RawReg, _offset: u32) -> Self::ReturnTy {
         self.cost += 1;
     }
 
     #[inline(always)]
-    fn load_i8(&mut self, _dst: Reg, _offset: u32) -> Self::ReturnTy {
+    fn load_i8(&mut self, _dst: RawReg, _offset: u32) -> Self::ReturnTy {
         self.cost += 1;
     }
 
     #[inline(always)]
-    fn load_u16(&mut self, _dst: Reg, _offset: u32) -> Self::ReturnTy {
+    fn load_u16(&mut self, _dst: RawReg, _offset: u32) -> Self::ReturnTy {
         self.cost += 1;
     }
 
     #[inline(always)]
-    fn load_i16(&mut self, _dst: Reg, _offset: u32) -> Self::ReturnTy {
+    fn load_i16(&mut self, _dst: RawReg, _offset: u32) -> Self::ReturnTy {
         self.cost += 1;
     }
 
     #[inline(always)]
-    fn load_u32(&mut self, _dst: Reg, _offset: u32) -> Self::ReturnTy {
+    fn load_u32(&mut self, _dst: RawReg, _offset: u32) -> Self::ReturnTy {
         self.cost += 1;
     }
 
     #[inline(always)]
-    fn branch_less_unsigned(&mut self, _s1: Reg, _s2: Reg, _imm: u32) -> Self::ReturnTy {
+    fn branch_less_unsigned(&mut self, _s1: RawReg, _s2: RawReg, _imm: u32) -> Self::ReturnTy {
         self.cost += 1;
         self.start_new_basic_block();
     }
 
     #[inline(always)]
-    fn branch_less_signed(&mut self, _s1: Reg, _s2: Reg, _imm: u32) -> Self::ReturnTy {
+    fn branch_less_signed(&mut self, _s1: RawReg, _s2: RawReg, _imm: u32) -> Self::ReturnTy {
         self.cost += 1;
         self.start_new_basic_block();
     }
 
     #[inline(always)]
-    fn branch_greater_or_equal_unsigned(&mut self, _s1: Reg, _s2: Reg, _imm: u32) -> Self::ReturnTy {
+    fn branch_greater_or_equal_unsigned(&mut self, _s1: RawReg, _s2: RawReg, _imm: u32) -> Self::ReturnTy {
         self.cost += 1;
         self.start_new_basic_block();
     }
 
     #[inline(always)]
-    fn branch_greater_or_equal_signed(&mut self, _s1: Reg, _s2: Reg, _imm: u32) -> Self::ReturnTy {
+    fn branch_greater_or_equal_signed(&mut self, _s1: RawReg, _s2: RawReg, _imm: u32) -> Self::ReturnTy {
         self.cost += 1;
         self.start_new_basic_block();
     }
 
     #[inline(always)]
-    fn branch_eq(&mut self, _s1: Reg, _s2: Reg, _imm: u32) -> Self::ReturnTy {
+    fn branch_eq(&mut self, _s1: RawReg, _s2: RawReg, _imm: u32) -> Self::ReturnTy {
         self.cost += 1;
         self.start_new_basic_block();
     }
 
     #[inline(always)]
-    fn branch_not_eq(&mut self, _s1: Reg, _s2: Reg, _imm: u32) -> Self::ReturnTy {
+    fn branch_not_eq(&mut self, _s1: RawReg, _s2: RawReg, _imm: u32) -> Self::ReturnTy {
         self.cost += 1;
         self.start_new_basic_block();
     }
 
     #[inline(always)]
-    fn branch_eq_imm(&mut self, _s1: Reg, _s2: u32, _imm: u32) -> Self::ReturnTy {
+    fn branch_eq_imm(&mut self, _s1: RawReg, _s2: u32, _imm: u32) -> Self::ReturnTy {
         self.cost += 1;
         self.start_new_basic_block();
     }
 
     #[inline(always)]
-    fn branch_not_eq_imm(&mut self, _s1: Reg, _s2: u32, _imm: u32) -> Self::ReturnTy {
+    fn branch_not_eq_imm(&mut self, _s1: RawReg, _s2: u32, _imm: u32) -> Self::ReturnTy {
         self.cost += 1;
         self.start_new_basic_block();
     }
 
     #[inline(always)]
-    fn branch_less_unsigned_imm(&mut self, _s1: Reg, _s2: u32, _imm: u32) -> Self::ReturnTy {
+    fn branch_less_unsigned_imm(&mut self, _s1: RawReg, _s2: u32, _imm: u32) -> Self::ReturnTy {
         self.cost += 1;
         self.start_new_basic_block();
     }
 
     #[inline(always)]
-    fn branch_less_signed_imm(&mut self, _s1: Reg, _s2: u32, _imm: u32) -> Self::ReturnTy {
+    fn branch_less_signed_imm(&mut self, _s1: RawReg, _s2: u32, _imm: u32) -> Self::ReturnTy {
         self.cost += 1;
         self.start_new_basic_block();
     }
 
     #[inline(always)]
-    fn branch_greater_or_equal_unsigned_imm(&mut self, _s1: Reg, _s2: u32, _imm: u32) -> Self::ReturnTy {
+    fn branch_greater_or_equal_unsigned_imm(&mut self, _s1: RawReg, _s2: u32, _imm: u32) -> Self::ReturnTy {
         self.cost += 1;
         self.start_new_basic_block();
     }
 
     #[inline(always)]
-    fn branch_greater_or_equal_signed_imm(&mut self, _s1: Reg, _s2: u32, _imm: u32) -> Self::ReturnTy {
+    fn branch_greater_or_equal_signed_imm(&mut self, _s1: RawReg, _s2: u32, _imm: u32) -> Self::ReturnTy {
         self.cost += 1;
         self.start_new_basic_block();
     }
 
     #[inline(always)]
-    fn branch_less_or_equal_unsigned_imm(&mut self, _s1: Reg, _s2: u32, _imm: u32) -> Self::ReturnTy {
+    fn branch_less_or_equal_unsigned_imm(&mut self, _s1: RawReg, _s2: u32, _imm: u32) -> Self::ReturnTy {
         self.cost += 1;
         self.start_new_basic_block();
     }
 
     #[inline(always)]
-    fn branch_less_or_equal_signed_imm(&mut self, _s1: Reg, _s2: u32, _imm: u32) -> Self::ReturnTy {
+    fn branch_less_or_equal_signed_imm(&mut self, _s1: RawReg, _s2: u32, _imm: u32) -> Self::ReturnTy {
         self.cost += 1;
         self.start_new_basic_block();
     }
 
     #[inline(always)]
-    fn branch_greater_unsigned_imm(&mut self, _s1: Reg, _s2: u32, _imm: u32) -> Self::ReturnTy {
+    fn branch_greater_unsigned_imm(&mut self, _s1: RawReg, _s2: u32, _imm: u32) -> Self::ReturnTy {
         self.cost += 1;
         self.start_new_basic_block();
     }
 
     #[inline(always)]
-    fn branch_greater_signed_imm(&mut self, _s1: Reg, _s2: u32, _imm: u32) -> Self::ReturnTy {
+    fn branch_greater_signed_imm(&mut self, _s1: RawReg, _s2: u32, _imm: u32) -> Self::ReturnTy {
         self.cost += 1;
         self.start_new_basic_block();
     }
 
     #[inline(always)]
-    fn load_imm(&mut self, _dst: Reg, _value: u32) -> Self::ReturnTy {
+    fn load_imm(&mut self, _dst: RawReg, _value: u32) -> Self::ReturnTy {
         self.cost += 1;
     }
 
     #[inline(always)]
-    fn load_imm_and_jump(&mut self, _ra: Reg, _value: u32, _target: u32) -> Self::ReturnTy {
+    fn load_imm_and_jump(&mut self, _ra: RawReg, _value: u32, _target: u32) -> Self::ReturnTy {
         self.cost += 1;
         self.start_new_basic_block();
     }
 
     #[inline(always)]
-    fn load_imm_and_jump_indirect(&mut self, _ra: Reg, _base: Reg, _value: u32, _offset: u32) -> Self::ReturnTy {
+    fn load_imm_and_jump_indirect(&mut self, _ra: RawReg, _base: RawReg, _value: u32, _offset: u32) -> Self::ReturnTy {
         self.cost += 1;
         self.start_new_basic_block();
     }
@@ -480,7 +480,7 @@ impl polkavm_common::program::InstructionVisitor for GasVisitor {
     }
 
     #[inline(always)]
-    fn jump_indirect(&mut self, _base: Reg, _offset: u32) -> Self::ReturnTy {
+    fn jump_indirect(&mut self, _base: RawReg, _offset: u32) -> Self::ReturnTy {
         self.cost += 1;
         self.start_new_basic_block();
     }
diff --git a/crates/polkavm/src/interpreter.rs b/crates/polkavm/src/interpreter.rs
index 1f0806ff..4d54e4b8 100644
--- a/crates/polkavm/src/interpreter.rs
+++ b/crates/polkavm/src/interpreter.rs
@@ -8,7 +8,7 @@ use core::num::NonZeroU32;
 use polkavm_common::abi::VM_ADDR_RETURN_TO_HOST;
 use polkavm_common::error::Trap;
 use polkavm_common::operation::*;
-use polkavm_common::program::{Instruction, InstructionVisitor, ParsedInstruction, Reg};
+use polkavm_common::program::{Instruction, InstructionVisitor, ParsedInstruction, RawReg, Reg};
 use polkavm_common::utils::{align_to_next_page_usize, byte_slice_init, Access, AsUninitSliceMut, Gas};
 use polkavm_common::{
     VM_RPC_FLAG_CLEAR_PROGRAM_AFTER_EXECUTION, VM_RPC_FLAG_RESET_MEMORY_AFTER_EXECUTION, VM_RPC_FLAG_RESET_MEMORY_BEFORE_EXECUTION,
@@ -21,7 +21,7 @@ trait IntoRegImm {
     fn into(self) -> RegImm;
 }
 
-impl IntoRegImm for Reg {
+impl IntoRegImm for RawReg {
     #[inline(always)]
     fn into(self) -> RegImm {
         RegImm::Reg(self)
@@ -424,7 +424,7 @@ impl InterpretedInstance {
             self.compiled_offset = (compiled_offset.get() - 1) as usize;
         } else {
             self.compiled_offset = self.compiled_instructions.len();
-            self.compile_block::<DEBUG>();
+            self.compile_block::<DEBUG>()?;
         }
 
         if self.gas_remaining.is_some() {
@@ -469,10 +469,11 @@ impl InterpretedInstance {
 
     #[inline(never)]
     #[cold]
-    fn compile_block<const DEBUG: bool>(&mut self) {
+    fn compile_block<const DEBUG: bool>(&mut self) -> Result<(), ExecutionError> {
         let Some(instructions) = self.module.blob().instructions_at(self.instruction_offset) else {
-            return;
+            return Err(ExecutionError::Trap(Default::default()));
         };
+
         if DEBUG {
             log::debug!("Compiling block at {}:", self.instruction_offset);
         }
@@ -496,6 +497,8 @@ impl InterpretedInstance {
 
         self.compiled_offset_for_block
             .insert(self.instruction_offset, NonZeroU32::new((starting_offset + 1) as u32).unwrap());
+
+        Ok(())
     }
 
     fn check_gas(&mut self) -> Result<(), ExecutionError> {
@@ -598,13 +601,14 @@ impl<'a, 'b, const DEBUG: bool> Visitor<'a, 'b, DEBUG> {
     #[inline(always)]
     fn get(&self, regimm: impl IntoRegImm) -> u32 {
         match regimm.into() {
-            RegImm::Reg(reg) => self.inner.regs[reg as usize],
+            RegImm::Reg(reg) => self.inner.regs[reg.get() as usize],
             RegImm::Imm(value) => value,
         }
     }
 
     #[inline(always)]
-    fn set(&mut self, dst: Reg, value: u32) -> Result<(), ExecutionError> {
+    fn set(&mut self, dst: RawReg, value: u32) -> Result<(), ExecutionError> {
+        let dst = dst.get();
         self.inner.regs[dst as usize] = value;
 
         if DEBUG {
@@ -622,7 +626,7 @@ impl<'a, 'b, const DEBUG: bool> Visitor<'a, 'b, DEBUG> {
     #[inline(always)]
     fn set3(
         &mut self,
-        dst: Reg,
+        dst: RawReg,
         s1: impl IntoRegImm,
         s2: impl IntoRegImm,
         callback: impl Fn(u32, u32) -> u32,
@@ -652,10 +656,10 @@ impl<'a, 'b, const DEBUG: bool> Visitor<'a, 'b, DEBUG> {
     }
 
     #[cfg_attr(not(debug_assertions), inline(always))]
-    fn load<T: LoadTy>(&mut self, dst: Reg, base: Option<Reg>, offset: u32) -> Result<(), ExecutionError> {
+    fn load<T: LoadTy>(&mut self, dst: RawReg, base: Option<RawReg>, offset: u32) -> Result<(), ExecutionError> {
         assert!(core::mem::size_of::<T>() >= 1);
 
-        let address = base.map_or(0, |base| self.inner.regs[base as usize]).wrapping_add(offset);
+        let address = base.map_or(0, |base| self.inner.regs[base.get() as usize]).wrapping_add(offset);
         let length = core::mem::size_of::<T>() as u32;
         let Some(slice) = self.inner.memory.get_memory_slice(&self.inner.module, address, length) else {
             if DEBUG {
@@ -682,12 +686,13 @@ impl<'a, 'b, const DEBUG: bool> Visitor<'a, 'b, DEBUG> {
         Ok(())
     }
 
-    fn store<T: StoreTy>(&mut self, src: impl IntoRegImm, base: Option<Reg>, offset: u32) -> Result<(), ExecutionError> {
+    fn store<T: StoreTy>(&mut self, src: impl IntoRegImm, base: Option<RawReg>, offset: u32) -> Result<(), ExecutionError> {
         assert!(core::mem::size_of::<T>() >= 1);
 
-        let address = base.map_or(0, |base| self.inner.regs[base as usize]).wrapping_add(offset);
+        let address = base.map_or(0, |base| self.inner.regs[base.get() as usize]).wrapping_add(offset);
         let value = match src.into() {
             RegImm::Reg(src) => {
+                let src = src.get();
                 let value = self.inner.regs[src as usize];
                 if DEBUG {
                     log::trace!("{kind} [0x{address:x}] = {src} = 0x{value:x}", kind = core::any::type_name::<T>());
@@ -732,7 +737,8 @@ impl<'a, 'b, const DEBUG: bool> Visitor<'a, 'b, DEBUG> {
     }
 
     #[cfg_attr(not(debug_assertions), inline(always))]
-    fn jump_indirect_impl(&mut self, call: Option<(Reg, u32)>, base: Reg, offset: u32) -> Result<(), ExecutionError> {
+    fn jump_indirect_impl(&mut self, call: Option<(RawReg, u32)>, base: RawReg, offset: u32) -> Result<(), ExecutionError> {
+        let base = base.get();
         let target = self.inner.regs[base as usize].wrapping_add(offset);
         if let Some((ra, return_address)) = call {
             self.set(ra, return_address)?;
@@ -845,7 +851,7 @@ impl<'a, 'b, const DEBUG: bool> InstructionVisitor for Visitor<'a, 'b, DEBUG> {
         self.inner.on_start_new_basic_block::<DEBUG>()
     }
 
-    fn sbrk(&mut self, dst: Reg, size: Reg) -> Self::ReturnTy {
+    fn sbrk(&mut self, dst: RawReg, size: RawReg) -> Self::ReturnTy {
         let size = self.get(size);
         let result = self.inner.sbrk(size).unwrap_or(0);
         self.set(dst, result)?;
@@ -864,158 +870,158 @@ impl<'a, 'b, const DEBUG: bool> InstructionVisitor for Visitor<'a, 'b, DEBUG> {
         }
     }
 
-    fn set_less_than_unsigned(&mut self, d: Reg, s1: Reg, s2: Reg) -> Self::ReturnTy {
+    fn set_less_than_unsigned(&mut self, d: RawReg, s1: RawReg, s2: RawReg) -> Self::ReturnTy {
         self.set3(d, s1, s2, |s1, s2| u32::from(s1 < s2))
     }
 
-    fn set_less_than_signed(&mut self, d: Reg, s1: Reg, s2: Reg) -> Self::ReturnTy {
+    fn set_less_than_signed(&mut self, d: RawReg, s1: RawReg, s2: RawReg) -> Self::ReturnTy {
         self.set3(d, s1, s2, |s1, s2| u32::from((s1 as i32) < (s2 as i32)))
     }
 
-    fn shift_logical_right(&mut self, d: Reg, s1: Reg, s2: Reg) -> Self::ReturnTy {
+    fn shift_logical_right(&mut self, d: RawReg, s1: RawReg, s2: RawReg) -> Self::ReturnTy {
         self.set3(d, s1, s2, u32::wrapping_shr)
     }
 
-    fn shift_arithmetic_right(&mut self, d: Reg, s1: Reg, s2: Reg) -> Self::ReturnTy {
+    fn shift_arithmetic_right(&mut self, d: RawReg, s1: RawReg, s2: RawReg) -> Self::ReturnTy {
         self.set3(d, s1, s2, |s1, s2| ((s1 as i32).wrapping_shr(s2)) as u32)
     }
 
-    fn shift_logical_left(&mut self, d: Reg, s1: Reg, s2: Reg) -> Self::ReturnTy {
+    fn shift_logical_left(&mut self, d: RawReg, s1: RawReg, s2: RawReg) -> Self::ReturnTy {
         self.set3(d, s1, s2, u32::wrapping_shl)
     }
 
-    fn xor(&mut self, d: Reg, s1: Reg, s2: Reg) -> Self::ReturnTy {
+    fn xor(&mut self, d: RawReg, s1: RawReg, s2: RawReg) -> Self::ReturnTy {
         self.set3(d, s1, s2, |s1, s2| s1 ^ s2)
     }
 
-    fn and(&mut self, d: Reg, s1: Reg, s2: Reg) -> Self::ReturnTy {
+    fn and(&mut self, d: RawReg, s1: RawReg, s2: RawReg) -> Self::ReturnTy {
         self.set3(d, s1, s2, |s1, s2| s1 & s2)
     }
 
-    fn or(&mut self, d: Reg, s1: Reg, s2: Reg) -> Self::ReturnTy {
+    fn or(&mut self, d: RawReg, s1: RawReg, s2: RawReg) -> Self::ReturnTy {
         self.set3(d, s1, s2, |s1, s2| s1 | s2)
     }
 
-    fn add(&mut self, d: Reg, s1: Reg, s2: Reg) -> Self::ReturnTy {
+    fn add(&mut self, d: RawReg, s1: RawReg, s2: RawReg) -> Self::ReturnTy {
         self.set3(d, s1, s2, u32::wrapping_add)
     }
 
-    fn sub(&mut self, d: Reg, s1: Reg, s2: Reg) -> Self::ReturnTy {
+    fn sub(&mut self, d: RawReg, s1: RawReg, s2: RawReg) -> Self::ReturnTy {
         self.set3(d, s1, s2, u32::wrapping_sub)
     }
 
-    fn negate_and_add_imm(&mut self, d: Reg, s1: Reg, s2: u32) -> Self::ReturnTy {
+    fn negate_and_add_imm(&mut self, d: RawReg, s1: RawReg, s2: u32) -> Self::ReturnTy {
         self.set3(d, s1, s2, |s1, s2| s2.wrapping_sub(s1))
     }
 
-    fn mul(&mut self, d: Reg, s1: Reg, s2: Reg) -> Self::ReturnTy {
+    fn mul(&mut self, d: RawReg, s1: RawReg, s2: RawReg) -> Self::ReturnTy {
         self.set3(d, s1, s2, u32::wrapping_mul)
     }
 
-    fn mul_imm(&mut self, d: Reg, s1: Reg, s2: u32) -> Self::ReturnTy {
+    fn mul_imm(&mut self, d: RawReg, s1: RawReg, s2: u32) -> Self::ReturnTy {
         self.set3(d, s1, s2, u32::wrapping_mul)
     }
 
-    fn mul_upper_signed_signed(&mut self, d: Reg, s1: Reg, s2: Reg) -> Self::ReturnTy {
+    fn mul_upper_signed_signed(&mut self, d: RawReg, s1: RawReg, s2: RawReg) -> Self::ReturnTy {
         self.set3(d, s1, s2, |s1, s2| mulh(s1 as i32, s2 as i32) as u32)
     }
 
-    fn mul_upper_signed_signed_imm(&mut self, d: Reg, s1: Reg, s2: u32) -> Self::ReturnTy {
+    fn mul_upper_signed_signed_imm(&mut self, d: RawReg, s1: RawReg, s2: u32) -> Self::ReturnTy {
         self.set3(d, s1, s2, |s1, s2| mulh(s1 as i32, s2 as i32) as u32)
     }
 
-    fn mul_upper_unsigned_unsigned(&mut self, d: Reg, s1: Reg, s2: Reg) -> Self::ReturnTy {
+    fn mul_upper_unsigned_unsigned(&mut self, d: RawReg, s1: RawReg, s2: RawReg) -> Self::ReturnTy {
         self.set3(d, s1, s2, mulhu)
     }
 
-    fn mul_upper_unsigned_unsigned_imm(&mut self, d: Reg, s1: Reg, s2: u32) -> Self::ReturnTy {
+    fn mul_upper_unsigned_unsigned_imm(&mut self, d: RawReg, s1: RawReg, s2: u32) -> Self::ReturnTy {
         self.set3(d, s1, s2, mulhu)
     }
 
-    fn mul_upper_signed_unsigned(&mut self, d: Reg, s1: Reg, s2: Reg) -> Self::ReturnTy {
+    fn mul_upper_signed_unsigned(&mut self, d: RawReg, s1: RawReg, s2: RawReg) -> Self::ReturnTy {
         self.set3(d, s1, s2, |s1, s2| mulhsu(s1 as i32, s2) as u32)
     }
 
-    fn div_unsigned(&mut self, d: Reg, s1: Reg, s2: Reg) -> Self::ReturnTy {
+    fn div_unsigned(&mut self, d: RawReg, s1: RawReg, s2: RawReg) -> Self::ReturnTy {
         self.set3(d, s1, s2, divu)
     }
 
-    fn div_signed(&mut self, d: Reg, s1: Reg, s2: Reg) -> Self::ReturnTy {
+    fn div_signed(&mut self, d: RawReg, s1: RawReg, s2: RawReg) -> Self::ReturnTy {
         self.set3(d, s1, s2, |s1, s2| div(s1 as i32, s2 as i32) as u32)
     }
 
-    fn rem_unsigned(&mut self, d: Reg, s1: Reg, s2: Reg) -> Self::ReturnTy {
+    fn rem_unsigned(&mut self, d: RawReg, s1: RawReg, s2: RawReg) -> Self::ReturnTy {
         self.set3(d, s1, s2, remu)
     }
 
-    fn rem_signed(&mut self, d: Reg, s1: Reg, s2: Reg) -> Self::ReturnTy {
+    fn rem_signed(&mut self, d: RawReg, s1: RawReg, s2: RawReg) -> Self::ReturnTy {
         self.set3(d, s1, s2, |s1, s2| rem(s1 as i32, s2 as i32) as u32)
     }
 
-    fn set_less_than_unsigned_imm(&mut self, d: Reg, s1: Reg, s2: u32) -> Self::ReturnTy {
+    fn set_less_than_unsigned_imm(&mut self, d: RawReg, s1: RawReg, s2: u32) -> Self::ReturnTy {
         self.set3(d, s1, s2, |s1, s2| u32::from(s1 < s2))
     }
 
-    fn set_greater_than_unsigned_imm(&mut self, d: Reg, s1: Reg, s2: u32) -> Self::ReturnTy {
+    fn set_greater_than_unsigned_imm(&mut self, d: RawReg, s1: RawReg, s2: u32) -> Self::ReturnTy {
         self.set3(d, s1, s2, |s1, s2| u32::from(s1 > s2))
     }
 
-    fn set_less_than_signed_imm(&mut self, d: Reg, s1: Reg, s2: u32) -> Self::ReturnTy {
+    fn set_less_than_signed_imm(&mut self, d: RawReg, s1: RawReg, s2: u32) -> Self::ReturnTy {
         self.set3(d, s1, s2, |s1, s2| u32::from((s1 as i32) < (s2 as i32)))
     }
 
-    fn set_greater_than_signed_imm(&mut self, d: Reg, s1: Reg, s2: u32) -> Self::ReturnTy {
+    fn set_greater_than_signed_imm(&mut self, d: RawReg, s1: RawReg, s2: u32) -> Self::ReturnTy {
         self.set3(d, s1, s2, |s1, s2| u32::from((s1 as i32) > (s2 as i32)))
     }
 
-    fn shift_logical_right_imm(&mut self, d: Reg, s1: Reg, s2: u32) -> Self::ReturnTy {
+    fn shift_logical_right_imm(&mut self, d: RawReg, s1: RawReg, s2: u32) -> Self::ReturnTy {
         self.set3(d, s1, s2, u32::wrapping_shr)
     }
 
-    fn shift_logical_right_imm_alt(&mut self, d: Reg, s2: Reg, s1: u32) -> Self::ReturnTy {
+    fn shift_logical_right_imm_alt(&mut self, d: RawReg, s2: RawReg, s1: u32) -> Self::ReturnTy {
         self.set3(d, s1, s2, u32::wrapping_shr)
     }
 
-    fn shift_arithmetic_right_imm(&mut self, d: Reg, s1: Reg, s2: u32) -> Self::ReturnTy {
+    fn shift_arithmetic_right_imm(&mut self, d: RawReg, s1: RawReg, s2: u32) -> Self::ReturnTy {
         self.set3(d, s1, s2, |s1, s2| ((s1 as i32) >> s2) as u32)
     }
 
-    fn shift_arithmetic_right_imm_alt(&mut self, d: Reg, s2: Reg, s1: u32) -> Self::ReturnTy {
+    fn shift_arithmetic_right_imm_alt(&mut self, d: RawReg, s2: RawReg, s1: u32) -> Self::ReturnTy {
         self.set3(d, s1, s2, |s1, s2| ((s1 as i32) >> s2) as u32)
     }
 
-    fn shift_logical_left_imm(&mut self, d: Reg, s1: Reg, s2: u32) -> Self::ReturnTy {
+    fn shift_logical_left_imm(&mut self, d: RawReg, s1: RawReg, s2: u32) -> Self::ReturnTy {
         self.set3(d, s1, s2, u32::wrapping_shl)
     }
 
-    fn shift_logical_left_imm_alt(&mut self, d: Reg, s2: Reg, s1: u32) -> Self::ReturnTy {
+    fn shift_logical_left_imm_alt(&mut self, d: RawReg, s2: RawReg, s1: u32) -> Self::ReturnTy {
         self.set3(d, s1, s2, u32::wrapping_shl)
     }
 
-    fn or_imm(&mut self, d: Reg, s1: Reg, s2: u32) -> Self::ReturnTy {
+    fn or_imm(&mut self, d: RawReg, s1: RawReg, s2: u32) -> Self::ReturnTy {
         self.set3(d, s1, s2, |s1, s2| s1 | s2)
     }
 
-    fn and_imm(&mut self, d: Reg, s1: Reg, s2: u32) -> Self::ReturnTy {
+    fn and_imm(&mut self, d: RawReg, s1: RawReg, s2: u32) -> Self::ReturnTy {
         self.set3(d, s1, s2, |s1, s2| s1 & s2)
     }
 
-    fn xor_imm(&mut self, d: Reg, s1: Reg, s2: u32) -> Self::ReturnTy {
+    fn xor_imm(&mut self, d: RawReg, s1: RawReg, s2: u32) -> Self::ReturnTy {
         self.set3(d, s1, s2, |s1, s2| s1 ^ s2)
     }
 
-    fn load_imm(&mut self, dst: Reg, imm: u32) -> Self::ReturnTy {
+    fn load_imm(&mut self, dst: RawReg, imm: u32) -> Self::ReturnTy {
         self.set(dst, imm)?;
         Ok(())
     }
 
-    fn move_reg(&mut self, d: Reg, s: Reg) -> Self::ReturnTy {
+    fn move_reg(&mut self, d: RawReg, s: RawReg) -> Self::ReturnTy {
         let imm = self.get(s);
         self.set(d, imm)?;
         Ok(())
     }
 
-    fn cmov_if_zero(&mut self, d: Reg, s: Reg, c: Reg) -> Self::ReturnTy {
+    fn cmov_if_zero(&mut self, d: RawReg, s: RawReg, c: RawReg) -> Self::ReturnTy {
         if self.get(c) == 0 {
             let value = self.get(s);
             self.set(d, value)?;
@@ -1024,7 +1030,7 @@ impl<'a, 'b, const DEBUG: bool> InstructionVisitor for Visitor<'a, 'b, DEBUG> {
         Ok(())
     }
 
-    fn cmov_if_zero_imm(&mut self, d: Reg, c: Reg, s: u32) -> Self::ReturnTy {
+    fn cmov_if_zero_imm(&mut self, d: RawReg, c: RawReg, s: u32) -> Self::ReturnTy {
         if self.get(c) == 0 {
             self.set(d, s)?;
         }
@@ -1032,7 +1038,7 @@ impl<'a, 'b, const DEBUG: bool> InstructionVisitor for Visitor<'a, 'b, DEBUG> {
         Ok(())
     }
 
-    fn cmov_if_not_zero(&mut self, d: Reg, s: Reg, c: Reg) -> Self::ReturnTy {
+    fn cmov_if_not_zero(&mut self, d: RawReg, s: RawReg, c: RawReg) -> Self::ReturnTy {
         if self.get(c) != 0 {
             let value = self.get(s);
             self.set(d, value)?;
@@ -1041,7 +1047,7 @@ impl<'a, 'b, const DEBUG: bool> InstructionVisitor for Visitor<'a, 'b, DEBUG> {
         Ok(())
     }
 
-    fn cmov_if_not_zero_imm(&mut self, d: Reg, c: Reg, s: u32) -> Self::ReturnTy {
+    fn cmov_if_not_zero_imm(&mut self, d: RawReg, c: RawReg, s: u32) -> Self::ReturnTy {
         if self.get(c) != 0 {
             self.set(d, s)?;
         }
@@ -1049,159 +1055,159 @@ impl<'a, 'b, const DEBUG: bool> InstructionVisitor for Visitor<'a, 'b, DEBUG> {
         Ok(())
     }
 
-    fn add_imm(&mut self, d: Reg, s1: Reg, s2: u32) -> Self::ReturnTy {
+    fn add_imm(&mut self, d: RawReg, s1: RawReg, s2: u32) -> Self::ReturnTy {
         self.set3(d, s1, s2, u32::wrapping_add)
     }
 
-    fn store_imm_u8(&mut self, value: u32, offset: u32) -> Self::ReturnTy {
+    fn store_imm_u8(&mut self, offset: u32, value: u32) -> Self::ReturnTy {
         self.store::<u8>(value, None, offset)
     }
 
-    fn store_imm_u16(&mut self, value: u32, offset: u32) -> Self::ReturnTy {
+    fn store_imm_u16(&mut self, offset: u32, value: u32) -> Self::ReturnTy {
         self.store::<u16>(value, None, offset)
     }
 
-    fn store_imm_u32(&mut self, value: u32, offset: u32) -> Self::ReturnTy {
+    fn store_imm_u32(&mut self, offset: u32, value: u32) -> Self::ReturnTy {
         self.store::<u32>(value, None, offset)
     }
 
-    fn store_imm_indirect_u8(&mut self, base: Reg, offset: u32, value: u32) -> Self::ReturnTy {
+    fn store_imm_indirect_u8(&mut self, base: RawReg, offset: u32, value: u32) -> Self::ReturnTy {
         self.store::<u8>(value, Some(base), offset)
     }
 
-    fn store_imm_indirect_u16(&mut self, base: Reg, offset: u32, value: u32) -> Self::ReturnTy {
+    fn store_imm_indirect_u16(&mut self, base: RawReg, offset: u32, value: u32) -> Self::ReturnTy {
         self.store::<u16>(value, Some(base), offset)
     }
 
-    fn store_imm_indirect_u32(&mut self, base: Reg, offset: u32, value: u32) -> Self::ReturnTy {
+    fn store_imm_indirect_u32(&mut self, base: RawReg, offset: u32, value: u32) -> Self::ReturnTy {
         self.store::<u32>(value, Some(base), offset)
     }
 
-    fn store_indirect_u8(&mut self, src: Reg, base: Reg, offset: u32) -> Self::ReturnTy {
+    fn store_indirect_u8(&mut self, src: RawReg, base: RawReg, offset: u32) -> Self::ReturnTy {
         self.store::<u8>(src, Some(base), offset)
     }
 
-    fn store_indirect_u16(&mut self, src: Reg, base: Reg, offset: u32) -> Self::ReturnTy {
+    fn store_indirect_u16(&mut self, src: RawReg, base: RawReg, offset: u32) -> Self::ReturnTy {
         self.store::<u16>(src, Some(base), offset)
     }
 
-    fn store_indirect_u32(&mut self, src: Reg, base: Reg, offset: u32) -> Self::ReturnTy {
+    fn store_indirect_u32(&mut self, src: RawReg, base: RawReg, offset: u32) -> Self::ReturnTy {
         self.store::<u32>(src, Some(base), offset)
     }
 
-    fn store_u8(&mut self, src: Reg, offset: u32) -> Self::ReturnTy {
+    fn store_u8(&mut self, src: RawReg, offset: u32) -> Self::ReturnTy {
         self.store::<u8>(src, None, offset)
     }
 
-    fn store_u16(&mut self, src: Reg, offset: u32) -> Self::ReturnTy {
+    fn store_u16(&mut self, src: RawReg, offset: u32) -> Self::ReturnTy {
         self.store::<u16>(src, None, offset)
     }
 
-    fn store_u32(&mut self, src: Reg, offset: u32) -> Self::ReturnTy {
+    fn store_u32(&mut self, src: RawReg, offset: u32) -> Self::ReturnTy {
         self.store::<u32>(src, None, offset)
     }
 
-    fn load_u8(&mut self, dst: Reg, offset: u32) -> Self::ReturnTy {
+    fn load_u8(&mut self, dst: RawReg, offset: u32) -> Self::ReturnTy {
         self.load::<u8>(dst, None, offset)
     }
 
-    fn load_i8(&mut self, dst: Reg, offset: u32) -> Self::ReturnTy {
+    fn load_i8(&mut self, dst: RawReg, offset: u32) -> Self::ReturnTy {
         self.load::<i8>(dst, None, offset)
     }
 
-    fn load_u16(&mut self, dst: Reg, offset: u32) -> Self::ReturnTy {
+    fn load_u16(&mut self, dst: RawReg, offset: u32) -> Self::ReturnTy {
         self.load::<u16>(dst, None, offset)
     }
 
-    fn load_i16(&mut self, dst: Reg, offset: u32) -> Self::ReturnTy {
+    fn load_i16(&mut self, dst: RawReg, offset: u32) -> Self::ReturnTy {
         self.load::<i16>(dst, None, offset)
     }
 
-    fn load_u32(&mut self, dst: Reg, offset: u32) -> Self::ReturnTy {
+    fn load_u32(&mut self, dst: RawReg, offset: u32) -> Self::ReturnTy {
         self.load::<u32>(dst, None, offset)
     }
 
-    fn load_indirect_u8(&mut self, dst: Reg, base: Reg, offset: u32) -> Self::ReturnTy {
+    fn load_indirect_u8(&mut self, dst: RawReg, base: RawReg, offset: u32) -> Self::ReturnTy {
         self.load::<u8>(dst, Some(base), offset)
     }
 
-    fn load_indirect_i8(&mut self, dst: Reg, base: Reg, offset: u32) -> Self::ReturnTy {
+    fn load_indirect_i8(&mut self, dst: RawReg, base: RawReg, offset: u32) -> Self::ReturnTy {
         self.load::<i8>(dst, Some(base), offset)
     }
 
-    fn load_indirect_u16(&mut self, dst: Reg, base: Reg, offset: u32) -> Self::ReturnTy {
+    fn load_indirect_u16(&mut self, dst: RawReg, base: RawReg, offset: u32) -> Self::ReturnTy {
         self.load::<u16>(dst, Some(base), offset)
     }
 
-    fn load_indirect_i16(&mut self, dst: Reg, base: Reg, offset: u32) -> Self::ReturnTy {
+    fn load_indirect_i16(&mut self, dst: RawReg, base: RawReg, offset: u32) -> Self::ReturnTy {
         self.load::<i16>(dst, Some(base), offset)
     }
 
-    fn load_indirect_u32(&mut self, dst: Reg, base: Reg, offset: u32) -> Self::ReturnTy {
+    fn load_indirect_u32(&mut self, dst: RawReg, base: RawReg, offset: u32) -> Self::ReturnTy {
         self.load::<u32>(dst, Some(base), offset)
     }
 
-    fn branch_less_unsigned(&mut self, s1: Reg, s2: Reg, i: u32) -> Self::ReturnTy {
+    fn branch_less_unsigned(&mut self, s1: RawReg, s2: RawReg, i: u32) -> Self::ReturnTy {
         self.branch(s1, s2, i, |s1, s2| s1 < s2)
     }
 
-    fn branch_less_unsigned_imm(&mut self, s1: Reg, s2: u32, i: u32) -> Self::ReturnTy {
+    fn branch_less_unsigned_imm(&mut self, s1: RawReg, s2: u32, i: u32) -> Self::ReturnTy {
         self.branch(s1, s2, i, |s1, s2| s1 < s2)
     }
 
-    fn branch_less_signed(&mut self, s1: Reg, s2: Reg, i: u32) -> Self::ReturnTy {
+    fn branch_less_signed(&mut self, s1: RawReg, s2: RawReg, i: u32) -> Self::ReturnTy {
         self.branch(s1, s2, i, |s1, s2| (s1 as i32) < (s2 as i32))
     }
 
-    fn branch_less_signed_imm(&mut self, s1: Reg, s2: u32, i: u32) -> Self::ReturnTy {
+    fn branch_less_signed_imm(&mut self, s1: RawReg, s2: u32, i: u32) -> Self::ReturnTy {
         self.branch(s1, s2, i, |s1, s2| (s1 as i32) < (s2 as i32))
     }
 
-    fn branch_eq(&mut self, s1: Reg, s2: Reg, i: u32) -> Self::ReturnTy {
+    fn branch_eq(&mut self, s1: RawReg, s2: RawReg, i: u32) -> Self::ReturnTy {
         self.branch(s1, s2, i, |s1, s2| s1 == s2)
     }
 
-    fn branch_eq_imm(&mut self, s1: Reg, s2: u32, i: u32) -> Self::ReturnTy {
+    fn branch_eq_imm(&mut self, s1: RawReg, s2: u32, i: u32) -> Self::ReturnTy {
         self.branch(s1, s2, i, |s1, s2| s1 == s2)
     }
 
-    fn branch_not_eq(&mut self, s1: Reg, s2: Reg, i: u32) -> Self::ReturnTy {
+    fn branch_not_eq(&mut self, s1: RawReg, s2: RawReg, i: u32) -> Self::ReturnTy {
         self.branch(s1, s2, i, |s1, s2| s1 != s2)
     }
 
-    fn branch_not_eq_imm(&mut self, s1: Reg, s2: u32, i: u32) -> Self::ReturnTy {
+    fn branch_not_eq_imm(&mut self, s1: RawReg, s2: u32, i: u32) -> Self::ReturnTy {
         self.branch(s1, s2, i, |s1, s2| s1 != s2)
     }
 
-    fn branch_greater_or_equal_unsigned(&mut self, s1: Reg, s2: Reg, i: u32) -> Self::ReturnTy {
+    fn branch_greater_or_equal_unsigned(&mut self, s1: RawReg, s2: RawReg, i: u32) -> Self::ReturnTy {
         self.branch(s1, s2, i, |s1, s2| s1 >= s2)
     }
 
-    fn branch_greater_or_equal_unsigned_imm(&mut self, s1: Reg, s2: u32, i: u32) -> Self::ReturnTy {
+    fn branch_greater_or_equal_unsigned_imm(&mut self, s1: RawReg, s2: u32, i: u32) -> Self::ReturnTy {
         self.branch(s1, s2, i, |s1, s2| s1 >= s2)
     }
 
-    fn branch_greater_or_equal_signed(&mut self, s1: Reg, s2: Reg, i: u32) -> Self::ReturnTy {
+    fn branch_greater_or_equal_signed(&mut self, s1: RawReg, s2: RawReg, i: u32) -> Self::ReturnTy {
         self.branch(s1, s2, i, |s1, s2| (s1 as i32) >= (s2 as i32))
     }
 
-    fn branch_greater_or_equal_signed_imm(&mut self, s1: Reg, s2: u32, i: u32) -> Self::ReturnTy {
+    fn branch_greater_or_equal_signed_imm(&mut self, s1: RawReg, s2: u32, i: u32) -> Self::ReturnTy {
         self.branch(s1, s2, i, |s1, s2| (s1 as i32) >= (s2 as i32))
     }
 
-    fn branch_less_or_equal_unsigned_imm(&mut self, s1: Reg, s2: u32, i: u32) -> Self::ReturnTy {
+    fn branch_less_or_equal_unsigned_imm(&mut self, s1: RawReg, s2: u32, i: u32) -> Self::ReturnTy {
         self.branch(s1, s2, i, |s1, s2| s1 <= s2)
     }
 
-    fn branch_less_or_equal_signed_imm(&mut self, s1: Reg, s2: u32, i: u32) -> Self::ReturnTy {
+    fn branch_less_or_equal_signed_imm(&mut self, s1: RawReg, s2: u32, i: u32) -> Self::ReturnTy {
         self.branch(s1, s2, i, |s1, s2| (s1 as i32) <= (s2 as i32))
     }
 
-    fn branch_greater_unsigned_imm(&mut self, s1: Reg, s2: u32, i: u32) -> Self::ReturnTy {
+    fn branch_greater_unsigned_imm(&mut self, s1: RawReg, s2: u32, i: u32) -> Self::ReturnTy {
         self.branch(s1, s2, i, |s1, s2| s1 > s2)
     }
 
-    fn branch_greater_signed_imm(&mut self, s1: Reg, s2: u32, i: u32) -> Self::ReturnTy {
+    fn branch_greater_signed_imm(&mut self, s1: RawReg, s2: u32, i: u32) -> Self::ReturnTy {
         self.branch(s1, s2, i, |s1, s2| (s1 as i32) > (s2 as i32))
     }
 
@@ -1214,16 +1220,16 @@ impl<'a, 'b, const DEBUG: bool> InstructionVisitor for Visitor<'a, 'b, DEBUG> {
         self.inner.on_start_new_basic_block::<DEBUG>()
     }
 
-    fn jump_indirect(&mut self, base: Reg, offset: u32) -> Self::ReturnTy {
+    fn jump_indirect(&mut self, base: RawReg, offset: u32) -> Self::ReturnTy {
         self.jump_indirect_impl(None, base, offset)
     }
 
-    fn load_imm_and_jump(&mut self, ra: Reg, value: u32, target: u32) -> Self::ReturnTy {
+    fn load_imm_and_jump(&mut self, ra: RawReg, value: u32, target: u32) -> Self::ReturnTy {
         self.load_imm(ra, value)?;
         self.jump(target)
     }
 
-    fn load_imm_and_jump_indirect(&mut self, ra: Reg, base: Reg, value: u32, offset: u32) -> Self::ReturnTy {
+    fn load_imm_and_jump_indirect(&mut self, ra: RawReg, base: RawReg, value: u32, offset: u32) -> Self::ReturnTy {
         self.jump_indirect_impl(Some((ra, value)), base, offset)
     }
 }
diff --git a/crates/polkavm/src/tests.rs b/crates/polkavm/src/tests.rs
index fbdc581e..d8bddcdf 100644
--- a/crates/polkavm/src/tests.rs
+++ b/crates/polkavm/src/tests.rs
@@ -96,7 +96,7 @@ fn basic_test_blob() -> ProgramBlob<'static> {
     builder.add_import("hostcall".into());
     builder.set_code(
         &[
-            asm::store_imm_u32(0x12345678, memory_map.rw_data_address()),
+            asm::store_imm_u32(memory_map.rw_data_address(), 0x12345678),
             asm::add(S0, A0, A1),
             asm::ecalli(0),
             asm::add(A0, A0, S0),
diff --git a/crates/polkavm/src/utils.rs b/crates/polkavm/src/utils.rs
index 0c78f716..4a0662ee 100644
--- a/crates/polkavm/src/utils.rs
+++ b/crates/polkavm/src/utils.rs
@@ -1,15 +1,15 @@
 use alloc::vec::Vec;
-use polkavm_common::program::Reg;
+use polkavm_common::program::RawReg;
 
 #[derive(Copy, Clone)]
 pub enum RegImm {
-    Reg(Reg),
+    Reg(RawReg),
     Imm(u32),
 }
 
-impl From<Reg> for RegImm {
+impl From<RawReg> for RegImm {
     #[inline]
-    fn from(reg: Reg) -> Self {
+    fn from(reg: RawReg) -> Self {
         RegImm::Reg(reg)
     }
 }
@@ -58,6 +58,11 @@ where
         self.inner.get(key as usize).and_then(|value| *value)
     }
 
+    #[inline]
+    pub fn len(&self) -> u32 {
+        self.inner.len() as u32
+    }
+
     #[inline]
     pub fn insert(&mut self, key: u32, value: T) {
         self.inner[key as usize] = Some(value);
diff --git a/examples/doom/roms/doom.polkavm b/examples/doom/roms/doom.polkavm
index c1f9707e..debdc353 100644
Binary files a/examples/doom/roms/doom.polkavm and b/examples/doom/roms/doom.polkavm differ