1 files changed, 126 insertions, 0 deletions
diff --git a/third_party/rust/cranelift-codegen/src/regalloc/affinity.rs b/third_party/rust/cranelift-codegen/src/regalloc/affinity.rs
new file mode 100644
index 0000000000..efcc4dabfa
--- /dev/null
+++ b/third_party/rust/cranelift-codegen/src/regalloc/affinity.rs
@@ -0,0 +1,126 @@
+//! Value affinity for register allocation.
+//!
+//! An SSA value's affinity is a hint used to guide the register allocator. It specifies the class
+//! of allocation that is likely to cause the least amount of fixup moves in order to satisfy
+//! instruction operand constraints.
+//!
+//! For values that want to be in registers, the affinity hint includes a register class or
+//! subclass. This is just a hint, and the register allocator is allowed to pick a register from a
+//! larger register class instead.
+
+use crate::ir::{AbiParam, ArgumentLoc};
+use crate::isa::{ConstraintKind, OperandConstraint, RegClassIndex, RegInfo, TargetIsa};
+use core::fmt;
+
+/// Preferred register allocation for an SSA value.
+#[derive(Clone, Copy, Debug)]
+pub enum Affinity {
+    /// No affinity.
+    ///
+    /// This indicates a value that is not defined or used by any real instructions. It is a ghost
+    /// value that won't appear in the final program.
+    Unassigned,
+
+    /// This value should be placed in a spill slot on the stack.
+    Stack,
+
+    /// This value prefers a register from the given register class.
+    Reg(RegClassIndex),
+}
+
+impl Default for Affinity {
+    fn default() -> Self {
+        Self::Unassigned
+    }
+}
+
+impl Affinity {
+    /// Create an affinity that satisfies a single constraint.
+    ///
+    /// This will never create an `Affinity::Unassigned`.
+    /// Use the `Default` implementation for that.
+    pub fn new(constraint: &OperandConstraint) -> Self {
+        if constraint.kind == ConstraintKind::Stack {
+            Self::Stack
+        } else {
+            Self::Reg(constraint.regclass.into())
+        }
+    }
+
+    /// Create an affinity that matches an ABI argument for `isa`.
+    pub fn abi(arg: &AbiParam, isa: &dyn TargetIsa) -> Self {
+        match arg.location {
+            ArgumentLoc::Unassigned => Self::Unassigned,
+            ArgumentLoc::Reg(_) => Self::Reg(isa.regclass_for_abi_type(arg.value_type).into()),
+            ArgumentLoc::Stack(_) => Self::Stack,
+        }
+    }
+
+    /// Is this the `Unassigned` affinity?
+    pub fn is_unassigned(self) -> bool {
+        match self {
+            Self::Unassigned => true,
+            _ => false,
+        }
+    }
+
+    /// Is this the `Reg` affinity?
+    pub fn is_reg(self) -> bool {
+        match self {
+            Self::Reg(_) => true,
+            _ => false,
+        }
+    }
+
+    /// Is this the `Stack` affinity?
+    pub fn is_stack(self) -> bool {
+        match self {
+            Self::Stack => true,
+            _ => false,
+        }
+    }
+
+    /// Merge an operand constraint into this affinity.
+    ///
+    /// Note that this does not guarantee that the register allocator will pick a register that
+    /// satisfies the constraint.
+    pub fn merge(&mut self, constraint: &OperandConstraint, reginfo: &RegInfo) {
+        match *self {
+            Self::Unassigned => *self = Self::new(constraint),
+            Self::Reg(rc) => {
+                // If the preferred register class is a subclass of the constraint, there's no need
+                // to change anything.
+                if constraint.kind != ConstraintKind::Stack && !constraint.regclass.has_subclass(rc)
+                {
+                    // If the register classes overlap, try to shrink our preferred register class.
+                    if let Some(subclass) = constraint.regclass.intersect_index(reginfo.rc(rc)) {
+                        *self = Self::Reg(subclass);
+                    }
+                }
+            }
+            Self::Stack => {}
+        }
+    }
+
+    /// Return an object that can display this value affinity, using the register info from the
+    /// target ISA.
+    pub fn display<'a, R: Into<Option<&'a RegInfo>>>(self, regs: R) -> DisplayAffinity<'a> {
+        DisplayAffinity(self, regs.into())
+    }
+}
+
+/// Displaying an `Affinity` correctly requires the associated `RegInfo` from the target ISA.
+pub struct DisplayAffinity<'a>(Affinity, Option<&'a RegInfo>);
+
+impl<'a> fmt::Display for DisplayAffinity<'a> {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        match self.0 {
+            Affinity::Unassigned => write!(f, "unassigned"),
+            Affinity::Stack => write!(f, "stack"),
+            Affinity::Reg(rci) => match self.1 {
+                Some(regs) => write!(f, "{}", regs.rc(rci)),
+                None => write!(f, "{}", rci),
+            },
+        }
+    }
+}