Fixed width immediates (https://github.com/Shopify/ruby/pull/437)

There are a lot of times when encoding AArch64 instructions that we
need to represent an integer value with a custom fixed width. For
example, the offset for a B instruction is 26 bits, so we store an
i32 on the instruction struct and then mask it when we encode.

We've been doing this masking everywhere, which has worked, but
it's getting a bit copy-pasty all over the place. This commit
centralizes that logic to make sure we stay consistent.

12 files changed, 110 insertions, 48 deletions

diff --git a/yjit/src/asm/arm64/arg/mod.rs b/yjit/src/asm/arm64/arg/mod.rs
index 30f3cc3dfe..9bf4a8ea13 100644
--- a/yjit/src/asm/arm64/arg/mod.rs
+++ b/yjit/src/asm/arm64/arg/mod.rs
@@ -6,9 +6,11 @@ mod condition;
 mod sf;
 mod shifted_imm;
 mod sys_reg;
+mod truncate;
 
 pub use bitmask_imm::BitmaskImmediate;
 pub use condition::Condition;
 pub use sf::Sf;
 pub use shifted_imm::ShiftedImmediate;
 pub use sys_reg::SystemRegister;
+pub use truncate::{truncate_imm, truncate_uimm};
diff --git a/yjit/src/asm/arm64/arg/truncate.rs b/yjit/src/asm/arm64/arg/truncate.rs
new file mode 100644
index 0000000000..52f2c012cb
--- /dev/null
+++ b/yjit/src/asm/arm64/arg/truncate.rs
@@ -0,0 +1,66 @@
+// There are many instances in AArch64 instruction encoding where you represent
+// an integer value with a particular bit width that isn't a power of 2. These
+// functions represent truncating those integer values down to the appropriate
+// number of bits.
+
+/// Truncate a signed immediate to fit into a compile-time known width. It is
+/// assumed before calling this function that the value fits into the correct
+/// size. If it doesn't, then this function will panic.
+/// 
+/// When the value is positive, this should effectively be a no-op since we're
+/// just dropping leading zeroes. When the value is negative we should only be
+/// dropping leading ones.
+pub fn truncate_imm<T: Into<i32>, const WIDTH: usize>(imm: T) -> u32 {
+    let value: i32 = imm.into();
+    let masked = (value as u32) & ((1 << WIDTH) - 1);
+
+    // Assert that we didn't drop any bits by truncating.
+    if value >= 0 {
+        assert_eq!(value as u32, masked);
+    } else {
+        assert_eq!(value as u32, masked | (u32::MAX << WIDTH));
+    }
+
+    masked
+}
+
+/// Truncate an unsigned immediate to fit into a compile-time known width. It is
+/// assumed before calling this function that the value fits into the correct
+/// size. If it doesn't, then this function will panic.
+/// 
+/// This should effectively be a no-op since we're just dropping leading zeroes.
+pub fn truncate_uimm<T: Into<u32>, const WIDTH: usize>(uimm: T) -> u32 {
+    let value: u32 = uimm.into();
+    let masked = (value & ((1 << WIDTH) - 1));
+
+    // Assert that we didn't drop any bits by truncating.
+    assert_eq!(value, masked);
+
+    masked
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_truncate_imm_positive() {
+        let inst = truncate_imm::<i32, 4>(5);
+        let result: u32 = inst.into();
+        assert_eq!(0b0101, result);
+    }
+
+    #[test]
+    fn test_truncate_imm_negative() {
+        let inst = truncate_imm::<i32, 4>(-5);
+        let result: u32 = inst.into();
+        assert_eq!(0b1011, result);
+    }
+
+    #[test]
+    fn test_truncate_uimm() {
+        let inst = truncate_uimm::<u32, 4>(5);
+        let result: u32 = inst.into();
+        assert_eq!(0b0101, result);
+    }
+}
diff --git a/yjit/src/asm/arm64/inst/branch_cond.rs b/yjit/src/asm/arm64/inst/branch_cond.rs
index 33cc9c3649..a6bc79dffe 100644
--- a/yjit/src/asm/arm64/inst/branch_cond.rs
+++ b/yjit/src/asm/arm64/inst/branch_cond.rs
@@ -1,4 +1,4 @@
-use super::super::arg::Condition;
+use super::super::arg::{Condition, truncate_imm};
 
 /// The struct that represents an A64 conditional branch instruction that can be
 /// encoded.
@@ -31,12 +31,10 @@ const FAMILY: u32 = 0b101;
 impl From<BranchCond> for u32 {
     /// Convert an instruction into a 32-bit value.
     fn from(inst: BranchCond) -> Self {
-        let imm19 = (inst.imm19 as u32) & ((1 << 19) - 1);
-
         0
         | (1 << 30)
         | (FAMILY << 26)
-        | (imm19 << 5)
+        | (truncate_imm::<_, 19>(inst.imm19) << 5)
         | (inst.cond as u32)
     }
 }
@@ -66,8 +64,14 @@ mod tests {
     }
 
     #[test]
-    fn test_b_ne_neg() {
-        let result: u32 = BranchCond::bcond(Condition::NE, -128).into();
-        assert_eq!(0x54fffc01, result);
+    fn test_b_eq_max() {
+        let result: u32 = BranchCond::bcond(Condition::EQ, (1 << 20) - 4).into();
+        assert_eq!(0x547fffe0, result);
+    }
+
+    #[test]
+    fn test_b_eq_min() {
+        let result: u32 = BranchCond::bcond(Condition::EQ, -(1 << 20)).into();
+        assert_eq!(0x54800000, result);
     }
 }
diff --git a/yjit/src/asm/arm64/inst/call.rs b/yjit/src/asm/arm64/inst/call.rs
index 8d65359f77..32d924f799 100644
--- a/yjit/src/asm/arm64/inst/call.rs
+++ b/yjit/src/asm/arm64/inst/call.rs
@@ -1,3 +1,5 @@
+use super::super::arg::truncate_imm;
+
 /// The operation to perform for this instruction.
 enum Op {
     /// Branch directly, with a hint that this is not a subroutine call or
@@ -45,12 +47,10 @@ const FAMILY: u32 = 0b101;
 impl From<Call> for u32 {
     /// Convert an instruction into a 32-bit value.
     fn from(inst: Call) -> Self {
-        let imm26 = (inst.imm26 as u32) & ((1 << 26) - 1);
-
         0
         | ((inst.op as u32) << 31)
         | (FAMILY << 26)
-        | imm26
+        | truncate_imm::<_, 26>(inst.imm26)
     }
 }
 
@@ -92,13 +92,13 @@ mod tests {
 
     #[test]
     fn test_b_positive() {
-        let result: u32 = Call::b(256).into();
-        assert_eq!(0x14000100, result);
+        let result: u32 = Call::b((1 << 25) - 1).into();
+        assert_eq!(0x15ffffff, result);
     }
 
     #[test]
     fn test_b_negative() {
-        let result: u32 = Call::b(-256).into();
-        assert_eq!(0x17ffff00, result);
+        let result: u32 = Call::b(-(1 << 25)).into();
+        assert_eq!(0x16000000, result);
     }
 }
diff --git a/yjit/src/asm/arm64/inst/data_reg.rs b/yjit/src/asm/arm64/inst/data_reg.rs
index e2c2723fcf..a742121f1f 100644
--- a/yjit/src/asm/arm64/inst/data_reg.rs
+++ b/yjit/src/asm/arm64/inst/data_reg.rs
@@ -1,4 +1,4 @@
-use super::super::arg::Sf;
+use super::super::arg::{Sf, truncate_uimm};
 
 /// The operation being performed by this instruction.
 enum Op {
@@ -129,8 +129,6 @@ const FAMILY: u32 = 0b0101;
 impl From<DataReg> for u32 {
     /// Convert an instruction into a 32-bit value.
     fn from(inst: DataReg) -> Self {
-        let imm6 = (inst.imm6 as u32) & ((1 << 6) - 1);
-
         0
         | ((inst.sf as u32) << 31)
         | ((inst.op as u32) << 30)
@@ -139,7 +137,7 @@ impl From<DataReg> for u32 {
         | (1 << 24)
         | ((inst.shift as u32) << 22)
         | ((inst.rm as u32) << 16)
-        | (imm6 << 10)
+        | (truncate_uimm::<_, 6>(inst.imm6) << 10)
         | ((inst.rn as u32) << 5)
         | inst.rd as u32
     }
diff --git a/yjit/src/asm/arm64/inst/halfword_imm.rs b/yjit/src/asm/arm64/inst/halfword_imm.rs
index 675e33d4a8..c31d1f8945 100644
--- a/yjit/src/asm/arm64/inst/halfword_imm.rs
+++ b/yjit/src/asm/arm64/inst/halfword_imm.rs
@@ -1,3 +1,5 @@
+use super::super::arg::truncate_imm;
+
 /// Whether this is a load or a store.
 enum Op {
     Load = 1,
@@ -95,11 +97,12 @@ impl From<HalfwordImm> for u32 {
     fn from(inst: HalfwordImm) -> Self {
         let (mut opc, imm) = match inst.index {
             Index::None => {
-                let mut imm12 = ((inst.imm / 2) as u32) & ((1 << 12) - 1);
+                assert_eq!(inst.imm & 1, 0, "immediate offset must be even");
+                let imm12 = truncate_imm::<_, 12>(inst.imm / 2);
                 (0b100, imm12)
             },
             Index::PreIndex | Index::PostIndex => {
-                let mut imm9 = (inst.imm as u32) & ((1 << 9) - 1);
+                let imm9 = truncate_imm::<_, 9>(inst.imm);
                 (0b000, (imm9 << 2) | (inst.index as u32))
             }
         };
diff --git a/yjit/src/asm/arm64/inst/load_literal.rs b/yjit/src/asm/arm64/inst/load_literal.rs
index d2a5d57eea..c5ab09713c 100644
--- a/yjit/src/asm/arm64/inst/load_literal.rs
+++ b/yjit/src/asm/arm64/inst/load_literal.rs
@@ -1,3 +1,5 @@
+use super::super::arg::truncate_imm;
+
 /// The size of the operands being operated on.
 enum Opc {
     Size32 = 0b00,
@@ -50,13 +52,11 @@ const FAMILY: u32 = 0b0100;
 impl From<LoadLiteral> for u32 {
     /// Convert an instruction into a 32-bit value.
     fn from(inst: LoadLiteral) -> Self {
-        let imm19 = (inst.imm19 as u32) & ((1 << 19) - 1);
-
         0
         | ((inst.opc as u32) << 30)
         | (1 << 28)
         | (FAMILY << 25)
-        | (imm19 << 5)
+        | (truncate_imm::<_, 19>(inst.imm19) << 5)
         | (inst.rt as u32)
     }
 }
diff --git a/yjit/src/asm/arm64/inst/load_store.rs b/yjit/src/asm/arm64/inst/load_store.rs
index 80a67c837e..ea42f2d17f 100644
--- a/yjit/src/asm/arm64/inst/load_store.rs
+++ b/yjit/src/asm/arm64/inst/load_store.rs
@@ -1,3 +1,5 @@
+use super::super::arg::truncate_imm;
+
 /// The size of the operands being operated on.
 enum Size {
     Size32 = 0b10,
@@ -110,14 +112,12 @@ const FAMILY: u32 = 0b0100;
 impl From<LoadStore> for u32 {
     /// Convert an instruction into a 32-bit value.
     fn from(inst: LoadStore) -> Self {
-        let imm9 = (inst.imm9 as u32) & ((1 << 9) - 1);
-
         0
         | ((inst.size as u32) << 30)
         | (0b11 << 28)
         | (FAMILY << 25)
         | ((inst.opc as u32) << 22)
-        | (imm9 << 12)
+        | (truncate_imm::<_, 9>(inst.imm9) << 12)
         | ((inst.idx as u32) << 10)
         | ((inst.rn as u32) << 5)
         | (inst.rt as u32)
diff --git a/yjit/src/asm/arm64/inst/logical_reg.rs b/yjit/src/asm/arm64/inst/logical_reg.rs
index 83230ac5b2..a96805c9f9 100644
--- a/yjit/src/asm/arm64/inst/logical_reg.rs
+++ b/yjit/src/asm/arm64/inst/logical_reg.rs
@@ -1,4 +1,4 @@
-use super::super::arg::Sf;
+use super::super::arg::{Sf, truncate_uimm};
 
 /// Whether or not this is a NOT instruction.
 enum N {
@@ -124,8 +124,6 @@ const FAMILY: u32 = 0b0101;
 impl From<LogicalReg> for u32 {
     /// Convert an instruction into a 32-bit value.
     fn from(inst: LogicalReg) -> Self {
-        let imm6 = (inst.imm6 as u32) & ((1 << 6) - 1);
-
         0
         | ((inst.sf as u32) << 31)
         | ((inst.opc as u32) << 29)
@@ -133,7 +131,7 @@ impl From<LogicalReg> for u32 {
         | ((inst.shift as u32) << 22)
         | ((inst.n as u32) << 21)
         | ((inst.rm as u32) << 16)
-        | (imm6 << 10)
+        | (truncate_uimm::<_, 6>(inst.imm6) << 10)
         | ((inst.rn as u32) << 5)
         | inst.rd as u32
     }
diff --git a/yjit/src/asm/arm64/inst/reg_pair.rs b/yjit/src/asm/arm64/inst/reg_pair.rs
index d8fece2ed6..87690e3b4a 100644
--- a/yjit/src/asm/arm64/inst/reg_pair.rs
+++ b/yjit/src/asm/arm64/inst/reg_pair.rs
@@ -1,3 +1,5 @@
+use super::super::arg::truncate_imm;
+
 /// The operation to perform for this instruction.
 enum Opc {
     /// When the registers are 32-bits wide.
@@ -114,18 +116,12 @@ const FAMILY: u32 = 0b0100;
 impl From<RegisterPair> for u32 {
     /// Convert an instruction into a 32-bit value.
     fn from(inst: RegisterPair) -> Self {
-        let mut imm7 = (inst.imm7 as u32) & ((1 << 7) - 1);
-
-        if inst.imm7 < 0 {
-            imm7 |= 1 << 6;
-        }
-
         0
         | ((inst.opc as u32) << 30)
         | (1 << 29)
         | (FAMILY << 25)
         | ((inst.index as u32) << 22)
-        | (imm7 << 15)
+        | (truncate_imm::<_, 7>(inst.imm7) << 15)
         | ((inst.rt2 as u32) << 10)
         | ((inst.rn as u32) << 5)
         | (inst.rt1 as u32)
diff --git a/yjit/src/asm/arm64/inst/sbfm.rs b/yjit/src/asm/arm64/inst/sbfm.rs
index 6f69e58043..8602998980 100644
--- a/yjit/src/asm/arm64/inst/sbfm.rs
+++ b/yjit/src/asm/arm64/inst/sbfm.rs
@@ -1,4 +1,4 @@
-use super::super::arg::Sf;
+use super::super::arg::{Sf, truncate_uimm};
 
 /// The struct that represents an A64 signed bitfield move instruction that can
 /// be encoded.
@@ -56,16 +56,13 @@ const FAMILY: u32 = 0b1001;
 impl From<SBFM> for u32 {
     /// Convert an instruction into a 32-bit value.
     fn from(inst: SBFM) -> Self {
-        let immr = (inst.immr as u32) & ((1 << 6) - 1);
-        let imms = (inst.imms as u32) & ((1 << 6) - 1);
-
         0
         | ((inst.sf as u32) << 31)
         | (FAMILY << 25)
         | (1 << 24)
         | ((inst.n as u32) << 22)
-        | (immr << 16)
-        | (imms << 10)
+        | (truncate_uimm::<_, 6>(inst.immr) << 16)
+        | (truncate_uimm::<_, 6>(inst.imms) << 10)
         | ((inst.rn as u32) << 5)
         | inst.rd as u32
     }
diff --git a/yjit/src/asm/arm64/inst/test_bit.rs b/yjit/src/asm/arm64/inst/test_bit.rs
index 1961e65949..c57a05ad2b 100644
--- a/yjit/src/asm/arm64/inst/test_bit.rs
+++ b/yjit/src/asm/arm64/inst/test_bit.rs
@@ -1,3 +1,5 @@
+use super::super::arg::truncate_imm;
+
 /// The upper bit of the bit number to test.
 #[derive(Debug)]
 enum B5 {
@@ -77,11 +79,7 @@ impl From<TestBit> for u32 {
     /// Convert an instruction into a 32-bit value.
     fn from(inst: TestBit) -> Self {
         let b40 = (inst.b40 & 0b11111) as u32;
-        let mut imm14 = (inst.imm14 & ((1 << 13) - 1)) as u32;
-
-        if inst.imm14 < 0 {
-            imm14 |= (1 << 13);
-        }
+        let imm14 = truncate_imm::<_, 14>(inst.imm14);
 
         0
         | ((inst.b5 as u32) << 31)