asm: refactor how REX suffixes are emitted (#10749)

This change simplifies the generated code for encoding REX suffixes
(ModR/M byte, SIB byte, displacement) but does not alter the underlying
logic. As a part of this refactoring, the displacement struct is now
appropriately named `Disp` and several public functions no longer need
to be so.

Author: abrown

cranelift/assembler-x64/meta/src/generate/format.rs

@@ -136,64 +136,18 @@ impl dsl::Format {
             [FixedReg(_), Imm(_)] => {
                 // No need to emit a ModRM byte: we know the register used.
             }
-            [Mem(dst), Imm(_)] => {
-                let digit = rex
-                    .digit
-                    .expect("REX digit must be set for operands: [RegMem, Imm]");
+            [Mem(mem), Imm(_)] | [RegMem(mem), Imm(_)] => {
+                let digit = rex.digit.unwrap();
                 fmtln!(f, "let digit = 0x{digit:x};");
-                fmtln!(
-                    f,
-                    "emit_modrm_sib_disp(buf, off, digit, &self.{dst}, 0, None);"
-                );
+                fmtln!(f, "self.{mem}.encode_rex_suffixes(buf, off, digit, 0);");
             }
-            [RegMem(dst), Imm(_)] => {
-                let digit = rex
-                    .digit
-                    .expect("REX digit must be set for operands: [RegMem, Imm]");
-                fmtln!(f, "let digit = 0x{digit:x};");
-                f.add_block(&format!("match &self.{dst}"), |f| {
-                    fmtln!(f, "GprMem::Gpr({dst}) => emit_modrm(buf, digit, {dst}.enc()),");
-                    fmtln!(f, "GprMem::Mem({dst}) => emit_modrm_sib_disp(buf, off, digit, {dst}, 0, None),");
-                });
-            }
-            [Reg(dst), RegMem(src)] => {
-                fmtln!(f, "let {dst} = self.{dst}.enc();");
-                f.add_block(&format!("match &self.{src}"), |f| {
-                    match dst.bits() {
-                        128 => {
-                            fmtln!(f, "XmmMem::Xmm({src}) => emit_modrm(buf, {dst}, {src}.enc()),");
-                            fmtln!(f, "XmmMem::Mem({src}) => emit_modrm_sib_disp(buf, off, {dst}, {src}, 0, None),");
-                        }
-                        _ => {
-                            fmtln!(f, "GprMem::Gpr({src}) => emit_modrm(buf, {dst}, {src}.enc()),");
-                            fmtln!(f, "GprMem::Mem({src}) => emit_modrm_sib_disp(buf, off, {dst}, {src}, 0, None),");
-                        }
-                    };
-                });
-            }
-            [Mem(dst), Reg(src)] => {
-                fmtln!(f, "let {src} = self.{src}.enc();");
-                fmtln!(
-                    f,
-                    "emit_modrm_sib_disp(buf, off, {src}, &self.{dst}, 0, None);"
-                );
-            }
-            [RegMem(dst), Reg(src)]
-            | [RegMem(dst), Reg(src), Imm(_)]
-            | [RegMem(dst), Reg(src), FixedReg(_)] => {
-                fmtln!(f, "let {src} = self.{src}.enc();");
-                f.add_block(&format!("match &self.{dst}"), |f| {
-                    match src.bits() {
-                        128 => {
-                            fmtln!(f, "XmmMem::Xmm({dst}) => emit_modrm(buf, {src}, {dst}.enc()),");
-                            fmtln!(f, "XmmMem::Mem({dst}) => emit_modrm_sib_disp(buf, off, {src}, {dst}, 0, None),");
-                        }
-                        _ => {
-                            fmtln!(f, "GprMem::Gpr({dst}) => emit_modrm(buf, {src}, {dst}.enc()),");
-                            fmtln!(f, "GprMem::Mem({dst}) => emit_modrm_sib_disp(buf, off, {src}, {dst}, 0, None),");
-                        }
-                    };
-                });
+            [Reg(reg), RegMem(mem)]
+            | [Mem(mem), Reg(reg)]
+            | [RegMem(mem), Reg(reg)]
+            | [RegMem(mem), Reg(reg), Imm(_)]
+            | [RegMem(mem), Reg(reg), FixedReg(_)] => {
+                fmtln!(f, "let reg = self.{reg}.enc();");
+                fmtln!(f, "self.{mem}.encode_rex_suffixes(buf, off, reg, 0);");
             }
             unknown => unimplemented!("unknown pattern: {unknown:?}"),
         }

cranelift/assembler-x64/src/inst.rs

@@ -6,16 +6,10 @@
 use crate::api::{AsReg, CodeSink, KnownOffsetTable, RegisterVisitor, Registers};
 use crate::gpr::{self, Gpr, Size};
 use crate::imm::{Extension, Imm16, Imm32, Imm8, Simm32, Simm8};
-use crate::mem::{emit_modrm_sib_disp, visit_amode, Amode, GprMem, XmmMem};
-use crate::rex::{self, RexPrefix};
+use crate::mem::{visit_amode, Amode, GprMem, XmmMem};
+use crate::rex::RexPrefix;
 use crate::xmm::Xmm;
 use crate::Fixed;
 
 // Include code generated by the `meta` crate.
 include!(concat!(env!("OUT_DIR"), "/assembler.rs"));
-
-/// Helper function to make code generation simpler.
-fn emit_modrm(buffer: &mut impl CodeSink, enc_reg_g: u8, rm_e: u8) {
-    let modrm = rex::encode_modrm(0b11, enc_reg_g & 7, rm_e & 7);
-    buffer.put1(modrm);
-}

cranelift/assembler-x64/src/mem.rs

@@ -2,7 +2,7 @@
 
 use crate::api::{AsReg, CodeSink, Constant, KnownOffset, KnownOffsetTable, Label, TrapCode};
 use crate::gpr::{self, NonRspGpr, Size};
-use crate::rex::{encode_modrm, encode_sib, Imm, RexPrefix};
+use crate::rex::{encode_modrm, encode_sib, Disp, RexPrefix};
 use crate::{RegisterVisitor, Registers};
 
 /// x64 memory addressing modes.
@@ -48,6 +48,18 @@ impl<R: AsReg> Amode<R> {
             Amode::RipRelative { .. } => RexPrefix::two_op(enc_reg, 0, has_w_bit, uses_8bit),
         }
     }
+
+    /// Emit the ModR/M, SIB, and displacement suffixes as neeeded for this
+    /// `Amode`.
+    pub(crate) fn encode_rex_suffixes(
+        &self,
+        sink: &mut impl CodeSink,
+        offsets: &impl KnownOffsetTable,
+        enc_reg: u8,
+        bytes_at_end: u8,
+    ) {
+        emit_modrm_sib_disp(sink, offsets, enc_reg, self, bytes_at_end, None);
+    }
 }
 
 /// Visit the registers in an [`Amode`].
@@ -268,6 +280,24 @@ impl<R: AsReg, M: AsReg> GprMem<R, M> {
             GprMem::Mem(amode) => amode.as_rex_prefix(enc_reg, has_w_bit, uses_8bit),
         }
     }
