Merge pull request RustPython#4269 from afonso360/jit-more-ops

Implement a few more ops in the JIT

Author: youknowone

jit/src/instructions.rs

@@ -8,6 +8,12 @@ use std::collections::HashMap;
 
 use super::{JitCompileError, JitSig, JitType};
 
+#[repr(u16)]
+enum CustomTrapCode {
+    /// Raised when shifting by a negative number
+    NegativeShiftCount = 0,
+}
+
 #[derive(Clone)]
 struct Local {
     var: Variable,
@@ -340,64 +346,75 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> {
                 // the rhs is popped off first
                 let b = self.stack.pop().ok_or(JitCompileError::BadBytecode)?;
                 let a = self.stack.pop().ok_or(JitCompileError::BadBytecode)?;
-                match (a.ty, b.ty) {
-                    (JitType::Int, JitType::Int) => match op {
-                        BinaryOperator::Add => {
-                            let (out, carry) = self.builder.ins().iadd_ifcout(a.val, b.val);
-                            self.builder.ins().trapif(
-                                IntCC::Overflow,
-                                carry,
-                                TrapCode::IntegerOverflow,
-                            );
-                            self.stack.push(JitValue {
-                                val: out,
-                                ty: JitType::Int,
-                            });
-                            Ok(())
-                        }
-                        BinaryOperator::Subtract => {
-                            let out = self.compile_sub(a.val, b.val);
-                            self.stack.push(JitValue {
-                                val: out,
-                                ty: JitType::Int,
-                            });
-                            Ok(())
-                        }
-                        _ => Err(JitCompileError::NotSupported),
-                    },
-                    (JitType::Float, JitType::Float) => match op {
-                        BinaryOperator::Add => {
-                            self.stack.push(JitValue {
-                                val: self.builder.ins().fadd(a.val, b.val),
-                                ty: JitType::Float,
-                            });
-                            Ok(())
-                        }
-                        BinaryOperator::Subtract => {
-                            self.stack.push(JitValue {
-                                val: self.builder.ins().fsub(a.val, b.val),
-                                ty: JitType::Float,
-                            });
-                            Ok(())
-                        }
-                        BinaryOperator::Multiply => {
-                            self.stack.push(JitValue {
-                                val: self.builder.ins().fmul(a.val, b.val),
-                                ty: JitType::Float,
-                            });
-                            Ok(())
-                        }
-                        BinaryOperator::Divide => {
-                            self.stack.push(JitValue {
-                                val: self.builder.ins().fdiv(a.val, b.val),
-                                ty: JitType::Float,
-                            });
-                            Ok(())
-                        }
-                        _ => Err(JitCompileError::NotSupported),
-                    },
-                    _ => Err(JitCompileError::NotSupported),
-                }
+                let (val, ty) = match (op, a.ty, b.ty) {
+                    (BinaryOperator::Add, JitType::Int, JitType::Int) => {
+                        let (out, carry) = self.builder.ins().iadd_ifcout(a.val, b.val);
+                        self.builder.ins().trapif(
+                            IntCC::Overflow,
+                            carry,
+                            TrapCode::IntegerOverflow,
+                        );
+                        (out, JitType::Int)
+                    }
+                    (BinaryOperator::Subtract, JitType::Int, JitType::Int) => {
+                        (self.compile_sub(a.val, b.val), JitType::Int)
+                    }
+                    (BinaryOperator::FloorDivide, JitType::Int, JitType::Int) => {
+                        (self.builder.ins().sdiv(a.val, b.val), JitType::Int)
+                    }
+                    (BinaryOperator::Modulo, JitType::Int, JitType::Int) => {
+                        (self.builder.ins().srem(a.val, b.val), JitType::Int)
+                    }
+                    (
+                        BinaryOperator::Lshift | BinaryOperator::Rshift,
+                        JitType::Int,
+                        JitType::Int,
+                    ) => {
+                        // Shifts throw an exception if we have a negative shift count
+                        // Remove all bits except the sign bit, and trap if its 1 (i.e. negative).
+                        let sign = self.builder.ins().ushr_imm(b.val, 63);
+                        self.builder.ins().trapnz(
+                            sign,
+                            TrapCode::User(CustomTrapCode::NegativeShiftCount as u16),
+                        );
+
+                        let out = if *op == BinaryOperator::Lshift {
+                            self.builder.ins().ishl(a.val, b.val)
+                        } else {
+                            self.builder.ins().sshr(a.val, b.val)
+                        };
+
+                        (out, JitType::Int)
+                    }
+                    (BinaryOperator::And, JitType::Int, JitType::Int) => {
+                        (self.builder.ins().band(a.val, b.val), JitType::Int)
+                    }
+                    (BinaryOperator::Or, JitType::Int, JitType::Int) => {
+                        (self.builder.ins().bor(a.val, b.val), JitType::Int)
+                    }
+                    (BinaryOperator::Xor, JitType::Int, JitType::Int) => {
+                        (self.builder.ins().bxor(a.val, b.val), JitType::Int)
+                    }
+
+                    // Floats
+                    (BinaryOperator::Add, JitType::Float, JitType::Float) => {
+                        (self.builder.ins().fadd(a.val, b.val), JitType::Float)
+                    }
+                    (BinaryOperator::Subtract, JitType::Float, JitType::Float) => {
+                        (self.builder.ins().fsub(a.val, b.val), JitType::Float)
+                    }
+                    (BinaryOperator::Multiply, JitType::Float, JitType::Float) => {
+                        (self.builder.ins().fmul(a.val, b.val), JitType::Float)
+                    }
+                    (BinaryOperator::Divide, JitType::Float, JitType::Float) => {
+                        (self.builder.ins().fdiv(a.val, b.val), JitType::Float)
+                    }
+                    _ => return Err(JitCompileError::NotSupported),
+                };
+
+                self.stack.push(JitValue { val, ty });
+
+                Ok(())
             }
             Instruction::SetupLoop { .. } | Instruction::PopBlock => {
                 // TODO: block support

jit/tests/int_tests.rs

@@ -23,6 +23,111 @@ fn test_sub() {
     assert_eq!(sub(-3, -10), Ok(7));
 }
 
+#[test]
+fn test_floor_div() {
+    let floor_div = jit_function! { floor_div(a:i64, b:i64) -> i64 => r##"
+        def floor_div(a: int, b: int):
+            return a // b
+    "## };
+
+    assert_eq!(floor_div(5, 10), Ok(0));
+    assert_eq!(floor_div(5, 2), Ok(2));
+    assert_eq!(floor_div(12, 10), Ok(1));
+    assert_eq!(floor_div(7, 10), Ok(0));
+    assert_eq!(floor_div(-3, -1), Ok(3));
+    assert_eq!(floor_div(-3, 1), Ok(-3));
+}
+
+#[test]
+fn test_mod() {
+    let modulo = jit_function! { modulo(a:i64, b:i64) -> i64 => r##"
+        def modulo(a: int, b: int):
+            return a % b
+    "## };
+
+    assert_eq!(modulo(5, 10), Ok(5));
+    assert_eq!(modulo(5, 2), Ok(1));
+    assert_eq!(modulo(12, 10), Ok(2));
+    assert_eq!(modulo(7, 10), Ok(7));
+    assert_eq!(modulo(-3, 1), Ok(0));
+    assert_eq!(modulo(-5, 10), Ok(-5));
+}
+
+#[test]
+fn test_lshift() {
+    let lshift = jit_function! { lshift(a:i64, b:i64) -> i64 => r##"
+        def lshift(a: int, b: int):
+            return a << b
+    "## };
+
+    assert_eq!(lshift(5, 10), Ok(5120));
+    assert_eq!(lshift(5, 2), Ok(20));
+    assert_eq!(lshift(12, 10), Ok(12288));
+    assert_eq!(lshift(7, 10), Ok(7168));
+    assert_eq!(lshift(-3, 1), Ok(-6));
+    assert_eq!(lshift(-10, 2), Ok(-40));
+}
+
+#[test]
+fn test_rshift() {
+    let rshift = jit_function! { rshift(a:i64, b:i64) -> i64 => r##"
+        def rshift(a: int, b: int):
+            return a >> b
+    "## };
+
+    assert_eq!(rshift(5120, 10), Ok(5));
+    assert_eq!(rshift(20, 2), Ok(5));
+    assert_eq!(rshift(12288, 10), Ok(12));
+    assert_eq!(rshift(7168, 10), Ok(7));
+    assert_eq!(rshift(-3, 1), Ok(-2));
+    assert_eq!(rshift(-10, 2), Ok(-3));
+}
+
+#[test]
+fn test_and() {
+    let bitand = jit_function! { bitand(a:i64, b:i64) -> i64 => r##"
+        def bitand(a: int, b: int):
+            return a & b
+    "## };
+
+    assert_eq!(bitand(5120, 10), Ok(0));
+    assert_eq!(bitand(20, 16), Ok(16));
+    assert_eq!(bitand(12488, 4249), Ok(4232));
+    assert_eq!(bitand(7168, 2), Ok(0));
+    assert_eq!(bitand(-3, 1), Ok(1));
+    assert_eq!(bitand(-10, 2), Ok(2));
+}
+
+#[test]
+fn test_or() {
+    let bitor = jit_function! { bitor(a:i64, b:i64) -> i64 => r##"
+        def bitor(a: int, b: int):
+            return a | b
+    "## };
+
+    assert_eq!(bitor(5120, 10), Ok(5130));
+    assert_eq!(bitor(20, 16), Ok(20));
+    assert_eq!(bitor(12488, 4249), Ok(12505));
+    assert_eq!(bitor(7168, 2), Ok(7170));
+    assert_eq!(bitor(-3, 1), Ok(-3));
+    assert_eq!(bitor(-10, 2), Ok(-10));
+}
+
+#[test]
+fn test_xor() {
+    let bitxor = jit_function! { bitxor(a:i64, b:i64) -> i64 => r##"
+        def bitxor(a: int, b: int):
+            return a ^ b
+    "## };
+
+    assert_eq!(bitxor(5120, 10), Ok(5130));
+    assert_eq!(bitxor(20, 16), Ok(4));
+    assert_eq!(bitxor(12488, 4249), Ok(8273));
+    assert_eq!(bitxor(7168, 2), Ok(7170));
+    assert_eq!(bitxor(-3, 1), Ok(-4));
+    assert_eq!(bitxor(-10, 2), Ok(-12));
+}
+
 #[test]
 fn test_eq() {
     let eq = jit_function! { eq(a:i64, b:i64) -> i64 => r##"