Repository: bytecodealliance/wasmtime
Pull request #8308 (2024-04-05 21:02:38)

cranelift/aarch64: Remove types from address-modes
Previously, the precondition on every use of AMode::*Offset was that the type specified in the address mode must be the same size as the access type of the memory instruction where the address mode was used.

However, at the point where we replace these magic address modes with real modes that the target architecture understands, we know what instruction is using the address, so we can get the type from there.

This simplifies constructing addresses while also ensuring that the types are always correct.

The changes here center around threading the access type through `mem_finalize` and `AMode::pretty_print`, then deleting it from all `AMode` variants.

In the process I introduced a helper method on `Inst` to centralize the mapping from instruction to access type, and used that everywhere.


Patch:
diff --git a/cranelift/codegen/src/isa/aarch64/abi.rs b/cranelift/codegen/src/isa/aarch64/abi.rs
index 7f5903dac..94b5897b6 100644
--- a/cranelift/codegen/src/isa/aarch64/abi.rs
+++ b/cranelift/codegen/src/isa/aarch64/abi.rs
@@ -36,9 +36,9 @@ impl Into<AMode> for StackAMode {
     fn into(self) -> AMode {
         match self {
             // Argument area begins after saved frame pointer + return address.
-            StackAMode::ArgOffset(off, ty) => AMode::FPOffset { off: off + 16, ty },
-            StackAMode::NominalSPOffset(off, ty) => AMode::NominalSPOffset { off, ty },
-            StackAMode::SPOffset(off, ty) => AMode::SPOffset { off, ty },
+            StackAMode::ArgOffset(off, _ty) => AMode::FPOffset { off: off + 16 },
+            StackAMode::NominalSPOffset(off, _ty) => AMode::NominalSPOffset { off },
+            StackAMode::SPOffset(off, _ty) => AMode::SPOffset { off },
         }
     }
 }
@@ -471,7 +471,6 @@ impl ABIMachineSpec for AArch64MachineDeps {
         let mem = AMode::RegOffset {
             rn: base,
             off: offset as i64,
-            ty,
         };
         Inst::gen_load(into_reg, mem, ty, MemFlags::trusted())
     }
@@ -480,7 +479,6 @@ impl ABIMachineSpec for AArch64MachineDeps {
         let mem = AMode::RegOffset {
             rn: base,
             off: offset as i64,
-            ty,
         };
         Inst::gen_store(mem, from_reg, ty, MemFlags::trusted())
     }
@@ -1190,7 +1188,7 @@ impl AArch64MachineDeps {
             insts.extend(Self::gen_sp_reg_adjust(-(guard_size as i32)));
 
             insts.push(Inst::gen_store(
-                AMode::SPOffset { off: 0, ty: I8 },
+                AMode::SPOffset { off: 0 },
                 zero_reg(),
                 I32,
                 MemFlags::trusted(),
diff --git a/cranelift/codegen/src/isa/aarch64/inst.isle b/cranelift/codegen/src/isa/aarch64/inst.isle
index 51107ab36..f706f8048 100644
--- a/cranelift/codegen/src/isa/aarch64/inst.isle
+++ b/cranelift/codegen/src/isa/aarch64/inst.isle
@@ -1119,15 +1119,13 @@
         ;; Register plus register offset, scaled by type's size.
         (RegScaled
          (rn Reg)
-         (rm Reg)
-         (ty Type))
+         (rm Reg))
 
         ;; Register plus register offset, scaled by type's size, with index
         ;; sign- or zero-extended first.
         (RegScaledExtended
          (rn Reg)
          (rm Reg)
-         (ty Type)
          (extendop ExtendOp))
 
         ;; Register plus register offset, with index sign- or zero-extended
@@ -1157,18 +1155,15 @@
         ;; offsets with multiple instructions as necessary during code emission.
         (RegOffset
          (rn Reg)
-         (off i64)
-         (ty Type))
+         (off i64))
 
         ;; Offset from the stack pointer.
         (SPOffset
-         (off i64)
-         (ty Type))
+         (off i64))
 
         ;; Offset from the frame pointer.
         (FPOffset
-         (off i64)
-         (ty Type))
+         (off i64))
 
         ;; A reference to a constant which is placed outside of the function's
         ;; body, typically at the end.
@@ -1188,8 +1183,7 @@
         ;; clobber pushes. See the diagram in the documentation for
         ;; [crate::isa::aarch64::abi](the ABI module) for more details.
         (NominalSPOffset
-         (off i64)
-         (ty Type))))
+         (off i64))))
 
 ;; A memory argument to a load/store-pair.
 (type PairAMode (enum
@@ -3169,18 +3163,18 @@
 ;; first and then if that succeeds afterwards try to find an extend.
 (rule 6 (amode_no_more_iconst ty (iadd x (ishl y (iconst (u64_from_imm64 n)))) offset)
         (if-let $true (u64_eq (ty_bytes ty) (u64_shl 1 n)))
-        (amode_reg_scaled (amode_add x offset) y ty))
+        (amode_reg_scaled (amode_add x offset) y))
 (rule 7 (amode_no_more_iconst ty (iadd (ishl y (iconst (u64_from_imm64 n))) x) offset)
         (if-let $true (u64_eq (ty_bytes ty) (u64_shl 1 n)))
-        (amode_reg_scaled (amode_add x offset) y ty))
+        (amode_reg_scaled (amode_add x offset) y))
 
