cranelift/riscv64: Idiomatic AMode operand collection (#8454)

Other backends have functions for passing any registers in an AMode to
the operand collector. Use the same pattern in this backend too: it's
simpler to use, and also sets up for more changes I'm planning.

cranelift/codegen/src/isa/riscv64/inst/args.rs

@@ -125,17 +125,21 @@ impl AMode {
         }
     }
 
-    /// Returns the registers that known to the register allocator.
+    /// Add the registers referenced by this AMode to `collector`.
     /// Keep this in sync with `with_allocs`.
-    pub(crate) fn get_allocatable_register(&self) -> Option<Reg> {
+    pub(crate) fn get_operands<F: Fn(regalloc2::VReg) -> regalloc2::VReg>(
+        &self,
+        collector: &mut OperandCollector<'_, F>,
+    ) {
         match self {
-            AMode::RegOffset(reg, ..) => Some(*reg),
+            &AMode::RegOffset(reg, ..) => collector.reg_use(reg),
+            // Registers used in these modes aren't allocatable.
             AMode::SPOffset(..)
             | AMode::FPOffset(..)
             | AMode::NominalSPOffset(..)
             | AMode::IncomingArg(..)
             | AMode::Const(..)
-            | AMode::Label(..) => None,
+            | AMode::Label(..) => {}
         }
     }
 

cranelift/codegen/src/isa/riscv64/inst/mod.rs

@@ -376,15 +376,11 @@ fn riscv64_get_operands<F: Fn(VReg) -> VReg>(inst: &Inst, collector: &mut Operan
             collector.reg_def(rd);
         }
         &Inst::Load { rd, from, .. } => {
-            if let Some(r) = from.get_allocatable_register() {
-                collector.reg_use(r);
-            }
+            from.get_operands(collector);
             collector.reg_def(rd);
         }
         &Inst::Store { to, src, .. } => {
-            if let Some(r) = to.get_allocatable_register() {
-                collector.reg_use(r);
-            }
+            to.get_operands(collector);
             collector.reg_use(src);
         }
 
@@ -459,9 +455,7 @@ fn riscv64_get_operands<F: Fn(VReg) -> VReg>(inst: &Inst, collector: &mut Operan
             collector.reg_clobbers(clobbers);
         }
         &Inst::LoadAddr { rd, mem } => {
-            if let Some(r) = mem.get_allocatable_register() {
-                collector.reg_use(r);
-            }
+            mem.get_operands(collector);
             collector.reg_early_def(rd);
         }
 
@@ -752,9 +746,7 @@ fn riscv64_get_operands<F: Fn(VReg) -> VReg>(inst: &Inst, collector: &mut Operan
             ref mask,
             ..
         } => {
-            if let Some(r) = from.get_allocatable_register() {
-                collector.reg_use(r);
-            }
+            from.get_operands(collector);
             collector.reg_def(to);
             vec_mask_operands(mask, collector);
         }
@@ -764,9 +756,7 @@ fn riscv64_get_operands<F: Fn(VReg) -> VReg>(inst: &Inst, collector: &mut Operan
             ref mask,
             ..
         } => {
-            if let Some(r) = to.get_allocatable_register() {
-                collector.reg_use(r);
-            }
+            to.get_operands(collector);
             collector.reg_use(from);
             vec_mask_operands(mask, collector);
         }

cranelift/codegen/src/isa/riscv64/inst/vector.rs

@@ -4,7 +4,7 @@ use crate::isa::riscv64::lower::isle::generated_code::{
     VecAMode, VecAluOpRImm5, VecAluOpRR, VecAluOpRRImm5, VecAluOpRRR, VecAluOpRRRImm5, VecAvl,
     VecElementWidth, VecLmul, VecMaskMode, VecOpCategory, VecOpMasking, VecTailMode,
 };
-use crate::machinst::RegClass;
+use crate::machinst::{OperandCollector, RegClass};
 use crate::Reg;
 use core::fmt;
 
@@ -1070,9 +1070,12 @@ impl VecAMode {
         }
     }
 
-    pub fn get_allocatable_register(&self) -> Option<Reg> {
+    pub fn get_operands<F: Fn(regalloc2::VReg) -> regalloc2::VReg>(
+        &self,
+        collector: &mut OperandCollector<'_, F>,
+    ) {
         match self {
-            VecAMode::UnitStride { base, .. } => base.get_allocatable_register(),
+            VecAMode::UnitStride { base, .. } => base.get_operands(collector),
         }
     }