Merge pull request RustPython#5086 from MannarAmuthan/float-int-ops

youknowone · web-flow · commit f1b36233b6a2 · 2023-10-07T21:15:04.000+09:00
[JIT]  Binary operations for integers and floating-point numbers allowing mixed type calculations
diff --git a/jit/src/instructions.rs b/jit/src/instructions.rs
@@ -381,6 +381,10 @@ 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)?;
+
+                let a_type = a.to_jit_type();
+                let b_type = b.to_jit_type();
+
                 let val = match (op, a, b) {
                     (BinaryOperator::Add, JitValue::Int(a), JitValue::Int(b)) => {
                         let (out, carry) = self.builder.ins().iadd_ifcout(a, b);
@@ -443,6 +447,36 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> {
                     (BinaryOperator::Divide, JitValue::Float(a), JitValue::Float(b)) => {
                         JitValue::Float(self.builder.ins().fdiv(a, b))
                     }
+
+                    // Floats and Integers
+                    (_, JitValue::Int(a), JitValue::Float(b))
+                    | (_, JitValue::Float(a), JitValue::Int(b)) => {
+                        let operand_one = match a_type.unwrap() {
+                            JitType::Int => self.builder.ins().fcvt_from_sint(types::F64, a),
+                            _ => a,
+                        };
+
+                        let operand_two = match b_type.unwrap() {
+                            JitType::Int => self.builder.ins().fcvt_from_sint(types::F64, b),
+                            _ => b,
+                        };
+
+                        match op {
+                            BinaryOperator::Add => {
+                                JitValue::Float(self.builder.ins().fadd(operand_one, operand_two))
+                            }
+                            BinaryOperator::Subtract => {
+                                JitValue::Float(self.builder.ins().fsub(operand_one, operand_two))
+                            }
+                            BinaryOperator::Multiply => {
+                                JitValue::Float(self.builder.ins().fmul(operand_one, operand_two))
+                            }
+                            BinaryOperator::Divide => {
+                                JitValue::Float(self.builder.ins().fdiv(operand_one, operand_two))
+                            }
+                            _ => return Err(JitCompileError::NotSupported),
+                        }
+                    }
                     _ => return Err(JitCompileError::NotSupported),
                 };
                 self.stack.push(val);
diff --git a/jit/tests/float_tests.rs b/jit/tests/float_tests.rs
@@ -32,6 +32,18 @@ fn test_add() {
     assert_eq!(add(1.0, f64::NEG_INFINITY), Ok(f64::NEG_INFINITY));
 }
 
+#[test]
+fn test_add_with_integer() {
+    let add = jit_function! { add(a:f64, b:i64) -> f64 => r##"
+        def add(a: float, b: int):
+            return a + b
+    "## };
+
+    assert_approx_eq!(add(5.5, 10), Ok(15.5));
+    assert_approx_eq!(add(-4.6, 7), Ok(2.4));
+    assert_approx_eq!(add(-5.2, -3), Ok(-8.2));
+}
+
 #[test]
 fn test_sub() {
     let sub = jit_function! { sub(a:f64, b:f64) -> f64 => r##"
@@ -49,6 +61,19 @@ fn test_sub() {
     assert_eq!(sub(1.0, f64::INFINITY), Ok(f64::NEG_INFINITY));
 }
 
+#[test]
+fn test_sub_with_integer() {
+    let sub = jit_function! { sub(a:i64, b:f64) -> f64 => r##"
+        def sub(a: int, b: float):
+            return a - b
+    "## };
+
+    assert_approx_eq!(sub(5, 3.6), Ok(1.4));
+    assert_approx_eq!(sub(3, -4.2), Ok(7.2));
+    assert_approx_eq!(sub(-2, 1.3), Ok(-3.3));
+    assert_approx_eq!(sub(-3, -1.3), Ok(-1.7));
+}
+
 #[test]
 fn test_mul() {
     let mul = jit_function! { mul(a:f64, b:f64) -> f64 => r##"
@@ -70,6 +95,21 @@ fn test_mul() {
     assert_eq!(mul(f64::NEG_INFINITY, f64::INFINITY), Ok(f64::NEG_INFINITY));
 }
 
+#[test]
+fn test_mul_with_integer() {
+    let mul = jit_function! { mul(a:f64, b:i64) -> f64 => r##"
+        def mul(a: float, b: int):
+            return a * b
+    "## };
+
+    assert_approx_eq!(mul(5.2, 2), Ok(10.4));
+    assert_approx_eq!(mul(3.4, -1), Ok(-3.4));
+    assert_bits_eq!(mul(1.0, 0), Ok(0.0f64));
+    assert_bits_eq!(mul(-0.0, 1), Ok(-0.0f64));
+    assert_bits_eq!(mul(0.0, -1), Ok(-0.0f64));
+    assert_bits_eq!(mul(-0.0, -1), Ok(0.0f64));
+}
+
 #[test]
 fn test_div() {
     let div = jit_function! { div(a:f64, b:f64) -> f64 => r##"
@@ -91,6 +131,23 @@ fn test_div() {
     assert_bits_eq!(div(-1.0, f64::INFINITY), Ok(-0.0f64));
 }
 
+#[test]
+fn test_div_with_integer() {
+    let div = jit_function! { div(a:f64, b:i64) -> f64 => r##"
+        def div(a: float, b: int):
+            return a / b
+    "## };
+
+    assert_approx_eq!(div(5.2, 2), Ok(2.6));
+    assert_approx_eq!(div(3.4, -1), Ok(-3.4));
+    assert_eq!(div(1.0, 0), Ok(f64::INFINITY));
+    assert_eq!(div(1.0, -0), Ok(f64::INFINITY));
+    assert_eq!(div(-1.0, 0), Ok(f64::NEG_INFINITY));
+    assert_eq!(div(-1.0, -0), Ok(f64::NEG_INFINITY));
+    assert_eq!(div(f64::INFINITY, 2), Ok(f64::INFINITY));
+    assert_eq!(div(f64::NEG_INFINITY, 3), Ok(f64::NEG_INFINITY));
+}
+
 #[test]
 fn test_if_bool() {
     let if_bool = jit_function! { if_bool(a:f64) -> i64 => r##"
