Cranelift: Make ir::{Constant,Immediate} considered entities (#11207)

fitzgen · web-flow · commit 8a23cc748fda · 2025-07-09T19:24:48.000Z
* Cranelift: Make `ir::{Constant,Immediate}` considered entities

They reference data in out-of-line pools rather than storing their data inline
in the instruction, and when an instruction containing them is moved from one
`ir::Function` to another, they need their indices updated
accordingly. Therefore, they really are entities rather than immediates.

This recategorization means that they will now be properly mapped in
`ir::InstructionData::map` calls.

* fix tests
diff --git a/cranelift/codegen/meta/src/shared/entities.rs b/cranelift/codegen/meta/src/shared/entities.rs
@@ -48,6 +48,19 @@ pub(crate) struct EntityRefs {
 
     /// A variable-sized list of value operands. Use for Block and function call arguments.
     pub(crate) varargs: OperandKind,
+
+    /// A constant stored in the constant pool.
+    ///
+    /// This operand is used to pass constants to instructions like `vconst`
+    /// while storing the actual bytes in the constant pool.
+    pub(crate) pool_constant: OperandKind,
+
+    /// An unsigned 128-bit immediate integer operand, stored out-of-line in the
+    /// `DataFlowGraph::immediates` pool.
+    ///
+    /// This operand is used to pass entire 128-bit vectors as immediates to instructions like
+    /// `shuffle` and `mask`.
+    pub(crate) uimm128: OperandKind,
 }
 
 impl EntityRefs {
@@ -101,6 +114,18 @@ impl EntityRefs {
                         returned from an instruction.
                     "#,
             ),
+
+            pool_constant: new(
+                "constant_handle",
+                "ir::Constant",
+                "A constant stored in the constant pool.",
+            ),
+
+            uimm128: new(
+                "imm",
+                "ir::Immediate",
+                "A 128-bit immediate unsigned integer.",
+            ),
         }
     }
 }
