Merge pull request #900 from youknowone/int-pow

Fix int.__pow__

tests/snippets/ints.py

@@ -37,6 +37,9 @@ assert (2).__mul__(1) == 2
 assert (2).__rmul__(1) == 2
 assert (2).__truediv__(1) == 2.0
 assert (2).__rtruediv__(1) == 0.5
+assert (2).__pow__(3) == 8
+assert (10).__pow__(-1) == 0.1
+assert (2).__rpow__(3) == 9
 
 # real/imag attributes
 assert (1).real == 1
@@ -58,6 +61,8 @@ assert (2).__mul__(1.0) == NotImplemented
 assert (2).__rmul__(1.0) == NotImplemented
 assert (2).__truediv__(1.0) == NotImplemented
 assert (2).__rtruediv__(1.0) == NotImplemented
+assert (2).__pow__(3.0) == NotImplemented
+assert (2).__rpow__(3.0) == NotImplemented
 
 
 assert int() == 0

tests/snippets/math_basics.py

@@ -14,7 +14,7 @@ assert a ** 3 == 64
 assert a * 3 == 12
 assert a / 2 == 2
 assert 2 == a / 2
-# assert a % 3 == 1
+assert a % 3 == 1
 assert a - 3 == 1
 assert -a == -4
 assert +a == 4
@@ -38,3 +38,34 @@ assert_raises(
     OverflowError,
     lambda: round(-float('inf')),
     'OverflowError: cannot convert float NaN to integer')
+
+assert pow(0, 0) == 1
+assert pow(2, 2) == 4
+assert pow(1, 2.0) == 1.0
+assert pow(2.0, 1) == 2.0
+assert pow(0, 10**1000) == 0
+assert pow(1, 10**1000) == 1
+assert pow(-1, 10**1000+1) == -1
+assert pow(-1, 10**1000) == 1
+
+assert pow(2, 4, 5) == 1
+assert_raises(
+    TypeError,
+    lambda: pow(2, 4, 5.0),
+    'pow() 3rd argument not allowed unless all arguments are integers')
+assert_raises(
+    TypeError,
+    lambda: pow(2, 4.0, 5),
+    'pow() 3rd argument not allowed unless all arguments are integers')
+assert_raises(
+    TypeError,
+    lambda: pow(2.0, 4, 5),
+    'pow() 3rd argument not allowed unless all arguments are integers')
+assert_raises(
+    ValueError,
+    lambda: pow(2, -1, 5),
+    'pow() 2nd argument cannot be negative when 3rd argument specified')
+assert_raises(
+    ValueError,
+    lambda: pow(2, 2, 0),
+    'pow() 3rd argument cannot be 0')

vm/src/builtins.rs

@@ -6,7 +6,8 @@ use std::char;
 use std::io::{self, Write};
 use std::path::PathBuf;
 
-use num_traits::Signed;
+use num_bigint::Sign;
+use num_traits::{Signed, Zero};
 
 use crate::compile;
 use crate::import::import_module;
