Merge pull request #3958 from roc-lang/dev-backend-bitwise-shifts

Dev backend bitwise shifts

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

@@ -899,6 +899,45 @@ impl Assembler<AArch64GeneralReg, AArch64FloatReg> for AArch64Assembler {
     ) {
         todo!("bitwise xor for AArch64")
     }
+
+    fn shl_reg64_reg64_reg64<'a, 'r, ASM, CC>(
+        _buf: &mut Vec<'a, u8>,
+        _storage_manager: &mut StorageManager<'a, 'r, AArch64GeneralReg, AArch64FloatReg, ASM, CC>,
+        _dst: AArch64GeneralReg,
+        _src1: AArch64GeneralReg,
+        _src2: AArch64GeneralReg,
+    ) where
+        ASM: Assembler<AArch64GeneralReg, AArch64FloatReg>,
+        CC: CallConv<AArch64GeneralReg, AArch64FloatReg, ASM>,
+    {
+        todo!("shl for AArch64")
+    }
+
+    fn shr_reg64_reg64_reg64<'a, 'r, ASM, CC>(
+        _buf: &mut Vec<'a, u8>,
+        _storage_manager: &mut StorageManager<'a, 'r, AArch64GeneralReg, AArch64FloatReg, ASM, CC>,
+        _dst: AArch64GeneralReg,
+        _src1: AArch64GeneralReg,
+        _src2: AArch64GeneralReg,
+    ) where
+        ASM: Assembler<AArch64GeneralReg, AArch64FloatReg>,
+        CC: CallConv<AArch64GeneralReg, AArch64FloatReg, ASM>,
+    {
+        todo!("shr for AArch64")
+    }
+
+    fn sar_reg64_reg64_reg64<'a, 'r, ASM, CC>(
+        _buf: &mut Vec<'a, u8>,
+        _storage_manager: &mut StorageManager<'a, 'r, AArch64GeneralReg, AArch64FloatReg, ASM, CC>,
+        _dst: AArch64GeneralReg,
+        _src1: AArch64GeneralReg,
+        _src2: AArch64GeneralReg,
+    ) where
+        ASM: Assembler<AArch64GeneralReg, AArch64FloatReg>,
+        CC: CallConv<AArch64GeneralReg, AArch64FloatReg, ASM>,
+    {
+        todo!("sar for AArch64")
+    }
 }
 
 impl AArch64Assembler {}

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

@@ -170,6 +170,36 @@ pub trait Assembler<GeneralReg: RegTrait, FloatReg: RegTrait>: Sized + Copy {
         src2: GeneralReg,
     );
 
+    fn shl_reg64_reg64_reg64<'a, 'r, ASM, CC>(
+        buf: &mut Vec<'a, u8>,
+        storage_manager: &mut StorageManager<'a, 'r, GeneralReg, FloatReg, ASM, CC>,
+        dst: GeneralReg,
+        src1: GeneralReg,
+        src2: GeneralReg,
+    ) where
+        ASM: Assembler<GeneralReg, FloatReg>,
+        CC: CallConv<GeneralReg, FloatReg, ASM>;
+
+    fn shr_reg64_reg64_reg64<'a, 'r, ASM, CC>(
+        buf: &mut Vec<'a, u8>,
+        storage_manager: &mut StorageManager<'a, 'r, GeneralReg, FloatReg, ASM, CC>,
+        dst: GeneralReg,
+        src1: GeneralReg,
+        src2: GeneralReg,
+    ) where
+        ASM: Assembler<GeneralReg, FloatReg>,
+        CC: CallConv<GeneralReg, FloatReg, ASM>;
+
+    fn sar_reg64_reg64_reg64<'a, 'r, ASM, CC>(
+        buf: &mut Vec<'a, u8>,
+        storage_manager: &mut StorageManager<'a, 'r, GeneralReg, FloatReg, ASM, CC>,
+        dst: GeneralReg,
+        src1: GeneralReg,
+        src2: GeneralReg,
+    ) where
+        ASM: Assembler<GeneralReg, FloatReg>,
+        CC: CallConv<GeneralReg, FloatReg, ASM>;
+
     fn call(buf: &mut Vec<'_, u8>, relocs: &mut Vec<'_, Relocation>, fn_name: String);
 
     /// Jumps by an offset of offset bytes unconditionally.