-(decl amode_reg_scaled (Reg Value Type) AMode)
-(rule 0 (amode_reg_scaled base index ty)
-        (AMode.RegScaled base index ty))
-(rule 1 (amode_reg_scaled base (uextend index @ (value_type $I32)) ty)
-        (AMode.RegScaledExtended base index ty (ExtendOp.UXTW)))
-(rule 2 (amode_reg_scaled base (sextend index @ (value_type $I32)) ty)
-        (AMode.RegScaledExtended base index ty (ExtendOp.SXTW)))
+(decl amode_reg_scaled (Reg Value) AMode)
+(rule 0 (amode_reg_scaled base index)
+        (AMode.RegScaled base index))
+(rule 1 (amode_reg_scaled base (uextend index @ (value_type $I32)))
+        (AMode.RegScaledExtended base index (ExtendOp.UXTW)))
+(rule 1 (amode_reg_scaled base (sextend index @ (value_type $I32)))
+        (AMode.RegScaledExtended base index (ExtendOp.SXTW)))
 
 ;; Helper to add a 32-bit signed immediate to the register provided. This will
 ;; select an appropriate `add` instruction to use.
@@ -3519,7 +3513,7 @@
             ;; the frame record that corresponds to the current subroutine on
             ;; the stack; the presence of the record is guaranteed by the
             ;; `preserve_frame_pointers` setting.
-            (addr AMode (AMode.FPOffset 8 $I64))
+            (addr AMode (AMode.FPOffset 8))
             (_ Unit (emit (MInst.ULoad64 dst addr (mem_flags_trusted)))))
            dst))
 
@@ -3529,7 +3523,7 @@
       ;; the frame record. Furthermore, we can use LR as a scratch register
       ;; because the function will set it to the return address immediately
       ;; before returning.
