asm: allow printing of virtual registers (bytecodealliance#10389)

* Rename module: `reg` to `gpr`

* Override `to_string` implementation for `cranelift-codegen` types

* Add pretty-printing tests

Author: abrown

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

@@ -72,8 +72,8 @@ impl dsl::Location {
     pub fn generate_fixed_reg(&self) -> Option<&str> {
         use dsl::Location::*;
         match self {
-            al | ax | eax | rax => Some("reg::enc::RAX"),
-            cl => Some("reg::enc::RCX"),
+            al | ax | eax | rax => Some("gpr::enc::RAX"),
+            cl => Some("gpr::enc::RCX"),
             imm8 | imm16 | imm32 | r8 | r16 | r32 | r64 | xmm | rm8 | rm16 | rm32 | rm64 | rm128 => None,
         }
     }

cranelift/assembler-x64/src/api.rs

@@ -1,6 +1,6 @@
 //! Contains traits that a user of this assembler must implement.
 
-use crate::reg;
+use crate::gpr;
 use crate::xmm;
 use std::{num::NonZeroU8, ops::Index, vec::Vec};
 
@@ -130,10 +130,10 @@ pub trait AsReg: Clone + std::fmt::Debug {
     fn enc(&self) -> u8;
 
     /// Return the register name.
-    fn to_string(&self, size: Option<reg::Size>) -> &str {
+    fn to_string(&self, size: Option<gpr::Size>) -> String {
         match size {
-            Some(size) => reg::enc::to_string(self.enc(), size),
-            None => xmm::enc::to_string(self.enc()),
+            Some(size) => gpr::enc::to_string(self.enc(), size).into(),
+            None => xmm::enc::to_string(self.enc()).into(),
         }
     }
 }

cranelift/assembler-x64/src/fuzz.rs

@@ -198,7 +198,7 @@ impl Arbitrary<'_> for AmodeOffsetPlusKnownOffset {
 }
 impl<R: AsReg> Arbitrary<'_> for NonRspGpr<R> {
     fn arbitrary(u: &mut Unstructured<'_>) -> Result<Self> {
-        use crate::reg::enc::*;
+        use crate::gpr::enc::*;
         let gpr = u.choose(&[
             RAX, RCX, RDX, RBX, RBP, RSI, RDI, R8, R9, R10, R11, R12, R13, R14, R15,
         ])?;

cranelift/assembler-x64/src/reg.rs cranelift/assembler-x64/src/gpr.rscranelift/assembler-x64/src/reg.rs renamed to cranelift/assembler-x64/src/gpr.rs

@@ -30,8 +30,8 @@ impl<R: AsReg> Gpr<R> {
     }
 
     /// Return the register name at the given `size`.
-    pub fn to_string(&self, size: Size) -> &str {
-        enc::to_string(self.enc(), size)
+    pub fn to_string(&self, size: Size) -> String {
+        self.0.to_string(Some(size))
     }
 
     /// Proxy on the 8-bit REX flag emission; helpful for simplifying generated

cranelift/assembler-x64/src/inst.rs

@@ -4,9 +4,9 @@
 //! See also: [`Inst`], an `enum` containing all these instructions.
 
 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, GprMem, XmmMem};
-use crate::reg::{self, Gpr, Size};
 use crate::rex::{self, RexFlags};
 use crate::xmm::Xmm;
 

cranelift/assembler-x64/src/lib.rs

@@ -44,12 +44,12 @@
 )]
 
 mod api;
+pub mod gpr;
 mod imm;
 pub mod inst;
 mod mem;
-mod reg;
 mod rex;
-mod xmm;
+pub mod xmm;
 
 #[cfg(any(test, feature = "fuzz"))]
 pub mod fuzz;
@@ -76,11 +76,11 @@ pub use api::{
     AsReg, CodeSink, Constant, KnownOffset, KnownOffsetTable, Label, RegisterVisitor, Registers,
     TrapCode,
 };
+pub use gpr::{Gpr, NonRspGpr, Size};
 pub use imm::{Extension, Imm16, Imm32, Imm8, Simm16, Simm32, Simm8};
 pub use mem::{
     Amode, AmodeOffset, AmodeOffsetPlusKnownOffset, DeferredTarget, GprMem, Scale, XmmMem,
 };
-pub use reg::{Gpr, NonRspGpr, Size};
 pub use rex::RexFlags;
 pub use xmm::Xmm;
 

cranelift/assembler-x64/src/mem.rs

@@ -1,7 +1,7 @@
 //! Memory operands to instructions.
 
 use crate::api::{AsReg, CodeSink, Constant, KnownOffset, KnownOffsetTable, Label, TrapCode};
-use crate::reg::{self, NonRspGpr, Size};
+use crate::gpr::{self, NonRspGpr, Size};
 use crate::rex::{encode_modrm, encode_sib, Imm, RexFlags};
 use crate::xmm;
 
@@ -169,7 +169,7 @@ impl<R: AsReg> std::fmt::Display for Amode<R> {
             Amode::ImmReg { simm32, base, .. } => {
                 // Note: size is always 8; the address is 64 bits,
                 // even if the addressed operand is smaller.
-                let base = reg::enc::to_string(base.enc(), Size::Quadword);
+                let base = gpr::enc::to_string(base.enc(), Size::Quadword);
                 write!(f, "{simm32:x}({base})")
             }
             Amode::ImmRegRegShift {
@@ -179,8 +179,8 @@ impl<R: AsReg> std::fmt::Display for Amode<R> {
                 scale,
                 ..
             } => {
-                let base = reg::enc::to_string(base.enc(), Size::Quadword);
-                let index = reg::enc::to_string(index.enc(), Size::Quadword);
+                let base = gpr::enc::to_string(base.enc(), Size::Quadword);
+                let index = gpr::enc::to_string(index.enc(), Size::Quadword);
                 let shift = scale.shift();
                 if shift > 1 {
                     write!(f, "{simm32:x}({base}, {index}, {shift})")
@@ -256,7 +256,7 @@ impl<R: AsReg, M: AsReg> GprMem<R, M> {
     /// Pretty-print the operand.
     pub fn to_string(&self, size: Size) -> String {
         match self {
-            GprMem::Gpr(gpr) => reg::enc::to_string(gpr.enc(), size).to_owned(),
+            GprMem::Gpr(gpr) => gpr.to_string(Some(size)),
             GprMem::Mem(amode) => amode.to_string(),
         }
     }
@@ -313,7 +313,7 @@ pub fn emit_modrm_sib_disp<R: AsReg>(
             // optional immediate. If rsp is the base register, however, then a
             // SIB byte must be used.
             let enc_e_low3 = enc_e & 7;
-            if enc_e_low3 == reg::enc::RSP {
+            if enc_e_low3 == gpr::enc::RSP {
                 // Displacement from RSP is encoded with a SIB byte where
                 // the index and base are both encoded as RSP's encoding of
                 // 0b100. This special encoding means that the index register
@@ -326,7 +326,7 @@ pub fn emit_modrm_sib_disp<R: AsReg>(
                 // If the base register is rbp and there's no offset then force
                 // a 1-byte zero offset since otherwise the encoding would be
                 // invalid.
-                if enc_e_low3 == reg::enc::RBP {
+                if enc_e_low3 == gpr::enc::RBP {
                     imm.force_immediate();
                 }
                 sink.put1(encode_modrm(imm.m0d(), enc_g & 7, enc_e & 7));
@@ -348,14 +348,14 @@ pub fn emit_modrm_sib_disp<R: AsReg>(
             // ever be rsp. Note, though, that the encoding of r12, whose three
             // lower bits match the encoding of rsp, is explicitly allowed with
             // REX bytes so only rsp is disallowed.
-            assert!(enc_index != reg::enc::RSP);
+            assert!(enc_index != gpr::enc::RSP);
 
             // If the offset is zero then there is no immediate. Note, though,
             // 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);
-            if enc_base & 7 == reg::enc::RBP {
+            if enc_base & 7 == gpr::enc::RBP {
                 imm.force_immediate();
             }
 

cranelift/codegen/src/isa/x64/inst/external.rs

@@ -6,6 +6,7 @@ use super::{
 };
 use crate::ir::TrapCode;
 use cranelift_assembler_x64 as asm;
+use std::string::String;
 
 /// Define the types of registers Cranelift will use.
 #[derive(Clone, Debug)]
@@ -38,6 +39,16 @@ impl asm::AsReg for PairedGpr {
         write
     }
 
+    fn to_string(&self, size: Option<asm::Size>) -> String {
+        if self.read.is_real() {
+            asm::gpr::enc::to_string(self.enc(), size.unwrap()).into()
+        } else {
+            let read = self.read.to_reg();
+            let write = self.write.to_reg().to_reg();
+            format!("(%{write:?} <- %{read:?})")
+        }
+    }
+
     fn new(_: u8) -> Self {
         panic!("disallow creation of new assembler registers")
     }
@@ -59,6 +70,17 @@ impl asm::AsReg for PairedXmm {
         write
     }
 
+    fn to_string(&self, size: Option<asm::Size>) -> String {
+        assert!(size.is_none(), "XMM registers do not have size variants");
+        if self.read.is_real() {
+            asm::xmm::enc::to_string(self.enc()).into()
+        } else {
+            let read = self.read.to_reg();
+            let write = self.write.to_reg().to_reg();
+            format!("(%{write:?} <- %{read:?})")
+        }
+    }
+
     fn new(_: u8) -> Self {
         panic!("disallow creation of new assembler registers")
     }
@@ -70,6 +92,14 @@ impl asm::AsReg for Gpr {
         enc_gpr(self)
     }
 
+    fn to_string(&self, size: Option<asm::Size>) -> String {
+        if self.is_real() {
+            asm::gpr::enc::to_string(self.enc(), size.unwrap()).into()
+        } else {
+            format!("%{:?}", self.to_reg())
+        }
+    }
+
     fn new(_: u8) -> Self {
         panic!("disallow creation of new assembler registers")
     }
@@ -81,6 +111,15 @@ impl asm::AsReg for Xmm {
         enc_xmm(self)
     }
 
+    fn to_string(&self, size: Option<asm::Size>) -> String {
+        assert!(size.is_none(), "XMM registers do not have size variants");
+        if self.is_real() {
+            asm::xmm::enc::to_string(self.enc()).into()
+        } else {
+            format!("%{:?}", self.to_reg())
+        }
+    }
+
     fn new(_: u8) -> Self {
         panic!("disallow creation of new assembler registers")
     }
@@ -286,3 +325,45 @@ impl From<asm::Constant> for VCodeConstant {
 // the assembler instructions exposed to ISLE.
 include!(concat!(env!("OUT_DIR"), "/assembler-isle-macro.rs"));
 pub(crate) use isle_assembler_methods;
+
+#[cfg(test)]
+mod tests {
+    use super::asm::{AsReg, Size};
+    use super::PairedGpr;
+    use crate::isa::x64::args::{FromWritableReg, Gpr, WritableGpr, WritableXmm, Xmm};
+    use crate::isa::x64::inst::external::PairedXmm;
+    use crate::{Reg, Writable};
+    use regalloc2::{RegClass, VReg};
+
+    #[test]
+    fn pretty_print_registers() {
+        // For logging, we need to be able to pretty-print the virtual registers
+        // that Cranelift uses before register allocation. This test ensures
+        // that these remain printable using the `AsReg::to_string` interface
+        // (see issue #10631).
+
+        let v200: Reg = VReg::new(200, RegClass::Int).into();
+        let gpr200 = Gpr::new(v200).unwrap();
+        assert_eq!(gpr200.to_string(Some(Size::Quadword)), "%v200");
+
+        let v300: Reg = VReg::new(300, RegClass::Int).into();
+        let wgpr300 = WritableGpr::from_writable_reg(Writable::from_reg(v300).into()).unwrap();
+        let pair = PairedGpr {
+            read: gpr200,
+            write: wgpr300,
+        };
+        assert_eq!(pair.to_string(Some(Size::Quadword)), "(%v300 <- %v200)");
+
+        let v400: Reg = VReg::new(400, RegClass::Float).into();
+        let xmm400 = Xmm::new(v400).unwrap();
+        assert_eq!(xmm400.to_string(None), "%v400");
+
+        let v500: Reg = VReg::new(500, RegClass::Float).into();
+        let wxmm500 = WritableXmm::from_writable_reg(Writable::from_reg(v500).into()).unwrap();
+        let pair = PairedXmm {
+            read: xmm400,
+            write: wxmm500,
+        };
+        assert_eq!(pair.to_string(None), "(%v500 <- %v400)");
+    }
+}