+
+    /// Emit the ModR/M, SIB, and displacement suffixes for this [`GprMem`].
+    pub(crate) fn encode_rex_suffixes(
+        &self,
+        sink: &mut impl CodeSink,
+        offsets: &impl KnownOffsetTable,
+        enc_reg: u8,
+        bytes_at_end: u8,
+    ) {
+        match self {
+            GprMem::Gpr(gpr) => {
+                sink.put1(encode_modrm(0b11, enc_reg & 0b111, gpr.enc() & 0b111));
+            }
+            GprMem::Mem(amode) => {
+                amode.encode_rex_suffixes(sink, offsets, enc_reg, bytes_at_end);
+            }
+        }
+    }
 }
 
 /// An XMM register or memory operand.
@@ -299,10 +329,28 @@ impl<R: AsReg, M: AsReg> XmmMem<R, M> {
             XmmMem::Mem(amode) => amode.as_rex_prefix(enc_reg, has_w_bit, uses_8bit),
         }
     }
+
+    /// Emit the ModR/M, SIB, and displacement suffixes for this [`XmmMem`].
+    pub(crate) fn encode_rex_suffixes(
+        &self,
+        sink: &mut impl CodeSink,
+        offsets: &impl KnownOffsetTable,
+        enc_reg: u8,
+        bytes_at_end: u8,
+    ) {
+        match self {
+            XmmMem::Xmm(xmm) => {
+                sink.put1(encode_modrm(0b11, enc_reg & 0b111, xmm.enc() & 0b111));
+            }
+            XmmMem::Mem(amode) => {
+                amode.encode_rex_suffixes(sink, offsets, enc_reg, bytes_at_end);
+            }
+        }
+    }
 }
 
 /// Emit the ModRM/SIB/displacement sequence for a memory operand.