@@ -2032,6 +2062,127 @@ impl<
             }
         }
     }
+
+    fn build_int_shift_left(
+        &mut self,
+        dst: &Symbol,
+        src1: &Symbol,
+        src2: &Symbol,
+        int_width: IntWidth,
+    ) {
+        let buf = &mut self.buf;
+
+        match int_width {
+            IntWidth::U128 | IntWidth::I128 => todo!(),
+            _ => {
+                let dst_reg = self.storage_manager.claim_general_reg(buf, dst);
+                let src1_reg = self.storage_manager.load_to_general_reg(buf, src1);
+                let src2_reg = self.storage_manager.load_to_general_reg(buf, src2);
+
+                ASM::shl_reg64_reg64_reg64(
+                    buf,
+                    &mut self.storage_manager,
+                    dst_reg,
+                    src1_reg,
+                    src2_reg,
+                );
+            }
+        }
+    }
+
+    fn build_int_shift_right(
+        &mut self,
+        dst: &Symbol,
+        src1: &Symbol,
+        src2: &Symbol,
+        int_width: IntWidth,
+    ) {
+        let buf = &mut self.buf;
+
+        match int_width {
+            IntWidth::U128 | IntWidth::I128 => todo!(),
+            _ => {
+                let dst_reg = self.storage_manager.claim_general_reg(buf, dst);
+                let src1_reg = self.storage_manager.load_to_general_reg(buf, src1);
+                let src2_reg = self.storage_manager.load_to_general_reg(buf, src2);
+
+                // to get sign extension "for free", we move our bits to the left
+                // so the integers sign bit is stored in the register's sign bit.
+                // Then we arithmetic shift right, getting the correct sign extension behavior,
+                // then shift logical right to get the bits back into the position they should
+                // be for our particular integer width
+                let sign_extend_shift_amount = 64 - (int_width.stack_size() as i64 * 8);
+
+                if sign_extend_shift_amount > 0 {
+                    self.storage_manager.with_tmp_general_reg(
+                        buf,
+                        |storage_manager, buf, tmp_reg| {
+                            ASM::mov_reg64_imm64(buf, tmp_reg, sign_extend_shift_amount);
+                            ASM::shl_reg64_reg64_reg64(
+                                buf,
+                                storage_manager,
+                                src1_reg,
+                                src1_reg,
+                                tmp_reg,
+                            );
+                        },
+                    )
+                }
+
+                ASM::sar_reg64_reg64_reg64(
+                    buf,
+                    &mut self.storage_manager,
+                    dst_reg,
+                    src1_reg,
+                    src2_reg,
+                );
+
+                if sign_extend_shift_amount > 0 {
+                    // shift back if needed
+                    self.storage_manager.with_tmp_general_reg(
+                        &mut self.buf,
+                        |storage_manager, buf, tmp_reg| {
+                            ASM::mov_reg64_imm64(buf, tmp_reg, sign_extend_shift_amount);
+                            ASM::shr_reg64_reg64_reg64(
+                                buf,
+                                storage_manager,
+                                dst_reg,
+                                dst_reg,
+                                tmp_reg,
+                            );
+                        },
+                    )
+                }
+            }
+        }
+    }
+
+    fn build_int_shift_right_zero_fill(
+        &mut self,
+        dst: &Symbol,
+        src1: &Symbol,
+        src2: &Symbol,
+        int_width: IntWidth,
+    ) {
+        let buf = &mut self.buf;
+
+        match int_width {
+            IntWidth::U128 | IntWidth::I128 => todo!(),
+            _ => {
+                let dst_reg = self.storage_manager.claim_general_reg(buf, dst);
+                let src1_reg = self.storage_manager.load_to_general_reg(buf, src1);
+                let src2_reg = self.storage_manager.load_to_general_reg(buf, src2);
+
+                ASM::shr_reg64_reg64_reg64(
+                    buf,
+                    &mut self.storage_manager,
+                    dst_reg,
+                    src1_reg,
+                    src2_reg,
+                );
+            }
+        }
+    }
 }
 
 /// This impl block is for ir related instructions that need backend specific information.

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