-      (let ((addr AMode (AMode.FPOffset 8 $I64))
+      (let ((addr AMode (AMode.FPOffset 8))
             (lr WritableReg (writable_link_reg))
             (_ Unit (emit (MInst.ULoad64 lr addr (mem_flags_trusted))))
             (_ Unit (emit (MInst.Xpaclri))))
diff --git a/cranelift/codegen/src/isa/aarch64/inst/args.rs b/cranelift/codegen/src/isa/aarch64/inst/args.rs
index a9f185bf6..3375215d2 100644
--- a/cranelift/codegen/src/isa/aarch64/inst/args.rs
+++ b/cranelift/codegen/src/isa/aarch64/inst/args.rs
@@ -138,11 +138,10 @@ impl AMode {
 
     /// Memory reference using `reg1 + sizeof(ty) * reg2` as an address, with `reg2` sign- or
     /// zero-extended as per `op`.
-    pub fn reg_plus_reg_scaled_extended(reg1: Reg, reg2: Reg, ty: Type, op: ExtendOp) -> AMode {
+    pub fn reg_plus_reg_scaled_extended(reg1: Reg, reg2: Reg, op: ExtendOp) -> AMode {
         AMode::RegScaledExtended {
             rn: reg1,
             rm: reg2,
-            ty,
             extendop: op,
         }
     }
@@ -162,20 +161,13 @@ impl AMode {
                 rn: allocs.next(rn),
                 rm: allocs.next(rm),
             },
-            &AMode::RegScaled { rn, rm, ty } => AMode::RegScaled {
+            &AMode::RegScaled { rn, rm } => AMode::RegScaled {
                 rn: allocs.next(rn),
                 rm: allocs.next(rm),
-                ty,
             },
-            &AMode::RegScaledExtended {
-                rn,
-                rm,
-                ty,
-                extendop,
-            } => AMode::RegScaledExtended {
+            &AMode::RegScaledExtended { rn, rm, extendop } => AMode::RegScaledExtended {
                 rn: allocs.next(rn),
                 rm: allocs.next(rm),
-                ty,
                 extendop,
             },
             &AMode::RegExtended { rn, rm, extendop } => AMode::RegExtended {
@@ -183,10 +175,9 @@ impl AMode {
                 rm: allocs.next(rm),
                 extendop,
             },
-            &AMode::RegOffset { rn, off, ty } => AMode::RegOffset {
+            &AMode::RegOffset { rn, off } => AMode::RegOffset {
                 rn: allocs.next(rn),
                 off,
-                ty,
             },
             &AMode::SPPreIndexed { .. }
             | &AMode::SPPostIndexed { .. }
@@ -385,19 +376,20 @@ impl PrettyPrint for MemLabel {
     }
 }
 
-fn shift_for_type(ty: Type) -> usize {
-    match ty.bytes() {
+fn shift_for_type(size_bytes: u8) -> usize {
+    match size_bytes {
         1 => 0,
         2 => 1,
         4 => 2,
         8 => 3,
         16 => 4,
-        _ => panic!("unknown type: {}", ty),
+        _ => panic!("unknown type size: {size_bytes}"),
     }
 }
 
 impl PrettyPrint for AMode {
-    fn pretty_print(&self, _: u8, allocs: &mut AllocationConsumer<'_>) -> String {
+    fn pretty_print(&self, size_bytes: u8, allocs: &mut AllocationConsumer<'_>) -> String {
+        debug_assert!(size_bytes != 0);
         match self {
             &AMode::Unscaled { rn, simm9 } => {
                 let reg = pretty_print_reg(rn, allocs);
@@ -410,7 +402,7 @@ impl PrettyPrint for AMode {
             }
             &AMode::UnsignedOffset { rn, uimm12 } => {
                 let reg = pretty_print_reg(rn, allocs);
-                if uimm12.value != 0 {
+                if uimm12.value() != 0 {
                     let uimm12 = uimm12.pretty_print(8, allocs);
                     format!("[{}, {}]", reg, uimm12)
                 } else {
@@ -422,19 +414,14 @@ impl PrettyPrint for AMode {
                 let r2 = pretty_print_reg(rm, allocs);
                 format!("[{}, {}]", r1, r2)
             }
-            &AMode::RegScaled { rn, rm, ty } => {
+            &AMode::RegScaled { rn, rm } => {
                 let r1 = pretty_print_reg(rn, allocs);
                 let r2 = pretty_print_reg(rm, allocs);
-                let shift = shift_for_type(ty);
+                let shift = shift_for_type(size_bytes);
                 format!("[{}, {}, LSL #{}]", r1, r2, shift)
             }
-            &AMode::RegScaledExtended {
-                rn,
-                rm,
-                ty,
-                extendop,
-            } => {
-                let shift = shift_for_type(ty);
+            &AMode::RegScaledExtended { rn, rm, extendop } => {
+                let shift = shift_for_type(size_bytes);
                 let size = match extendop {
                     ExtendOp::SXTW | ExtendOp::UXTW => OperandSize::Size32,
                     _ => OperandSize::Size64,
diff --git a/cranelift/codegen/src/isa/aarch64/inst/emit.rs b/cranelift/codegen/src/isa/aarch64/inst/emit.rs
index b02fe200e..df840e5b1 100644
--- a/cranelift/codegen/src/isa/aarch64/inst/emit.rs
+++ b/cranelift/codegen/src/isa/aarch64/inst/emit.rs
@@ -15,13 +15,14 @@ use crate::trace;
 pub fn mem_finalize(
     sink: Option<&mut MachBuffer<Inst>>,
     mem: &AMode,
+    access_ty: Type,
     state: &EmitState,
 ) -> (SmallVec<[Inst; 4]>, AMode) {
     match mem {
-        &AMode::RegOffset { off, ty, .. }
-        | &AMode::SPOffset { off, ty }
-        | &AMode::FPOffset { off, ty }
-        | &AMode::NominalSPOffset { off, ty } => {
+        &AMode::RegOffset { off, .. }
+        | &AMode::SPOffset { off }
+        | &AMode::FPOffset { off }
+        | &AMode::NominalSPOffset { off } => {
             let basereg = match mem {
                 &AMode::RegOffset { rn, .. } => rn,
                 &AMode::SPOffset { .. } | &AMode::NominalSPOffset { .. } => stack_reg(),
@@ -45,7 +46,7 @@ pub fn mem_finalize(
             if let Some(simm9) = SImm9::maybe_from_i64(off) {
                 let mem = AMode::Unscaled { rn: basereg, simm9 };
                 (smallvec![], mem)
-            } else if let Some(uimm12) = UImm12Scaled::maybe_from_i64(off, ty) {
+            } else if let Some(uimm12) = UImm12Scaled::maybe_from_i64(off, access_ty) {
                 let mem = AMode::UnsignedOffset {
                     rn: basereg,
                     uimm12,
@@ -979,7 +980,8 @@ impl MachInstEmit for Inst {
             | &Inst::FpuLoad128 { rd, ref mem, flags } => {
                 let rd = allocs.next_writable(rd);
                 let mem = mem.with_allocs(&mut allocs);
-                let (mem_insts, mem) = mem_finalize(Some(sink), &mem, state);
+                let access_ty = self.mem_type().unwrap();
+                let (mem_insts, mem) = mem_finalize(Some(sink), &mem, access_ty, state);
 
                 for inst in mem_insts.into_iter() {
                     inst.emit(&[], sink, emit_info, state);
@@ -991,17 +993,17 @@ impl MachInstEmit for Inst {
                 // This is the base opcode (top 10 bits) for the "unscaled
                 // immediate" form (Unscaled). Other addressing modes will OR in
                 // other values for bits 24/25 (bits 1/2 of this constant).
-                let (op, bits) = match self {
-                    &Inst::ULoad8 { .. } => (0b0011100001, 8),
-                    &Inst::SLoad8 { .. } => (0b0011100010, 8),
-                    &Inst::ULoad16 { .. } => (0b0111100001, 16),
-                    &Inst::SLoad16 { .. } => (0b0111100010, 16),
-                    &Inst::ULoad32 { .. } => (0b1011100001, 32),
-                    &Inst::SLoad32 { .. } => (0b1011100010, 32),
-                    &Inst::ULoad64 { .. } => (0b1111100001, 64),
-                    &Inst::FpuLoad32 { .. } => (0b1011110001, 32),
-                    &Inst::FpuLoad64 { .. } => (0b1111110001, 64),
-                    &Inst::FpuLoad128 { .. } => (0b0011110011, 128),
+                let op = match self {
+                    Inst::ULoad8 { .. } => 0b0011100001,
+                    Inst::SLoad8 { .. } => 0b0011100010,
+                    Inst::ULoad16 { .. } => 0b0111100001,
+                    Inst::SLoad16 { .. } => 0b0111100010,
+                    Inst::ULoad32 { .. } => 0b1011100001,
+                    Inst::SLoad32 { .. } => 0b1011100010,
+                    Inst::ULoad64 { .. } => 0b1111100001,
+                    Inst::FpuLoad32 { .. } => 0b1011110001,
+                    Inst::FpuLoad64 { .. } => 0b1111110001,
+                    Inst::FpuLoad128 { .. } => 0b0011110011,
                     _ => unreachable!(),
                 };
 
@@ -1017,9 +1019,6 @@ impl MachInstEmit for Inst {
                     }
                     &AMode::UnsignedOffset { rn, uimm12 } => {
                         let reg = allocs.next(rn);
-                        if uimm12.value() != 0 {
-                            assert_eq!(bits, ty_bits(uimm12.scale_ty()));
-                        }
                         sink.put4(enc_ldst_uimm12(op, uimm12, reg, rd));
                     }
                     &AMode::RegReg { rn, rm } => {
@@ -1029,11 +1028,9 @@ impl MachInstEmit for Inst {
                             op, r1, r2, /* scaled = */ false, /* extendop = */ None, rd,
                         ));
                     }
-                    &AMode::RegScaled { rn, rm, ty }
-                    | &AMode::RegScaledExtended { rn, rm, ty, .. } => {
+                    &AMode::RegScaled { rn, rm } | &AMode::RegScaledExtended { rn, rm, .. } => {
                         let r1 = allocs.next(rn);
                         let r2 = allocs.next(rm);
-                        assert_eq!(bits, ty_bits(ty));
                         let extendop = match &mem {
                             &AMode::RegScaled { .. } => None,
                             &AMode::RegScaledExtended { extendop, .. } => Some(extendop),
@@ -1123,20 +1120,21 @@ impl MachInstEmit for Inst {
             | &Inst::FpuStore128 { rd, ref mem, flags } => {
                 let rd = allocs.next(rd);
                 let mem = mem.with_allocs(&mut allocs);
-                let (mem_insts, mem) = mem_finalize(Some(sink), &mem, state);
+                let access_ty = self.mem_type().unwrap();
+                let (mem_insts, mem) = mem_finalize(Some(sink), &mem, access_ty, state);
 
                 for inst in mem_insts.into_iter() {
                     inst.emit(&[], sink, emit_info, state);
                 }
 
-                let (op, bits) = match self {
-                    &Inst::Store8 { .. } => (0b0011100000, 8),
-                    &Inst::Store16 { .. } => (0b0111100000, 16),
-                    &Inst::Store32 { .. } => (0b1011100000, 32),
-                    &Inst::Store64 { .. } => (0b1111100000, 64),
-                    &Inst::FpuStore32 { .. } => (0b1011110000, 32),
-                    &Inst::FpuStore64 { .. } => (0b1111110000, 64),
-                    &Inst::FpuStore128 { .. } => (0b0011110010, 128),
+                let op = match self {
+                    Inst::Store8 { .. } => 0b0011100000,
+                    Inst::Store16 { .. } => 0b0111100000,
+                    Inst::Store32 { .. } => 0b1011100000,
+                    Inst::Store64 { .. } => 0b1111100000,
+                    Inst::FpuStore32 { .. } => 0b1011110000,
+                    Inst::FpuStore64 { .. } => 0b1111110000,
+                    Inst::FpuStore128 { .. } => 0b0011110010,
                     _ => unreachable!(),
                 };
 
@@ -1152,9 +1150,6 @@ impl MachInstEmit for Inst {
                     }
                     &AMode::UnsignedOffset { rn, uimm12 } => {
                         let reg = allocs.next(rn);
-                        if uimm12.value() != 0 {
-                            assert_eq!(bits, ty_bits(uimm12.scale_ty()));
-                        }
                         sink.put4(enc_ldst_uimm12(op, uimm12, reg, rd));
                     }
                     &AMode::RegReg { rn, rm } => {
@@ -1164,7 +1159,7 @@ impl MachInstEmit for Inst {
                             op, r1, r2, /* scaled = */ false, /* extendop = */ None, rd,
                         ));
                     }
-                    &AMode::RegScaled { rn, rm, .. } | &AMode::RegScaledExtended { rn, rm, .. } => {
+                    &AMode::RegScaled { rn, rm } | &AMode::RegScaledExtended { rn, rm, .. } => {
                         let r1 = allocs.next(rn);
                         let r2 = allocs.next(rm);
                         let extendop = match &mem {
@@ -3341,7 +3336,6 @@ impl MachInstEmit for Inst {
                     mem: AMode::reg_plus_reg_scaled_extended(
                         rtmp1.to_reg(),
                         rtmp2.to_reg(),
-                        I32,
                         ExtendOp::UXTW,
                     ),
                     flags: MemFlags::trusted(),
@@ -3436,7 +3430,7 @@ impl MachInstEmit for Inst {
             &Inst::LoadAddr { rd, ref mem } => {
                 let rd = allocs.next_writable(rd);
                 let mem = mem.with_allocs(&mut allocs);
-                let (mem_insts, mem) = mem_finalize(Some(sink), &mem, state);
+                let (mem_insts, mem) = mem_finalize(Some(sink), &mem, I8, state);
                 for inst in mem_insts.into_iter() {
                     inst.emit(&[], sink, emit_info, state);
                 }
@@ -3874,7 +3868,6 @@ fn emit_return_call_common_sequence(
             rd: tmp2,
             mem: AMode::SPOffset {
                 off: i64::from(i * 8),
-                ty: types::I64,
             },
             flags: ir::MemFlags::trusted(),
         }
@@ -3886,7 +3879,6 @@ fn emit_return_call_common_sequence(
             rd: tmp2.to_reg(),
             mem: AMode::FPOffset {
                 off: fp_to_callee_sp + i64::from(i * 8),
-                ty: types::I64,
             },
             flags: ir::MemFlags::trusted(),
         }
diff --git a/cranelift/codegen/src/isa/aarch64/inst/imms.rs b/cranelift/codegen/src/isa/aarch64/inst/imms.rs
index cd8a2aa76..b87209dda 100644
--- a/cranelift/codegen/src/isa/aarch64/inst/imms.rs
+++ b/cranelift/codegen/src/isa/aarch64/inst/imms.rs
@@ -220,9 +220,9 @@ impl SImm9 {
 #[derive(Clone, Copy, Debug)]
 pub struct UImm12Scaled {
     /// The value.
-    pub value: u16,
+    value: u16,
     /// multiplied by the size of this type
-    pub scale_ty: Type,
+    scale_ty: Type,
 }
 
 impl UImm12Scaled {
@@ -257,11 +257,6 @@ impl UImm12Scaled {
     pub fn value(&self) -> u32 {
         self.value as u32
     }
-
-    /// The value type which is the scaling base.
-    pub fn scale_ty(&self) -> Type {
-        self.scale_ty
-    }
 }
 
 /// A shifted immediate value in 'imm12' format: supports 12 bits, shifted
diff --git a/cranelift/codegen/src/isa/aarch64/inst/mod.rs b/cranelift/codegen/src/isa/aarch64/inst/mod.rs
index 82fdccab8..e1447d60a 100644
--- a/cranelift/codegen/src/isa/aarch64/inst/mod.rs
+++ b/cranelift/codegen/src/isa/aarch64/inst/mod.rs
@@ -367,6 +367,32 @@ impl Inst {
             }
         }
     }
+
+    /// What type does this load or store instruction access in memory? When
+    /// uimm12 encoding is used, the size of this type is the amount that
+    /// immediate offsets are scaled by.
+    pub fn mem_type(&self) -> Option<Type> {
+        match self {
+            Inst::ULoad8 { .. } => Some(I8),
+            Inst::SLoad8 { .. } => Some(I8),
+            Inst::ULoad16 { .. } => Some(I16),
+            Inst::SLoad16 { .. } => Some(I16),
+            Inst::ULoad32 { .. } => Some(I32),
+            Inst::SLoad32 { .. } => Some(I32),
+            Inst::ULoad64 { .. } => Some(I64),
+            Inst::FpuLoad32 { .. } => Some(F32),
+            Inst::FpuLoad64 { .. } => Some(F64),
+            Inst::FpuLoad128 { .. } => Some(I8X16),
+            Inst::Store8 { .. } => Some(I8),
+            Inst::Store16 { .. } => Some(I16),
+            Inst::Store32 { .. } => Some(I32),
+            Inst::Store64 { .. } => Some(I64),
+            Inst::FpuStore32 { .. } => Some(F32),
+            Inst::FpuStore64 { .. } => Some(F64),
+            Inst::FpuStore128 { .. } => Some(I8X16),
+            _ => None,
+        }
+    }
 }
 
 //=============================================================================
@@ -1211,8 +1237,8 @@ impl MachInst for Inst {
 //=============================================================================
 // Pretty-printing of instructions.
 
-fn mem_finalize_for_show(mem: &AMode, state: &EmitState) -> (String, AMode) {
-    let (mem_insts, mem) = mem_finalize(None, mem, state);
+fn mem_finalize_for_show(mem: &AMode, access_ty: Type, state: &EmitState) -> (String, String) {
+    let (mem_insts, mem) = mem_finalize(None, mem, access_ty, state);
     let mut mem_str = mem_insts
         .into_iter()
         .map(|inst| {
@@ -1224,6 +1250,7 @@ fn mem_finalize_for_show(mem: &AMode, state: &EmitState) -> (String, AMode) {
         mem_str += " ; ";
     }
 
+    let mem = mem.pretty_print(access_ty.bytes() as u8, &mut AllocationConsumer::new(&[]));
     (mem_str, mem)
 }
 
@@ -1411,8 +1438,8 @@ impl Inst {
 
                 let rd = pretty_print_ireg(rd.to_reg(), size, allocs);
                 let mem = mem.with_allocs(allocs);
-                let (mem_str, mem) = mem_finalize_for_show(&mem, state);
-                let mem = mem.pretty_print_default();
+                let access_ty = self.mem_type().unwrap();
+                let (mem_str, mem) = mem_finalize_for_show(&mem, access_ty, state);
 
                 format!("{}{} {}, {}", mem_str, op, rd, mem)
             }
@@ -1438,8 +1465,8 @@ impl Inst {
 
                 let rd = pretty_print_ireg(rd, size, allocs);
                 let mem = mem.with_allocs(allocs);
-                let (mem_str, mem) = mem_finalize_for_show(&mem, state);
-                let mem = mem.pretty_print_default();
+                let access_ty = self.mem_type().unwrap();
+                let (mem_str, mem) = mem_finalize_for_show(&mem, access_ty, state);
 
                 format!("{}{} {}, {}", mem_str, op, rd, mem)
             }
@@ -1822,45 +1849,45 @@ impl Inst {
             &Inst::FpuLoad32 { rd, ref mem, .. } => {
                 let rd = pretty_print_vreg_scalar(rd.to_reg(), ScalarSize::Size32, allocs);
                 let mem = mem.with_allocs(allocs);
-                let (mem_str, mem) = mem_finalize_for_show(&mem, state);
-                let mem = mem.pretty_print_default();
+                let access_ty = self.mem_type().unwrap();
+                let (mem_str, mem) = mem_finalize_for_show(&mem, access_ty, state);
                 format!("{}ldr {}, {}", mem_str, rd, mem)
             }
             &Inst::FpuLoad64 { rd, ref mem, .. } => {
                 let rd = pretty_print_vreg_scalar(rd.to_reg(), ScalarSize::Size64, allocs);
                 let mem = mem.with_allocs(allocs);
-                let (mem_str, mem) = mem_finalize_for_show(&mem, state);
-                let mem = mem.pretty_print_default();
+                let access_ty = self.mem_type().unwrap();
+                let (mem_str, mem) = mem_finalize_for_show(&mem, access_ty, state);
                 format!("{}ldr {}, {}", mem_str, rd, mem)
             }
             &Inst::FpuLoad128 { rd, ref mem, .. } => {
                 let rd = pretty_print_reg(rd.to_reg(), allocs);
                 let rd = "q".to_string() + &rd[1..];
                 let mem = mem.with_allocs(allocs);
-                let (mem_str, mem) = mem_finalize_for_show(&mem, state);
-                let mem = mem.pretty_print_default();
+                let access_ty = self.mem_type().unwrap();
+                let (mem_str, mem) = mem_finalize_for_show(&mem, access_ty, state);
                 format!("{}ldr {}, {}", mem_str, rd, mem)
             }
             &Inst::FpuStore32 { rd, ref mem, .. } => {
                 let rd = pretty_print_vreg_scalar(rd, ScalarSize::Size32, allocs);
                 let mem = mem.with_allocs(allocs);
-                let (mem_str, mem) = mem_finalize_for_show(&mem, state);
-                let mem = mem.pretty_print_default();
+                let access_ty = self.mem_type().unwrap();
+                let (mem_str, mem) = mem_finalize_for_show(&mem, access_ty, state);
                 format!("{}str {}, {}", mem_str, rd, mem)
             }
             &Inst::FpuStore64 { rd, ref mem, .. } => {
                 let rd = pretty_print_vreg_scalar(rd, ScalarSize::Size64, allocs);
                 let mem = mem.with_allocs(allocs);
-                let (mem_str, mem) = mem_finalize_for_show(&mem, state);
-                let mem = mem.pretty_print_default();
+                let access_ty = self.mem_type().unwrap();
+                let (mem_str, mem) = mem_finalize_for_show(&mem, access_ty, state);
                 format!("{}str {}, {}", mem_str, rd, mem)
             }
             &Inst::FpuStore128 { rd, ref mem, .. } => {
                 let rd = pretty_print_reg(rd, allocs);
                 let rd = "q".to_string() + &rd[1..];
                 let mem = mem.with_allocs(allocs);
-                let (mem_str, mem) = mem_finalize_for_show(&mem, state);
-                let mem = mem.pretty_print_default();
+                let access_ty = self.mem_type().unwrap();
+                let (mem_str, mem) = mem_finalize_for_show(&mem, access_ty, state);
                 format!("{}str {}, {}", mem_str, rd, mem)
             }
             &Inst::FpuLoadP64 {
@@ -2771,7 +2798,7 @@ impl Inst {
                 // of the existing legalization framework).
                 let rd = allocs.next_writable(rd);
                 let mem = mem.with_allocs(allocs);
-                let (mem_insts, mem) = mem_finalize(None, &mem, state);
+                let (mem_insts, mem) = mem_finalize(None, &mem, I8, state);
                 let mut ret = String::new();
                 for inst in mem_insts.into_iter() {
                     ret.push_str(
diff --git a/cranelift/codegen/src/isa/aarch64/pcc.rs b/cranelift/codegen/src/isa/aarch64/pcc.rs
index e17daffd4..cdc6ca42e 100644
--- a/cranelift/codegen/src/isa/aarch64/pcc.rs
+++ b/cranelift/codegen/src/isa/aarch64/pcc.rs
@@ -38,7 +38,8 @@ pub(crate) fn check(
     inst_idx: InsnIndex,
     state: &mut FactFlowState,
 ) -> PccResult<()> {
-    trace!("Checking facts on inst: {:?}", vcode[inst_idx]);
+    let inst = &vcode[inst_idx];
+    trace!("Checking facts on inst: {:?}", inst);
 
     // We only persist flag state for one instruction, because we
     // can't exhaustively enumerate all flags-effecting ops; so take
@@ -47,28 +48,29 @@ pub(crate) fn check(
     let cmp_flags = state.cmp_flags.take();
     trace!(" * with cmp_flags = {cmp_flags:?}");
 
-    match vcode[inst_idx] {
+    match *inst {
         Inst::Args { .. } => {
             // Defs on the args have "axiomatic facts": we trust the
             // ABI code to pass through the values unharmed, so the
             // facts given to us in the CLIF should still be true.
             Ok(())
         }
-        Inst::ULoad8 { rd, ref mem, flags } | Inst::SLoad8 { rd, ref mem, flags } => {
-            check_load(ctx, Some(rd.to_reg()), flags, mem, vcode, I8)
+        Inst::ULoad8 { rd, ref mem, flags }
+        | Inst::SLoad8 { rd, ref mem, flags }
+        | Inst::ULoad16 { rd, ref mem, flags }
+        | Inst::SLoad16 { rd, ref mem, flags }
+        | Inst::ULoad32 { rd, ref mem, flags }
+        | Inst::SLoad32 { rd, ref mem, flags }
+        | Inst::ULoad64 { rd, ref mem, flags } => {
+            let access_ty = inst.mem_type().unwrap();
+            check_load(ctx, Some(rd.to_reg()), flags, mem, vcode, access_ty)
         }
-        Inst::ULoad16 { rd, ref mem, flags } | Inst::SLoad16 { rd, ref mem, flags } => {
-            check_load(ctx, Some(rd.to_reg()), flags, mem, vcode, I16)
+        Inst::FpuLoad32 { ref mem, flags, .. }
+        | Inst::FpuLoad64 { ref mem, flags, .. }
+        | Inst::FpuLoad128 { ref mem, flags, .. } => {
+            let access_ty = inst.mem_type().unwrap();
+            check_load(ctx, None, flags, mem, vcode, access_ty)
         }
-        Inst::ULoad32 { rd, ref mem, flags } | Inst::SLoad32 { rd, ref mem, flags } => {
-            check_load(ctx, Some(rd.to_reg()), flags, mem, vcode, I32)
-        }
-        Inst::ULoad64 { rd, ref mem, flags } => {
-            check_load(ctx, Some(rd.to_reg()), flags, mem, vcode, I64)
-        }
-        Inst::FpuLoad32 { ref mem, flags, .. } => check_load(ctx, None, flags, mem, vcode, F32),
-        Inst::FpuLoad64 { ref mem, flags, .. } => check_load(ctx, None, flags, mem, vcode, F64),
-        Inst::FpuLoad128 { ref mem, flags, .. } => check_load(ctx, None, flags, mem, vcode, I8X16),
         Inst::LoadP64 { ref mem, flags, .. } => check_load_pair(ctx, flags, mem, vcode, 16),
         Inst::FpuLoadP64 { ref mem, flags, .. } => check_load_pair(ctx, flags, mem, vcode, 16),
         Inst::FpuLoadP128 { ref mem, flags, .. } => check_load_pair(ctx, flags, mem, vcode, 32),
@@ -82,14 +84,18 @@ pub(crate) fn check(
             ..
         } => check_load_addr(ctx, flags, rn, vcode, access_ty),
 
-        Inst::Store8 { rd, ref mem, flags } => check_store(ctx, Some(rd), flags, mem, vcode, I8),
-        Inst::Store16 { rd, ref mem, flags } => check_store(ctx, Some(rd), flags, mem, vcode, I16),
-        Inst::Store32 { rd, ref mem, flags } => check_store(ctx, Some(rd), flags, mem, vcode, I32),
-        Inst::Store64 { rd, ref mem, flags } => check_store(ctx, Some(rd), flags, mem, vcode, I64),
-        Inst::FpuStore32 { ref mem, flags, .. } => check_store(ctx, None, flags, mem, vcode, F32),
-        Inst::FpuStore64 { ref mem, flags, .. } => check_store(ctx, None, flags, mem, vcode, F64),
-        Inst::FpuStore128 { ref mem, flags, .. } => {
-            check_store(ctx, None, flags, mem, vcode, I8X16)
+        Inst::Store8 { rd, ref mem, flags }
+        | Inst::Store16 { rd, ref mem, flags }
+        | Inst::Store32 { rd, ref mem, flags }
+        | Inst::Store64 { rd, ref mem, flags } => {
+            let access_ty = inst.mem_type().unwrap();
+            check_store(ctx, Some(rd), flags, mem, vcode, access_ty)
+        }
+        Inst::FpuStore32 { ref mem, flags, .. }
+        | Inst::FpuStore64 { ref mem, flags, .. }
+        | Inst::FpuStore128 { ref mem, flags, .. } => {
+            let access_ty = inst.mem_type().unwrap();
+            check_store(ctx, None, flags, mem, vcode, access_ty)
         }
         Inst::StoreP64 { ref mem, flags, .. } => check_store_pair(ctx, flags, mem, vcode, 16),
         Inst::FpuStoreP64 { ref mem, flags, .. } => check_store_pair(ctx, flags, mem, vcode, 16),
@@ -444,19 +450,14 @@ fn check_addr<'a>(
             trace!("rn = {rn:?} rm = {rm:?} sum = {sum:?}");
             check(&sum, ty)
         }
-        &AMode::RegScaled { rn, rm, ty } => {
+        &AMode::RegScaled { rn, rm } => {
             let rn = get_fact_or_default(vcode, rn, 64);
             let rm = get_fact_or_default(vcode, rm, 64);
             let rm_scaled = fail_if_missing(ctx.scale(&rm, 64, ty.bytes()))?;
             let sum = fail_if_missing(ctx.add(&rn, &rm_scaled, 64))?;
             check(&sum, ty)
         }
-        &AMode::RegScaledExtended {
-            rn,
-            rm,
-            ty,
-            extendop,
-        } => {
+        &AMode::RegScaledExtended { rn, rm, extendop } => {
             let rn = get_fact_or_default(vcode, rn, 64);
             let rm = get_fact_or_default(vcode, rm, 64);
             let rm_extended = fail_if_missing(extend_fact(ctx, &rm, extendop))?;
@@ -485,7 +486,7 @@ fn check_addr<'a>(
             // 32760 (= 4095 * 8) for I64s. The `UImm12Scaled` type
             // stores the *scaled* value, so we don't need to multiply
             // (again) by the type's size here.
-            let offset: u64 = uimm12.value.into();
+            let offset: u64 = uimm12.value().into();
             // This `unwrap()` will always succeed because the value
             // will always be positive and much smaller than
             // `i64::MAX` (see above).
diff --git a/winch/codegen/src/isa/aarch64/address.rs b/winch/codegen/src/isa/aarch64/address.rs
index 11521e5ce..02b1680d6 100644
--- a/winch/codegen/src/isa/aarch64/address.rs
+++ b/winch/codegen/src/isa/aarch64/address.rs
@@ -137,7 +137,6 @@ impl TryFrom<Address> for AMode {
             Offset { base, offset } => Ok(AMode::RegOffset {
                 rn: base.into(),
                 off: offset,
-                ty: types::I64,
             }),
         }
     }


Test patch:
diff --git a/cranelift/codegen/src/isa/aarch64/inst/emit_tests.rs b/cranelift/codegen/src/isa/aarch64/inst/emit_tests.rs
index 3a0a81324..d470ed901 100644
--- a/cranelift/codegen/src/isa/aarch64/inst/emit_tests.rs
+++ b/cranelift/codegen/src/isa/aarch64/inst/emit_tests.rs
@@ -1551,7 +1551,6 @@ fn test_aarch64_binemit() {
             mem: AMode::RegScaled {
                 rn: xreg(2),
                 rm: xreg(3),
-                ty: I16,
             },
             flags: MemFlags::trusted(),
         },
@@ -1588,7 +1587,6 @@ fn test_aarch64_binemit() {
             mem: AMode::RegScaled {
                 rn: xreg(20),
                 rm: xreg(20),
-                ty: I16,
             },
             flags: MemFlags::trusted(),
         },
@@ -1625,7 +1623,6 @@ fn test_aarch64_binemit() {
             mem: AMode::RegScaled {
                 rn: xreg(2),
                 rm: xreg(12),
-                ty: I32,
             },
             flags: MemFlags::trusted(),
         },
@@ -1662,7 +1659,6 @@ fn test_aarch64_binemit() {
             mem: AMode::RegScaled {
                 rn: xreg(5),
                 rm: xreg(1),
-                ty: I32,
             },
             flags: MemFlags::trusted(),
         },
@@ -1735,7 +1731,6 @@ fn test_aarch64_binemit() {
             mem: AMode::RegScaled {
                 rn: xreg(2),
                 rm: xreg(3),
-                ty: I64,
             },
             flags: MemFlags::trusted(),
         },
@@ -1748,7 +1743,6 @@ fn test_aarch64_binemit() {
             mem: AMode::RegScaledExtended {
                 rn: xreg(2),
                 rm: xreg(3),
-                ty: I64,
                 extendop: ExtendOp::SXTW,
             },
             flags: MemFlags::trusted(),
@@ -1805,7 +1799,7 @@ fn test_aarch64_binemit() {
     insns.push((
         Inst::ULoad64 {
             rd: writable_xreg(1),
-            mem: AMode::FPOffset { off: 32768, ty: I8 },
+            mem: AMode::FPOffset { off: 32768 },
             flags: MemFlags::trusted(),
         },
         "100090D2A1EB70F8",
@@ -1814,10 +1808,7 @@ fn test_aarch64_binemit() {
     insns.push((
         Inst::ULoad64 {
             rd: writable_xreg(1),
-            mem: AMode::FPOffset {
-                off: -32768,
-                ty: I8,
-            },
+            mem: AMode::FPOffset { off: -32768 },
             flags: MemFlags::trusted(),
         },
         "F0FF8F92A1EB70F8",
@@ -1826,10 +1817,7 @@ fn test_aarch64_binemit() {
     insns.push((
         Inst::ULoad64 {
             rd: writable_xreg(1),
-            mem: AMode::FPOffset {
-                off: 1048576,
-                ty: I8,
-            }, // 2^20
+            mem: AMode::FPOffset { off: 1048576 }, // 2^20
             flags: MemFlags::trusted(),
         },
         "1002A0D2A1EB70F8",
@@ -1838,10 +1826,7 @@ fn test_aarch64_binemit() {
     insns.push((
         Inst::ULoad64 {
             rd: writable_xreg(1),
-            mem: AMode::FPOffset {
-                off: 1048576 + 1,
-                ty: I8,
-            }, // 2^20 + 1
+            mem: AMode::FPOffset { off: 1048576 + 1 }, // 2^20 + 1
             flags: MemFlags::trusted(),
         },
         "300080521002A072A1EB70F8",
@@ -1854,7 +1839,6 @@ fn test_aarch64_binemit() {
             mem: AMode::RegOffset {
                 rn: xreg(7),
                 off: 8,
-                ty: I64,
             },
             flags: MemFlags::trusted(),
         },
@@ -1868,7 +1852,6 @@ fn test_aarch64_binemit() {
             mem: AMode::RegOffset {
                 rn: xreg(7),
                 off: 1024,
-                ty: I64,
             },
             flags: MemFlags::trusted(),
         },
@@ -1882,7 +1865,6 @@ fn test_aarch64_binemit() {
             mem: AMode::RegOffset {
                 rn: xreg(7),
                 off: 1048576,
-                ty: I64,
             },
             flags: MemFlags::trusted(),
         },
@@ -2004,7 +1986,6 @@ fn test_aarch64_binemit() {
             mem: AMode::RegScaled {
                 rn: xreg(2),
                 rm: xreg(3),
-                ty: I64,
             },
             flags: MemFlags::trusted(),
         },
@@ -2017,7 +1998,6 @@ fn test_aarch64_binemit() {
             mem: AMode::RegScaledExtended {
                 rn: xreg(2),
                 rm: xreg(3),
-                ty: I64,
                 extendop: ExtendOp::UXTW,
             },
             flags: MemFlags::trusted(),
@@ -6692,7 +6672,6 @@ fn test_aarch64_binemit() {
             mem: AMode::RegScaled {
                 rn: xreg(8),
                 rm: xreg(9),
-                ty: F32,
             },
             flags: MemFlags::trusted(),
         },
@@ -6706,7 +6685,6 @@ fn test_aarch64_binemit() {
             mem: AMode::RegScaled {
                 rn: xreg(8),
                 rm: xreg(9),
-                ty: F64,
             },
             flags: MemFlags::trusted(),
         },
@@ -6720,7 +6698,6 @@ fn test_aarch64_binemit() {
             mem: AMode::RegScaled {
                 rn: xreg(8),
                 rm: xreg(9),
-                ty: I128,
             },
             flags: MemFlags::trusted(),
         },
@@ -6770,7 +6747,6 @@ fn test_aarch64_binemit() {
             mem: AMode::RegScaled {
                 rn: xreg(8),
                 rm: xreg(9),
-                ty: F32,
             },
             flags: MemFlags::trusted(),
         },
@@ -6784,7 +6760,6 @@ fn test_aarch64_binemit() {
             mem: AMode::RegScaled {
                 rn: xreg(8),
                 rm: xreg(9),
-                ty: F64,
             },
             flags: MemFlags::trusted(),
         },
@@ -6798,7 +6773,6 @@ fn test_aarch64_binemit() {
             mem: AMode::RegScaled {
                 rn: xreg(8),
                 rm: xreg(9),
-                ty: I128,
             },
             flags: MemFlags::trusted(),
         },