-pub fn emit_modrm_sib_disp<R: AsReg>(
+fn emit_modrm_sib_disp<R: AsReg>(
     sink: &mut impl CodeSink,
     offsets: &impl KnownOffsetTable,
     enc_g: u8,
@@ -313,7 +361,7 @@ pub fn emit_modrm_sib_disp<R: AsReg>(
     match mem_e.clone() {
         Amode::ImmReg { simm32, base, .. } => {
             let enc_e = base.enc();
-            let mut imm = Imm::new(simm32.value(offsets), evex_scaling);
+            let mut imm = Disp::new(simm32.value(offsets), evex_scaling);
 
             // Most base registers allow for a single ModRM byte plus an
             // optional immediate. If rsp is the base register, however, then a
@@ -360,7 +408,7 @@ pub fn emit_modrm_sib_disp<R: AsReg>(
             // that if the base register's lower three bits are `101` then an
             // offset must be present. This is a special case in the encoding of
             // the SIB byte and requires an explicit displacement with rbp/r13.
-            let mut imm = Imm::new(simm32.value(), evex_scaling);
+            let mut imm = Disp::new(simm32.value(), evex_scaling);
             if enc_base & 7 == gpr::enc::RBP {
                 imm.force_immediate();
             }

cranelift/assembler-x64/src/rex.rs

@@ -2,13 +2,13 @@
 
 use crate::api::CodeSink;
 
-pub(crate) fn low8_will_sign_extend_to_32(xs: i32) -> bool {
+fn low8_will_sign_extend_to_32(xs: i32) -> bool {
     xs == ((xs << 24) >> 24)
 }
 
 /// Encode the ModR/M byte.
 #[inline]
-pub fn encode_modrm(m0d: u8, enc_reg_g: u8, rm_e: u8) -> u8 {
+pub(crate) fn encode_modrm(m0d: u8, enc_reg_g: u8, rm_e: u8) -> u8 {
     debug_assert!(m0d < 4);
     debug_assert!(enc_reg_g < 8);
     debug_assert!(rm_e < 8);
@@ -17,7 +17,7 @@ pub fn encode_modrm(m0d: u8, enc_reg_g: u8, rm_e: u8) -> u8 {
 
 /// Encode the SIB byte (scale-index-base).
 #[inline]
-pub fn encode_sib(scale: u8, enc_index: u8, enc_base: u8) -> u8 {
+pub(crate) fn encode_sib(scale: u8, enc_index: u8, enc_base: u8) -> u8 {
     debug_assert!(scale < 4);
     debug_assert!(enc_index < 8);
     debug_assert!(enc_base < 8);
@@ -136,14 +136,15 @@ impl RexPrefix {
     }
 }
 
+/// The displacement bytes used after the ModR/M and SIB bytes.
 #[derive(Copy, Clone)]
-pub enum Imm {
+pub enum Disp {
     None,
     Imm8(i8),
     Imm32(i32),
 }
 
-impl Imm {
+impl Disp {
     /// Classifies the 32-bit immediate `val` as how this can be encoded
     /// with ModRM/SIB bytes.
     ///
@@ -157,51 +158,51 @@ impl Imm {
     /// The `evex_scaling` factor provided here is `Some(N)` for EVEX
     /// instructions.  This is taken into account where the `Imm` value
     /// contained is the raw byte offset.
-    pub fn new(val: i32, evex_scaling: Option<i8>) -> Imm {
+    pub fn new(val: i32, evex_scaling: Option<i8>) -> Disp {
         if val == 0 {
-            return Imm::None;
+            return Disp::None;
         }
         match evex_scaling {
             Some(scaling) => {
                 if val % i32::from(scaling) == 0 {
                     let scaled = val / i32::from(scaling);
                     if low8_will_sign_extend_to_32(scaled) {
-                        return Imm::Imm8(scaled as i8);
+                        return Disp::Imm8(scaled as i8);
                     }
                 }
-                Imm::Imm32(val)
+                Disp::Imm32(val)
             }
             None => match i8::try_from(val) {
-                Ok(val) => Imm::Imm8(val),
-                Err(_) => Imm::Imm32(val),
+                Ok(val) => Disp::Imm8(val),
+                Err(_) => Disp::Imm32(val),
             },
         }
     }
 
     /// Forces `Imm::None` to become `Imm::Imm8(0)`, used for special cases
     /// where some base registers require an immediate.
     pub fn force_immediate(&mut self) {
-        if let Imm::None = self {
-            *self = Imm::Imm8(0);
+        if let Disp::None = self {
+            *self = Disp::Imm8(0);
         }
     }
 
     /// Returns the two "mod" bits present at the upper bits of the mod/rm
     /// byte.
     pub fn m0d(self) -> u8 {
         match self {
-            Imm::None => 0b00,
-            Imm::Imm8(_) => 0b01,
-            Imm::Imm32(_) => 0b10,
+            Disp::None => 0b00,
+            Disp::Imm8(_) => 0b01,
+            Disp::Imm32(_) => 0b10,
         }
     }
 
     /// Emit the truncated immediate into the code sink.
     pub fn emit(self, sink: &mut impl CodeSink) {
         match self {
-            Imm::None => {}
-            Imm::Imm8(n) => sink.put1(n as u8),
-            Imm::Imm32(n) => sink.put4(n as u32),
+            Disp::None => {}
+            Disp::Imm8(n) => sink.put1(n as u8),
+            Disp::Imm32(n) => sink.put4(n as u32),
         }
     }
 }