movzx function return values, to clear upper bits in zig function results

crates/compiler/gen_dev/src/generic64/aarch64.rs

@@ -795,6 +795,16 @@ impl Assembler<AArch64GeneralReg, AArch64FloatReg> for AArch64Assembler {
         todo!("move with sign extension");
     }
 
+    #[inline(always)]
+    fn movzx_reg_reg(
+        _buf: &mut Vec<'_, u8>,
+        _input_width: RegisterWidth,
+        _dst: AArch64GeneralReg,
+        _src: AArch64GeneralReg,
+    ) {
+        todo!("move with zero extension");
+    }
+
     #[inline(always)]
     fn mov_freg64_base32(_buf: &mut Vec<'_, u8>, _dst: AArch64FloatReg, _offset: i32) {
         todo!("loading floating point reg from base offset for AArch64");

crates/compiler/gen_dev/src/generic64/mod.rs

@@ -306,6 +306,14 @@ pub trait Assembler<GeneralReg: RegTrait, FloatReg: RegTrait>: Sized + Copy {
         src: GeneralReg,
     );
 
+    // move with zero extension
+    fn movzx_reg_reg(
+        buf: &mut Vec<'_, u8>,
+        input_width: RegisterWidth,
+        dst: GeneralReg,
+        src: GeneralReg,
+    );
+
     // base32 is similar to stack based instructions but they reference the base/frame pointer.
     fn mov_freg64_base32(buf: &mut Vec<'_, u8>, dst: FloatReg, offset: i32);
 
@@ -914,7 +922,7 @@ impl<
                 let width = RegisterWidth::try_from_layout(ret_repr).unwrap();
 
                 let dst_reg = self.storage_manager.claim_general_reg(&mut self.buf, dst);
-                ASM::mov_reg_reg(&mut self.buf, width, dst_reg, CC::GENERAL_RETURN_REGS[0]);
+                ASM::movzx_reg_reg(&mut self.buf, width, dst_reg, CC::GENERAL_RETURN_REGS[0]);
             }
             single_register_floats!() => {
                 let dst_reg = self.storage_manager.claim_float_reg(&mut self.buf, dst);

crates/compiler/gen_dev/src/generic64/x86_64.rs

@@ -1511,6 +1511,7 @@ impl Assembler<X86_64GeneralReg, X86_64FloatReg> for X86_64Assembler {
     ) {
         mov_reg_reg(buf, register_width, dst, src);
     }
+
     #[inline(always)]
     fn movsx_reg_reg(
         buf: &mut Vec<'_, u8>,
@@ -1518,7 +1519,27 @@ impl Assembler<X86_64GeneralReg, X86_64FloatReg> for X86_64Assembler {
         dst: X86_64GeneralReg,
         src: X86_64GeneralReg,
     ) {
-        raw_movsx_reg_reg(buf, input_width, dst, src);
+        use RegisterWidth::*;
+
+        match input_width {
+            W8 | W16 | W32 => raw_movsx_reg_reg(buf, input_width, dst, src),
+            W64 => mov_reg_reg(buf, input_width, dst, src),
+        }
+    }
+
+    #[inline(always)]
+    fn movzx_reg_reg(
+        buf: &mut Vec<'_, u8>,
+        input_width: RegisterWidth,
+        dst: X86_64GeneralReg,
+        src: X86_64GeneralReg,
+    ) {
+        use RegisterWidth::*;
+
+        match input_width {
+            W8 | W16 => raw_movzx_reg_reg(buf, input_width, dst, src),
+            W32 | W64 => mov_reg_reg(buf, input_width, dst, src),
+        }
     }
 
     #[inline(always)]
@@ -2677,6 +2698,35 @@ fn raw_movsx_reg_reg(
     }
 }
 
+#[allow(unused)]
+fn raw_movzx_reg_reg(
+    buf: &mut Vec<u8>,
+    input_width: RegisterWidth,
+    dst: X86_64GeneralReg,
+    src: X86_64GeneralReg,
+) {
+    let dst_high = dst as u8 > 7;
+    let dst_mod = dst as u8 % 8;
+    let src_high = src as u8 > 7;
+    let src_mod = src as u8 % 8;
+
+    // NOTE src and dst seem to be flipped here. It works this way though
+    let mod_rm = 0xC0 | (dst_mod << 3) | src_mod;
+
+    let rex = add_rm_extension(src, REX_W);
+    let rex = add_reg_extension(dst, rex);
+
+    match input_width {
+        RegisterWidth::W8 => {
+            buf.extend([rex, 0x0f, 0xb6, mod_rm]);
+        }
+        RegisterWidth::W16 => {
+            buf.extend([rex, 0x0f, 0xb7, mod_rm]);
+        }
+        RegisterWidth::W32 | RegisterWidth::W64 => { /* do nothing */ }
+    }
+}
+
 /// `MOV r/m64,r64` -> Move r64 to r/m64.
 /// This will not generate anything if dst and src are the same.
 #[inline(always)]
@@ -3852,6 +3902,32 @@ mod tests {
         );
     }
 
+    #[test]
+    fn test_movzx_reg64_reg64() {
+        disassembler_test!(
+            raw_movzx_reg_reg,
+            |w, reg1, reg2| {
+                match w {
+                    RegisterWidth::W8 => format!(
+                        "movzx {}, {}",
+                        reg1,
+                        X86_64GeneralReg::low_8bits_string(&reg2)
+                    ),
+                    RegisterWidth::W16 => format!(
+                        "movzx {}, {}",
+                        reg1,
+                        X86_64GeneralReg::low_16bits_string(&reg2)
+                    ),
+                    RegisterWidth::W32 => String::new(),
+                    RegisterWidth::W64 => String::new(),
+                }
+            },
+            ALL_REGISTER_WIDTHS,
+            ALL_GENERAL_REGS,
+            ALL_GENERAL_REGS
+        );
+    }
+
     #[test]
     fn test_movsd_freg64_base64_offset32() {
         disassembler_test!(