diff --git a/cranelift/codegen/meta/src/shared/formats.rs b/cranelift/codegen/meta/src/shared/formats.rs
@@ -57,7 +57,9 @@ impl Formats {
 
             unary_ieee64: Builder::new("UnaryIeee64").imm(&imm.ieee64).build(),
 
-            unary_const: Builder::new("UnaryConst").imm(&imm.pool_constant).build(),
+            unary_const: Builder::new("UnaryConst")
+                .imm(&entities.pool_constant)
+                .build(),
 
             unary_global_value: Builder::new("UnaryGlobalValue")
                 .imm(&entities.global_value)
@@ -94,7 +96,7 @@ impl Formats {
             shuffle: Builder::new("Shuffle")
                 .value()
                 .value()
-                .imm(&imm.uimm128)
+                .imm(&entities.uimm128)
                 .build(),
 
             int_compare: Builder::new("IntCompare")
diff --git a/cranelift/codegen/meta/src/shared/immediates.rs b/cranelift/codegen/meta/src/shared/immediates.rs
@@ -14,18 +14,6 @@ pub(crate) struct Immediates {
     /// counts on shift instructions.
     pub uimm8: OperandKind,
 
-    /// An unsigned 128-bit immediate integer operand.
-    ///
-    /// This operand is used to pass entire 128-bit vectors as immediates to instructions like
-    /// const.
-    pub uimm128: OperandKind,
-
-    /// A constant stored in the constant pool.
-    ///
-    /// This operand is used to pass constants to instructions like vconst while storing the
-    /// actual bytes in the constant pool.
-    pub pool_constant: OperandKind,
-
     /// A 32-bit immediate signed offset.
     ///
     /// This is used to represent an immediate address offset in load/store instructions.
@@ -109,16 +97,6 @@ impl Immediates {
                 "ir::immediates::Uimm8",
                 "An 8-bit immediate unsigned integer.",
             ),
-            uimm128: new_imm(
-                "imm",
-                "ir::Immediate",
-                "A 128-bit immediate unsigned integer.",
-            ),
-            pool_constant: new_imm(
-                "constant_handle",
-                "ir::Constant",
-                "A constant stored in the constant pool.",
-            ),
             offset32: new_imm(
                 "offset",
                 "ir::immediates::Offset32",
diff --git a/cranelift/codegen/meta/src/shared/instructions.rs b/cranelift/codegen/meta/src/shared/instructions.rs
@@ -1478,7 +1478,7 @@ pub(crate) fn define(
         "#,
             &formats.unary_const,
         )
-        .operands_in(vec![Operand::new("N", &imm.pool_constant)])
+        .operands_in(vec![Operand::new("N", &entities.pool_constant)])
         .operands_out(vec![
             Operand::new("a", f128_).with_doc("A constant f128 scalar value"),
         ]),
@@ -1495,7 +1495,7 @@ pub(crate) fn define(
             &formats.unary_const,
         )
         .operands_in(vec![
-            Operand::new("N", &imm.pool_constant)
+            Operand::new("N", &entities.pool_constant)
                 .with_doc("The 16 immediate bytes of a 128-bit vector"),
         ])
         .operands_out(vec![
@@ -1530,7 +1530,7 @@ pub(crate) fn define(
         .operands_in(vec![
             Operand::new("a", Tx16).with_doc("A vector value"),
             Operand::new("b", Tx16).with_doc("A vector value"),
-            Operand::new("mask", &imm.uimm128)
+            Operand::new("mask", &entities.uimm128)
                 .with_doc("The 16 immediate bytes used for selecting the elements to shuffle"),
         ])
         .operands_out(vec![Operand::new("a", Tx16).with_doc("A vector value")]),
diff --git a/cranelift/codegen/src/ir/instructions.rs b/cranelift/codegen/src/ir/instructions.rs
@@ -1062,6 +1062,12 @@ pub trait InstructionMapper {
         &mut self,
         dynamic_stack_slot: ir::DynamicStackSlot,
     ) -> ir::DynamicStackSlot;
+
+    /// Map a function over a `Constant`.
+    fn map_constant(&mut self, constant: ir::Constant) -> ir::Constant;
+
+    /// Map a function over an `Immediate`.
+    fn map_immediate(&mut self, immediate: ir::Immediate) -> ir::Immediate;
 }
 
 impl<'a, T> InstructionMapper for &'a mut T
@@ -1110,6 +1116,14 @@ where
     ) -> ir::DynamicStackSlot {
         (**self).map_dynamic_stack_slot(dynamic_stack_slot)
     }
+
+    fn map_constant(&mut self, constant: ir::Constant) -> ir::Constant {
+        (**self).map_constant(constant)
+    }
+
+    fn map_immediate(&mut self, immediate: ir::Immediate) -> ir::Immediate {
+        (**self).map_immediate(immediate)
+    }
 }
 
 #[cfg(test)]
@@ -1320,6 +1334,14 @@ mod tests {
             ) -> ir::DynamicStackSlot {
                 DynamicStackSlot::from_u32(dynamic_stack_slot.as_u32() + 1)
             }
+
+            fn map_constant(&mut self, constant: ir::Constant) -> ir::Constant {
+                ir::Constant::from_u32(constant.as_u32() + 1)
+            }
+
+            fn map_immediate(&mut self, immediate: ir::Immediate) -> ir::Immediate {
+                ir::Immediate::from_u32(immediate.as_u32() + 1)
+            }
         }
 
         let mut pool = ValueListPool::new();
@@ -1450,5 +1472,31 @@ mod tests {
                 dynamic_stack_slot: DynamicStackSlot::from_u32(1),
             },
         );
+
+        // Mapping `Constant`s
+        assert_eq!(
+            map(InstructionData::UnaryConst {
+                opcode: ir::Opcode::Vconst,
+                constant_handle: ir::Constant::from_u32(2)
+            }),
+            InstructionData::UnaryConst {
+                opcode: ir::Opcode::Vconst,
+                constant_handle: ir::Constant::from_u32(3)
+            },
+        );
+
+        // Mapping `Immediate`s
+        assert_eq!(
+            map(InstructionData::Shuffle {
+                opcode: ir::Opcode::Shuffle,
+                args: [Value::from_u32(0), Value::from_u32(1)],
+                imm: ir::Immediate::from_u32(41),
+            }),
+            InstructionData::Shuffle {
+                opcode: ir::Opcode::Shuffle,
+                args: [Value::from_u32(1), Value::from_u32(2)],
+                imm: ir::Immediate::from_u32(42),
+            },
+        );
     }
 }