@@ -1452,6 +1452,83 @@ impl Assembler<X86_64GeneralReg, X86_64FloatReg> for X86_64Assembler {
     fn xor_reg64_reg64_reg64(buf: &mut Vec<'_, u8>, dst: Reg64, src1: Reg64, src2: Reg64) {
         binop_move_src_to_dst_reg64(buf, xor_reg64_reg64, dst, src1, src2)
     }
+
+    fn shl_reg64_reg64_reg64<'a, 'r, ASM, CC>(
+        buf: &mut Vec<'a, u8>,
+        storage_manager: &mut StorageManager<'a, 'r, X86_64GeneralReg, X86_64FloatReg, ASM, CC>,
+        dst: X86_64GeneralReg,
+        src1: X86_64GeneralReg,
+        src2: X86_64GeneralReg,
+    ) where
+        ASM: Assembler<X86_64GeneralReg, X86_64FloatReg>,
+        CC: CallConv<X86_64GeneralReg, X86_64FloatReg, ASM>,
+    {
+        shift_reg64_reg64_reg64(buf, storage_manager, shl_reg64_reg64, dst, src1, src2)
+    }
+
+    fn shr_reg64_reg64_reg64<'a, 'r, ASM, CC>(
+        buf: &mut Vec<'a, u8>,
+        storage_manager: &mut StorageManager<'a, 'r, X86_64GeneralReg, X86_64FloatReg, ASM, CC>,
+        dst: X86_64GeneralReg,
+        src1: X86_64GeneralReg,
+        src2: X86_64GeneralReg,
+    ) where
+        ASM: Assembler<X86_64GeneralReg, X86_64FloatReg>,
+        CC: CallConv<X86_64GeneralReg, X86_64FloatReg, ASM>,
+    {
+        shift_reg64_reg64_reg64(buf, storage_manager, shr_reg64_reg64, dst, src1, src2)
+    }
+
+    fn sar_reg64_reg64_reg64<'a, 'r, ASM, CC>(
+        buf: &mut Vec<'a, u8>,
+        storage_manager: &mut StorageManager<'a, 'r, X86_64GeneralReg, X86_64FloatReg, ASM, CC>,
+        dst: X86_64GeneralReg,
+        src1: X86_64GeneralReg,
+        src2: X86_64GeneralReg,
+    ) where
+        ASM: Assembler<X86_64GeneralReg, X86_64FloatReg>,
+        CC: CallConv<X86_64GeneralReg, X86_64FloatReg, ASM>,
+    {
+        shift_reg64_reg64_reg64(buf, storage_manager, sar_reg64_reg64, dst, src1, src2)
+    }
+}
+
+fn shift_reg64_reg64_reg64<'a, 'r, ASM, CC>(
+    buf: &mut Vec<'a, u8>,
+    storage_manager: &mut StorageManager<'a, 'r, X86_64GeneralReg, X86_64FloatReg, ASM, CC>,
+    shift_function: fn(buf: &mut Vec<'_, u8>, X86_64GeneralReg),
+    dst: X86_64GeneralReg,
+    src1: X86_64GeneralReg,
+    src2: X86_64GeneralReg,
+) where
+    ASM: Assembler<X86_64GeneralReg, X86_64FloatReg>,
+    CC: CallConv<X86_64GeneralReg, X86_64FloatReg, ASM>,
+{
+    macro_rules! helper {
+        ($buf:expr, $dst:expr, $src1:expr, $src2:expr) => {{
+            mov_reg64_reg64($buf, $dst, $src1);
+            mov_reg64_reg64($buf, X86_64GeneralReg::RCX, $src2);
+
+            shift_function($buf, $dst)
+        }};
+    }
+
+    // if RCX is one of our input registers, we need to move some stuff around
+    if let X86_64GeneralReg::RCX = dst {
+        storage_manager.with_tmp_general_reg(buf, |_, buf, tmp| {
+            helper!(buf, tmp, src1, src2);
+
+            mov_reg64_reg64(buf, dst, tmp);
+        })
+    } else if let X86_64GeneralReg::RCX = src2 {
+        storage_manager.with_tmp_general_reg(buf, |_, buf, tmp| {
+            mov_reg64_reg64(buf, tmp, src2);
+
+            helper!(buf, dst, src1, tmp);
+        })
+    } else {
+        helper!(buf, dst, src1, src2)
+    }
 }
 
 impl X86_64Assembler {
@@ -1576,6 +1653,36 @@ fn xor_reg64_reg64(buf: &mut Vec<'_, u8>, dst: X86_64GeneralReg, src: X86_64Gene
     binop_reg64_reg64(0x33, buf, src, dst);
 }
 
+/// `SHL r/m64, CL` -> Multiply r/m64 by 2, CL times.
+#[inline(always)]
+fn shl_reg64_reg64(buf: &mut Vec<'_, u8>, dst: X86_64GeneralReg) {
+    let rex = add_rm_extension(dst, REX_W);
+    let rex = add_reg_extension(dst, rex);
+
+    let dst_mod = dst as u8 % 8;
+    buf.extend([rex, 0xD3, 0xC0 | (4 << 3) | dst_mod]);
+}
+
+/// `SHR r/m64, CL` -> Unsigned divide r/m64 by 2, CL times.
+#[inline(always)]
+fn shr_reg64_reg64(buf: &mut Vec<'_, u8>, dst: X86_64GeneralReg) {
+    let rex = add_rm_extension(dst, REX_W);
+    let rex = add_reg_extension(dst, rex);
+
+    let dst_mod = dst as u8 % 8;
+    buf.extend([rex, 0xD3, 0xC0 | (5 << 3) | dst_mod]);
+}
+
+/// `SAR r/m64, CL` -> Signed divide r/m64 by 2, CL times.
+#[inline(always)]
+fn sar_reg64_reg64(buf: &mut Vec<'_, u8>, dst: X86_64GeneralReg) {
+    let rex = add_rm_extension(dst, REX_W);
+    let rex = add_reg_extension(dst, rex);
+
+    let dst_mod = dst as u8 % 8;
+    buf.extend([rex, 0xD3, 0xC0 | (7 << 3) | dst_mod]);
+}
+
 /// `ADDSD xmm1,xmm2/m64` -> Add the low double-precision floating-point value from xmm2/mem to xmm1 and store the result in xmm1.
 #[inline(always)]
 fn addsd_freg64_freg64(buf: &mut Vec<'_, u8>, dst: X86_64FloatReg, src: X86_64FloatReg) {
@@ -2447,6 +2554,33 @@ mod tests {
         );
     }
 
+    #[test]
+    fn test_shl_reg64_reg64() {
+        disassembler_test!(
+            shl_reg64_reg64,
+            |reg| format!("shl {reg}, cl"),
+            ALL_GENERAL_REGS
+        );
+    }
+
+    #[test]
+    fn test_shr_reg64_reg64() {
+        disassembler_test!(
+            shr_reg64_reg64,
+            |reg| format!("shr {reg}, cl"),
+            ALL_GENERAL_REGS
+        );
+    }
+
+    #[test]
+    fn test_sar_reg64_reg64() {
+        disassembler_test!(
+            sar_reg64_reg64,
+            |reg| format!("sar {reg}, cl"),
+            ALL_GENERAL_REGS
+        );
+    }
+
     #[test]
     fn test_cmovl_reg64_reg64() {
         disassembler_test!(