refactor: change LowerCtx::get_immediate to return a DataValue

This change abstracts away (from the perspective of the new backend) how immediate values are stored in InstructionData. It gathers large immediates from necessary places (e.g. constant pool) and delegates to `InstructionData::imm_value` for the rest. This refactor only touches original users of `LowerCtx::get_immediate` but a future change could do the same for any place the new backend is accessing InstructionData directly to retrieve immediates.

cranelift/codegen/src/ir/constant.rs

@@ -63,6 +63,11 @@ impl ConstantData {
         self.0.is_empty()
     }
 
+    /// Return the data as a slice.
+    pub fn as_slice(&self) -> &[u8] {
+        self.0.as_slice()
+    }
+
     /// Convert the data to a vector.
     pub fn into_vec(self) -> Vec<u8> {
         self.0

cranelift/codegen/src/isa/aarch64/lower.rs

@@ -10,7 +10,7 @@
 use crate::ir::condcodes::{FloatCC, IntCC};
 use crate::ir::types::*;
 use crate::ir::Inst as IRInst;
-use crate::ir::{InstructionData, Opcode, Type};
+use crate::ir::{Opcode, Type};
 use crate::machinst::lower::*;
 use crate::machinst::*;
 use crate::CodegenResult;
@@ -20,6 +20,7 @@ use crate::isa::aarch64::AArch64Backend;
 
 use super::lower_inst;
 
+use crate::data_value::DataValue;
 use log::{debug, trace};
 use regalloc::{Reg, RegClass, Writable};
 use smallvec::SmallVec;
@@ -126,22 +127,9 @@ pub(crate) fn const_param_to_u128<C: LowerCtx<I = Inst>>(
     ctx: &mut C,
     inst: IRInst,
 ) -> Option<u128> {
-    let data = match ctx.data(inst) {
-        &InstructionData::Shuffle { mask, .. } => ctx.get_immediate(mask),
-        &InstructionData::UnaryConst {
-            constant_handle, ..
-        } => ctx.get_constant_data(constant_handle),
-        _ => return None,
-    };
-    let data = data.clone().into_vec();
-
-    if data.len() == 16 {
-        let mut bytes = [0u8; 16];
-
-        bytes.copy_from_slice(&data);
-        Some(u128::from_le_bytes(bytes))
-    } else {
-        None
+    match ctx.get_immediate(inst) {
+        Some(DataValue::V128(bytes)) => Some(u128::from_le_bytes(bytes)),
+        _ => None,
     }
 }
 

cranelift/codegen/src/isa/x64/lower.rs

@@ -2,6 +2,7 @@
 
 #![allow(non_snake_case)]
 
+use crate::data_value::DataValue;
 use crate::ir::{
     condcodes::FloatCC, types, AbiParam, ArgumentPurpose, ExternalName, Inst as IRInst,
     InstructionData, LibCall, Opcode, Signature, Type,
@@ -2642,10 +2643,9 @@ fn lower_insn_to_regs<C: LowerCtx<I = Inst>>(
             let lhs_ty = ctx.input_ty(insn, 0);
             let lhs = put_input_in_reg(ctx, inputs[0]);
             let rhs = put_input_in_reg(ctx, inputs[1]);
-            let mask = if let &InstructionData::Shuffle { mask, .. } = ctx.data(insn) {
-                ctx.get_immediate(mask).clone()
-            } else {
-                unreachable!("shuffle should always have the shuffle format")
+            let mask = match ctx.get_immediate(insn) {
+                Some(DataValue::V128(bytes)) => bytes.to_vec(),
+                _ => unreachable!("shuffle should always have a 16-byte immediate"),
             };
 
             // A mask-building helper: in 128-bit SIMD, 0-15 indicate which lane to read from and a

cranelift/codegen/src/machinst/lower.rs

@@ -8,9 +8,8 @@ use crate::inst_predicates::{has_lowering_side_effect, is_constant_64bit};
 use crate::ir::instructions::BranchInfo;
 use crate::ir::types::I64;
 use crate::ir::{
-    ArgumentPurpose, Block, Constant, ConstantData, ExternalName, Function, GlobalValueData,
-    Immediate, Inst, InstructionData, MemFlags, Opcode, Signature, SourceLoc, Type, Value,
-    ValueDef,
+    ArgumentPurpose, Block, Constant, ConstantData, ExternalName, Function, GlobalValueData, Inst,
+    InstructionData, MemFlags, Opcode, Signature, SourceLoc, Type, Value, ValueDef,
 };
 use crate::machinst::{
     ABICallee, BlockIndex, BlockLoweringOrder, LoweredBlock, MachLabel, VCode, VCodeBuilder,
@@ -20,8 +19,10 @@ use crate::CodegenResult;
 
 use regalloc::{Reg, RegClass, StackmapRequestInfo, VirtualReg, Writable};
 
+use crate::data_value::DataValue;
 use alloc::boxed::Box;
 use alloc::vec::Vec;
+use core::convert::TryInto;
 use log::debug;
 use smallvec::SmallVec;
 
@@ -161,8 +162,9 @@ pub trait LowerCtx {
     fn is_reg_needed(&self, ir_inst: Inst, reg: Reg) -> bool;
     /// Retrieve constant data given a handle.
     fn get_constant_data(&self, constant_handle: Constant) -> &ConstantData;
-    /// Retrieve an immediate given a reference.
-    fn get_immediate(&self, imm: Immediate) -> &ConstantData;
+    /// Retrieve the value immediate from an instruction. This will perform necessary lookups on the
+    /// `DataFlowGraph` to retrieve even large immediates.
+    fn get_immediate(&self, ir_inst: Inst) -> Option<DataValue>;
     /// Cause the value in `reg` to be in a virtual reg, by copying it into a new virtual reg
     /// if `reg` is a real reg.  `ty` describes the type of the value in `reg`.
     fn ensure_in_vreg(&mut self, reg: Reg, ty: Type) -> Reg;
@@ -1007,8 +1009,23 @@ impl<'func, I: VCodeInst> LowerCtx for Lower<'func, I> {
         self.f.dfg.constants.get(constant_handle)
     }
 
-    fn get_immediate(&self, imm: Immediate) -> &ConstantData {
-        self.f.dfg.immediates.get(imm).unwrap()
+    fn get_immediate(&self, ir_inst: Inst) -> Option<DataValue> {
+        let inst_data = self.data(ir_inst);
+        match inst_data {
+            InstructionData::Shuffle { mask, .. } => {
+                let buffer = self.f.dfg.immediates.get(mask.clone()).unwrap().as_slice();
+                let value = DataValue::V128(buffer.try_into().expect("a 16-byte data buffer"));
+                Some(value)
+            }
+            InstructionData::UnaryConst {
+                constant_handle, ..
+            } => {
+                let buffer = self.f.dfg.constants.get(constant_handle.clone()).as_slice();
+                let value = DataValue::V128(buffer.try_into().expect("a 16-byte data buffer"));
+                Some(value)
+            }
+            _ => inst_data.imm_value(),
+        }
     }
 
     fn ensure_in_vreg(&mut self, reg: Reg, ty: Type) -> Reg {