@@ -529,25 +530,33 @@ fn builtin_pow(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult {
         required = [(x, None), (y, None)],
         optional = [(mod_value, Some(vm.ctx.int_type()))]
     );
-    let pow_method_name = "__pow__";
-    let result = match vm.get_method(x.clone(), pow_method_name) {
-        Ok(attrib) => vm.invoke(attrib, vec![y.clone()]),
-        Err(..) => Err(vm.new_type_error("unsupported operand type(s) for pow".to_string())),
-    };
-    //Check if the 3rd argument is defined and perform modulus on the result
-    //this should be optimized in the future to perform a "power-mod" algorithm in
-    //order to improve performance
+
     match mod_value {
-        Some(mod_value) => {
-            let mod_method_name = "__mod__";
-            match vm.get_method(result.expect("result not defined").clone(), mod_method_name) {
-                Ok(value) => vm.invoke(value, vec![mod_value.clone()]),
-                Err(..) => {
-                    Err(vm.new_type_error("unsupported operand type(s) for mod".to_string()))
-                }
+        None => vm.call_or_reflection(x.clone(), y.clone(), "__pow__", "__rpow__", |vm, x, y| {
+            Err(vm.new_unsupported_operand_error(x, y, "pow"))
+        }),
+        Some(m) => {
+            // Check if the 3rd argument is defined and perform modulus on the result
+            if !(objtype::isinstance(x, &vm.ctx.int_type())
+                && objtype::isinstance(y, &vm.ctx.int_type()))
+            {
+                return Err(vm.new_type_error(
+                    "pow() 3rd argument not allowed unless all arguments are integers".to_string(),
+                ));
             }
+            let y = objint::get_value(y);
+            if y.sign() == Sign::Minus {
+                return Err(vm.new_value_error(
+                    "pow() 2nd argument cannot be negative when 3rd argument specified".to_string(),
+                ));
+            }
+            let m = objint::get_value(m);
+            if m.is_zero() {
+                return Err(vm.new_value_error("pow() 3rd argument cannot be 0".to_string()));
+            }
+            let x = objint::get_value(x);
+            Ok(vm.new_int(x.modpow(&y, &m)))
         }
-        None => result,
     }
 }
 

vm/src/obj/objint.rs

@@ -3,7 +3,7 @@ use std::hash::{Hash, Hasher};
 
 use num_bigint::BigInt;
 use num_integer::Integer;
-use num_traits::{Pow, Signed, ToPrimitive, Zero};
+use num_traits::{One, Pow, Signed, ToPrimitive, Zero};
 
 use crate::format::FormatSpec;
 use crate::function::{OptionalArg, PyFuncArgs};
@@ -13,7 +13,6 @@ use crate::pyobject::{
 };
 use crate::vm::VirtualMachine;
 
-use super::objfloat::{self, PyFloat};
 use super::objstr::{PyString, PyStringRef};
 use super::objtype;
 use crate::obj::objtype::PyClassRef;
@@ -111,6 +110,30 @@ impl_try_from_object_int!(
     (u64, to_u64),
 );
 
+fn inner_pow(int1: &PyInt, int2: &PyInt, vm: &VirtualMachine) -> PyResult {
+    Ok(if int2.value.is_negative() {
+        let v1 = int1.float(vm)?;
+        let v2 = int2.float(vm)?;
+        vm.ctx.new_float(v1.pow(v2))
+    } else {
+        if let Some(v2) = int2.value.to_u64() {
+            vm.ctx.new_int(int1.value.pow(v2))
+        } else if int1.value.is_one() || int1.value.is_zero() {
+            vm.ctx.new_int(int1.value.clone())
+        } else if int1.value == BigInt::from(-1) {
+            if int2.value.is_odd() {
+                vm.ctx.new_int(-1)
+            } else {
+                vm.ctx.new_int(1)
+            }
+        } else {
+            // missing feature: BigInt exp
+            // practically, exp over u64 is not possible to calculate anyway
+            vm.ctx.not_implemented()
+        }
+    })
+}
+
 #[pyimpl]
 impl PyInt {
     #[pymethod(name = "__eq__")]
@@ -323,15 +346,22 @@ impl PyInt {
     }
 
     #[pymethod(name = "__pow__")]
-    fn pow(&self, other: PyObjectRef, vm: &VirtualMachine) -> PyObjectRef {
+    fn pow(&self, other: PyObjectRef, vm: &VirtualMachine) -> PyResult {
         if objtype::isinstance(&other, &vm.ctx.int_type()) {
-            let v2 = get_value(&other).to_u32().unwrap();
-            vm.ctx.new_int(self.value.pow(v2))
-        } else if objtype::isinstance(&other, &vm.ctx.float_type()) {
-            let v2 = objfloat::get_value(&other);
-            vm.ctx.new_float((self.value.to_f64().unwrap()).powf(v2))
+            let other = other.payload::<PyInt>().unwrap();
+            inner_pow(self, &other, vm)
         } else {
-            vm.ctx.not_implemented()
+            Ok(vm.ctx.not_implemented())
+        }
+    }
+
+    #[pymethod(name = "__rpow__")]
+    fn rpow(&self, other: PyObjectRef, vm: &VirtualMachine) -> PyResult {
+        if objtype::isinstance(&other, &vm.ctx.int_type()) {
+            let other = other.payload::<PyInt>().unwrap();
+            inner_pow(&other, self, vm)
+        } else {
+            Ok(vm.ctx.not_implemented())
         }
     }
 
@@ -403,10 +433,9 @@ impl PyInt {
     }
 
     #[pymethod(name = "__float__")]
-    fn float(&self, vm: &VirtualMachine) -> PyResult<PyFloat> {
+    fn float(&self, vm: &VirtualMachine) -> PyResult<f64> {
         self.value
             .to_f64()
-            .map(PyFloat::from)
             .ok_or_else(|| vm.new_overflow_error("int too large to convert to float".to_string()))
     }
 
@@ -544,9 +573,7 @@ pub fn get_value(obj: &PyObjectRef) -> &BigInt {
 }
 
 pub fn get_float_value(obj: &PyObjectRef, vm: &VirtualMachine) -> PyResult<f64> {
-    get_value(obj).to_f64().ok_or_else(|| {
-        vm.new_overflow_error("OverflowError: int too large to convert to float".to_string())
-    })
+    obj.payload::<PyInt>().unwrap().float(vm)
 }
 
 #[inline]